用完整高斯-牛顿法加速大模型训练,效果远超现有优化器。
The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
- 采用完整高斯-牛顿预处理,提升训练收敛速度。
- 相比SOAP和Muon,训练迭代次数减少5.4倍。
- 分层近似已接近最优效果,适合研究高效优化器。
近期大模型预训练加速工作多依赖计算高效的二阶结构近似,但其性能损失尚不明确。本文通过在最大150M参数的Transformer模型上应用完整高斯-牛顿(GN)预处理,建立实际迭代复杂度的上限。实验表明,完整GN更新相比强基线如SOAP和Muon,可实现5.4倍的训练迭代减少。此外,忽略跨层信息的分层GN预处理器几乎达到完整GN的性能。结果表明:(1) GN近似对预处理高度有效,说明高阶损失项对加速并非关键;(2) 分层海森结构已包含实现大部分加速所需信息;(3) 当前近似方法与理想分层最优解之间存在显著性能差距。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。