Résumé Technique : Une Défense du Modèle Quadratique

Énoncé du Problème

La théorie de l'optimisation pour les grands modèles de langage (LLM) repose souvent sur des modèles idéalisés pour naviguer dans la complexité des paysages de perte non convexes. Cependant, un écart important existe entre la tractabilité théorique et la précision empirique : les modèles mathématiquement simples échouent souvent à décrire les dynamiques réelles, tandis que les modèles capturant des structures de problèmes spécifiques sont fréquemment intraitables sur le plan théorique. La question centrale abordée par ce travail est la suivante : Dans quelle mesure le modèle d'optimisation le plus simple possible — le modèle quadratique — capture-t-il réellement la dynamique d'entraînement des LLM ?

Méthodologie

Les auteurs mènent un « test de résistance » rigoureux du modèle quadratique en utilisant un transformeur de 150 millions de paramètres entraîné sur 3 milliards de tokens de FineWeb. La méthodologie comprend trois composantes principales :

1. Approximation Quadratique Locale (Section 3)

Les auteurs évaluent le pouvoir prédictif des expansions de Taylor locales à des points de contrôle intermédiaires (tous les 10 % de l'entraînement). Ils définissent deux pertes de substitution (proxy losses) :

• Prox (Linéarisé) : Linéarise le modèle $f_\theta$ autour d'un point de référence $\theta_{ref}$ tout en conservant l'objectif de cross-entropy original. Cela équivaut à une régression logistique multi-classe sous le régime du Noyau de Tangente Neurale (NTK).
• Quad (Quadratique) : Développe en série de Taylor la perte Prox au second ordre, résultant en un objectif quadratique convexe. Crucialement, ils utilisent la matrice de Gauss-Newton ($G$) plutôt que l'Hessien complet ($H$) pour garantir que l'objectif quadratique reste convexe (semi-défini positif), évitant ainsi l'indéterminance dans les directions de courbure négative.

Ils entraînent ces modèles de substitution pendant 10 % du budget total de tokens à partir de chaque point de contrôle et comparent les trajectoires d'optimisation résultantes avec les dynamiques d'entraînement réelles du LLM.

2. Estimation du Spectre de l'Hessien (Section 4)

Pour analyser la structure de ces problèmes quadratiques locaux, les auteurs estiment les spectres de l'Hessien et des matrices de Gauss-Newton (à la fois bruts et préconditionnés par Adam) le long de la trajectoire d'entraînement.

• Technique : Ils emploient la Quadrature de Lanczos Stochastique (SLQ) avec des sondes extrêmement profondes ($m=1200$) et des bandes d'erreur rigoureuses (bornes de Gauss-Radau).
• Échelle : Les sondes sont exécutées sur des points de contrôle à 10 %, 50 % et 100 % de l'entraînement à travers des tailles de lots $B \in \{1, 64, 1024\}$.
• Analyse : Ils examinent la décroissance spectrale (lois de puissance), la distribution des vecteurs propres à travers les couches du réseau, et l'alignement entre le gradient et les vecteurs propres (conditions de source).

3. Analyse de la Stabilité Limite (Edge of Stability - EOS) (Section 5)

Les auteurs testent empiriquement si l'entraînement des LLM opère à la « limite de stabilité ».

• Configuration : Ils figent le préconditionneur Adam et exécutent un SGD préconditionné sur le substitut quadratique local.
• Perturbation : Ils font varier les multiplicateurs du taux d'apprentissage ($\eta$) et de la taille du lot ($B$) sur une grille pour déterminer la frontière de stabilité.
• Métrique : Une exécution est considérée comme divergente si la perte dépasse un seuil (100). Ils distinguent la Stabilité de Moyenne Carrée (sensible aux queues lourdes) et la Stabilité de Lyapunov (comportement typique).

Résultats Clés

1. Pouvoir Prédictif des Modèles Quadratiques

• Fenêtre d'Accord : Les modèles quadratiques et linéarisés locaux prédisent avec précision la dynamique d'optimisation réelle pour des fenêtres durant jusqu'à 10 % du budget total d'entraînement.
• Précision Temporelle : La qualité de l'approximation est nettement meilleure dans les phases milieu et tardives de l'entraînement par rapport aux phases initiales.
• Prox vs Quad : Le modèle linéarisé (Prox) correspond généralement mieux à la perte réelle que le modèle quadratique (Quad) en début d'entraînement, tandis que Quad ne converge vers une haute précision qu'après 60 % de l'entraînement.

2. Structure du Spectre de l'Hessien

• Division Tête-Queue : Le spectre présente une transition nette autour de l'indice $i \approx 8192$ (la taille du vocabulaire). La « tête » (valeurs propres supérieures) est dominée par la couche d'unembedding, ce qui est cohérent avec le phénomène de « effondrement neural » (neural collapse).
• Queue Universelle : La « queue » du spectre (indices $> 10^5$) suit une loi de puissance qui est universelle : elle est identique pour différentes tailles de lots et n'est pas affectée par le préconditionneur Adam.
• Exposants :
  • Exposant de Capacité ($\alpha$) : Estimé à $\alpha \approx 0,96$ pour la queue de Gauss-Newton et $\alpha \approx 0,65$ pour la queue de l'Hessien.
  • Exposant de Source ($\beta$) : La condition de source (alignement du gradient) donne $\beta \in [0,1, 0,2]$ pour Gauss-Newton et $\beta \in [0,3, 0,6]$ pour l'Hessien. Cela contredit l'hypothèse courante de la littérature $\beta > 1$, suggérant que le problème se comporte comme s'il était de dimension finie avec une décroissance exponentielle du biais après un temps critique.
• Valeurs Propres Négatives : L'Hessien préconditionné présente de grandes valeurs propres négatives, particulièrement à de grandes tailles de lots, localisées dans la couche de projection ascendante (up-projection) des MLP. Les auteurs émettent l'hypothèse qu'elles sont liées aux « neurones morts » et à l'effet des tokens rares.

3. Dynamiques de la Limite de Stabilité (Edge of Stability)

• Stochastique vs Déterministe :
  • À de petites et moyennes tailles de lots ($B=1, 64$), l'entraînement opère à une limite de stabilité stochastique. La frontière de stabilité suit une ligne diagonale dans l'espace $(\eta, B)$, indiquant que l'instabilité est pilotée par la stochasticité (le bruit).
  • À de grandes tailles de lots ($B=1024$), l'entraînement opère à la limite de stabilité déterministe, caractérisée par un arrêt plat dans la région de stabilité, ce qui est cohérent avec le comportement de la descente de gradient en lot complet (full-batch).
• Proximité de la Frontière : Dans tous les réglages, l'entraînement opère à un facteur de 2 de la frontière de stabilité.

Signification et Revendications

L'article soutient que le modèle quadratique n'est pas seulement un jouet théorique, mais un substitut valide et expérimentalement vérifié pour les dynamiques de pré-entraînement des LLM.

• Tractabilité Théorique : Les résultats suggèrent que l'optimisation complexe des LLM peut être comprise localement comme une séquence de problèmes quadratiques structurés. Cela permet aux chercheurs de tirer parti de la riche boîte à outils de la théorie de l'optimisation quadratique (ex: analyse de stabilité, planification du taux d'apprentissage) pour raisonner sur l'entraînement des LLM.
• Révision des Hypothèses : Les résultats remettent en question les hypothèses standards de la littérature, spécifiquement la sommabilité des conditions de source ($\beta > 1$) et l'universalité de la queue de l'Hessien sous préconditionnement.
• Implications Pratiques : L'identification de la « limite de stabilité stochastique » pour les petits lots offre une compréhension affinée de la raison pour laquelle l'entraînement diverge sous certaines configurations d'hyperparamètres, distinguant les erreurs de discrétisation du bruit stochastique.

Les auteurs concluent que, bien que les paysages de perte globaux soient non convexes, l'approximation quadratique locale offre un cadre puissant et tractable pour analyser et concevoir des pipelines d'entraînement pour les modèles à grande échelle.