2025-11-16T03:28:12.300331

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

Informations de Base

  • ID de l'article: 2510.09378
  • Titre: The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
  • Auteurs: Natalie Abreu (Harvard), Nikhil Vyas (Harvard/OpenAI), Sham Kakade (Harvard), Depen Morwani (Harvard)
  • Classification: cs.LG cs.AI
  • Date de publication: 10 octobre 2025 (prépublication arXiv)
  • Lien de l'article: https://arxiv.org/abs/2510.09378

Résumé

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.

Contexte et Motivation de la Recherche

Définition du Problème

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.

Motivation de la Recherche

  1. 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.
  2. É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.
  3. 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 ? »

Contributions Principales

  1. É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.
  2. 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.
  3. 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.
  4. Scalabilité de la taille de lot: Élargit considérablement la taille de lot critique, démontrant une performance de scalabilité quasi-optimale.

Détails de la Méthode

Définition de la Tâche

É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).

Principes de la Méthode Gauss-Newton

Fondements Mathématiques

La matrice Hessienne complète peut être décomposée comme :

∇²θL(θ) = ∇θf(θ)ᵀ∇²zL(θ)∇θf(θ) + Σₐ(δL/δzₐ)∇²θ[f(θ)]ₐ

où le premier terme est la matrice Gauss-Newton G et le second terme est la courbure du modèle.

Implémentation de l'Algorithme

Algorithme 1 : Méthode Gauss-Newton

  1. Effectuer une expansion de Taylor du premier ordre du modèle : f⁽¹⁾θₜ(θ,x) := f(θₜ,x) + ∇f(θₜ,x)ᵀ(θ-θₜ)
  2. Convexifier la perte : L̃θₜ(θ) := (1/b)Σ₍ₓ,ᵧ₎∈B ℓ(f⁽¹⁾θₜ(θ,x), y)
  3. Construire une approximation de Taylor du second ordre : L̃⁽²⁾θₜ(θ)
  4. Résoudre le problème des moindres carrés : θ̂ = argminθ L̃⁽²⁾θₜ(θ)
  5. Recherche linéaire : θₜ₊₁ ← θₜ + α*(θ̂ - θₜ)

Implémentation Réalisable en Mémoire

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.

Variantes de Méthode

Méthode GN-prox-linear

Minimiser directement la perte sur le modèle linéarisé : θ* = argminθ L̃θₜ(θ), utilisée pour étudier l'impact des termes de perte d'ordre supérieur.

Gauss-Newton par Couche

Pour chaque couche l indépendamment :

  1. Calculer l'expansion de Taylor du premier ordre de cette couche f⁽¹⁾θₗ,ₜ(θₗ)
  2. Résoudre : θₗ,ₜ₊₁ = argminθₗ L̃⁽²⁾θₗ,ₜ(θₗ)
  3. Fusionner les mises à jour de toutes les couches et appliquer la recherche linéaire

Configuration Expérimentale

Ensemble de Données et Modèle

  • Modèle: Architecture LLaMA avec 45M et 150M paramètres
  • Ensemble de données: Ensemble de données C4
  • Longueur de séquence: 1024

Méthodes de Référence

  • AdamW: L'optimiseur LLM le plus largement utilisé
  • Muon: Méthode utilisant l'orthogonalisation Newton-Schulz
  • SOAP: Variante récente de Shampoo

Configuration Expérimentale

  • Optimiseur interne: Utilisation de Muon pour résoudre le problème des moindres carrés
  • Taille de lot: Contrôlée par accumulation de gradients, bᵢₙₜₑᵣₙₑ = 32(45M) / 128(150M)
  • Calendrier d'apprentissage: Trois stratégies - cosinus global, cosinus global + interne, constant + interne
  • Régularisation: Décroissance des poids, recherche linéaire et autres stratégies multiples

Résultats Expérimentaux

Résultats Principaux

Complexité Itérative

Dans l'expérience atteignant une perte de 3,25 :

  • Gauss-Newton: 54 étapes
  • SOAP: 292 étapes (écart de 5,4 fois)
  • Muon: Écart d'environ 16 fois
  • GN par couche: 78 étapes (écart de seulement 1,4 fois)

Scalabilité de la Taille de Lot

Lors d'un entraînement fixe de 3B tokens :

  • Gauss-Newton maintient une bonne performance même à une taille de lot de 120M (perte 3,45)
  • AdamW se dégrade sévèrement à la même taille de lot (perte >4,4)
  • La taille de lot critique s'élargit considérablement, approchant la tendance de scalabilité quasi-optimale

Études d'Ablation

GN vs GN-prox-linear

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.

GN Complet vs GN par Couche

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.

Découvertes Clés

  1. Importance du calendrier d'apprentissage: Le calendrier cosinus global montre les meilleures performances dans les configurations de lots petits à moyens
  2. Nécessité de la recherche linéaire: Critique pour la convergence stable de la méthode GN
  3. Choix de l'optimiseur interne: Muon surpasse AdamW en tant qu'optimiseur interne

Travaux Connexes

Méthodes d'Optimisation du Second Ordre

  • Optimisation sans Hessienne: Méthode de gradient conjugué proposée par Martens (2010)
  • Approximations Hessiennes diagonales: Méthodes légères telles que AdaHessian et Sophia
  • Approximations par couche: Idée centrale de la série de méthodes Shampoo

Évolution des Optimiseurs LLM

  • Méthodes traditionnelles: Séries SGD et Adam
  • Méthodes modernes du second ordre: Shampoo a surpassé les concurrents de 28% dans la compétition AlgoPerf
  • Méthodes spécialisées: Muon et SOAP conçus spécifiquement pour les LLM

Conclusions et Discussion

Conclusions Principales

  1. Établissement de bornes de performance: La méthode GN complète fournit un objectif de performance clair pour l'optimisation du second ordre
  2. Importance structurelle: La structure Hessienne par couche contient des informations suffisantes pour réaliser la plupart des gains
  3. Efficacité d'approximation: Les méthodes d'approximation actuelles présentent un écart de performance significatif par rapport à l'oracle par couche idéalisé

Limitations

  1. Coûts de calcul: L'implémentation actuelle est 4 à 5 fois plus lente que l'entraînement standard
  2. Restrictions d'échelle: Les expériences sont limitées aux modèles de 150M paramètres
  3. Praticité: Principalement un outil d'analyse plutôt qu'un optimiseur directement pratique

Directions Futures

  1. Implémentations efficaces: Développer des méthodes du second ordre exactes et computationnellement efficaces
  2. Meilleures approximations: Améliorer les méthodes d'approximation Hessienne par couche
  3. Scalabilité: Valider les découvertes sur des modèles plus volumineux

Évaluation Approfondie

Avantages

  1. Profondeur théorique: Fournit des intuitions théoriques importantes sur les limites de performance de l'optimisation du second ordre
  2. Rigueur expérimentale: Recherche extensive d'hyperparamètres et stratégies de régularisation multiples
  3. Valeur pratique: Fournit des objectifs clairs pour améliorer les méthodes du second ordre existantes
  4. Innovation méthodologique: Utilisation ingénieuse des JVP pour éviter le stockage explicite de la Hessienne

Insuffisances

  1. Coûts de calcul: Les frais de calcul élevés limitent les applications pratiques
  2. Limitations d'échelle: Non validé sur de véritables LLM à grande échelle
  3. Analyse théorique: Manque d'explication théorique approfondie sur l'efficacité des approximations par couche

Impact

  1. Contribution académique: Fournit un point de référence important pour la recherche en optimisation du second ordre
  2. Orientation pratique: Indique les directions pour améliorer les méthodes existantes
  3. Valeur méthodologique: Établit un nouveau cadre pour évaluer les méthodes du second ordre

Scénarios Applicables

  • Analyse théorique des méthodes d'optimisation du second ordre
  • Points de référence de performance pour les nouveaux algorithmes d'optimisation
  • Choix d'optimisation pour les scénarios d'entraînement en lots volumineux

Références

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.