Deconstructing What Makes a Good Optimizer for Autoregressive Language Models
Authors: Rosie Zhao
Depen Morwani
David Brandfonbrener
Nikhil Vyas
Sham M. Kakade
Venue: ICLR 2025 Poster (2025)
Keywords: optimization
LLMs
language models
Adam
Primary area: foundation or frontier models, including LLMs

Abstract: Training language models becomes increasingly expensive with scale, prompting numerous attempts to improve optimization efficiency. Despite these efforts, the Adam optimizer remains the most widely used, due to a prevailing view that it is the most effective approach. We aim to compare several optimization algorithms, including SGD, Adafactor, Adam, Lion, and Sophia in the context of autoregressive language modeling across a range of model sizes, hyperparameters, and architecture variants. Our findings indicate that, except for SGD, these algorithms all perform comparably both in their optimal performance and also in terms of how they fare across a wide range of hyperparameter choices. Our results suggest to practitioners that the choice of optimizer can be guided by practical considerations like memory constraints and ease of implementation, as no single algorithm emerged as a clear winner in terms of performance or stability to hyperparameter misspecification. Given our findings, we further dissect these approaches, examining two simplified versions of Adam: a) signed momentum (Signum)  which we see recovers both the performance and hyperparameter stability of Adam and b) Adalayer, a layerwise variant of Adam which we introduce to study the impact on Adam's preconditioning for different layers of the network. Examining Adalayer leads us to the conclusion that, perhaps surprisingly, adaptivity on *both* the last layer and LayerNorm parameters in particular are necessary for retaining performance and stability to learning rate.

Published as a conference paper at ICLR 2025
DECONSTRUCTING WHAT MAKES A GOOD OPTIMIZER FOR AUTOREGRESSIVE LANGUAGE MODELS
Rosie Zhao" Harvard University
Depen Morwani"
David Brandfonbrener'
Harvard University
Kempner Institute at Harvard University
Nikhil Vyas" Harvard University
Sham Kakade
Kempner Institute at Harvard University
ABSTRACT
Training language models becomes increasingly expensive with scale, prompting numerous attempts to improve optimization efficiency. Despite these efforts, the Adam optimizer remains the most widely used, due to a prevailing view that it is the most effective approach. We aim compare several optimization algorithms, including SGD, Adafactor, Adam, Lion, and Sophia in the context of autoregressive language modeling across a range of model sizes, hyperparameters, and architecture variants. Our findings indicate that, except for SGD, these algorithms all perform comparably both in their optimal performance and also terms of how they fare across wide range of hyperparameter choices. Our results suggest to practitioners that the choice of optimizer can be guided by practical considerations like memory constraints and ease of implementation, as no single algorithm emerged as a clear winner in terms of performance or stability to hyperparameter misspecification. Given our findings, fied versions of Adam: a) signed momentum (Signum) which we see recovers both the performance and hyperparameter stability of Adam and b) Adalayer, a layerwise variant of Adam which we introduce to study the impact on Adam's 's preconditioning for different layers of the network. Examining Adalayer leads us to the conclusion that, perhaps surprisingly, adaptivity on both the last layer and LayerNorm parameters particular are necessary for retaining performance and stability to learning rate.
further dissect these approaches, examining two simpli-
INTRODUCTION
As language model architectures increase in scale, pretraining becomes more expensive. In response. numerous efforts have been made to design efficient optimizers to mitigate these costs, and yet Adam (Kingma Ba 2013) remains the primary optimizer used for training language models. This persistent preference for Adam is rooted in an underlying belief that Adam generally outperforms alternative optimization algorithms. Although newly proposed optimizers run ablations to demonstrate superior performance to Adam for select architectures and tasks (Liu et al. 2024| Chen et al. (223), there no consensus among the literature about the relative performance of these optimizers. In fact, to the best of our knowledge, Kaddour et al. (2024) is the only work comparing these optimizers but in the context of masked language modeling and at a single model scale.
In this work, we aim to rigorously evaluate whether the widespread reliance on Adam is justified and to further explore the characteristics that make adaptive optimizers particularly effective for training autoregressive language models. As such, perform comprehensive sweep across different optimizers, hyperparameters, architectures, and scale. Along with looking at optimal performance, we argue that due to the difficulty of hyperparameter tuning with increasing scale (Yang et al.) [021] the stability of performance with respect to hyperparameter choices is equally important. Prior work has explored the learning rate stability of Adam (Wortsman al.) 2024). We extend this investigation to include the stability of multiple optimizers with respect to various hyperparameter
"qqual contribution, randomized author ordering. Correspondence rosi ezhacog harva rd. edu.

---

Published as a conference paper at ICLR 2025
150m (standard)
300m (standard)
3.4
3.2
Final Valitation Looa
3.3
3.1
Validalion
3.2
3.0
10"
10"
10
10"
10"
10"
10'
Mutppi of opoiaal LR
Multipi at optimal LR 1.2b (standard)
800m (standard)
3.0
Fintt Valldatoon Loos
2.8
2.9
Validafion
2.7
2.6
10"
10"
10
10
10"
10'
10'
10²
Muttile of optimal
Multiple optimal
SED
Adam
Adafactor
Lion
Signum
Figure 1: Final validation loss when training language models with 150m, 300m, 600m, and 1.2b parameters, sweeping across learning rates for five standard optimizers (SGD, Adam, Adafactor Lion, and Signum). Plots have been shifted to align the optimal learning rates for each optimizer. Except for SGD, other optimizers seem comparable in their optimal performance and stability with respect to learning rate tuning.
choices. Surprisingly, we find that multiple optimizers introduced in the literature after Adam such Lion (Chen et al. 2023) and Adafactor (with momentum) (Shazeer & Stern 2018 Zhai al. 2022) demonstrate robustness comparable to Adam and significantly superior to SGD. Figure | illustrates the remarkable similarity in performance and robustness of these optimizers across different learning rates and across multiple model scales (150m, 300m, 600m, and 1.2b parameters). This challenges the prevailing notion that Adam should be the default optimizer, where we see no single algorithm emerged as a clear winner in terms of performance or hyperparameter stability.
Following our initial ablations, wish identify the essential components of these optimizers that facilitate performance and stability. Thus, we conduct a series of investigations of simplified versions of these algorithms. We study signed momentum (Signum) (Bernstein et al. 2018 2019), a special case of Lion. Prior works have also studied its similarities to Adam (Balles & Hennig 2018). We find that Signum also recovers the stability and performance exhibited by Adam. This finding aligns with recent work (Kunstner et al., 2023). suggesting that the primary distinction between SGD and Adam is driven by Adam's ' resemblance to signSGD.
To further understand the role of preconditioning on various network parameters, we study Adalayer, which performs preconditioning on a per-layer basis. We empirically demonstrate that this variant nearly recovers the stability and performance of the other optimizers in previous ablations. Through empirical studies of Adalayer and its variants, we show that adapting the parameters of the last layer and LayerNorm parameters in transformer is necessary achieve stability and performance.
Main contributions
We empirically study the stability to hyperparameters of various optimization algorithms including SGD, Adam, Lion and Adafactor, showing that with the exception of SGD, these optimizers are comparable in terms of both performance and hyperparameter stability. This holds across multiple scales (150m, 300m, 600m, 1. 1.2b) and across two transformer architecture variants (Section().
Our implementation of Adafactor adds hack momentum described in Section 2.2
2

---

Published as a conference paper at ICLR 2025
To dive into the reasons behind the stability and performance of these algorithms. empirically examine signed momentum (Signum), and show that it recovers their stability and performance (Section 2.6
We study a coarser variant of Adam called Adalayer, that does per-layer preconditioning and recovers much the stability and performance exhibited by Adam (Section 3.1
Through an empirical study of Adalayer and its variants, study the importance of adaptivity with respect to different layers of the network. (Section 3.2
2 COMPARING OPTIMIZERS ACROSS HYPERPARAMETERS, ES AND SCALE
2.1 METHODOLOGY
To conduct our experiments, we start with hyperparametens recommended by previous work (e.g.. 31 = 0.9). We initially perform learning rate sweep identify the optimal learning rate. After determining the optimal learning rate for each algorithm, we conduct one-dimensional sweeps for each of the other hyperparameters.
A limitation of this methodology is the potential neglect of "higher-dimensional" interactions between hyperparameters. This an important direction for future work, but beyond the computational budget of this project. For example, some parameters like batch size and learning rate indeed are likely to exhibit 2D interactions (Shallue al.) 2019] Porian al. 2024). However, we argue that the sweeps provide tractable methodology that gives us useful signal about the hyperparameter stability of a variety of algorithms around the parameters that are common in practice.
2.2 SETUP
We train decoder-only language models on C4 tokenized with the T5 tokenizer (Raffel et al. 2020) and report results in terms of validation loss. Due the computational resources required to sweep numerous hyperparameters, optimizers, and scales, we restrict the scope of our work to autoregressive language models, leaving other domains and language model architectures future work. As discussed in the introduction, we argue that it is best to evaluate algorithms both in terms of the loss achieved by the best hyperparameters (performance) as well as the robustness across values of the hyperparameters (stability).
In this section we first present the results sweeps across the two most sensitive hyperparameters: learning rate and momentum (31). We sweep across five algorithms: Adam, Adafactor, Lion, Signum, and SGD, Further ablations of weight decay, warmup. $2, and e for the 150m standard model can be found in Section 2.5 Before diving into the results, we provide some more details about the setup.
Algorithms. We use the standard Pytorch implementation AdamW (Paszke et al. 2019), the timm implementation of SGDW (Wightman) 2019), and the OLMo implementation of Lion (Groeneveld et al. 2024). Following Zhai et al. (2022) we implement ourselves a modified version Adafactor which maintains the factored estimates of second moments but has momentum i.e. it is equivalent to Adam with factored second moment estimates. Since Signum is equivalent to Lion with 31 32 we reuse the OLMo implementation of Lion (Groeneveld et al., 2024) for it. We conducted experiments with the Sophia optimizer (Liu et al. 2024) in Appendix G However, since it does not outperform Signum (which can be achieved by setting p in Sophia), we did not include in other plots.
Models. We start from the OLMo codebase (Groeneveld et al., 2024) and train decoder-only
transformer models of four sizes: 150m, 300m, 600m, and 1.2b, where the parameter count refers to non-embedding parameters. The models have widths of 1024, 1024, and 1408 and depths of 12, 24, 24. The MLP hidden dimension is 4x the width. The activation function is GeLU (Hendrycks) & Gimpel| 2016). We use RoPE positional encodings (Su et al., 2024). Attention heads always dimension 64. We use Py Torch default LayerNorm. Following Wortsman et al. (2024) we do not learn biases for the linear layers LayerNorms. We train in mixed precision with b6
3

---

Published as a conference paper at ICLR 2025
150m (no QK norm or 2-loss)
300m (no QK norm or z-loss)
Findl Valdaton Loss
Final Valiatton
10"
10"
10
10"1
10" Muttple optmal Sigrum
10'
10*
optiml
SGD
Adam
Adafactor
Lion
Figure 2: Sweeping learning rate without QK norm or z-loss for (Left) the 150m model, and (Right) the 300m model. These models are less stable than the standard model, but the same general trend across algorithms hold here.
150m-long (standard)
150m-long OK
3.25
3.55
3.20
3.20
Final Valdation Loss
3.16
Finil Validation Loss
3:15
3.10
3.10
3.06
3.06
3.00
3.00
2.96
2.96
2.90
2.90
10"1
Multple of optmal LR
Multple of optimal LR
30D
Adam
Adafactor
Signum
Figure 3: Sweeping learning rate 150m models trained for 4x longer (100k steps) than in the base runs for (Left) the standard model, and (Right) the model without QK norm or z-loss. Compared to the shorter runs, these models achieve better performance and increased stability across learning rates.
Training variants. We note that Wortsman et al. (2024) observe that QK LayerNorm (Dehghani let al. 2023) and z-loss (Chowdhery et al.] (023) can have substantial effects on the stability of model training. As such, consider with QK LayerNorms and 7-loss with coefficient le-4, and no QK norm or z-loss which refers to the same model without the QK norm layers or the z-loss.
variants in experiments: standard which refers model
Token counts. For all models, we use a batch size of 256 and sequence length of 512 (as in Wortsman al. (024)). We default training models for the approximately "chinchilla optimal" number of tokens that is =20 times the number of parameters. Explicitly, this means for the 150m models we train for 25k steps or x3.3b tokens. The 300m models are trained 50k steps, the 600m models are trained for 100k steps, the 1.2b models are trained for 200k steps, and the 150m-long models are also trained for 100k steps.
Other hy peperaarreeeers. We default to using 0 weight decay. We default to using learning rate schedule with 10% of the training steps for warmup and then cosine decay with minimum that is 10% of the maximum learning rate. We default to 3n 0.95 and 1e-15 following Wortsman et al. 2024). These parameters are ablated Section 2.5
2.3 SWEEPING LEARNING RATES
First, we sweep over the most important hyperparameter: learning rate. Note, in all these sweeps learning rate set 31 0.9 for all algorithms except for SGD, where set B1 0.98. As
4

---

Published as a conference paper at ICLR 2025
(standard)
100m (no 2-1050)
300m (standorc))
3.40
3.40
9.20
5.90
2.30
0.10
5.25
3.25
3.05
2.20
3.00
3.15
3.16
2:50
305
3.06
2.58
0.5 0.5
0.9
Figure 4: Sweeping momentum for fixed learning rate across three settings: (Left) 150m standard, (Middle) 150m with no QK norm or z-loss, (Right) 300m standard. Adam and Adafactor are similarly robust to 31 while Lion and Signum are slightly more sensitive low values and SGD is substantially more sensitive.
we will see in the following subsection, SGD is more sensitive to the momentum hyperparameters and requires more momentum to be competitive with the other optimizers.
Main results for our standard architecture across four scales are presented Figure [ Note that the x-axis shifts the learning rates to align the optimal learning rates across algorithms. In terms of absolute learning rates, sweep in multiples of v10 from le-4 to I for Adam and Adafactor, from le-5 to le-1 for Lion and Signum, and from 1e-3 to 10 for SGD. We report on the optimal learning rates for each optimizer in Appendix B
The key takeaway that not only do the algorithms achieve similar performance at the optimal learning rate, but the learning rate stability itself is similar across algorithms and scales. The one exception is SGD, which worse both terms of optimal performance and terms of stability.
Further ablations presented in Figure and Figure 3 illustrating performance for models with no QK norm or z-loss and 4x longer training time respectively. While we find that the architecture choices clearly impact the amount of stability learning rate, the cross-algorithm comparisons remain the same: Adafactor and Lion are competitive with Adam, while SGD is worse both in terms of performance and stability learning rate. Similarly, training for longer can improve performance and stability to learning rate, but does not change the high-level cross- algorithm comparisons.
Takeaway: performance and stability learning rate are comparable across the non-SGD algo- rithms that we tested.
2.4 SWEEPING MOMENTUM
Now we also sweep across momentum values (i.e. 31 To do this sweep we fix the per-algorithm learning rate to be the optimal learning rate from the corresponding learning rate sweep.
Results are presented Figure 4 We observe that across various settings, the robustness to 31 is similar across the non-SGD algorithms when we stay in the range of momentums between 0.8 and 0.98. However, for high 31 Lion is better and low B3 Adam and Adafactor are better. Again we observe SGD being very sensitive to momentum.
Takeaway:
performance and stability to momentum are comparable across the non-SGD algorithms that we tested if we stay within the usual range of momentum values.
2.5 ADDITIONAL HYYERPARAMETTER SWEEPS
We also sweep learning rate and momentum. We observe that SGD is less stable with respect to weight decay and
variety of other hyperparameters in Figure 5
using the best per-algorithm
Note that in Lion, both 31 and 82 can be thought of as different types of "momentum" with 31 being the "one-step" momentum and 82 the "long-term" momentum. For consistency, we only sweep 88 here and sweep 3g in Section 2.5

---

Published as a conference paper at ICLR 2025
Finel Voldetion LOS
Einil Voldatoon LOSS
-.3
3.2
322
3.2
311
3.3
10
10.4
10"1
10"1
1014
10-4
105
3
5
10
20
40
60
27
20
350
705
324
Weight
Percent Thaining
Batch Size
Fini Weldebion Loss
SGD Adafactor
3.2
32
Signum
0,919900999 0.295 0.99 0.95
0.9
10
10"33
13-11
10"
10
beta,
Figure 5: Sweeps over other hyperparameters. Top: weight decay, warmup duration, and batch size. Bottom: and 32. We generally find little effect for the non-SGD algorithms, however there parameters that differ from our defaults that can offer to 0.02 improvements in validation loss
warmup length. And while it is possible to get small benefits from higher weight decay, longer warmup, and higher B2 than our defaults, generally the algorithms are much more stable these parameters than learning rate and momentum. We observe two exceptions: Lion shows poorer performance at extreme values of B2. However, it is important note that the (32 parameter in Lion functions more like form of "momentum", whereas in Adam and Adafactor, 83 regulates the moving average of the squared gradients. We also see that the performance across optimizers worsens at high batch size. This is expected as optimization algorithms are known to exhibit diminishing performance with increasing batch size (Shallue et al.. 2019; Zhang et al. 2024a).
stable with respect other hyperparameters- with the exception of batch size- and the possible gains in performance are relatively small compared to learning rate and momentum.
Takeaway:
generally algorithms
2.6
SIGNUM RECOVERS THE PERFORMANCE AND STABILITY ADAM
150m (standard)
150m (na OK norm 2-0050
300m (slandarc)
3.35
3.25
3.15
3.50
3.05
3.20
3.20
3.00
3.16
3.90
3.10
2:00
5.05
3.06
2.56
0.99
0.90
0.55
0.5
0.5
0.5
0.0
0.99
0.96
0.35
0.9
0.8
0.5
0.99
0.90
0.90
0.9
0.0
0.5
bet]_
be5a
teto_1
Ad2rt
Adam (30d bO3M))
Sgguu
Figure 6: Sweeping momentum with 31 = 32 tied together for Adam (dashed) and compared to Signum and Adam with fixed 83 0.95 (solid) across three settings: (Left) 150m standard, (Middle) 150m with no QK norm or z-loss, (Right) 300m standard. When 31 B2, Adam behaves very similarly to Signum.
In Figure I we observed that Adam and Signum have similar performance and stability for language modeling, even at scale. The following lemma from prior (alkeBalles & Hennig 2018) shows that Adam performs variance-adjusted sign gradient descent.
Lemma Balles & Hennig (2018). Consider a parameter with a history of gradients 9t- Let be the random variable that is equal st-r with probability (1 f3r }&} and be the random

---

Published as a conference paper at ICLR 2025
variable that is equal to 9t-r with probability (1 - -2) update ōsignum are related by
The Adam update dAdam and the Signum
| mm]
SAdam =
dsigmuu
If 31 B2 then m in Lemma and hence the parameter-wise ratio of Adam and Signum updates is equal to the ratio of the mean and the square root of second moment of m. Intuitively, this holds because when 31 32. the first moment estimates of Signum and Adam, and second moment estimates of Adam, average the previous gradients with same coefficients ((1 (3)8") This
intuitively suggests that when 31 = 82, Adam and Signum may behave similarly [ This motivates the conjecture that the main benefit of Adam over Signum is the fact that in Adam, .2 can be varied
independently of 31. In Figure [ we have 32 = 0.95 and 31 0.9 which are close, and as pointed out carlier, both optimizers have similar performance and stability.
We examine this hypothesis further in Figure varying 31 and setting 32 = 31, and again find that Signum and Adam behave very similarly. However, we also note that when we vary B1 for Adam while fixing 32 we get more stability for 31 as compared to Signum.
Takeaway: With B2 31 Adam and Signum behave similarly and the standard setting for training language models (32 = 0.95, 31 = 0.9) is close to this.
3 INVESTIGATING THE IMPACT OF ADAPTIVITY FOR OPTIMIZER STABILITY AND PERFORMANCE
Ablations in the previous section revealed the striking similarity in performance and stability across multiple optimizers compared to Adam. Adam and its other variants are designed to have a high degree of adaptivity at a fine granularity (per-parameter learning rates) throughout the training process. This adaptivity is often credited with the stability and robust performance observed in these optimizers. However, a critical question arises: to what extent is this adaptivity needed for different parameters of the network? By identifying the necessity of adaptivity for different network components to ensure both performance and stability, we aim to discern whether simpler optimizers like SGD can achieve similar benefits with minimal modifications. Since higher momentum often play the same role as a better preconditioner and to have all algorithms on an equal footing, we will fix 31 = 0.9 for all optimizers in this section.
The main optimizer we study in this section is a "layer-wise" variant of Adam, which we coin as 'Adalayer'. We use Adalayer for our investigations because it lends a greater ease of understanding compared to full-fledged Adam in identifying parts of the network which may be particularly critical for optimizer performance and stability. Note that this layerwise variant is a special case of previously known optimizer called Blockwise Adaptive Gradient with Momentum (BAGM) (Zheng & Kwok 2019)
3.1 ADALAYER
Algorithm 1: Adalayer
Parameters: Learning rate 1). exponential decay rates for the moment estimates 31, B2, number of steps T.c
To study the behavior of adaptive opti- mizers like Adam, we begin with describ- ing layer-wise version of Adam which we refer to as Adalayer. Adam, Adafac- tor and Adalayer all (approximately) store the diagonal second moment matrix, but with coarser and coarser granularity; for a Adam explicitly
while S do for each layer I with p parameters do
g +- ViL(w,)
of + B2 o/ 1 + (1 B2)
1/2 11/2
11
m B1 m 1 + (1 33 5))gt
layer of dimension maintains the second moment matrix using mn parameters in the shape of matrix. Adafactor stores row and column averages of the second moment matrix which serve as a rank-1 approximation to the second
w/+1 + wt 17- v 1/tt
mi
end end
Thii is not true in theory, the ratio can vary across parameters.
7

---

Published as a conference paper at ICLR 2025
Efllctive learning rate quantils for layer blocks
Effactiva learning quantilles other marri blocks
laamiigg quantiies LayenNorm blocks
210
207
Uhedee lemminy 44h
10P
010
200
105
107
'40
5000
10000
15603
30300
7,000
12000
1000
$0.00
5300
10300
1900
3000
Slep
Ovardi
0.1 02
05
0.8
Figure 7: Quantiles of effective learning rates (n/((v + +e) for each layer () for the last layer blocks (Left), the LayerNorm blocks (Right), and the other matrix blocks (Middle) for 150m model trained using Adalayer*. Unlike the other matrix blocks and LayerNorm parameters, the effective learning rates across logits vary across multiple orders of magnitude, providing evidence for the need to precondition them separately.
moment matrix. Finally, Adalayer stores a single scalar which is the average of the second moment matrix. We will later consider a generalization of Adalayer where instead of averaging second moment over layer we will average it over a "block" of parameters which can be subset of a layer. We note that similar algorithms have been studied before (Ginsburg et al.. 2019; Agarwal et al., (2020) but we choose to study this variant since it is direct analogue of Adam and Adafactor. A simplified version of Adalayer optimizer given Algorithm [] other details such as bias correction are kept same as that for Adam. In Appendix C. we provide more details about our Adalayer implementation and correction we make to how Adalayer treats the last layer in order to achieve similar performance and stability to Adam. Specifically, Adalayer when naively applied for each layer is neither performant nor stable to learning rate (Figure 10); however, if we additionally treat the set of weights in the last layer feeding into each logit its own block, this recovers most of the performance and stability of Adam (see the dotted blue lines in Figure 8). We henceforth refer to Adalayer with this correction as Adalayer
To study how Adalayer* preconditions the network, we plot effective learning rates used for different logits by Adalayer* in Figure (eeft). Here, the effective learning rate for layer I in the network is ne/(v o. + c). We find that Adalayer* indeed uses vastly different learning rates for different logits, supporting our hypothesis that preconditioning weight in different logits separately is important for performance and stability.
3.2 BOTH THE LAST LAYER AND LAYERNORM PARAMETERS NEED ADAPTIVITY
The results using Adalayer* in the previous section suggest that all layers except the last layer only need iteration-dependent scalar correction their learning rate. We now ask a stronger question: do we need these scales at all? Or can we train the remaining layers with SGD? This hypothesis supported by looking at Figure 7 (middle) where we observe that the learning rates for different matrix layers (except the last layer) assigned by Adalayer* are remarkably similar.
To test this, we train the last layer with Adalayer* (fixing a learning rate of 3.16e - 3) and the rest of the layers with SGD, both with 81 = 0.9. In & (Left) we show the results while sweeping over SGD learning rates from 0.1 to 3160. While this improves upon the performance of SGD, we do not recover stability of the Adalayer and Adam baselines. We trace this instability to LayerNorm blocks: Figure (Left) shows that the effective learning rates for the LayerNorm blocks are much smaller, which suggests that they may destabilize higher SGD learning rates. To ameliorate this, in Figure g (Right) we add LayerNorm parameters to those being trained with Adalayer* and find that this is sufficient to recover both performance and stability of Adalayer". In Figure 11 in the Appendix, see this trend continue to hold for 300m and 600m parameter models.
We conduct additional experiments investigating these 'hybrid' variants of SGD in the Appendix. Despite most language model parameters being in matrix layers, we show that we only apply * on the matrix layers and train the last layer and LayerNorm layers with SGD, this does not
8

---

Published as a conference paper at ICLR 2025
150m
150m
3.7
3.7
3.6
3.6
Valldation
3.4
3.4
3.2
3.2
3.1
3.1
10-1
10
10a
10a
10-a
10'
100-
10a
Multiple of optimal LR
Multiple of optimal
SGD
Adalayer
SCD
Adalayer
Adam
Adalayer* ILL, 3.16e-3) + SGD ILN Matrix)
Adam
Adalayer [LL + LN, 3.16e-3) + SGD (Matrid)
Figure& 8: (Left): Training the last layer using Adalayer* with a fixed learning rate of 3.16e 3 and other LayerNorm and matrix blocks using SGD achieves better performance than SGD, but does not recover stability. (Right): Training both the last layer and LayerNorm blocks using Adalayer* and the other matrix blocks using SGD nearly recovers or exceeds performance of Adalayer*, and achieves stability across learning rates. Dotted lines are baselines from optimizers previously given in Sections 2 and 3.1
recover performance and stability. Another plausible remedy to address the small effective learning
rates of the LayerNorm blocks is to simply not train these LayerNorm parameters. In Appendix D
Figure [3] we show that this improves performance and stability relative to SGD but does not recover
* performance. Finally in Appendix [E, we replace SGD with AdaSGD (Wang & Wiens) 2020) which still uses global learning rate but is now adaptive through training; in Figure [66 show AdaSGD + Adalayer* can recover full Adalayer* performance, but a small gap remains to recover Adam performance.
We acknowledge that while adaptivity for the last layer and LayerNorm parameters seems necessary retain optimizer performance and stability, there remains a small gap which is fulfilled by adaptivity on the other network parameters; in Figure 14 we also try training 150m and 300m models using SGD on the matrix blocks and Adafactor on the last layer and LayerNorm blocks; we find that performance and stability is comparable, but does not exceed that of Adafactor on the entire network.
Further, note that a caveat of the above results that we have introduced an additional hyperparameter- SGD learning rate, which we are sweeping over- while keeping the Adalayer* learning rate in the last layer and layer norm layers fixed. While decoupling the learning rates here is needed (due to SGD's performant learning rates being orders of magnitude higher than that of Adalayer*), this may be responsible for the observed stability. To address this, we perform the following experiment: stop updating the second moment estimates for all layers except the last layer and LayerNorm blocks after initialization. This implies that these layers are effectively being trained by SGD with fixed learning rate, though unlike the above results, these learning rates are different for different layers. We implement this by passing 1000 batches to initialized 150m and 300m models obtain second moment estimates for all layers without letting the model take a gradient step, and then allowing the model to train as normal under the same settings as all of our ablations. As in our previous investigation, we fix other hyperparameters to be the same: 31 = 0.9, 3n = 0.95 and e = 1e 15.
train all the layers with Adalayer* (sweeping over learning rate 7) but
In Figure [ we show the resulting learning rate sweep for freezing Adalayer* learning rate scales at initialization (with the exception of the last layer and LayerNorm). Surprisingly, we find for the 1 150m model that we can almost entirely recover the stability and performance of Adalayer*, For the 300m model we also match or exceed the performance of Adalayer", and even nearly match the peak performance of Adam. Note again that this sweeps learning rate across all network parameters. This provides further evidence for the importance of adaptivity in the last layer and LayerNorm, where in contrast we could have used fixed ratios from initialization for all other parameters to recover the performance and stability of Adalayer*. In Appendix F we present additional experiments which include allowing Adalayer to update for only one of the last layer or LayerNorm, or turning off LayerNorm training. None of these partial modifications fully recover performance and stability.

---

Published as a conference paper at ICLR 2025
3.7
150m
3.20
300m
3.6
3.15
Final Validation Loss
Phal Valddttio
3.5
3,05
3.3
3.00
3.2
2.95
3.1
10:4
10
10
10
10-3
10-2
10
Loaming Rate
Leariing Rate
Frozen Adalayer* (all but and LN params)
Frozen Adalayer* (all but LL and LN params)
Adalaayer*
AdaLayer*
Adam
Adam
Figure 9: Training 150m (Left) and 300m (Right) models using fixed Adalayer* learning rate ratios from initialization, with the exception of last layer and LayerNorm parameters. This almost entirely matches the performance and stability of Adalayer in the 150m model, and exceeds Adalayer**s peak performance to be comparable with Adam.
4 DISCUSSION AND LIMITATIONS
After a comprehensive comparison of variety of optimizers for language modeling, we have found that many optimizers seem to be roughly equivalent both in terms of optimal performance and hyperparameter stability. The only differing axis for practitioners to consider is thus memory constraints, but based on our results, memory considerations for optimizer choice should not affect the performance of their training run, least across the optimizers tested (i.e. optimizers using diagonal preconditioning). For tuning guidelines, we have reported the optimal learning rates in Appendix B found in sweeps starting point for practitioners in tuning. We also observed that most other parameters are very stable for the optimal learning rate (with exception to momentum) and thus believe tuning learning rate and momentum should be prioritized if one has limited compute budget for hyperparameter sweeping.
Diving deeper, have shown that the treatment of the last layer and LayerNorm parameters is crucial for realizing the benefits of adaptive optimizers. This highlights the following research directions the most promising for developing optimizers for language model training. Firstly, our results suggest that within the realm of diagonal preconditioning optimizers, there may be room for more memory-effiiien alternatives AdamW be designed beyond using factorization techniques (eg. Adafactor with momentum and Lion from our initial sweeps). We have provided one direction for cutting down on memory overhead, where we have shown that the need adaptivity is distributed unevenly across different parameters of the network (with similar ideas already explored in concurrent work (Zhang et al.) 2024c)). Secondly, terms of improving the performance of these models, our work suggests that optimizer choice is unlikely be promising approach- if we are limited to diagonal preconditioning optimizers. Thus, is likely that we need turn to non-diagonal methods current methods; this is indeed where previous work has demonstrated improved performance to Adam (Gupta et al. 2018a).
to improve
Of course, there are several limitations our study including the fact that due to computational constraints we only ablate few architecture decisions, that we only consider one dimensional hyperparameter sweeps, fix batch size, and that limit our study to autoregressive language modeling with a single dataset. Despite these limitations, we believe that the study sheds new light on the fundamentals of optimization for language modeling.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
SK, DM and RZ acknowledges support from the Office of Naval Research under award N00014- 22-1-2377 and the National Science Foundation Grant under award #IIS 2229881. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. NV, DM and RZ are supported by a Simons Investigator Fellowship. NSF grant DMS-2134157, DARPA grant W91 INF2010021 Land DOE grant DE-SC0022199, RZ and DM are supported by Kempner Institute Graduate Research Fellowships.
REFERENCES
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. CoRR, abs/2002. 1100002.1100... 1803, 2020. https:ts:://rr
orq/abs 2002. 11803
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). In The Twelfth International Conference on Learning Representations, 2024. URL/attpa://opeeer///
net http://ouumnii-uuu5s5.
Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients, 2018. htLhttt:: pttttp::://ppeneevviee wwtp://wwwwwweeeeeeeeeeeeeeeeeee Sttp:/oooommmid-d_ww-eeee
Lukas Balles, Fabian Pedregosa, and Nicolas Le Roux. The geometry of sign gradient descent, 2020.
Thomas Bartz-Beielstein, Carola Doerr, Daan den Berg, Jakob Bossek, Sowmya Chandrasekaran, Tome Eftimov, Andreas Fischbach, Pascal Kerschke, William La Cava, Manuel Lopez-lbanez, et al. Benchmarking in optimization: Best practice and open issues. arXiv preprint a0X 2020.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Arizzadenesheli, and Animashree Anandkumar. SignSGD: Compressed optimisation for non-convex problems. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 560 569. PMLR, 15-15 Jul 2018. URL http://ccceeiings......eeeee///a h 8a. html
https::/prroeeddnngg .mt.ss..l..ppp..... www.pprssssmmmeemmmeeeppeeerreee
Jeremy Bernstein, Jiawei Zhao, Kamyar Azizzadenesheli, and Anima Anandkumar. signSGD with majority vote is communication efficient and fault tolerant. In International Conference on Learning Representations, 2019. URL httpt:::///.. ptattts:://ppeeevviiii
net/ wew./eeoormiddBBae.e.
7ACY7
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe Le. Symbolic discovery optimization algorithms. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL wwt.:/outtteootemeeeemm
net.e/wweeeeeeeeeeeeeeeeeeeeeeee forum? id-idddeee neeeneeeeeeee 62eqLFC2
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski. Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to billion parameters. In International Conference on Machine Learning. pp. 7480-7512. PMLR, 2023.
Boris Ginsburg. Patrice Castonguay, Oleksii Hrinchuk, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Huyen Nguyen, Yang Zhang, and Jonathan M Cohen. Stochastic gradient methods with layer-wise adaptive moments for training of deep networks. arXiv preprint 1r0iv:1905.1 11286, 2019.
11

---

Published as a conference paper at ICLR 2025
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et al. Olmo: Accelerat- ing the science of language models. arXiv preprint arXiv:2402.008388 2024.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018a.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Leurning, ICML, 2018, Stockholmsmüssun, Stockholm, Sweden, July 10- 15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1837-1845. PMLR, 2018b. URL http:ttpppppcceeiinnggg .an.gssmmmml mir. press/v80 http://ppaaaaaaaaaaaaaaaa html
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415. 2016.
Samy Jelassi, David Dobre, Arthur Mensch, Yuanzhi Li, and Gauthier Gidel. Dissecting adaptive methods in gans, 2022.
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural network geometry? In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
net. www.:eetttooomr wew.rrummmmiiiiii
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt J Kusner. No train no gain: Revisiting efficient training algorithms for transformer- based language models. Advances in Neural Information Processing Systems, 36, 2024.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. URLhttp //arxiv.org/ abs 11221 6980
Ananya Kumar, Ruoqi Shen, Sebastien Bubeck, and Suriya Gunasekar. How to fine-tune vision models with SGD. In The Twelfth International Conference on Learning Representations, 2024.
URL htttps:://ppeeeeeeaeeteee net cenest/forunng forum? id- 2TssMmhC2X
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In The Eleventh International Conference on Learning Representations, 2023. httpt //openreview. net wtt.foooommmimddomm0a1
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. CoRR, 2024.
Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelfth International Conference Learning Representations, 2024. URLhttpa httpss:///ppeeeeeeeeee www.wviwwww..eeeeeee.e formm?
id 3x==3xDeA8No11
Yang Liu, Jeremy Bernstein, Markus Meister, and Yisong Yue. Learning by turning: Neural architecture aware optimisation. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 6748-6758. PMLR, 18-24 Jul 2021. WRLLhttps: htttts::/prroceeding
r. wttpprass///3///a21cc...
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017. URL https://
LU..E. Nesterov. Introductory Lectures on Convex Optimization: A Basic Course. Mathematics and its applications. Kluwer Academic Publishers, 2004. ISBN 9780004501444, URL https?
7/books. goog le. htt.e/c///ooooo.c/cllc//ooo// v wwtp
12

---

Published as a conference paper at ICLR 2025
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than SGD for transformers.
In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop), 2022. URL httpp:://penzeevee net/ Forum
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, sigh-performanc deep learning library. Advances in neural information processing systems, 32, 2019.
Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving discrepancies in compute-optimal scaling of language models. arXiv preprint 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21((40):1-611 2020.
Pranav Rajpurkar, Jian Zhang. Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text, 2016.
Herbert Robbins and Sutton Monro. A Stochastic Approximation Method. The Annals of Math- ematical Statistics, 22(3):400 407, 1951. doi: 10.1214/omm/111722556.. URL https?
7/doi. /00 / aww. 1777729588
Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through a crowded valley- benchmarking deep learning optimizers. In International Conference Machine Learning, pp. 9367-9376. PMLR, 2021.
Frank Schneider, Lukas Balles, and Philipp Hennig. Deepobs: A deep learning optimizer benchmark suite. In International Conference on Learning Representations, 2019.
Christopher J. Shallue, Jachoon Lee, Joseph Antognini, Jascha Sicl-Dickstein, Roy Frostig, and George E. Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(112):1 1949, 2019. URL http://////// jmlr. arg/papers/v20.
18-788 789.78888. html
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Jennifer G. Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 4603-4611 PMLR, 2018. URL http::/proccediini
mlr, press vB0 https////////eeeea3/e38 .tmll
shazeer
Prabhu Teja Sivaprasad, Florian Mai, Thijs Vogels, Martin Jaggi, and François Fleuret. Optimizer benchmarking needs account for hyperparameter tuning. In International conference on machine PMLR, 2020.
learning, pp.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Jiaxuan Wang and Jenna Wiens. Adasgd: Bridging the gap between sgd and adam. arXiv preprint arXiv:2006. 16541, 2020.
Ross Wightman. Pytorch image models.
https://github. com wwt.:wwggt/aaaannnm.mmmmmmmom
pytorch-image-mooe
2019.
Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie E Everett, Alexander A Alemi, Ben Adlam, John Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl- Dickstein, Kelvin Xu, Jachoon Lee, Justin Gilmer, and Simon Komblith. Small-scale proxies for large-scale transformer training instabilities. In The Twelfth International Conference on Learning Representations, 2024. WRl.hhttss///peeeeee htt.:/iwww..eeeeeeemeeeeeeee forumm wtt.:ennweeoroormmmmddod////////
13

---

Published as a conference paper at ICLR 2025
Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperpa- rameter transfer. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17084-17097. Cur- Associates, Inc., 2021. http:/tttpps:/prrceeeennoe neur pp wwwppppppppppppppppp. files/ Paper.rr.pp
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In IEEE/CVE Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp- 1213-121.. IEEE, 2022. doi: 10.1:0001111. 11109/CVPP57266822222210111 URL https: /
doi org 101110..///11/1/22886888 2022. 2ttp:////...2222222222221222221. 01179
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham Kakade. How does critical batch size scale in pre-training?, 2024a. URL https
//axiiv.orr//bb///110222.
Yushun Zhang, Congliang Chen, Tian Ding. Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why trans- formers need adam: hessian perspective, 2024b.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more, 2024c. URL https
//arxiv. /arxxii.org/aaa///44..111
Shuai Zheng and James T. Kwok. Blockwise adaptivity: Faster training and better generalization in deep learning. CoRR, abs/1905.098999 2019. Utbpcpe////rxttrr.rrr/aa//////
14

---

Published as a conference paper at ICLR 2025
A RELATED WORK
One closely related work to ours is Wortsman et al. (2024), which explores the stability of Adam with respect to learning rate. We extend the comparison to other optimizers including SGD, Lion and Adafactor, as well as other hyperparameters including momentum and weight decay.
Optimizers: SGD (Robbins & Monro 1951) had been the workhorse optimizer for deep learning until 2015, when Adam (Kingma & Ba 2013) was introduced. Adam is diagonal preconditionig algorithm that maintains a per- parameter learming rate. Over time, coarser variants of Adam have been proposed, which do not explicitly maintain a learning rate per parameter. Adafactor (Shazeer & Stern 2018 ZZai et al.) 2022) maintains a rank-1 1 approximation of the preconditioner matrix of Adam. Previous works have also explored Signum (Bernstein al., 2018 2019) and have observed its benefits in terms of communication efficiency and fault tolerance. Other works have also explored the similarity of Adam with variants of Signum (Balles & Hennig 2018), and recently, a close variant of Signum, called Lion (Chen et al., 2023), was discovered using symbolic search over algorithms. Some other optimizers that have recently gained increasing attention from the community include Shampoo (Gupta et al. 2018b) and Sophia (Liu et al. 2024).
Similar to us, prior work has explicitly studied optimizers comparisons across different domains and architectures. For instance, Schmiitt et al. (202T) performs a sweep over optimizers, schedulers, and seeds over various vision tasks, As mentioned in the introduction, Kaddour et al. (2024) compares optimizers for masked language modeling and at fixed model scale. Other works have called to attention better benchmarking practices for optimizers (Schneider et al.
2019 Sivaprasad et al. 2020
Bartz-Beielsten et al. 2020).
Adam and Signum: Many works have explored the relationship between Adam and variants of Signum (Balles & Hennig 2018 Balles et al.) 2020 Kunstner et al. 2023) and empirically demonstrated that Signum (or its close variants) generally performs comparably ly to Adam. [alles et al. (2020) also argued that signSGD generally performs better when the Hessian is close diagonal, however, it is unclear if this holds for practical settings. Kunstner el al. (2023) recently demonstrated that Adam and a close variant of Signum exhibit similar performance on a variety of datasets including WikiTexx-2 (Merity et al. (017)) ) and SQuAD (Rajpurkar et al. 2016). However, in contrast with our work, all of these are restricted to the setting of vision or masked language modeling, and generally do not sweep over multiple hyperparameters.
Layerwise or blockwise Adam:
We study Adalayer, a layerwise version of Adam. This a special
case of the BAGM optimizer (Zheng & Kwok (2019), specifically BAGM B.1. Similar algorithms
have also been studied by previous works (Ginsburg et al. | Agarwal al. 2020 Liu et al.. [202T]
Zhang et al. (204c)). In particular, concurrent to our work, Zhang et al. (202cc propose an algorithm termed Adam-mini, which closely tracks a modified version of Adalayer (called Adalayer*), and
demonstrate comparable performance to Adam W. Note that, in our work, Adalayer is introduced to understand the role played by preconditioning in Adam, and we do not specifically focus on the final performance. /Zhag et al. (2024b) empirically study the Hessian spectrum of transformers at initialization and find it to be more heterogeneous across layers as compared to ResNets. They argue that this heterogeneity is evidence towards the importance of Adam in training transformers. In contrast our results (Section 3.2) show that Adam's 's preconditioning is particularly important for the last layer and LayerNorm parameters achieve performance and learning rate stability.
For vision transformers, in the fine-tuning phase, Kumar et al. (2024) show that using SGD with frozen embedding parameters leads competitive performance with Adam. Jelassi et al. (2022) explore the similarity between Adam and normalized gradient descent (Nesterov) 2004) and show that normalized gradient descent on GANs does not suffer from mode collapse, while SGD does. Diang et al. (2023) empirically demonstrate that Adam steers the parameter trajectory towards better-conditioned regions than SGD. & Li (2022) also show that the parameter trajectory of Adam exhibits much higher directional smoothness than that of SGD. Ahn et al. (2024) show that the performance gap between Adam and SGD exacerbates with depth of the network. In a similar vein to us. Kunstner et al. (2024) show that Adam is less sensitive than gradient descent to class-ime flassiimallance present in language tasks; we provide further evidence for the importance of preconditioning the last layer, well the LayerNorm parameters.
Other related works:
15

---

Published as a conference paper at ICLR 2025
B ADDITIONAL MAIN SWEEP RESULTS
Optimizer
Optimal Learning Rate
Adam
3.16e-3 (150m), 1c-3 (300m). Ic-3 (600m). 1e-3 (1.2b)
Adafactor
3.16e-3 (150m). 1e-3 (300m). le-3 (600m), le-3 (1.2b)
Lion
3.16e-4 (150m), 3. 16 (300m), 3.16e-4 (600m), le-4 (1.2b)
Signum
3.16e-4 (150m), 3. 16e-4 (300m), 3.16 (600m), 36144 (1.2b)
Table 1: Optimal learning rates for various optimizers from Figure
report the optimal learning rate found for each optimizer in Ta In general, we find that the optimal learning rate for Adam and Adafactor are similar, with the optimal learning rate of Lion and Signum order of magnitude smaller.
For ablations in Figure I
C ADALAYER
As mentioned in Section 3 to investigate the role of preconditioning on language models for optimizers like Adam, we introduce the Adalayer optimizer for ease of analysis. In this section, first establish the performance and stability of Adalayer as reasonable proxy for Adam by making modification how Adalayer treats the last layer of the network.
In Figure I0 we study the behavior of Adalayer across learning rates. To preserve the correspondence with Adam we fix other hyperparameters to be the same: 31 0.9, B2 0.95 and le 15. We find that Adalayer has better performance than SGD, but it performs worse than Adam and also lacks Adam's stability across learning rates. The major difference between Adam and Adalayer the preconditioning done by Adam within a layer. Intuitively, this preconditioning will have large effects in layers where we expect different weights within a layer to have different gradient scales. The first candidate for such layer is the last layer, since different tokens have widely different frequencies leading to different gradient scales. To test this hypothesis, we run a corrected version of Adalayer where treat the set of weights feeding into logit separate block. We henceforth refer to Adalayer with this correction as Adalayer*. This is plotted in Figure 10 and we observe that Adalayer* almost recovers the performance well large fraction of the stability of Adam.
3.7
3.6
3.5
3.4
3.3
3.2
3.1
10-4
10a
10a
10a
Multiple optimal LR
Adalayer AdaLayer* (last layer correction)
Adam
SGD
Figure 10: Modifying Adalayer with the last layer correction improves performance and stability across learning rates.
*We note that this reasoning also applies to the first layer, but in our ablations applying this correction the first layer did make significant difference.
16

---

Published as a conference paper at ICLR 2025
D ADDITIONAL EXPERIMENTS. SGD + ADAPTIVE VARIANTS (ADALAYER*, ADAFACTOR)
In this section, we report additional experiments involving training language models with SGD on a fraction of the models' parameters and an adaptive optimizer on the remaining parameters. Firstly, we show that our results from Section [32 hold even when training 600m parameter models with Adalayer* applied only on the last layer and LayerNorm parameters in Figure II
600m
3.05
3.00
2.95
Validation
2.85
2.80
2.75
2.70
10:1
10a
10
10a
Multipic of optimal LR
Adalayer* (LL 1.0e3) (Matrix) AdaLayer*
SGD
Figure 11: As in Section 3.2 we train 300m (Left) and 600m models (Right) with Adalayer* on the last layer and LayerNorm parameters, and train the remaining model parameters with SGD. We that performance and stability continues to match that of Adalayer* or outperform Adalayer* even at these larger scales..
We also provide further ablations supporting our claim that the largest impact of the adaptivity of Adalayer* is concentrated on the last layer and LayerNorm parameters. Firstly, we train 150m models using Adalayer* on only the matrix parameters, while training the last layer and LayerNorm parameters with SGD. In Figure we see performance improves relative to SGD but we see similar instability at larger learning rates.
Secondly, given that the effective learning rates of the LayerNorm blocks were observed to be small reasonable to whether training the LayerNorm parameters is necessary at all; in Figure [3 we show results for training 150m and 300m models using Adalayer* only on the last layer, using SGD on all other matrix blocks, and turning off training for the LayerNorm parameters. This indeed yields greater stability in comparison to Figure (Left) but does not fully recover the performance of Adalayer*, indicating that training LayerNorm parameters helps with performance, which seems more pronounced in the larger model.
Figure 7 (Right),
150m
3.7
3.6
3.5
Validation
3.4
3.3
10
Multiple of optimal LR
SGD
AdaLayer
Adam
Adalayer (Matrix,
(LL UN)
Figure 12: We train 150m models using Adalayer* on the matrix layers with a fixed learning rate of le and using SGD on the last layer and LayerNorm parameters. Compared to the results in Figure we do not recover the same stability nor do we reach the optimal performance of Adalayer*.
17

---

Published as a conference paper at ICLR 2025
150m
300m
3.6
3.20
3.15
3.5
Validation
3.4
3.10
3.05
3.3
3.2
3.00
2.95
3.1
10
10a
10
10a
2,90
Multiple optimal
10-4
SGD
Multiple optimall LR
AdaLayer
Adalayer (U,, 3.16e-3) (Matrid) training off
Adalayer* ILL. 3.16e 3) SGD (Matric)+ LN training
Figure 13:
Training 150m (Left)) and 300m (Right) models using Adalayer* on the last layer with
fixed learning rate of 3.16e 3 and using SGD on other matrix blocks, while turning off the option to train LayerNorm parameters. We that while the performance and stability is improved compared to SGD, it is still not as performant as Adalayer*, This indicates a degree of importance of training LayerNorm parameters for these models.
We saw in Figure (Middle, Right) that using Adalayer* on only the last layer and LayerNorm parameters sufficed to recover or exceed the performance of Adalayer*. In Figure 14 we report learning rate sweep
the analogous experiment but using Adafactor the last layer and LayerNorm parameters with a fixed learning rate. For the 150m model, using a learning rate of 3. 16e with Adafactor yielded better performance than Adafactor for low learning rates, and is comparable in terms of performance and stability. For the 300m model, the difference between Adafactor and our 'hybrid' optimizer more distinct at higher learning rates for fixed Adafactor leaminng rate 1.0e 3 and 3.16e 3, but is comparable until the peak validation loss.
150m
300m
3.6
Rinal Valitation
Final Validation
3.4
3.00
3.2
2.95
2.90
2.85
10a
10a
10
10
10a
10a
Multiple of optimal LR
Multiple
AdaFactor (LL, + LN, 3.160-3) SGD (Matrix)
Adaliactor (LL LN, 3,16e: (Matrix)
AdeFactor (LL LN, 1.0e-3) SGD (Matrix)
AdaFactor [LL LN, 1.0e-3) SGD (Matri))
Figure 14: Training 150m (Left) and 300m (Right) models using Adafactor on the last layer with fixed learning rate and using SGD on other matrix blocks. We see that performance and stability is comparable Adafactor, but does not exceed it, particularly at higher learning rates.
E ADDITIONAL EXPERIMENTS. ADASGD + ADALAYER*
The previous experiments on SGD Adalayer* show that although using fixed learning rate for the matrix parameters largely recovers performance full Adalayer*, there remains a small gap to full Adalayer* or Adam performance. Can we recover what is remaining? We further investigate the importance of adaptivity at different levels of granularity by training the matrix parameters with AdaSGD (Wang& & Wiens 2020) instead of SGD. AdaSGD still uses global learning rate but allows this scalar to be adaptive across training- other words, this a "global" version of Adalayer* where all matrix parameters are treated block.
18

---

Published as a conference paper at ICLR 2025
In Figure 15 conduct analogous learning rate sweep for AdaSGD for Bg E {0.9. 0.98)
did in Section 2 The performance of AdaSGD is marginally better than SGD, but still suffers from learning rate instability; this is consistent with our previous findings in Figure 7 where certain parameters had effective learning rates which were multiple orders of magnitude different from other parameters, and thus single adaptive learning rate across all parameters likely would not suffice for achieving performance and stability.
that
In Figure 16]we we replace SGD with AdaSGD the matrix parameters and Adalayer train the last layer and LayerNorm parameters. The performance now essentially matches that of full Adalayer* compared SGD Adalayer*, showing that the single adaptive learning rate on matrix parameters is largely sufficient to recover Adalayer* performance. However, gap still remains between Adam performance even with an adaptive learning rate across matrix parameters; likely that this small gap is attributed to the need for increased adaptivity granularity for certain matrix parameters.
150m (standard)
3.40
3.35
Final Validation Loss
3.30
3.25
3.20
3.15
3.05
10-1
10
10a
10a
Multiple of optimal LR
(beta 0=0.9)
AdaSGD (beta 0=0.9)
SGD (beta_0=0.90 0-0.98)
AdaSGD (beta (beta_0-0.0 0-0988)
Figure 15: Learning rate sweep for AdaSGD with two 8o values. We observe that AdaSGD performs marginally better than SGD but still lacks learning rate stability.
150m
3.7
3.6
3.5
Final Validation
3.4
3.3
3.2
3.1
10-1
10a
10a
10a
Multiple of optimal LR
SGD
Adalayer* (LL + LN, 3.16e-3) SGD (Matrix)
Adam
Adalayer* (LL LN, 1.0e-3) AdaSGD (Matrix)
AdaLayer*
When replacing SGD with AdaSGD on the matrix parameters and using Adalayer* to train the last layer and LayerNorm parameters, the performance and stability has improved match that of Adalayer* but gap still remains between Adam performance even with an adaptive learning rate across matrix parameters.
Figure 16:
19

---

Published as a conference paper at ICLR 2025
F ADDITIONAL EXPERIMENTS: FREEZING ADALAYER LEARNING RATE RATIOS
In this section, report additional experiments exploring whether the results involving frozen * Section need both last layer and LayerNorm adaptivity. We show that this is indeed the case by conducting the same sweep for frozen Adalayer* while trying also freeze the learning rate ratios for last layer or LayerNorm parameters as well. In Figure 17 we show that fixing initialized learning rate ratios for all layers does not reach peak performance of Adalayer*, nor does it exhibit stability. In Figure [8] we show that either continuing to update the LayerNorm parameters or the last layer parameters can achieve the peak performance of Adalayer* but is still unstable. Finally, we show results for turning off LayerNorm training while fixing learning rate ratios (with the exception of the last layer) in Figure 19 We conclude that is necessary to maintain adaptivity for both the last layer and LayerNorm parameters, but understanding why the fixed ratios do not suffice would be an interesting question for future work.
150m
300m
3.20
3.6
3.15
3.5
3.4
3.05
3.3
3.00
3.2
2.95
2.90
10
10-3
10-2
10-4
Learning Rate
Learning
Frozen Adalayer* (all)
Frozen Adalayer* (all) Adalayer*
Adalayer*
Figure 17: Training 150m (Left) and 300m (Right) models using fixed Adalayer* learning rate ratios from initialization for all layers. We observe this quickly diverges, achieving neither peak performance nor stability.
3.6
3.6
Vaaliattion Loss
Volication
3.3
3.3
3.2
10
10
10
10
Learning Rate
Learning Rate
Frozen Adaloyer" (all
Frozen Adalayer (all LN parameters)
Adai ayer*
Adal.ayer*
Figure 18: Training 150m models using fixed Adalayer* learning rate ratios from initialization while either excluding only the last layer (Left) excluding only the LayerNorm parameters (Right). We observe both modifications reach peak performance but fails to be stable at higher learning rates.
20

---

Published as a conference paper at ICLR 2025
150m
3.7
300m
3.20
3.5
Validation
3.10
3.4
3.00
3.2
Fiail
3.1
10
10
10
10
2.90
Learning
Frozen Adalayer* fall but UL LN training off)
Learning Rate
AdaLayer
Adalayer" AdaLayer*
off)
Adalayer" (LN training off)
Figure 19: Training 150m (Left) and 300m (Right) models using fixed Adalayer* learning rate ratios
from initialization, while letting the last layer continue to update, and turning LayerNorm training off. Stability across learning rates has improved but is less performant; for the 150m model we also plot the sweep for regular Adalayer* with LayerNorm training off, and that it is worse in performance compared to Adalayer* with LayerNorm training.
G SOPHIA
In this section, we compare Sophia (Liu et al. 2024 to Signum. Note that Signum is special case of Sophia, achieved by setting p= = We find that Sophia does outperform Signum. No significant change in performance was observed when transferring the hyperparameters suggested by Liu et al. (2024) (eg. 31. 32, weight decay), nor when additionally scaling attention by the inverse of layer index which was used in the original Sophia implementation.
Scat
10
Figure 20: Comparing Sophia (Liu et al. 2024 and Signum for the 150M model in our default setup.
21


Reviews:
[{"review_id": "0DeuJKaNf1", "reviewer": "Reviewer_J9Wq", "summary": "The study compares optimization algorithms, including SGD, Adafactor, Adam, Lion, and Sophia, for training language models of different sizes and architectures. Results show that, apart from SGD, all algorithms achieve similar performance and robustness to hyperparameter variations, suggesting practical considerations like memory and simplicity may guide optimizer choice. The researchers further dissect Signum and Adalayer to explore Adam’s adaptivity effects. They find that adaptivity on the last layer and LayerNorm parameters is essential for maintaining performance and stability, highlighting these elements as crucial for reliable training.", "questions": "Despite the issues mentioned in the weakness, I have the following additional questions:\n\n1. why the loss curve for signum are different in Figure 4, $\\beta_1$ and Figure 5, $\\beta_2$? Signum is defined as lion with $\\beta_1=\\beta_2$, if I remember correctly.\n\n2. It would be interesting to add AdaSGD (Wang et al., 2020) into comparison, which only uses one more register than normal SGD, but has much better convergence, stability and tolerence to larger learning rate. AdaSGD can be viewed as a global version of adalayer, i.e., by viewing all the parameters as one layer. A very recent work by [Xie et al.,2024] shows that AdaSGD closes a large portion of the gap between SGD and Adam on GPT2. Replacing the SGD part in adalayer by AdaSGD might help us better understand the importance of adaptivity at different levels of granularity, under the unified framework of AdaLayer, or more generally Adam for arbitrary partition of parameter blocks.\n\n3. typo. \"in perplexity\" in line 317 should be \"in validation loss\"?\n\nI am willing to improve my score if the authors could address my concerns.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well-written and easy to understand. The experimental setups, discussion and interpretations of the results, and limitations of the methods are spelled out clearly. Though it might not have great novelty, I think the extensive study on hyperparameters on different optimizers is a valuable contribution to the LLM optimization community.  I enjoyed reading this paper.\n2. Soundness of the result: As the authors mentioned in the paper, due to limits of computational resources, they only do 1D hyperparameter sweep, instead of 2D or even higher-order sweeps. This is fine for the major finding of the paper, which is a positive result, that the non-SGD mainstread optimizers are robust to choice of wide range of hyperparameters. Tuning remaining hyperparameters will definitely make the best performance of the current hyperparameter better. \n3. The finding that last layer and layernorm parameters need more  fine-grained (than layerwise) adaptivity is interesting. It shed lights on future research towards how Adam works and memory-efficient variant of Adam.\n\n=============\nI increased my score to 6 in reflection of the authors' additional experiments and clarifications.", "weaknesses": "Though I appreciate the extensive abalation on the LLM training experiments with different hyperparameters, the interpretation and the takeways of the experiments are vague and imprecise. I will elaborate it below.\n\n1. The authors often use \"Comparable\" to describe the performance of models trained with different hyperparameters, which I could not find a precise definition in the paper. A validation loss of 3.07 may look close to 3.10 in terms of the relative ratio, but because the minimum of loss is not zero, the suboptimality between 3.07 and 3.10 maybe very significant, say if the minimum of population loss is 3.0 for 150M models. \n\n    As a result of imprecise definition, it is not clear what is the consequence of having comparable validation loss. Does it imply such gap between comparable losses are negligible so in practice if people use any hyper-parameters which lead to losses comparable to optimal validation loss in the first attempt, they are satisfied and do not need to spend more compute to rerun the experiments for the optimal optimal validation loss? If they would still like to rerun the experiment, no matter how small the gap in validation loss is, having a wide range of hyperparameters lead to comparable performance does not imply ease of hyperparameter tuning.\n\n2. The authors write in line 327 \"Takeaway: generally algorithms are more stable with respect to other hyperparameters and the possible gains in performance are relatively small compared to learning rate and momentum.\"\n\n    I do not fully agree with this claim take-home message, given the experiments in the current draft. It seems that the authors get this conclusion just from the shape of the loss curves in those 1D hyperparameter sweep plots, instead of based on some quantative metrics. This could be problematic, because the range of some hyperparameter sweep in section 3.5 for hyperparameters other than learning rate and momentum seems to be smaller than that of learning rate and momentum. I am convinced that for sufficiently small WD and epsilon, the final validation losses are nearly the same. But for warmup, batch size, and $beta_2$, the range are smaller and thus I am not convinced. For example, the ratio between maximal and minimal batch sizes tried in the experiments are just 8, while the ratio is more than a thousand for learning rate. But given that the authors are training in the \"chinchilla optimal\" regime, which means the flops are fixed for different batch sizes, we should expect changing batch size can have similar effect to changing learning rates given the linear scaling rule for SGD [Goyal et al., 2017, Li er al.,2021] or square-root scaling rules for Adam [Malladi et al., 2022]. Therefore I suspect if the authors also vary batch size in a range of 1000 times, the impact of batch size would be much larger. (The authors also mention the 2D effect between batch size and learning rate in line 166)\n\n    For pretraining percentage and $beta_2$, I encourage the authors to include more extreme values to support the claim that these hyperparameters do not matter. Instead showing them they do not matter in the current small range, it is more informative to show to the readers at what extreme values the loss starts to increase significantly. It is possible for warmup that we can completely get rid of it, and in this case it is useful to include that in the plot, just like weight decay.\n\n\n**Soundness of finding in Section 3.6**: First the citation of Lemma 1 from [Balles\\&Henni, 2018] seems to be quite different from what that is in the cited paper. Second, given the assumption in the Lemma 1, $\\delta_{adam}$ seems just defined as $\\frac{\\mathbb{E}[m]}{\\sqrt{\\mathbb{E}[v^2]}}$, and the lemma implies that $\\delta_{signum}$ must be $1$, which is not necessarily true, when sign of $m$ is negative. Maybe the authors are missing an absolute sign on $\\mathbb{E}[m]$.\n\nI also do not understand the argument between line 357-line 360 when $\\beta_1=\\beta_2$. The authors are essentially claiming for all coordinates and all time steps, the first and square root of second moment have the same ratio, which I do not see why it should be true. This prediction could be easily checked by just plotting the update magnitude per coordinate for Adam and see if they are roughly the same. If the authors allow a different scaling factor for every coordinates, then this claim becomes trivial because all first-order optimization algorithm using diagonal preconditioning have the same sign for each coordinate in its update and only differ by scaling.", "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": "2oqyjKziXh", "reviewer": "Reviewer_9bTN", "summary": "This paper conducts a thorough experimental evaluation of the performance and stability of common optimizers used in pre-training language models and uncovers interesting and previously unknown parameters in the transformer architecture that seemingly benefit more from the adaptivity of some of these optimizers.", "questions": "* It is intuitive to me that adaptivity can help with optimization of LayerNorm as the affine parameters are applied in a coordinate-wise way, but I can't understand why the last layer (which I think is the embedding layer in the case of these models?) needs adaptivity. Can the authors please elaborate on this?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The thorough empirical evaluation fulfills a missing piece in the current literature of optimizers used in training language models, as most works in this literature stick to Adam.\n* The experimental grids used for different experiments seem sufficient to draw robust conclusions. The authors incorporate models of different scales and explore 1 dimensional grid searches which is reasonable given the intense compute needs of such studies.\n* This work uncovers an interesting phenomenon which attributes most of the adaptivity gains from Adam and similar algorithms to normalization and last layer's parameters.\n* Considering stability of the training algorithms wrt hyperaparameters as opposed to the final performance is an important aspect since eventually a practitioner wouldn't be able to find the best possible hyperparameters but the hope would be to be able to easily food \"good\" set of hyperparameters.", "weaknesses": "* One major drawback of this study could be the lack of a search over min learning rate at the end of cosine scheduling. As the authors have mentioned, learning rate plays the most important role among all these parameters. Recent studies (https://arxiv.org/abs/2410.05192) have shown that the final stage of training where the learning rate is decreased down to the minimum learning rate plays a crucial role in determining the final performance. Based on this, I'd be curious to know if certain algorithms would outperform others if the min learning rates were chosen carefully, since in that case individual momentums and how the history of gradients is dealt with could significantly differ among different algorithms. I would predict that using some optimizers, some larger learning rates would result in drastically better performances given that the minimum learning rate is small enough.\n\n* I think implementing AdaFactor and AdaLayer would be more fruitful if parameters corresponding to different attention heads would be disentangled. Since the authors don't mention anything in this regard, I assume that they haven't implemented these algorithms to do so, but the nature of computations in attention heads would maybe benefit from adaptivity isolated to the attention head over all of the attention parameters. \n\n* Recently shampoo has gained significant popularity in training language models. It would be nice if the authors could incorporate shampoo among these optimizers to have a more complete story.\n\n* It would be nice if the authors could elaborate on how optimality of different hyperparameters changes as the scale increases.", "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": "Ghji1Lyif9", "reviewer": "Reviewer_iK1k", "summary": "The authors argue that the Adam optimizer is the default choice for most practitioners training language models, and they seek to evaluate this belief by comparing Adam against other popular optimizers. Their findings show that, except for SGD, other optimizers perform on par with Adam. Following this, they explore what components contribute to Adam's success by evaluating two simplified versions of it: Signed Momentum and Adalayer.", "questions": "- line 165, are there only two hyperparameters in the problem? Because if there are more, then the optimal solution is only achievable by solving the N-D problem, with N being the number of hyperparameters.\n\n- While I agree with the authors that cross evaluating all combinations of hyperparameters is intractable, why not use Bayesian optimization to find the best set of hyperparameters for each experiment?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "The Authors show that most optimizers, other than SGD, can match Adam in stability and performance across varied model sizes and hyperparameter settings. Furthermore, it reveals that, to recover stability concerning learning rate and to maintain strong performance, adaptivity is essential, particularly in the last layer and LayerNorm parameters. This insight challenges the notion that Adam is uniquely suited for language model training.", "weaknesses": "- Limited contribution: The main contribution of the paper is the finding that adaptivity on both the last layer and LayerNorm provides the necessary conditions for retaining performance and stability with respect to the learning rate.\n\n- The paper could serve as a guide for those looking to choose an optimizer by comparing performance versus computational cost. However, after reading it, I didn’t feel I gained a clear insight to help decide on a better optimizer. In this sense, the paper has potential but feels incomplete.", "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": "TyNsakPU5J", "reviewer": "Reviewer_4oK2", "summary": "This paper evaluates several optimizers, including SGD, Adam, Adafactor, Lion, and Sophia, in the context of large language models (LLMs). The results suggest that all optimizers, except for SGD, perform similarly regarding optimal performance and sensitivity to hyperparameters. The authors then further investigated two simplified versions of Adam: Signed Momentum and AdaLayer. They found that adaptivity in the LayerNorm and final layers is essential for achieving stable performance and learning rates.", "questions": "- Many follow-up works aim to improve Adam's performance or to develop memory- or compute-efficient versions of it. What is the rationale behind selecting these specific optimizers for comparison? Was the choice based on practical popularity, conceptual connections to Adam, or some other criteria?\n\n- The plots comparing final validation loss (e.g., Figure 1) are presented so that each optimizer’s optimal learning rate aligns, with the x-axis showing multiples of this optimal learning rate. However, why should different optimizers be compared over the same scale of learning rate values? For example, SGD appears the most sensitive to changes in learning rate, but could this be due to it requiring a finer-grained learning rate grid for a fair comparison?\n\n- A concurrent work, [Zhang et al. 2024: Adam-mini: Use Fewer Learning Rates To Gain More], also explores the concept of a coarser, \"parameter-group\" -wise learning rate for Adam, proposing it as the minimal level of adaptivity required. While it is of course not required to cite or evaluate concurrent work, a comparison with AdaLayer would be interesting given the similarity in approach. It would also be useful to see whether *Adam-mini* aligns with this paper’s findings on layer-specific adaptivity. This paper shows that adaptivity in the last layer and LayerNorm is necessary, yet limiting adaptivity to these layers alone still underperforms compared to Adam. So, what degree of adaptivity is \"sufficient\" to achieve Adam's full performance?\n\n- Minor changes are needed on some of the later plots, as they are difficult to read due to small font sizes in the legends and axes, and they lose resolution when zoomed in. Please adjust these for clarity.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The evaluation of hyperparameters is thorough, and the documentation of experimental details is complete. Although it is limited to one-dimensional changes and does not capture the interplay between hyperparameters, the authors clearly state these limitations.\n\n- The observation that adaptivity in the LayerNorm and final layers is necessary for LLMs is interesting.", "weaknesses": "- My main concern with this paper is the significance of its contribution. The paper does not introduce new algorithms or provide insights into why some evaluated algorithms work or fail; it is limited to an empirical comparison. While this can still be valuable, I’m unsure if the paper adequately addresses the question posed in its title, “What Makes a Good Optimizer for Autoregressive Language Models.” It also does not seem to explain how the characteristics of “autoregressive” or “language” models interact with optimization or why the results might differ in other tasks and architectures. Additionally, the organization is somewhat confusing, and it’s unclear how the two parts of the paper relate.\n\n- There are quite a few observations from each experiment, but it’s unclear what the main message or takeaway of the paper is. For example, it doesn’t clearly outline what practitioners should do, what future researchers in algorithm design should focus on, or provide any conceptual or theoretical insights that explain these observations. The paper could be significantly improved by clarifying the main questions it actually addresses and having a better discussion paragraph on what others could do with this information. \n\n- Citations of other optimizer-comparison papers could be more comprehensive. For instance, [Schmidt et al., 2020: *Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers*] is an example that could be included.", "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)