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

2025-10-10

AI Summary

This work investigates the performance gap between practical second-order optimization approximations and ideal full second-order methods in large language model (LLM) pretraining. Method: To establish a theoretical upper bound, the authors implement the first scalable, full Gauss–Newton (GN) preconditioned update on Transformer architectures and conduct a systematic analysis of inter-layer Hessian structural properties. Contribution/Results: Full GN preconditioning substantially reduces iteration counts; remarkably, layer-wise exact preconditioning—without cross-layer coupling—suffices to closely approximate full GN performance, exposing substantial untapped optimization potential in mainstream approximations (e.g., SOAP, Muon). On a 150M-parameter model, the proposed method achieves 5.4× speedup over baseline first-order training. This work provides the first scalable, full-GN benchmark for LLM optimization and identifies a principled, structured path toward efficient second-order training.

Technology Category

• Natural Language Processing: Learning & Optimization for NLP
• Machine Learning: Optimization
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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.

Problem

• Evaluating performance loss from second-order optimization approximations
• Establishing practical upper bound with full Gauss-Newton preconditioning
• Comparing layerwise versus full Gauss-Newton method effectiveness

Innovation

• Full Gauss-Newton preconditioning for transformer models
• Layerwise GN preconditioner ignoring cross-layer information
• Substantial iteration reduction compared to existing optimizers