Google DeepMind identifie un défaut du softmax : le gradient s'annule dès un écart de 16,6

Le 22 octobre 2026, Jonathan Barron, Dan Gnanapragasam et Michael Tschannen, de Google DeepMind, ont mis en ligne le preprint A Fix for Softmax and Cross-Entropy. Les auteurs analysent deux défauts numériques des implémentations standard du softmax et de la cross-entropy utilisées par PyTorch et JAX. Le premier, Softmax Collapse, annule en float32 le gradient de la bonne classe lorsque l'écart avec les concurrents dépasse environ 16,6. Le second, Softmax Cancellation, ajoute un gradient non nul en distillation, même lorsque l'élève reproduit exactement le professeur. Les deux défauts se corrigent en réarrangeant les termes de la formule, sans modifier les mathématiques.

Où float32 perd le petit nombre

Le schéma standard est déjà protégé contre le dépassement de capacité : on soustrait le maximum des logits, et l'exponentielle n'explose pas. Mais cette protection contre l'overflow ne garantit pas la précision des petites différences. Prenons les logits [20, 0] en float32. La probabilité de la seconde classe vaut 2,06·10⁻⁹, un nombre parfaitement représentable. La probabilité de la première est arrondie à exactement 1, car un ajout de 10⁻⁹ à l'unité ne tient pas dans sept chiffres significatifs. Après soustraction de 1, le gradient de la bonne classe est nul, alors qu'il devrait valoir environ −2,06·10⁻⁹. Les gradients des concurrents, eux, subsistent, et la somme des dérivées cesse d'être nulle.

Si l'on note m l'écart de la bonne réponse et r = e^(−m), la formule habituelle calcule 1/(1+r) − 1. La version corrigée calcule −r/(1+r). C'est la même fonction, mais la petite quantité figure au numérateur et n'est plus arrondie avec l'unité. Vérification en float32 : pour m = 17, le schéma habituel donne 0, le schéma corrigé donne −4,14·10⁻⁸. Pour m = 20, on obtient −2,06·10⁻⁹, et pour m = 30, −9,36·10⁻¹⁴.

Les seuils dépendent du nombre de bits de la partie fractionnaire du format de calcul :

• bfloat16, 7 bits : seuil d'écart ≈ 5,54 ;
• float16, 10 bits : ≈ 7,62 ;
• float32, 23 bits : ≈ 16,63 ;
• float64, 52 bits : ≈ 36,74.

Ces seuils sont ceux de l'arithmétique de l'opération concernée. Des poids en bfloat16 ne signifient pas que le softmax et la loss sont aussi calculés en bfloat16 : en précision mixte, les opérations sensibles s'exécutent souvent en float32.

L'élève a rejoint le professeur, mais le gradient demeure

Le second défaut apparaît en distillation. L'élève apprend à reproduire les probabilités d'un professeur fixe, et lorsqu'elles coïncident parfaitement, le gradient par rapport aux logits devrait être nul. La cross-entropy reste alors égale à l'entropie du professeur, tandis que le gradient et la divergence KL tombent à zéro.

En pratique, les probabilités du professeur et de l'élève sont souvent calculées par des chemins différents : l'une par division des exponentielles, l'autre par l'exponentielle de log_softmax. Pour les logits [2, 1, 0] en float32, l'écart est d'environ 1,49·10⁻⁸ sur deux composantes, avec une norme de bruit de 2,1·10⁻⁸. Les logits sont identiques, le gradient ne l'est pas.

Les auteurs proposent de comparer d'abord les logits, puis de passer aux probabilités. On calcule la différence entre les logits de l'élève et du professeur, on en soustrait la moyenne pondérée par les probabilités du professeur, puis on évalue la petite correction avec expm1 et log1p. Ces fonctions conservent leur précision pour un argument faible. Lorsque les logits coïncident, on obtient un zéro strict : expm1(0) renvoie 0.

Pourquoi 10⁻⁸ finit par peser sur l'entraînement

Un gradient de l'ordre de 10⁻⁹ paraît inoffensif, mais Adam normalise le pas par une estimation de l'échelle du gradient : la petitesse absolue du signal ne protège donc pas de ses conséquences. Avec Softmax Collapse, le gradient du logit cible a disparu, alors que celui des concurrents subsiste. Ce signal déformé traverse le réseau et parvient à l'optimiseur adaptatif. Les auteurs relient cette asymétrie à l'effet Slingshot : après une période de loss basse, des pics brutaux apparaissent.

Les expériences du preprint ont été menées en float32, sur cinq exécutions avec des seeds différents. Dans la première, une addition modulo 31 avec weight decay λ = 1 et Adam ε = 10⁻⁸. Dans la seconde, (a² + b²) mod 97 avec λ = 0 et ε = 10⁻¹⁶. Dans les deux cas, la correction du schéma numérique supprime les pics et modifie la dynamique d'entraînement qui suit. Ce résultat porte sur des modèles et des tâches précis. Il ne prouve pas que tout saut de loss dans un LLM soit causé par le softmax.

Pour calculer de façon stable la loss avec hard labels, les auteurs utilisent logsumexp sur les concurrents et softplus au lieu d'une addition directe avec 1. L'exemple pédagogique en NumPy de l'analyse ne fonctionne que lorsque la bonne classe est en tête, et ne remplace pas l'autograd.

Que vérifier dans votre code

Le correctif modifie l'ordre des opérations, pas l'architecture : il n'exige ni nouveaux paramètres ni mémoire supplémentaire. Pour les équipes qui entraînent des modèles à loss basse, c'est une vérification peu coûteuse.

• Fournissez au noyau de loss des logits avec un écart de 17 à 30 en float32 et vérifiez si le gradient de la bonne classe s'annule.
• Assurez-vous que, dans votre mode de précision mixte, le softmax et la loss sont bien calculés en float32, et non en bfloat16 ou float16.
• En distillation, donnez à l'élève les logits du professeur. Le gradient doit être exactement nul, et non de l'ordre de 10⁻⁸.
• Comparez le comportement pour différentes valeurs de ε dans Adam : ce paramètre détermine dans quelle mesure le bruit se transforme en pas de mise à jour.
• Si vous observez des pics de loss après un long plateau, vérifiez le schéma numérique avant de modifier les hyperparamètres.

L'implémentation complète pour les grands écarts de logits doit être reprise du preprint des auteurs : les démonstrations succinctes de l'analyse conviennent aux petites différences et ne remplacent pas une fonction prête à l'emploi.