Résumé technique : Grokking ou Dysfonctionnement ? Comment la faible précision provoque des pics de perte en mode fronde

1. Énoncé du problème

Les réseaux de neurones profonds entraînés sans régularisation explicite exhibent souvent un phénomène appelé « Grokking », où les modèles atteignent une généralisation soudaine bien après avoir atteint une précision d'entraînement parfaite. Ce processus est fréquemment accompagné du Mécanisme de Fronde : des instabilités périodiques caractérisées par une croissance rapide, quasi exponentielle, de la norme des paramètres de la dernière couche, suivie de pics abrupts dans la perte d'entraînement.

Alors que la littérature existante attribue ces pics à des dynamiques d'optimisation intrinsèques — telles que le Bord de Stabilité (EOS) ou les interactions entre les échelles de gradient et les optimiseurs adaptatifs — le mécanisme de déclenchement précis reste flou. Cet article examine si ces instabilités sont inhérentes au paysage d'optimisation ou des artefacts de l'arithmétique en précision finie.

2. Méthodologie et cadre théorique

Les auteurs proposent que le Mécanisme de Fronde n'est pas un phénomène d'optimisation intrinsèque, mais le résultat des limites de précision de l'arithmétique à virgule flottante lors du calcul de la perte d'Entropie Croisée (CE).

2.1 Effondrement Softmax (SC)

Le cœur du mécanisme réside dans l'erreur d'absorption inhérente à l'addition en virgule flottante. À mesure que l'entraînement progresse et que le modèle entre dans un régime de haute confiance (précision d'entraînement parfaite), le logit de la classe correcte ($z_r$) devient significativement plus grand que les logits des classes incorrectes ($z_k$).

• Dans le calcul du Softmax, le dénominateur implique une opération log-sum-exp : $Z = \log(\sum \exp(z_k))$.
• Lorsque la marge $z_r - \max_{k \neq r} z_k$ dépasse un seuil déterminé par la précision de la mantisse (par exemple, $\approx 16$ pour float32), les termes plus petits sont absorbés dans le terme plus grand en raison de la précision limitée.
• Par conséquent, la probabilité calculée pour la classe correcte devient exactement 1, et le gradient par rapport au logit de la classe correcte, $\partial L / \partial z_r$, s'arrondit à exactement zéro. Ce phénomène est nommé Effondrement Softmax (SC).

2.2 Rupture de la contrainte somme nulle

En arithmétique exacte, la somme des gradients sur toutes les classes pour un échantillon unique est nulle ($\sum_k \nabla_{W_k} L = 0$), assurant que la moyenne globale des poids du classifieur ($W_G$) reste statique.

• Sous SC, le gradient pour la classe correcte s'annule, tandis que les gradients pour les classes incorrectes restent non nuls (bien que faibles).
• Cela brise la contrainte somme nulle, introduisant une dérive nette dans la mise à jour de la moyenne globale du classifieur $W_G$. Plus précisément, $W_G$ dérive dans la direction de $-\mu_G$, où $\mu_G$ est la moyenne globale des caractéristiques de la couche pénultième.

2.3 Inflation numérique des caractéristiques (NFI)

L'article identifie une boucle de rétroaction positive nommée Inflation Numérique des Caractéristiques (NFI) :

1. La dérive de $W_G$ modifie le flux de gradient vers la couche de caractéristiques.
2. La moyenne des caractéristiques $\mu_G$ reçoit une dérive dans la direction de $-W_G$.
3. Cela crée une dynamique couplée où $W_G$ et $\mu_G$ deviennent anti-parallèles, et leurs normes croissent exponentiellement.
4. À mesure que ces normes croissent, la marge pour certains échantillons (en particulier les valeurs aberrantes) finit par tomber en dessous du seuil d'absorption, provoquant la réapparition brutale du gradient de la classe correcte.

2.4 Le pic de la fronde

La réapparition du gradient de la classe correcte déclenche le pic de perte :

• Avant le pic, le gradient de la classe correcte est nul, et les gradients des classes incorrectes sont minuscules. L'accumulateur du second moment de l'optimiseur Adam ($v_t$) décroît, faisant en sorte que le taux d'apprentissage effectif ($\eta / \sqrt{v_t + \epsilon}$) devienne extrêmement grand.
• Lorsque le gradient réapparaît soudainement, ce taux d'apprentissage effectif massif entraîne une mise à jour de paramètres disproportionnément grande, ramenant la perte à des niveaux de hasard pur (la « Fronde »).

3. Contributions clés

1. Explication théorique de la Fronde : L'article fournit une preuve rigoureuse que le Mécanisme de Fronde est un artefact numérique causé par le calcul de la perte CE en précision finie, le reliant spécifiquement à l'Effondrement Softmax et à la rupture subséquente de la contrainte de gradient somme nulle.
2. Identification de la NFI : Les auteurs définissent l'Inflation Numérique des Caractéristiques (NFI), une boucle de rétroaction entre la moyenne globale du classifieur et la moyenne globale des caractéristiques, qui entraîne la croissance exponentielle des normes de paramètres observée avant les pics de perte.
3. Distinction par rapport aux dynamiques intrinsèques : Le travail démontre que la Fronde est distincte du Bord de Stabilité (EOS). Alors que l'EOS implique des oscillations dues à la courbure de la Hessienne, la Fronde se produit dans des régimes où les valeurs propres de la Hessienne s'annulent (régime d'interpolation) et est éliminée par l'augmentation de la précision.
4. Interventions pratiques : L'article propose et valide des méthodes pour supprimer la NFI, notamment :
   • L'utilisation de float64 pour le calcul de la perte (éliminant les erreurs d'absorption).
   • L'imposition de la contrainte somme nulle par projection de gradient.
   • L'application de Batch Normalization immédiatement avant le classifieur final pour supprimer la composante de moyenne globale.
   • L'augmentation du paramètre $\epsilon$ d'Adam pour plafonner le taux d'apprentissage effectif.

4. Résultats expérimentaux

Les auteurs valident leur théorie sur plusieurs architectures (MLP, VGG11, ResNet18, ViT, Transformers) et ensembles de données (Arithmétique Modulaire, CIFAR-10, FineWeb pour les LLM).

• Sensibilité à la précision : Le passage du calcul de la perte en float64 élimine complètement les pics de Fronde, même lorsque les paramètres du modèle restent en float32. Cela confirme que l'instabilité provient du calcul de la perte, et non des dynamiques du modèle.
• Dépendance architecturale : Le phénomène est observé dans la plupart des architectures mais est notablement absent dans ResNet18. Les auteurs attribuent cela à la capacité de ResNet18 à réduire la masse de probabilité résiduelle ($\epsilon$) suffisamment rapidement pour empêcher la boucle de rétroaction NFI d'entrer dans le régime de croissance exponentielle avant que le modèle ne se stabilise.
• Atténuation :
  • La Normalisation par Couche (LN) appliquée avant le classifieur échoue à atténuer l'effet et peut accélérer les pics en contraignant les normes des caractéristiques.
  • La Normalisation par Lots (BN) appliquée immédiatement avant la dernière couche linéaire élimine avec succès les pics en centrant les caractéristiques et en supprimant la composante de dérive.
  • Le Lissage des étiquettes élimine les pics induits par la précision mais introduit une autre instabilité liée au Bord de Stabilité due à des valeurs propres de Hessienne non nulles.
• Contexte LLM : Dans l'entraînement des Grands Modèles de Langage, bien que les pics de Fronde en lot complet soient rares en raison de la stochasticité des mini-lots, le mécanisme NFI continue de conduire à une croissance anormale des logits. Les auteurs notent que dans la modélisation du langage, les déséquilibres de fréquence des tokens (loi de Zipf) peuvent créer une grande moyenne globale $W_G$ même sans effondrement numérique, conduisant à une dynamique différente où la faible précision peut partiellement supprimer la divergence induite par la fréquence.

5. Signification et affirmations

L'article réinterprète le Mécanisme de Fronde non pas comme un comportement d'optimisation mystérieux, mais comme une dynamique numérique de l'entraînement en précision finie.

• Réinterprétation du flux de gradient : Les résultats mettent en évidence un fossé entre les analyses théoriques de flux de gradient (qui supposent une précision infinie et prédisent souvent une croissance logarithmique des normes) et l'entraînement pratique, où les erreurs numériques peuvent entraîner une croissance exponentielle.
• Stabilité dans l'entraînement en faible précision : Alors que le domaine se tourne vers des précisions inférieures (BF16, FP8, FP4) pour l'entraînement à grande échelle, la compréhension de la NFI est cruciale. Les auteurs soutiennent que les erreurs numériques dans le calcul des log-probabilités représentent un mode de défaillance distinct qui peut déstabiliser l'entraînement, séparé des erreurs de quantification dans la multiplication de matrices ou les puits d'attention.
• Portée modeste : Les auteurs reconnaissent des limites, notant que leur analyse repose sur l'hypothèse du Modèle de Caractéristiques Non Contraintes et examine principalement l'interaction entre la dernière couche et les caractéristiques pénultièmes. Ils ne prétendent pas expliquer toutes les formes d'instabilité d'entraînement, mais ciblent spécifiquement les pics de perte périodiques observés dans l'entraînement à long terme sans régularisation.

En conclusion, l'article affirme que la Fronde est un dysfonctionnement, et non une fonctionnalité, entraîné par les erreurs d'absorption de l'arithmétique à virgule flottante, et offre des explications et des interventions concrètes et testables pour la croissance anormale des paramètres dans les phases tardives de l'entraînement.