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

The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton

基本信息

  • 论文ID: 2510.09378
  • 标题: The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
  • 作者: Natalie Abreu (Harvard), Nikhil Vyas (Harvard/OpenAI), Sham Kakade (Harvard), Depen Morwani (Harvard)
  • 分类: cs.LG cs.AI
  • 发表时间: 2025年10月10日 (arXiv预印本)
  • 论文链接: https://arxiv.org/abs/2510.09378

摘要

本文研究了在大语言模型(LLM)预训练中，现有二阶优化方法的计算高效近似会损失多少性能。作者通过在150M参数的Transformer模型上应用完整的Gauss-Newton(GN)预条件，建立了迭代复杂度的实用上界。实验显示，完整GN更新相比SOAP和Muon等强基线实现了5.4倍的训练迭代减少。此外，忽略跨层信息的精确逐层GN预条件器几乎达到了完整GN方法的性能。

研究背景与动机

问题定义

随着LLM计算需求的不断增长，优化方法的改进已成为提高训练效率的核心策略。传统的一阶方法(如SGD和Adam)虽然广泛使用，但二阶方法在理论上具有更快的收敛速度和更好的大批量扩展能力。

研究动机

  1. 现有二阶方法的局限性: 当前的二阶优化器(如Shampoo、SOAP、Muon)为了保持计算可行性，都使用了Hessian的近似，但这些近似会损失多少性能尚不清楚。
  2. 理论与实践的差距: 虽然二阶方法在理论上优越，但由于完整Hessian的存储和计算代价过高，实际应用中必须使用近似方法。
  3. 核心研究问题: "二阶优化在LLM中的基本性能极限是什么？Hessian的哪些结构性质对实现这些极限是必要的？"

核心贡献

  1. 建立性能上界: 通过完整Gauss-Newton方法为二阶优化建立了实用的性能上界，在迭代复杂度上相比SOAP实现5.4倍提升。
  2. 揭示关键结构: 发现逐层Hessian结构包含了实现大部分性能增益的充分信息，跨层曲率信息的重要性有限。
  3. 理论洞察: 证明了GN近似对预条件高度有效，暗示高阶损失项对收敛速度可能不是关键的。
  4. 批量大小扩展: 显著扩展了临界批量大小，展现了接近最优的扩展性能。

方法详解

任务定义

给定模型参数θ、输入x和标签y，定义损失函数L(f(θ,x), y)。目标是最小化期望损失，重点关注迭代复杂度（达到目标损失所需的步数）。

Gauss-Newton方法原理

数学基础

完整的Hessian矩阵可以分解为：

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

其中第一项即为Gauss-Newton矩阵G，第二项为模型的曲率。

算法实现

Algorithm 1: Gauss-Newton方法

  1. 对模型进行一阶泰勒展开：f⁽¹⁾θₜ(θ,x) := f(θₜ,x) + ∇f(θₜ,x)ᵀ(θ-θₜ)
  2. 凸化损失：L̃θₜ(θ) := (1/b)Σ₍ₓ,ᵧ₎∈B ℓ(f⁽¹⁾θₜ(θ,x), y)
  3. 构建二阶泰勒近似：L̃⁽²⁾θₜ(θ)
  4. 求解最小二乘问题：θ̂ = argminθ L̃⁽²⁾θₜ(θ)
  5. 线搜索：θₜ₊₁ ← θₜ + α*(θ̂ - θₜ)

内存可行的实现

为避免显式存储Hessian矩阵，使用Jacobian-向量乘积(JVPs)来实现功能等价的方法。核心思想是优化损失函数L的二阶泰勒近似和模型f的一阶泰勒近似。

变体方法

GN-prox-linear方法

直接最小化线性化模型上的损失：θ* = argminθ L̃θₜ(θ)，用于研究高阶损失项的影响。

逐层Gauss-Newton

对每一层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.25的实验中：

  • Gauss-Newton: 54步
  • SOAP: 292步 (5.4倍差距)
  • Muon: 约16倍差距
  • 逐层GN: 78步 (仅1.4倍差距)

批量大小扩展

在固定3B token训练中：

  • Gauss-Newton在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. 方法创新: 巧妙地使用JVPs避免显式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训练中的性能上界，为该领域提供了重要的理论洞察和实践指导。尽管存在计算成本和规模限制，但其学术价值和对未来研究的指导意义是显著的。