The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
Authors: Natalie Abreu
Nikhil Vyas
Sham M. Kakade
Depen Morwani
Venue: ICLR 2026 Poster (2026)
Keywords: LLMs
optimization
Primary area: optimization

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.

THE POTENTIAL OF SECOND-ORDER OPTIMIZATION FOR LLMS: A STUDY WITH FULL GAUSS-NEWTON
Natalie Abreu
Nikhil Vyas*
Kempner Institute, Harvard University natalieabreuêg harvard. edu
Department of Computer Science, Harvard University vyasnikhil gmail.mal.. com.
Sham Kakade
Depen Morwani
Kempner Institute, Harvard University sham@seas
Kempner Institute, Harvard University dmorwan @g. waaraaaaree rd edu
ABSTRACT
Recent efforts to accelerate LLM pretraining have focused on compurationall--fficienn approxima- tions that exploit second-order structure. This raises 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) precondition- ing to transformer models of 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 precise layerwise GN preconditioner, which ignores cross-layer information, nearly matches the performance of the full GN method. Collectively, results suggest: (I) 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 achieve most of these potential gains; and (3) a significant performance gap exists between current approximate methods and an idealized layerwise oracle.
Code is available at: httpps://githhb..cm///ttaaeaeeee
-aaasgaasssss- -auss-newttonn
1 INTRODUCTION
With rising compute requirements for training large language models (LLMs), improving optimization methods has become central strategy for improving training efficiency. Better optimizers can directly reduce the serial runtime to train LLM,, which is crucial for large-scale models that train from days to months. Optimization for LLMs has traditionally leveraged first-order methods such as SGD and Adam (Kingma & Ba 2017). However, recent research in optimization has started exploring the use of second-order optimizers for large-scale models, motivated by the faster convergence rates known from theory (Nesterov) 2018) and potential to scale to larger batch sizes (Zhang et al. 2019 two ways of reducing serial runtime.
Some recent popular (Jordan al. 2024b).. Shampoo the recent optimization algorithms benchmark called AlgoPerf (Kasimbeg et al. 2025), outperforming Adam by margin of 28%. SOAP, a recent generalization of the Shampoo algorithm, has shown impressive performance language modeling benchmarks, and has been used for training inforredd neural networks (PINNs) (Wang et al. 2025). Muon has been extensively optimized on the nanoGPT benchmark (Jordan et al. 2024a), and was also recently scaled up 16B LLMs, showing 50% improvements over AdamW (Liu et al.
methods include Shampoo (Gupta et al., 2018), SOAP (Vyas et al. 2025) and Muon
2025
However, these methods do not use complete second-order information, instead focusing on memory- and computationally--ffiiiii approximations of the Hessian. Indeed, precisely storing or computing the Hessian required for second-order methods such as Newton's method is prohibitively expensive for modern LLMs that have billions
"Currently at OpenAl

---

of parameters. To remain practical, these methods leverage computationallly-eficiinnt estimators for the layerwise Hessian of neural networks.
The success of current methods motivates a better understanding of the fundamental potential of second-order opti- mizers. Our work driven by the following question:
What are the fundamental performance limits of second-order optimization for LLMs, and what structural properties of the Hessian are essential for achieving them?
To answer this, we first establish the performance limits of an idealized second-order method, full Gauss-Newton (GN), and measure its performance in terms of iteration complexity (the number of steps to reach target loss). This serves as practical lower bound for any second-order approach. We also analyze how this idcalized method affects the critical batch size (McCandlish et al. 2018: Shallue al. 2019: Jain al., 2018b). key measure data parallelism efficiency.
To isolate the essential structural properties of the Hes- sian, we then compare the full GN method to two vari- ants:
1. A prox-linear version of Gauss-Newer (GN- prox-linear) (Burke, 1983; Drusvyatskiy 2017p, which utilizes the higher order informa- tion of the loss function itself (see Algorithms and for comparison).
4.75
Loss
2. A purely layerwise version, which ignores all cross-layer curvature information.
3.75
Our main findings are three-fold. First, both idealized second-order methods provide a substantial improve- ment over existing optimizers at large batch sizes, with the full Gauss-Newton method achieving 5.4x reduc- tion in iteration complexity over the SOAP optimizer. Second, the Gauss-Newton method significantly extends the critical batch size beyond that of prior methods, displaying near-optimal scaling. Finally, we find that the layerwise approach, despite its structural limitations, still substantially outperforms both SOAP and Adam. This suggests that even layer-wise curvature information alone is sufficient to achieve major gains in compute ef- ficiency.
3.25
10
50
300
1000
3000
Step
AdamW
Muon
SOAP
Gauss-Newton
Layerwise GN
Training step versus validation loss until loss 3.25 when each method beyond its critical batch size, Gauss- Newton and Layerwise Gaaus-Neetton reach the target loss in 54 and 78 steps respectively, compared to 292 steps for SOAP.
Figure I:
We stress that this work is empirical study aimed at understanding performance limits with better higher order information, not at directly designing computationally cheaper optimizers themselves. Our implementation avoids materializing the Hessian by using Jacobian- vector products, though still has substantial computational overhead. We believe this work is best viewed as tool for analysis that provides a target for more practical second-order methods to aspire to, and we provide further discussion this point.
Paper organization In Section 2 we cover related work and in Section 3 we provide background on existing op- timization methods, In Section we introduce the setting for full second-order optimization and the Gauss-Newton matrix. In Section 53 we provide the setup for our main experiments and in Section 6.1 we discuss our results on itera- tion complexity and critical batch size of the full second-order method. Section 6.2 we compare the Gauss-Newton method to the GN-prox-lines method and in Section 6.33 we compare to a layerwise variation. Finally, in Section 7 we discuss the implications as well as limitations of our work.
2 RELATED WORK
We mention a few of the most related works here and provide additional related work in Appendix C work on spe- cific optimizers for LLMs is discussed in Section Most related to our work is Hessian-free optimization, which
2

---

Algorithm 1: Gauss-Newton method
Input: On training set T, training iterations T, batch size
for 1= 0, 0,
do
Linearize model:
(1) (0.x)) = f(O,, x) + v 0()
Convexify loss:
With T, Lo. (0) = Σ (B E(), (0, (0,x).
Gauss-Newton loss:
Co. (0) = Le, (0.) + Vo, (0.) 0,) + (0 V (0:) 0z)
Set to be an approximate solution to the least squares problem:
ming ( (0)
Line search: Set 0x+1 + 02 + a² (0- 0z), where:
a² f arg min, [(. + a(0 0z))
Algorithm 2: GN GN-prox-lineat
Input: 0o training set T, training iterations I', batch size
for = 0, 0,
1 do
Linearize model:
pl)) (6, x) = f(0,, v) (B 0₂)
Convexify loss:
With B T, Co. (0) ( (0,
Set 4.41 to be an approximate solution the convex problem:
ming Co
avoids explicit Hessian formation by leveraging Hessian vector products (Martens) 2010 This approach serves as an alternative to layerwise approximation methods of the Hessian as discussed in Section Specifically, prior work on Hessian-free optimizers use the conjugate gradient (CG) to solve an incomplete (unconverged) optimization of the Newton step rather than storing approximation to the Hessian. This is introduced by Martens] (2010) classifi- cation and auto-encoder tasks, and extended to additional settings such as recurrent neural networks by Martens & Sutskever (2011a) and Cho et al. (2015). Garcia et al. (2023) amortizes the CG steps in Hessian- -ree optimization for deep linear and auto-encoder models. In contrast, our work focuses on the setting of LLMs, and we leverage optimizers that specifically designed for L.MSs (e.g. Adam and Muon) rather than CG to apply the Gauss-Newton step.
3 BACKGROUND ON EXISTING OPTIMIZERS
We will denote the weight matrix of a model layer at timestep by W, E Rmxx and the corresponding gradient by Gt. We use 77 to denote the learning rate.
The most widely used optimizer for LLMs is Adam (Kingma & Ba 2017). Adam maintains matrices for the first and second moment of the gradient G,, denoted M, and V, respectively. Adam performs the element-wise update
M,
W41 = W.
VV.
AdaGrad (Duchi et al. 2011) maintains an accumulator over the vectorized gradient gr vec(Gr) gm The preconditioner II, and vectorized weights wi at timestep t are updated as
H, = H,-1 + 9:9 T
We = : H, 1/2
Shampoo Gupta et al. 2018) was originally motivated by AdaGrad, but can be viewed as an approximation of the
Gauss-Newton component of the Hessian (Anil al. 202 Osawa al. 2023 Morwani al. 2024). These meth-
3

---

ods leverage computationally efficient approximations of the layerwise Hessian to precondition the gradient update. Shampoo maintains separate preconditioner for each dimension of the weight matrix: For weight matrix W E R Shampoo maintains a left matrix Lt E Rm and a right matrix Rg E IR The update rule is as follows:
L == LL-1 + G.G.; := R G. Ga
We == Wi-1 - L, 1/4 G, G 1/4
SOAP Vyas et al.
2025) is a recent variant of Shampoo that runs Adam (Loshchilov & Hutter 2019) in the
eigenbasis provided by Shampoo.
Muon (Joddnn et al. 2024b) tracks the first moment of the gradient, denoted M, and performs an orthonormal update:
O (NM W++1 nn
where NS denotes Newton-Schull orthogonalization procedure. (See Jordan al. (2024b) for further description of the Newton-Schulz method). Muon without momentum can be seen as a version of Shampoo without preconditioner accumulation (Bernstein & Newhouse 2024).
These second-order methods have been shown to scale effectively to larger batch sizes (Zhang et al. 2019 Vyas et 2025). However, these are restricted by the need for computationally efficient per-layer preconditioners given matrix.
the high computational and memory requirements for computing the full
4 FULL SECOND-ORDER OPTIMIZATION
4.1 NOTATION
Let S [(0, .r) denote the model with parameters and input x. Let (0) (000,.)))0) y) denote the (convex) loss function which takes the model output and the true labels y. We will use either VaL or g or simply V C to denote the gradient with respect to V denote the derivative of L with respect and И == V denote the Hessian. We will use (0) -. (O( v) to denote the first-order Taylor expansion of f around 0r and the second- order Taylor expansion of L around 6, respectively. Similarly, we will also use L(O) == ((f xx), v) when f. x, and y are clear from context. For simplicity, we will assume that we are working with cross-entropy loss throughout this
f0 (0,x:) and ())
work, however, our presentation holds for any general convex loss function.
4.2 NEWTON'S S METHOO & THE
MATRIX
Full second-order optimization requires full access to the Hessian H, which can be used to precondition the gradient in each parameter update. This is known as Newton's 's method, and results in the following update rule:
0" 0 H-1
In practice, for neural networks, the Hessian is not guaranteed to be positive semi-definite (PSD), and therefore New- ton's method does not guarantee that the loss decreases in each iteration or even converges. As a consequence, it is common to instead use the Gauss-Newton matrix.
The Gauss-Newton matrix is defined to be the first term of the following decomposition of the Hessian, where = = f (x) denotes the pre-softmax outputs of the model of and a goes over the output dimensions of f:
5L v (f(0)).
ViC(()) = Vof(0))V2<(0)00000 +
Σ
sza
Gauss-Newton matrix
That is, the Gauss-Newton matrix is defined as G = V VV the curvature of the loss function, but drops the curvature of the model. Unlike the Hessian, the Gauss-Newton matrix is PSD for MSE and cross-entropy loss (Martens, 2020). This avoids untrustworthy updates as negative curvature implies unbounded decrease in loss (Martens 2020). Indeed, methods using the Gauss-Newton matrix rather than the full Hessian have been found to lead to better optimization (Martens) 2010, Martens & Sutskever, 2012, Vinyals & Povey 2012).
Intuitively, the Gauss-Newton captures
WWe refer to diagonal preconditioners as first-order and non non- diagonal as second-order.
4

---

4.3 MEMORY-FEASIBLE
IMPLEMENTATION
To test the limits of second-order optimizers, we want to apply the full Gauss-Newton term as the preconditioner. Formally, for gradient g, Gauss-Newton matrix G and current parameters 8, the Gauss- update is
0" 0 G
(1)
However, given that computing the Gauss-Newton matrix directly is infeasible, we instead run functionally equiv- alent method that leverages Jacobian- vector products (JVPs) to avoid explicitly storing the Hessian. Specifically, we optimize the second-order Taylor approximation of the loss function L with first-order Taylor approximation of the model f. The minimization of the loss this setting is equivalent to using the Gauss-Newton matrix as a precondi- tioner (Martens & Sutskever 2011b). The proof is provided in Appendix
We are now ready to define our Gauss-Newton method. Let Co, (0) C(fo) (0,,),)) be the loss function on the first-order Taylor expansion of f around current parameters Oz. Let ( denote the second-order Taylor expansion of C around Oz.
(0)
Given current parameters O,, we define the Gauss-Newton update (Algorithm D as
0* argming 2 (0)
With this definition, there remains the problem of finding the minimizing 0* As it is difficult to solve for the min- imum directly, we instead use a separate optimizer to minimize Z%² (0). In our experiments we use Muon (Jordan et al., 2024b) as this "inner optimizer" as we found it to outperform AdamW. More details on this inner optimization procedure are given in Scction
For the GN-prox-linea algorithm (Algorithm 2) instead define the updated iterate
argming (0).
The results for GN-prox-linear algorithm are provided in Section 6.2
5 EXPERIMENT DETAILS
We train 45M and 150M parameter LLaMA models (Touvron et al. 2023) on the C4 dataset [Raffe] al. 2020). Full details models and hyperparameter sweeps are given in Appendix and Appendix H respectively.
Training details
Baselines We run Adam W (Loshchilov & Hutter 2019), Muon (Jordan et al. 2024b), and SOAP (Vyas et al. 2025) as baselines. For 45M models, for each method at each batch size, we run hyperparameter sweep over learning rate, weight decay, and weight averaging decay if applicable. We additionally sweep the 3g parameter for Adam, the μ parameter for Muon, and (31, (3) for SOAP. For 150M parameter models we run a more limited hyperparameter sweep over learning rate and 8 and parameters. To make sure runs well-initialized, we start all runs after Adam W warmup consisting of 5% of the Chinchilla-optimal number of tokens (Hoffmann et al. 2022). More details hyperparameter tuning are given in Appendix H
Gauss-Newton For each training step, we take a first-order Taylor approximation of the model around the current parameters. We initialize the parameters of the Taylor approximation to be the pre-linesearch parameters from the previous iteration (see Section 5.1). We then take a second-order Taylor approximation around the cross-entropy loss on the Taylorized model, also around the current model parameters. We use Muon (Jordan et al., 2024b) with batch size Dinner to minimize the Taylorized loss. We take N steps of Muon and then update the model parameters using a line search. We refer to the global batch size (in number of sequences) as b = N x dinner since this is the total amount of data seen per parameter update on the model. We dinner binner 128 for the 150M models with sequence length 1024, and vary N to control the overall batch size b. We start all runs from the same AdamW post-wardn wost-warmup checkpoint. To compute the necessary Taylor approximations we use the neural- tangents s library from Novak et al. (2020). See Algorithm [ for details. In addition we provide an efficient implementation for the minimization of the least squares subproblem in Appendix
32 for the 45M models and
5

---

Upper bound for Gauss-Newton method In our experiments, we use Muon Jordan et al., 2024b) with batch size Dennr to solve the least squares problem in Algorithm T Therefore, Muon with batch size dinner trained on the true model and loss marks the upper bound for the Gauss-Newton method: since Muon in our method optimizes over the respective Taylor approximations, it upper bounded by the performance of Muon with the same batch size on the true model and loss. We include results for Muon with batch size binner in order to judge the relative performance of the Gauss-Newton method.
GN-prox-linear For each training step, minimize the loss of the first order Taylor expansion of the model around the current parameters (as mentioned Algorithm 2) This method evaluates whether incorporating higher-order terms in the loss function yields improvements compared the Mauss-Newton method. The results for this method are discussed in Section 6.2
5.1 OPTIMIZATION STRATEGIES
We perform extensive hyperparameter sweeps, learning rate scheduling strategies, and regularization strategies to test the limits of the Gauss-Newton method.
Learning rate schedules We experiment with three learning rate schedules for Gauss-Newton runs, which we refer to "global cosine," "global+inner cosine," and "constant+inner cosine." We depict each learning rate schedule in Figure
Regularization We experiment with several types of regularization strategies improve the stability of the Gauss- Newton runs at high learning rate. These fall into two categories of inner optimization and outer optimization regular- ization strategies. For inner optimization strategies (regularization involving the inner optimization loop solve the least squares problem), we add weight decay to the optimizer as well as a weight decay term to the loss, which adds regularization on the & norm of the magnitude of the parameter update. For outer optimization, we experiment with line search to control the size of the parameter update.
Inner optimizer We note that the choice of inner optimizer has significant impact on the upper bound of the Gauss-Newton method; consider our results with respect to the performance of the inner optimizer. Regardless, we find that Muon outperforms A as the inner optimizer for the Gauss-Newton runs.
Takeaways from optimization strategies We found that learning rafe schedule and line search had major impact on the stability of training for the Gauss-Newton method. As for learning rate schedules, we find that the global cosine schedule outperforms the global+inner cosine schedule small to medium batch size, but the constant-inner cosine schedule can be helpful for runs at large batch size for Gauss-Newton. Additionally, we found line search to be essential for stable convergence for Gauss-Newton runs. However, we found that when using line search, it helps to set the initial parameters for the next inner minimization to be the pre-linesearch parameters from the previous step. These findings coincide with those of Martens (2010), which finds that sharing information across iterations and backtracking improve performance of a conjugate gradient-based Hessian-free optimization strategy. The importance of inner optimizer and sharing information across iterations seem imply that we are not finding the precise Gauss- Newton update at each step it is possible that with further optimization the Newssn method could achieve even better performance.
6 EXPERIMENTS
6.1 GAUSS-NEWTH EXPERIMENTS
6.1.1
ITERATION COMPLEXITY
We measure the iteration complexity of each method by measuring the number of steps it takes to reach loss 3.25 with extremely large batch size. Specifically, for each method, we use batch size significantly beyond that method's 's critical batch size (McCandlish al. 2018 Shallue al. 2019) such that further increasing batch size does not

---

Batch Size vs Loss
Target Loss 3.4
12
4.4
11
4.2
log2 (# Steps)
10
Final Loss
4.0
9
3.8
8
3.6
3.4
6
3.2
120M
40M
12M
4M
1.2M
1.2M
4M
12M
40M
Batch size
Batch size
AdamW
Muon
GN Upper bound Optimal
Gauss-Newton
SOAP
Figure 2: Left: Batch size vs final validation loss for models trained for Chinchilla-optimal number of tokens. The dotted line marks the loss achieved by a model trained with Muon with batch size 128k. This represents the upper bound of performance for Gauss-Newton method. Right: Critical batch size scaling. The dotted line marks the optimal scaling trend, where no sample efficiency is lost as batch size increases.
reduce the number of training steps needed to achieve a given performance Following our critical batch size findings in Section 6.1.2. use batch size of 40M tokens for AdamW and Muon gains disappear almost completely beyond this amount, and a batch size of 240M tokens for SOAP and Gauss-Newton. We choose this threshold and these batch sizes to enter the regime in which additional batch size increases no longer reduce the required steps, while keeping runs feasible.
We find that the Gauss-Newton method can make fast progress in the large batch size regime, particularly in the first few steps of training. After 10 steps, the loss for the Gauss-Newton model is below 3.75, while other methods have made marginal progress from the starting loss. The optimal learning rates for AdamW, Muon, and SOAP lead to initial instability but faster convergence overall see Appendix H for details on choosing hyperparameters. The Gauss-Newton method able reach loss 3.25 in 54 steps, 5.4x gain over SOAP and 16x gain over Muon. The results are shown in Figure
6.1.2 BATCH SIZE SCALING
While iteration complexity captures a purely sequential perspective of training efficiency, it is also important sider the sample efficiency. In an ideal setting, the number of samples seen at each step would vary proportionally to the number of steps, such that there is always constant number of samples required overall. However, it is known that sample efficiency is lost once the batch size is scaled past a given method's critical batch size (McCandlish et al. 2018, Shallue et al. 2019). That is, the total number of samples needed to achieve given loss will grow once the critical batch size is exceeded. Therefore, we study the batch size scaling behavior of the Gauss-Newtar method to understand how much iterations. We run experiments in two settings: first, we train models for fixed amount of data over a range of batch sizes and measure the final loss. Second, we measure how the number of steps to achieve given loss changes with increasing batch size. We run experiments for 45M- and 150M 15rrmeerreet models; Appendix G for results at 4ar-mmeee scale.
lose in sample and computational efficiency when we minimize the number of sequential
"Since batch sizes increased using gradient accumulation in experiments, choose batch size based each method's critical batch size to save compute.
7

---

models for 3B tokens following Chinchilla-opimaal scal- ing laws (Hoffmann et al. 2022), ranging batch size from 1.2M to 120M tokens. We observe similar performance between SOAP and Gauss-Newton up to batch size 4M, while substantial gains are achieved by Gauss-Newton for larger batch size (Figure D. Of existing methods, SOAP performs best, followed by Muon and then AdamW. Espe- cially noteworthy is the performance of the Gauss-Newtor method at batch size 120M, which uses only 20 steps of optimization. Here we are able to achieve loss 3.45 with Gauss-Newton. For comparison, AdamW achieves loss 3.4 with batch size 1. 1.2M, and degrades to loss above 4.4 with batch size 120M.
Training for fixed token count We train
Training to reach a target validation loss Following the methodology of Zhang et al. (2025). ). we plot the number of steps required for each optimization method reach the target validation loss of 3.4 as a function of batch size. The point at which the curve for each model plateaus defines its respective critical batch size (McCandlish et al. 2018 [Shallue al.] 2019 Jain et al. 2018b). We find that AdamW levels off near batch size of 4M with little further reduction. SOAP and Muon continue to decrease up to batch size of 12M but with diminishing reductions, and show little additional decrease by 40M. Meanwhile, the Gauss-Newton method continues to decrease through 40M, indicating better sample efficiency at large batch sizes.
6.2 ON-PROX-LINEAT METHOD
We define variation of our method that corresponds to another convex problem that retains the full loss function on the linearized model instead of using the second-order approximation. This method follows the same procedure as the Gauss-Newton method in Section but directly minimizes the loss on the linearized model, denoted by Cs, (0):
0 = argming Lo, (0)
See Algorithm 2 for details. Note that for any convex loss function (including cross-entropy loss), the above optimiza- tion problem is still convex. Moreover, this problem is related to the richly studied literature of kernelized classification "Shalev-Shward
& Ben- BaviD 2014) (albeit with cross-entropy loss, instead of the max max-margin loss). This GN-prox linear method is notably not second-order method. Rather, it allows us to study the effect of the higher order terms of the loss as compared to the Gauss-Newton update rule.
We train in the same Chinchilla-optimaal setting on 150M parameter models and perform the same hyperparameter sweeps as for Gauss-Newton (See Appendix(H]. We find that the inclusion of higher order loss terms has little effect on performance as compared to
results are shown in Figure However, unlike Gauss-Newton, we found that the global cosine schedule for the inner optimizer outperformed the constant+inner cosine schedule. Additionally, line search was not necessary for the GN- linear method.
6.3 LAYERWISE GAUSS-NEWTON
Many existing second-order optimizers use layerwise approximations of the Hessian for computational and memory feasibility. This prompts a further study on whether the full Hessian is necessary to achieve the performance gains discussed in S5ction[5] Specifically, we want to understand the importance of the cross-layer Hessian information.
We define a layerwise version of our layer and optimize the second-order Taylor expansion of the loss separately for each layer.
method, in which we take a Taylor expansion around each model
Formally, let 0r, be the set of parameters at time for layer of the network. For layer define ((()) (0:) as the first- order Taylor expansion of f with respect to only the parameters f, expanded around the current parameters fr,,, while keeping the parameters of all other layers fixed at their current values.
(2) At each timestep f, take Le. (0t) to be the second-order Taylor expansion of the loss on (0t). Then for each layer, we solve for
7(2) 0r.++1 argming, 00, (222
(0t)
After independently computing updates for each layer, we merge the updated layer parameters. We then apply line search over the merged parameter set to obtain the final parameter update at step t.
Due to compute costs, we train 150M parameter layerwise Gauss-Newton models only for the largest three batch sizes. We follow the same training setting for fixed token count as specified in Section 6.2.2 For layerwise experiments we
8

---

Layerwise GN
GN-prox-linear
3.6
Final Loss
3.4
3.2
120M
40M
12M
120M
40M
12M
4M
1.2M
Batch size
Batch size
SOAP
GN-prox-linear Layerwise GN
Upper bound
Gauss-Newton
Figure 3: Left: Comparison of Gauss-Newton the layerwise implementation for Chinchilla-optimal token count for 150M parameter models. The layerwise method achieves almost matching performance to that of the full Gauss- Newton. Right: The Gauss-Newton update closely matches the GN-prox- method that has access higher order loss terms.
set hyperparameters to match those of the best Gauss-Newton configuration at each batch size. However, we include smaller step size options for the line search as we find this is necessary for stable convergence. We provide more details on line search step sizes in Appendix D We find that the layerwise Gauss-Newton method also achieves comparable performance through batch size of 40M tokens (Figure 3 We additionally train a layerwise Gauss-Newton model with batch size of 120M tokens to loss 3.25 to compare its iteration complexity to that of full Gauss-Newton (See See I.II.I. D). We find that the layerwise Gauss-Newton takes only 1.4x more steps to reach the target loss compared to the full Gauss-Newton method and provides a 3,4x gain over SOAP (Figure I.
7 DISCUSSION AND CONCLUSION
In this work, we study whether full second-order optimization specifically, using the full Gauss-Newtor matrix preconditioner can offer further benefits for training large language models as compared to existing methods. In particular, we focus on the large batch size regime, following Dain al. (2018a) and Zhang et al. (2019) which show that the benefits of preconditioning may not appear at small batch size. While our current implementation is roughly 1.5x slower than standard training (e.g. with AdamW or Muon), we view this as a proof of concept demonstrating the potential of exact second-order methods. Our results indicate that further development in second-order methods could lead to substantial benefits in convergence and ability to scale to larger batch size.
While we perform extensive hyperparameter sweeps and regularization strategies, we acknowledge that there could be other optimization strategies to further improve the performance of the Gauss-Newton method. In addition, our work is limited to applying the inverse of the Gauss-Newtor matrix the preconditioner (G-1). There may be better ways to apply full second-order optimization for large language models. We encourage future work in this area and hope our findings are informative.
We also compare
to the GN-prox- -ineea method to study whether there is benefit to including higher order loss terms beyond second-order. Our results suggest that Gauss-Newton can achieve performance similar this method, indicating that higher-order loss terms are not necessary to achieve gains performance over current methods. addition. our layerwise Gauss-Newton experiments suggest that better approximations the per-layer Hessian may be sufficient to achieve substantial performance benefits over current methods. We encourage future work in developing computationally efficient and practical optimization methods in this direction.

---

8 ACKNOWLEDGEMENTS
We thank Jonathan Frankle and Kwangjun Ahn for helpful discussions. SK, NA, and DM acknowledge support from the Office of Naval Research under award N0001422-1-2377 and the National Science Foundation Grant under award #IIS 2229881. This work has been made possible in part by gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. NA. NV, and DM are supported by a Simons Investigator Fellowship. NSF grant DMS-2134155 DARPA grant W911NF2010022,,add DOE grant DE-SC0022199
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Towards practical second order optimiza- tion for deep learning, 2021. URL / http://ew....tt.ooc.ooooooooaoao
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology, 2024. URL https://arxiy: org/ abs/ /lbs/2409.00 24322
James V. Burke. Descent methods for composite nondifferentiable optimization problems. Math. Program., 33(3); 260 279, December 1985. ISSN 0025-5610. doi: 10.1007//FO01544377 Utthttps:: http://to..///.///.////// 120070.00007 BF01584377
Minhyung Cho, Chandra Dhir, and Jachyung Lee. flesian-frree optimization for learning deep mul- tidimnnsional recurrent neural networks. In C. Cortes, N. Lawrence. D. Lee, M. Sugiyama, and Garnett (eds.), Advances Neural Information Processing Systems, volume 28. Curran Asso- ciates, Inc., 2015. URL https::/prrceedddigg neurips. 0 neerruuip......s............... cc/paper httppppprrflappaaaeeeaeaee fflles/paper22015//1ee
Dmitriy Drusvyatskiy. The proximal point method revisited, 2017. URL httpss//arxxv..crr//b///1 06038
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimiza- tion. Journal of Machine Learning Research, 12(61): 2011. URL http://tppittt www.rorgrrpp......... papers/
v12/duchilla. html
Jezabel R Garcia, Federica Freddi, Stathi Fotiadis, Maolin Li, Sattar Vakili, Alberto Bernacchia, and Guillaume Hen- nequin. Fisher-legendre (fishleg) optimization of deep neural networks. In The Eleventh International Conference on Learning Representations, 2023. URL http::::///ppeeeeiieeeeee wtwp::/wwwweeeeeeeeeeeeee
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization, 2018 URL
095688
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford. Diego
de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL https: Vlaxxiv.org. abs 2203.15555
Prateek Jain, Sham M. Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerating stochastic gradient descent for least squares regression, 2018a. httttp:://srxiivorr//aaatt
08227
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Parallelizing stochastic gradient descent for least squares regression: mini-batching, averaging, and model misspecification. Journal of machine learning research, (66223)):1422. 2018b.
Keller Jordan, Jeremy Bernstein, Brendan Brapadn @fembbear bskyeaarbbk.sscciiall Boza Vlado, You Jiacheng. Franz Cesista, Speedrunning the nanogpt baseline, 2024a. URL
Braden Koszarsky, and 10mu662204977
nanogpt
10

---

Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024b. URL. httpps://kklllerjordarn gitub. i
/ostts/muon/
Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, Feng Boyuan, Less Wright, Edward Z. Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig. Michael Rabbat, and George E. Dahl. Accelerating neural network training: An analysis of the AlgoPerf competition. In The Thirteenth International Conference on Learning Representations 2025. URL httpps://peerrevieet wttpeeewwwwwwwweeeeeeeeeeeeeeeee forum?
1d=CtM5x fR
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. 2017. URL https://arxiv.
org/ abs/ 1412. 6980
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order oppoodiiederootimiize 14442
for language model pre-training, 2024. URL
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai. Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for Ilm training, 2025. URL https://: /aapsii///rxvv..o http://axxvvv..aaaaaaa. /bb//
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https:://axivv.o/
abs/1711.0510
James Martens. Deep learning via hessian-free optimization. In Proceedings of the 27th International Conference on International Conference on Machine Learning, ICML' 10, pp. 735-742, Madison, WI, USA, 2010. Omnipress. ISBN 9781605589077.
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020. URL /attl:/tt:///t0 hwtp.//ppp.ppppp//////----777771
org/ papppss/v1111-61111/
html
James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian-free optimization. In Proceedings of the 28th International Conference on International Conference Machine Learning, ICML' 11, pp. 1033-1040, Madison, WI, USA, 201 1a. Omnipress. ISBN 978145030611
James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian-free optimization. In Proceedings of the 28th International Conference on International Conference on Machine Learning, ICML' 11, pp. 1033-1040, Madison, WI, USA, 201 Omnipress. ISBN 978 4503061155.
James Martens and lya Sutskever. Training deep and recurrent networks with hessian-free optimization. In Neural Networks: Tricks of the Trade: Second Edition, pp. 479-535, Springer, 2012.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training, 2018. httts:::::// / /ttp://awxv////////a//ee////at orgg abs aatrgg////////
8 0288812811 061 62
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective shampoo's 's preconditioner, 2024. URL
wwg. wwtp://.wwww...////.//..// 17748
Yurüi Nesterov. Lectures on Convex Optimization, volume 137 of Springer Optimization and Its Applications. Springer, 2018. ISBN 978-3-319-91577-7, doi:
Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. International Conference on Learning Representations, 2020. /ttt::::::///
h/tpeoolleeeuueee neura -tangents
Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler. AsdI: A unified interface for gradient preconditioning in pytorch, 2023. URL URL/ttts:://arx wwtp://..w...b/...aaao/oo
04684
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lec, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 2 21(140): 1-67, 2020. URL hoLLhtpp:///
org/ www.oorgggppaa.. pttpppppppppp///oepppa 074 html
11

---

Shai Shalev-Shwarde and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, USA, 2014. ISBN 1107057132.
Christopher J. Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl. Measuring the effects of data parallelism neural network training, 2019.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023. URL https::/arxiy
org b2g/abs/2332.115
Oriol Vinyals and Daniel Povey. Krylov subspace descent for deep learning. Artificial intelligence and statistics, pp. 1261-1268. PMLR, 2012.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam, 2025. /RL.atttss://xxx..oorr /bbs/
2409.11321
Sifan Wang, Ananyae Kumar Bhartari, Bowen Li, and Paris Perdikaris. Gradient aligmment in physics-informed neural networks: A second-order optimization perspective, 2025. /tLLtttss///xxxtt wtw.://ww...r//bb//52....0.. 00604
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael W. Mahoney, Adahessian: An adaptive second order optimizer for machine learning. 2021. URL httts:://xxxx: otgg/bo//00s 00719
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George E. Dahl, Christopher Shallue, and Roger Grosse. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model, 2019. URL htttttt/s:////iiii https/////aaxiiiooor////////////
04164
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham Kakade, How does critical batch size scale pre-trane gra-trainng?, 2025. URL
abs/2410. 2214 21676
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade, Deconstructing what makes a good optimizer for language models, 2025, URL https::///// //axxv/./iivvvv/aa wttgabbs/200..00023/ 07972
12

---

A PROOF OF EQUIVALENCE TO GAUSS-NEWTOR METHOD
In our method, we compute the Gauss-Newton update by minimizing the second-order Taylor approximation of the loss around the first-order Taylor approximation of the function.
Taking the first-order Taylor approximation of around O2.
(1) fo. (0, x)) = f(O,, V 0,)
Let (0) (0) = (0, y) be the loss function on the Taylor expansion of f around 0,.
Then the second-order Taylor approximation of Cs, (0) around 0g gives
()) = (0.) (0)) 01) + 2(0 01) (o) - 8x)
Denoting = 100 (0, x) and applying the chain rule for the gradients, we have
Vc, = (0)=
A Lo) 11 (0.
Then substituting and evaluating at O
(0)=0000 = =110
= 1.00001444-211000133332201000100
Then
(0) (11 + (0- 0.)
0c)
Let 9 denote the gradient of the loss at G2, iee (' x), V ((), Let G denote the Gauss-Newton term
Vof(br,)) VV Vo x). Then we can write
( [(f(0..)) g(0 0.) -0 0.) G 0.)
Since the Causs-Newton matrix is PSID, find 0r to minimize [/ by setting its gradient to zero:
g + (0 Oo) =
which results in the update rule
0" = 0o - G g
13

---

B IMPLEMENTATION OF GAUSS-NEWTON
Following notation from Appendix A
we have the Gauss-Newton objective:
(:22 (0) 01100::0000-CCC000.caa0000000000 1
(0 01) 00000//000.0.11112///00000000000 01)
Then the gradient is:
V V. ff(0,, (' 1110
0,)
We compute the gradient directly in the following manner, which results only one additional forward pass as compared standard backpropagation step:
Algorithm 3: Gauss-Newton Inner Optimization Step
Input:
train_stale (0),
reference params 9o. batch B = (input, target]
Define Sbunct logits f(O, innpuut
return logits
// on current batch B
Define logits (y):
loss CE.loss B.target)
return loss
Define VALUEANDGRADIENT(OO 0):
// Linearize f at 0o: f(0) = yo + Jo(0 - -0)
(Yo, jvp) +- linearize(faact 0u)
d0 - f On jvp(df)
11 v Jo d0
// Logit-space gradient and Hessian-vector
Vy + grad (ngiis)
(go, Hyu) - JVP((VV)))))) Vaggis (vo), (v)
11 H = 72 l at
// Pullback through the linearization: JJ (yo + H,v)
ji < linear timearttrnspposeijvp 0o)
V + jL(go Hy²)
Quadratic objective on the linear model loss t (iugin + 90 + 1 v. (Hyv)
return (loss, Vo)
loss, g VALURRANDGRAADIETTOO, 0 train state.params)
Apply update:
train state <- staaeeapply.grraieeee
return train state
14

---

C ADDITIONAL RELATED WORK
Also related to our work are optimizers that leverage diagonal approximations to the Hessian, such as AdaHessian (Yao et 2021) and Sophia (Liu al., 2024). These propose lightweight approximations the diagonal Hessian rather than layerwise approximations as in Shampoo and SOAP. However, Zhao et al. (2025) show that Sophia performs comparably to AdamW, suggesting the need to go beyond the diagonal Hessian.
D ADDITIONAL DETAILS ON OPTIMIZATION STRATEGIES
D.1 LEARNING RATE SCHEDULES
Global Cosine Schedule
Global + Inner Cosine
Constant + Inner Cosine
0.008
0.006
Leaminn
0.000
10111223141516111
1011121 611111617
Step
Stop
Step
Figure 4: Three learning rate schedules used for the Gauss-Newton and GN-prox-linear runs. From left to right: global cosine, global+inner cosine, and constant+inner cosine. Each inner cosine period lasts the duration of the optimization over the current Taylor expansion; outer step refers each parameter update on the model.
We experiment with three learning rate schedules: "global cosine," "global+inner cosine," and "constant-iineer co- sine." " From preliminary experiments we find that global+inner cosine did not generally outperform the global cosine and constant+inner cosine options, we do not use global+inner cosine in our main experiments.
D.2 LINE SEARCH STEP SIZES
1.0
Gauss-Newts Newton Layenwise GN
0.8
Step size
0.6
0.4
0.2
www
20
40
60
Global step
Figure 5: Resulting step sizes used from line search for Gauss-Newton and layerwise Gauss-Newton.
We find that line search is necessary for the performance of Gauss-Newton and layerwise Gauss- runs. For the line search, we evaluate the true model loss step sizes 2 + with {O...., 4] for full Gauss-Newton and
15

---

E {O..... 9) for layerwise Gauss-Newton. Figure 5 shows the step sizes at each step for iteration complexity runs (See Section 6.1.1
E MODEL DETAILS
Configuration
45M Model
150M Model
Hidden Size
512
768
Intermediate Size
2048
3072
Number of Layers
4
12
Attention Heads
8
16
Key/Value Heads
8
16
Table I: Model configurations for the 45M and 150M parameter LLaMA-based models.
F COMPUTE RESOURCES
All 45M parameter model runs are trained on Nvidia 80GB H100. 45M runs for Adam Muon, and SOAP trained for 2-3 hours and Gauss-Newton and GN-prox-linear trained for 1-2 days. 150M parameter model runs for batch size scaling experiments for AdamW, Muon, and SOAP are each trained using H100 for roughly 6-20 hours. 150M parameter model runs for batch size scaling experiments for Gauss Newton and GN-prox-lineat 1-3 days with 80GB Nvidia H100s using distributed data parallel (DDP). Layerwise Gauss-Newton Newton runs each trained for 3-7 days with 4 H100s. Iteration complexity runs trained for 2-3 days for AdamW and Muon and 15-30 days for SOAP, Gauss-Newton, and Layerwise Gauss-Newtt Newton.
Gauss-Newton
each trained for
G EXPERIMENTS ON 45M PARAMETER MODELS
Batch Size vs Loss
Target Loss 3.7
5.0
12
4.8
11
4.6
log2 Steps)
10
Final Loss
4.2
4.0
3.8
Optimal
1.2m
400k
1.2m
Batch size
Batch size
AdamW
Muon
SOAP
Gauss-Newton
GN-prox-linear
GN Upper bound
Figure 6: Left: Batch size vs final validation loss for 45M parameter models. All models are trained for 900M tokens on the C4 dataset (Raffel al. 2020) following Chinchilla scaling laws (Hoffmann al.) 2022). The dotted line marks the loss achieved by model trained with Muon with batch size of 32k tokens. This represents the upper bound of performance for our Gauss-Newton and GN-prox-lines methods as we use Muon with batch size of 32k tokens as the inner optimizer to compute the parameter update in each step. Right: Critical batch size scaling for 45M parameter models. The dotted line marks the optimal scaling trend, where no sample efficiency lost as batch size increases.
16

---

H DETAILS ON HYPERPARAMETER TUNING
Shared Hyperparameters
Model Size
45M
Batch Size
400k, 1.2m, 4m, 12m, 40m
Context Length
1024
Adam
Learning Rate
0.001, 0.003, 0.01, 0.03 0, (0.001, 0.01)
Weight Decay
Additional Warmup Fraction
0, (0.1) 0.9
Momentum B1
Adam
0.95, 0.99, 0.999
LR Scheduler
constant+EWA, cosine 0.9, 0.99
EWA Decay Rate
Muon
Learning Rate
0.03, 0.1, (0.3)
Additional Warmup Fraction
Momentum f
0.9, 0.95, 0.99 constant+EWA, cosine (0.3), 0.5, 0.7, (0.9)
LR Scheduler
EWA Decay Rate
SOAP
Preconditioning Frequency
Learning Rate Additional Warmup Fraction
0.01, 0.03, (0.1)
Momentum B1
0.7, 0.8, 0.9
Adam
0.7, 0.8, 0.9
LR Scheduler
constant+EWA, cosine (0.5), 0.7, 0.8, (0.9)
EWA Decay Rate
Gauss-Newton
Learning Rate
(0.001), 0.003, 0.01, 0.03, (0.1) 0, (0.2)
Inner Loop Warmup Fraction
Global Warmup Fraction
0, (0.2)
Momentum μ
0.95
Optimizer weight decay
(0), 0.001
Loss weight decay
0, 0.01, (0.1)
Parameter weight decay
0, (0.01, 0.03, 0.1)
LR Scheduler
constant+inner cosine, global cosine True, False (0.99), 0.999
Linesearch
EWA Decay Rate
Table 2: Hyperparameter configurations used for 45M models. Values in parentheses not used for every sweep. For the critical batch size plot (Figure riihht only the constant+EWA learning rate schedule was used. We start each run after an AdamW warmup of 5% for 45M parameter models; additional warmup refers to warmup starting from this checkpoint. For baselines, weight averaging is used only with the constant schedule. All Gauss-Newton runs without line search use weight averaging; runs with line search use no weight averaging. Inner loop warmup applies only the constant+inner cosine schedule.
17

---

Shared Hyperparameters
Model Size
150M
Batch Size
1.2m, 4m, 12m, 40m, 120m 1024
Context Length
Adam
Learning Rate
0.001, 0,003, 0.01
Weight Decay
0.001, 0.1
Additional Warmup Fraction
0, (0.1)
Momentum 3
0.9
Adam 32
0.95, 0.99 cosine
LR Scheduler
Muon
Learning Rate
0.01, 0.03, 0.1
Weight Decay
Additional Warmup Fraction
0
Momentum μ
0.9, 0.95, 0.99 cosine, (constant+EWA) 0.7, 0.9
LR Scheduler
EWA Decay Rate
SOAP
Preconditioning Frequency
Learning Rate
0.01, 0.03, 0.1
Weight Decay
Additional Warmup Fraction
Momentum B1
(0.7), 0.9, (0.95) 0.7, 0.9, 0.95 cosine, (constant+EWA)
Adam 32
LR Scheduler
EWA Decay Rate
0.9
Gauss-Newton
Inner Loop Learning Rate
0.003, 0.01, 0.03, 0.1
Inner Loop Warmup Fraction
0, 0.2
Global Warmup Fraction
0
Inner Loop Weight Decay
0, 0.01, 0.1
Optimizer Weight Decay
0.001
Momentum μ
0.95
LR Scheduler
constant+inner cosine, global cosine True, False
Linesearch
EWA Decay Rate
(0.9, 0.99), 0.999
Table 3: Hyperparameter configurations used for 150M models for batch size scaling experiments (Section 6.1.2) For Muon, EoN learning rate schedule was included in the sweep for the three largest batch sizes. For Gauss- Newton, due to high compute costs we hand-tune over the range of provided values rather than conducting the entire sweep at each batch size. Inner loop warmup applies only to the constant+inner cosine schedule. start each run after an AdamW W warmup of 5% for 150M parameter models; additional warmup refers to warmup starting from this checkpoint. For baselines, weight averaging used only with constant schedule. All Gauss-Newton runs without line search weight averaging with 0.999; runs with line search are default to weight averaging with lower values of 7. For iteration complexity experiments (Section 6.1.II: For AdamW and Muon we take the best hyperparameters from sweeps 40M batch size in the batch size experiments. For SOAP and Gauss-Newton we use the best hyperparameters from sweeps at 120M batch size and run a limited sweep over learning rate.
swept
18


Reviews:
[{"review_id": "Xe0iDWz1QN", "reviewer": "Reviewer_CWUm", "summary": "The paper test the performance of full 2nd-order method (Gauss-Newton) on LLM pretraining. The aim of the paper is to explore the potential of classical 2nd-order methods in mordern pretrain tasks.  Although the proposed implementation is still compute intensive, this trial is meaningful and provides valuable guidance for the future optimizer design.", "questions": "**Summary:** The paper test the performance of full 2nd-order method (Gauss-Newton) on LLM pretraining. The aim of the paper is to explore the potential of classical 2nd-order methods in mordern pretrain tasks.  Although the proposed implementation is still compute intensive, this trial is meaningful and provides valuable guidance for the future optimizer design. \n\n**Strength:**  The topic of \"exploring the 2nd-order methods in LLM pretrain\" is crucially important. The experiments are sound. The message is valuable and quite worth sharing with the community. \n\n**Disclaimer:** Many may criticize that the computational cost of this paper remains high, which is true. However, I do not consider it a critical issue. For me, the directional guidance provided by this work outweighs its computational limitations. The reasons are as follows:\n\n1. Although traditional second-order optimization methods have a rich theoretical foundation, there have been almost no meaningful attempts to apply them directly to modern neural networks or large language models, except a few early attempts on RNNs 10 years ago by Martins and Sustkever.  For modern LLM pretrain, we only know that several heavily approximated versions of Gauss-Newton can work in practice (e.g., Muon and Adam), but it remains unclear whether the non-approximated versions would be effective. This script gives a postive answer on a meaningful setting (nanoGPT speedrun).\n\n2. Approximated methods often develop new characteristics that lead to new theoretical interpretations—for instance, Muon exhibits properties related to matrix orthogonalization. Researchers tend to construct new theoretical frameworks to explain these methods rather than treating them as approximations of classical approaches. While this promotes academic diversity and inspire new thoughts, it also creates confusion in algorithm design: Should we follow classical theories or newer principles? This paper demonstrates that following classical theory can indeed lead to success. This directional guidance is highly valuable, and I believe this work has the potential to inspire many future optimizer designs.\n\nIn summary, I think this is an important paper that provides important directional guidance to future optimizer design. The computational overhead is a minor issue since the main contribution lies in \"guidance\" rather than \"an algorithm itself\".  I vote for clear acceptance. \n\n**Questions related to presentation:**  The paper presentation is not prefect and can be further improved.\n\n**Question 1:**\n I am quite confused by Section 7.2. LINEARIZED MODEL METHOD. \n\n1-1. What is the explicit form of this Algorithm 2? As an analogy to previous sections: the minimization over 2nd-order Taylor of the loss over the linearized model gives  w= w - G-1 grad, does algorithm 2 admit closed-form expression, and what is the closed-form if it exists?\n\n1-2. Is algorithm 2 still a 2nd-order algorithm? \n\n1-3. What is the implication on the loss landscape if Algorithm 2 works well?\n\n1-4. How do you solve the subproblem in Section 7.2. \n\n1-5. The name \"Algorithm 2: Linearized model method\" is quite misleading. Both Algorithm 1 and 2 are based on the Linearized model, an  if I understand correctly, Algorithm 2 actually incorporates higher-order loss function information than Algorithm 1.  The name \"Linearized model method\" may leave the impression that \"Algorithm 1 incorporates higher-order information than Algorithm 2\", but actually it is the other way around (correct me if wrong). \n\n\n**Question 2:** The algorithm presentation (Algorithm 1 and Algorithm 2) is confusing. It is impossible to understand the line \"Update\", e.g., in line 231. I suggest adding a hyperlink to tell readers which equation is used in defining this \"Update\" operator. Similarly, the line \"Line search\" is confusing as well. In optimization theory, line search usually refers to a particular procedure of choosing the stepsize of GD. But here you mean something different. Please add more detailed description.\n\n\n**Question 3:** \"We start all runs from the same AdamW post-warmup checkpoint.\" I dont understand this sentence. What do you mean by \"all runs\" and what is \"AdamW post-warmup checkpoint\"? Further, why not use random initialization?\n\n**Questions related to experiments:**\n\n**Question 4:** When computing L(θ_t + α(θ̂ − θ_t)) during line search, do you reuse the same data as in the inner loop, a held-out/proxy mini-batch, or freshly sampled data? Please specify for both full GN and layerwise GN\n\n**Question 5:** Ablation study on line search. I’m wondering whether other optimization algorithms also rely on line search. If they don’t, why does line search improve the performance of Gauss–Newton? In addition, how much of Gauss–Newton’s performance gain comes from the algorithm itself rather than from the use of line search?\n\n**Question 6:** The current script uses  GN^-1: w = w  - G^-1 grad. How about GN^-1/2: w = w  - G^-1/2 grad? Is there any feasible way to implement GN^-1/2? How would you expect it to perform?\n\n**Question 7:** How would GN perform if you keep training on more tokens? Will the benefit shrink as the token size reaches 1x or 2x Chinchilla, or even more tokens?\n\n**Question 8:**  In Figure 1, the x-axis is the number of outer steps, while Gauss–Newton uses a much larger batch size (240 M tokens vs. 40 M for AdamW/Muon) and each “step” includes several inner optimizer iterations and line-search evaluations. This makes the comparison on a per-step basis potentially misleading. I suggest reporting loss versus total tokens processed.\n\n**Code:** No code is provided in the submission. I strongly suggest the authors provide an anonymous code link.  I will check the code, and I believe it is the most efficient way to help address my questions above.\n\nI am more than happy to raise my score to 10 if my questions are addressed.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The topic of \"exploring the 2nd-order methods in LLM pretrain\" is crucially important. The presentation is clear. The experiments are sound. The message is valuable and quite worth sharing with the community.", "weaknesses": "The presentation can be further improved, as I explained below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5i78tR9NF9", "reviewer": "Reviewer_Hzbf", "summary": "The paper provides an empirical study of the application of the Gauss-Newton algorithm in LLM training. The paper studies a valid approximation method for lowering the memory and computational cost of the original Gauss-Newton method, making it feasible for LLM training. Empirical results show potential of this second-order approach, especially under large batch settings.", "questions": "1. In practice, I am curious about how we should tune the hyperparameter $ N $ (the inner loop length)? Schedule choices may also be a concerning point. There seems to be still a long way to go to put this second-order algorithm into practical training settings.\n2. What about the wall-clock training time of the algorithm? How is it compared to Muon or Adam?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The motivation of the paper is interesting and inspiring, bringing second-order optimization back to the regime of LLM training, and shows promising performance.\n2. The idea for approximating the full Gauss-Newton is interesting, making the algorithm practical and applicable to practical problems.", "weaknesses": "1. The notations and claims in the paper may need clearer explanations. I suggest the authors include more intuition on how this approximation is valid in approximating the updates of real Gauss-Newton and what the major differences are between the proposed approach and the real Gauss-Newton method.\n2. For the experiment part, I have two major concerns. Firstly, as an empirical study of the application of the Gauss-Newton algorithm, the experiment scale seems not large enough, since the paper only examines models of size 45M and 150M. Also, it seems that only when a very large batch size (compared to the full token number) is adopted, the benefits of the Gauss-Newton method are clear, which seems somewhat limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "v4UUziifFh", "reviewer": "Reviewer_hy1M", "summary": "This paper investigates the upper bound potential of second-order optimization for large language models (LLMs) using full Gauss-Newton (GN) preconditioning. By implementing a memory-feasible GN method via Jacobian-vector products, the authors estimate the theoretical limit of convergence speed relative to modern optimizers like SOAP, Muon, and AdamW.\nExperiments on 45M- and 150M-parameter LLaMA models show 5.4× faster convergence and larger critical batch sizes for GN compared to SOAP. The layerwise GN variant nearly matches full GN, implying that most curvature information is local and that higher-order loss terms are unnecessary.  \nOverall, the paper provides a rigorous and well-documented study that clarifies how far second-order optimization could improve LLM training efficiency, even if the proposed method itself is not yet practical at scale.", "questions": "* How much slower is each GN iteration in wall-clock time compared to SOAP or Muon?\n* Would the results generalize to models with stronger inter-layer dependencies (e.g., MoE architectures)?\n* How sensitive are the results to the choice of inner optimizer (Muon vs. AdamW)?\n* Can mixed-precision or blockwise GN approximations retain most of the benefits?\n* Does the GN method still outperform others in smaller-batch or lower-data regimes?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* Clearly defined and novel research question: measuring the upper bound of second-order optimization performance for LLMs.\n* Strong empirical methodology, including detailed hyperparameter sweeps and compute transparency.\n* Insightful findings: full GN drastically improves iteration efficiency, and layerwise GN nearly replicates its gains.\n* Excellent documentation and reproducibility, with appendices covering training setups, tuning, and resources.\n* Practical relevance: establishes a benchmark for how close current optimizers are to theoretical limits.", "weaknesses": "* Experiments are limited to 150M-parameter models, far smaller than modern LLMs.\n* The study focuses on iteration count, not wall-clock time or energy cost, so real-world efficiency gains remain unclear.\n* Comparisons may be slightly biased since the GN method uses Muon as an inner optimizer, giving it an inherent advantage.\n* No multiple-run variance reporting; results could benefit from confidence intervals.\n* Theoretical analysis of why layerwise GN performs nearly as well as full GN is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oRlbMiVVMZ", "reviewer": "Reviewer_Td6r", "summary": "The authors of this manuscript investigate the practical upper bound of second-order optimization methods for training Large Language Models (LLMs). They establish this bound by applying a full Gauss-Newton (GN) preconditioner to transformer models of up to 150M parameters, comparing its iteration complexity against strong baselines like SOAP and Muon. They report that the full GN update achieves a significant reduction in training iterations (e.g., 5.4x over SOAP). Furthermore, they find that a precise layerwise GN preconditioner, which ignores cross-layer information, achieves nearly identical performance to the full GN method.", "questions": "1- How does the performance and scaling of the full Gauss-Newton oracle compare to the KFAC family of optimizers [1], including scaled implementations like KAISA [2] and MKOR [3]?\n\n2- Could the authors provide an ablation that includes both Adam and the baseline AdamW [4] to clarify the impact of decoupled weight decay in this context?\n\n3- How do the observed iteration complexity gains from GN hold up in the later stages of training, after the initial rapid loss decrease and into the long-tail plateau phase?\n\n---\n\n[1] Martens, J., & Grosse, R. (2015). Optimizing Neural Networks with Kronecker-factored Approximate Curvature. arXiv:1503.05671.\n\n[2] Pauloski, J. G., et al. (2021). KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks. SC21.\n\n[3] Mozaffari, M., et al. (2023). MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates. NeurIPS 2023.\n\n[4] Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR 2019.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1- Studies showing the capability and limitation of current and ideal second-order optimization methods are very useful and insightful.\n\n2- The paper is well-written and addresses a clear and important question: what is the practical performance ceiling for second-order methods in LLM training?\n\n3- The finding that a precise layerwise Gauss-Newton preconditioner nearly matches the full GN method is a very strong and useful insight, as it suggests where future efforts in developing practical algorithms should be focused.", "weaknesses": "1- The current study lacks a comparison to an important family of second-order optimizers: the KFAC [1] family, which is orthogonal to the Shampoo-family methods (SOAP, Muon) that were analyzed. KFAC optimizers have been successfully scaled to large models, such as in KAISA [2] and MKOR [3]. It is important to compare the capabilities of such methods and understand how they relate to the Gauss-Newton \"oracle\" presented in this work.\n\n2- In line 117, the authors claim that the most common optimizer used for LLMs is Adam, while AdamW [4] is more commonly used and is the default for many libraries (e.g., HuggingFace Trainer). While the authors use AdamW as a baseline, a direct comparison including both Adam and AdamW as baselines would be insightful to clarify the impact of decoupled weight decay relative to the second-order gains.\n\n4- This study mainly focuses on the beginning part of the training process, where the loss decreases very rapidly (e.g., Figures 1 & 2 show convergence to a relatively high loss of 3.25-3.4). This is in contrast to the most dominant and expensive part of LLM pretraining, where the loss decrease is very gradual and almost plateaus. This focus is a limitation and leads to questioning the practicality of the findings for the full, end-to-end training process. I understand that scaling these optimizers (especially the full GN method) is very expensive, but without experiments that explore this later-stage convergence, the findings on iteration complexity gains remain inconclusive.\n\n---\n\n[1] Martens, J., & Grosse, R. (2015). Optimizing Neural Networks with Kronecker-factored Approximate Curvature. arXiv:1503.05671.\n\n[2] Pauloski, J. G., et al. (2021). KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks. SC21.\n\n[3] Mozaffari, M., et al. (2023). MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates. NeurIPS 2023.\n\n[4] Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)