The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
Abreu, Vyas, Kakade et al.
Recent efforts to accelerate LLM pretraining have focused on computationally-efficient approximations that exploit second-order structure. This raises a key question for large-scale training: how much performance is forfeited by these approximations? To probe this question, we establish a practical upper bound on iteration complexity by applying full Gauss-Newton (GN) preconditioning to transformer models of up to 150M parameters. Our experiments show that full GN updates yield substantial gains over existing optimizers, achieving a 5.4x reduction in training iterations compared to strong baselines like SOAP and Muon. Furthermore, we find that a precise layerwise GN preconditioner, which ignores cross-layer information, nearly matches the performance of the full GN method. Collectively, our results suggest: (1) the GN approximation is highly effective for preconditioning, implying higher-order loss terms may not be critical for convergence speed; (2) the layerwise Hessian structure contains sufficient information to achieve most of these potential gains; and (3) a significant performance gap exists between current approximate methods and an idealized layerwise oracle.
academic
Le Potentiel de l'Optimisation du Second Ordre pour les LLM : Une Étude avec Gauss-Newton Complet
Cet article examine la dégradation de performance des approximations efficaces en termes de calcul des méthodes d'optimisation du second ordre existantes dans l'entraînement préalable des grands modèles de langage (LLM). Les auteurs établissent des bornes supérieures pratiques de la complexité itérative en appliquant le préconditionnement Gauss-Newton (GN) complet sur un modèle Transformer de 150M paramètres. Les expériences montrent que les mises à jour GN complètes réalisent une réduction de 5,4 fois du nombre d'itérations d'entraînement par rapport aux lignes de base fortes telles que SOAP et Muon. De plus, le préconditionnement GN exact par couche, qui ignore les informations entre couches, atteint presque les performances de la méthode GN complète.
Avec la croissance continue des besoins de calcul des LLM, l'amélioration des méthodes d'optimisation est devenue une stratégie centrale pour augmenter l'efficacité de l'entraînement. Bien que les méthodes du premier ordre (telles que SGD et Adam) soient largement utilisées, les méthodes du second ordre possèdent théoriquement une convergence plus rapide et une meilleure scalabilité en lots volumineux.
Limitations des méthodes du second ordre existantes: Les optimiseurs du second ordre actuels (tels que Shampoo, SOAP, Muon) utilisent tous des approximations de la Hessienne pour maintenir la faisabilité computationnelle, mais il reste flou de savoir quelle performance ces approximations perdent.
Écart entre théorie et pratique: Bien que les méthodes du second ordre soient théoriquement supérieures, en raison du coût élevé du stockage et du calcul de la Hessienne complète, les applications pratiques doivent utiliser des méthodes d'approximation.
Question de recherche centrale: « Quelles sont les limites de performance fondamentales de l'optimisation du second ordre dans les LLM ? Quelles propriétés structurelles de la Hessienne sont nécessaires pour atteindre ces limites ? »
Établissement de bornes de performance: Établit des bornes de performance pratiques pour l'optimisation du second ordre via la méthode Gauss-Newton complète, réalisant une amélioration de 5,4 fois en complexité itérative par rapport à SOAP.
Révélation de structures clés: Découvre que la structure Hessienne par couche contient des informations suffisantes pour réaliser la plupart des gains de performance, avec une importance limitée des informations de courbure entre couches.
Intuitions théoriques: Démontre que l'approximation GN est hautement efficace pour le préconditionnement, suggérant que les termes de perte d'ordre supérieur pourraient ne pas être critiques pour la vitesse de convergence.
Scalabilité de la taille de lot: Élargit considérablement la taille de lot critique, démontrant une performance de scalabilité quasi-optimale.
Étant donné les paramètres du modèle θ, l'entrée x et l'étiquette y, on définit la fonction de perte L(f(θ,x), y). L'objectif est de minimiser la perte attendue, en mettant l'accent sur la complexité itérative (nombre d'étapes nécessaires pour atteindre la perte cible).
Pour éviter le stockage explicite de la matrice Hessienne, on utilise des produits Jacobien-vecteur (JVP) pour implémenter une méthode fonctionnellement équivalente. L'idée centrale est d'optimiser l'approximation de Taylor du second ordre de la perte L et l'approximation de Taylor du premier ordre du modèle f.
Les deux méthodes présentent des performances presque identiques, indiquant que les termes de perte d'ordre supérieur contribuent peu à l'amélioration des performances.
La méthode par couche approche les performances du GN complet dans la plupart des configurations, suggérant une importance limitée des informations de courbure entre couches.
Importance du calendrier d'apprentissage: Le calendrier cosinus global montre les meilleures performances dans les configurations de lots petits à moyens
Nécessité de la recherche linéaire: Critique pour la convergence stable de la méthode GN
Choix de l'optimiseur interne: Muon surpasse AdamW en tant qu'optimiseur interne
Établissement de bornes de performance: La méthode GN complète fournit un objectif de performance clair pour l'optimisation du second ordre
Importance structurelle: La structure Hessienne par couche contient des informations suffisantes pour réaliser la plupart des gains
Efficacité d'approximation: Les méthodes d'approximation actuelles présentent un écart de performance significatif par rapport à l'oracle par couche idéalisé
Cet article cite des travaux importants dans le domaine de l'optimisation, notamment :
Martens (2010): Travail fondateur en optimisation sans Hessienne
Gupta et al. (2018): Optimiseur Shampoo
Jordan et al. (2024): Optimiseur Muon
Vyas et al. (2025): Optimiseur SOAP
Évaluation Générale: Ceci est un article de recherche de haute qualité qui établit rigoureusement les bornes de performance de l'optimisation du second ordre dans l'entraînement des LLM par le biais d'expériences rigoureuses, fournissant des intuitions théoriques importantes et des orientations pratiques au domaine. Malgré les limitations en termes de coûts de calcul et d'échelle, sa valeur académique et son importance directrice pour les recherches futures sont significatifs.