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

大規模言語モデルにおける二階最適化の可能性：完全ガウス・ニュートン法による研究

基本情報

  • 論文ID: 2510.09378
  • タイトル: The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
  • 著者: Natalie Abreu (ハーバード大学)、Nikhil Vyas (ハーバード大学/OpenAI)、Sham Kakade (ハーバード大学)、Depen Morwani (ハーバード大学)
  • 分類: cs.LG cs.AI
  • 発表日: 2025年10月10日 (arXiv プレプリント)
  • 論文リンク: https://arxiv.org/abs/2510.09378

要約

本論文は、大規模言語モデル(LLM)の事前学習において、既存の二階最適化手法の計算効率的な近似がどの程度性能を損失するかを調査している。著者は150M パラメータの Transformer モデルに完全ガウス・ニュートン(GN)前処理を適用することにより、反復複雑度の実用的な上界を確立した。実験結果は、完全GN更新が SOAP および Muon などの強力なベースラインと比較して、訓練反復を5.4倍削減することを示している。さらに、層間情報を無視した正確な層ごとGN前処理器は、完全GN手法の性能にほぼ達している。

研究背景と動機

問題定義

LLMの計算要件が継続的に増加する中で、最適化手法の改善は訓練効率向上の中核戦略となっている。従来の一階手法(SGDおよびAdam など)は広く使用されているが、二階手法は理論的には高速な収束速度と優れた大バッチスケーリング能力を有している。

研究動機

  1. 既存二階手法の限界: 現在の二階オプティマイザ(Shampoo、SOAP、Muon など)は計算可行性を維持するため、Hessian の近似を使用しているが、これらの近似がどの程度性能を損失するかは不明確である。
  2. 理論と実践のギャップ: 二階手法は理論的に優れているが、完全Hessian の保存と計算コストが過度に高いため、実際の応用では近似手法を使用する必要がある。
  3. 中核的研究問題: 「LLMにおける二階最適化の基本的性能限界は何か？Hessian のどの構造的性質がこれらの限界を実現するために必要か？」

核心的貢献

  1. 性能上界の確立: 完全ガウス・ニュートン手法を通じて二階最適化の実用的な性能上界を確立し、反復複雑度において SOAP と比較して5.4倍の向上を実現した。
  2. 重要な構造の解明: 層ごとの Hessian 構造が大部分の性能向上を実現するための十分な情報を含んでおり、層間曲率情報の重要性は限定的であることを発見した。
  3. 理論的洞察: GN近似が前処理に対して高度に有効であることを証明し、高階損失項が収束速度にとって重要でない可能性を示唆している。
  4. バッチサイズスケーリング: 臨界バッチサイズを大幅に拡張し、ほぼ最適なスケーリング性能を実現した。

方法論の詳細

タスク定義

モデルパラメータ θ、入力 x、ラベル y が与えられた場合、損失関数 L(f(θ,x), y) を定義する。目標は期待損失を最小化することであり、反復複雑度(目標損失に到達するために必要なステップ数)に焦点を当てている。

ガウス・ニュートン法の原理

数学的基礎

完全Hessian 行列は以下のように分解できる：

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

ここで第1項がガウス・ニュートン行列 G であり、第2項がモデルの曲率である。

アルゴリズム実装

アルゴリズム1: ガウス・ニュートン法

  1. モデルに対して一階テイラー展開を実行：f⁽¹⁾θₜ(θ,x) := f(θₜ,x) + ∇f(θₜ,x)ᵀ(θ-θₜ)
  2. 損失を凸化：L̃θₜ(θ) := (1/b)Σ₍ₓ,ᵧ₎∈B ℓ(f⁽¹⁾θₜ(θ,x), y)
  3. 二階テイラー近似を構築：L̃⁽²⁾θₜ(θ)
  4. 最小二乗問題を求解：θ̂ = argminθ L̃⁽²⁾θₜ(θ)
  5. 直線探索：θₜ₊₁ ← θₜ + α*(θ̂ - θₜ)

メモリ実行可能な実装

Hessian 行列の明示的な保存を回避するため、Jacobian-ベクトル積(JVP)を使用して機能的に等価な方法を実装する。中核的な考え方は、損失関数 L の二階テイラー近似とモデル f の一階テイラー近似を最適化することである。

変種手法

GN-prox-linear 法

線形化モデル上の損失を直接最小化：θ* = argminθ L̃θₜ(θ)。高階損失項の影響を調査するために使用される。

層ごとガウス・ニュートン法

各層 l に対して独立的に：

  1. その層の一階テイラー展開を計算：f⁽¹⁾θₗ,ₜ(θₗ)
  2. 求解：θₗ,ₜ₊₁ = argminθₗ L̃⁽²⁾θₗ,ₜ(θₗ)
  3. すべての層の更新を統合し、直線探索を適用

実験設定

データセットとモデル

  • モデル: 45M および 150M パラメータの LLaMA アーキテクチャ
  • データセット: C4 データセット
  • シーケンス長: 1024

ベースライン手法

  • AdamW: 最も広く使用されている LLM オプティマイザ
  • Muon: Newton-Schulz 直交化を使用する手法
  • SOAP: Shampoo の最新変種

実験構成

  • 内部オプティマイザ: 最小二乗問題の求解に Muon を使用
  • バッチサイズ: 勾配累積により制御、bᵢₙₙₑᵣ = 32(45M) / 128(150M)
  • 学習率スケジュール: グローバルコサイン、グローバル+内部コサイン、定数+内部コサインの3つの戦略
  • 正則化: 重み減衰、直線探索など複数の戦略

実験結果

主要結果

反復複雑度

損失3.25に到達する実験において：

  • ガウス・ニュートン法: 54ステップ
  • SOAP: 292ステップ (5.4倍の差)
  • Muon: 約16倍の差
  • 層ごとGN: 78ステップ (わずか1.4倍の差)

バッチサイズスケーリング

固定3Bトークン訓練において：

  • ガウス・ニュートン法は120Mバッチサイズでも良好な性能を維持(損失3.45)
  • AdamW は同じバッチサイズで性能が大幅に低下(損失>4.4)
  • 臨界バッチサイズが大幅に拡張され、ほぼ最適なスケーリング傾向に近い

アブレーション実験

GN vs GN-prox-linear

両手法の性能はほぼ同一であり、高階損失項が性能向上に対する寄与は限定的であることを示している。

完全GN vs 層ごとGN

層ごとの手法はほとんどの設定で完全GNの性能に近く、層間曲率情報の重要性が限定的であることを示している。

重要な発見

  1. 学習率スケジュールの重要性: グローバルコサインスケジュールは中小バッチサイズで最良の性能を示す
  2. 直線探索の必要性: GN法の安定した収束に不可欠である
  3. 内部オプティマイザの選択: Muon は内部オプティマイザとして AdamW より優れている

関連研究

二階最適化手法

  • Hessian-free 最適化: Martens(2010)が提案した共役勾配法
  • 対角Hessian近似: AdaHessian、Sophia などの軽量手法
  • 層ごとの近似: Shampoo シリーズ手法の中核的思想

LLM オプティマイザの発展

  • 従来の手法: SGD、Adam シリーズ
  • 現代的二階手法: Shampoo が AlgoPerf コンペティションで28%勝利
  • 特化手法: LLM 用に設計された Muon、SOAP など

結論と議論

主要な結論

  1. 性能上界の確立: 完全GN法は二階最適化に対して明確な性能目標を提供する
  2. 構造の重要性: 層ごとの Hessian 構造は大部分の増益を実現するための十分な情報を含んでいる
  3. 近似効果: 現在の近似手法と理想化された層ごとのオラクルの間に顕著な性能ギャップが存在する

限界

  1. 計算オーバーヘッド: 現在の実装は標準訓練より4～5倍遅い
  2. 規模の制限: 実験は150M パラメータモデルに限定されている
  3. 実用性: 主に分析ツールとしての役割であり、直接的な実用的オプティマイザではない

将来の方向性

  1. 効率的な実装: 計算効率的な正確な二階手法の開発
  2. より良い近似: 層ごとの Hessian 近似手法の改善
  3. 規模の拡張: より大規模なモデルでの発見の検証

深層的評価

利点

  1. 理論的深さ: 二階最適化の性能限界に関する重要な理論的洞察を提供
  2. 実験の厳密性: 広範なハイパーパラメータ探索と複数の正則化戦略
  3. 実用的価値: 既存の二階手法を改善するための明確な目標を提供
  4. 方法論的革新: JVP を巧妙に使用して明示的な Hessian 保存を回避

不足点

  1. 計算コスト: 高い計算オーバーヘッドが実際の応用を制限
  2. 規模の限界: 真の大規模LLMでの検証がない
  3. 理論分析: 層ごとの近似がなぜそれほど有効であるかについての深い理論的説明が不足

影響力

  1. 学術的貢献: 二階最適化研究に重要なベンチマークを提供
  2. 実践的指導: 既存手法を改善するための方向性を示唆
  3. 方法論的価値: 二階手法を評価するための新しいフレームワークを確立

適用シーン

  • 二階最適化手法の理論分析
  • 新しい最適化アルゴリズムの性能ベンチマーク
  • 大バッチ訓練シナリオの最適化選択

参考文献

本論文は最適化分野の重要な研究を引用している：

  • Martens (2010): Hessian-free 最適化の先駆的研究
  • Gupta et al. (2018): Shampoo オプティマイザ
  • Jordan et al. (2024): Muon オプティマイザ
  • Vyas et al. (2025): SOAP オプティマイザ

総合評価: これは高品質な研究論文であり、厳密な実験を通じて LLM 訓練における二階最適化の性能上界を確立し、当該分野に重要な理論的洞察と実践的指導を提供している。計算コストと規模の制限は存在するが、その学術的価値と将来の研究に対する指導意義は顕著である。