The AdEMAMix Optimizer: Better, Faster, Older
Authors: Matteo Pagliardini
Pierre Ablin
David Grangier
Venue: ICLR 2025 Poster (2025)
Keywords: Optimization
LLM
Deep Learning
Momentum
Primary area: foundation or frontier models, including LLMs

Abstract: Momentum based optimizers are central to a wide range of machine learning applications. These typically rely on an Exponential Moving Average (EMA) of gradients, which decays exponentially the present contribution of older gradients. This accounts for gradients being local linear approximations which lose their relevance as the iterate moves along the loss landscape. This work questions the use of a single EMA to accumulate past gradients and empirically demonstrates how this choice can be sub-optimal: a single EMA cannot simultaneously give a high weight to the immediate past, and a non-negligible weight to older gradients. Building on this observation, we propose AdEMAMix, a simple modification of the Adam optimizer with a mixture of two EMAs to better take advantage of past gradients. Our experiments on language modeling and image classification show---quite surprisingly---that gradients can stay relevant for tens of thousands of steps. They help to converge faster, and often to lower minima: e.g., a $1.3$B parameter AdEMAMix LLM trained on $101$B tokens performs comparably to an AdamW model trained on $197$B tokens ($+95\%$). Moreover, our method significantly slows-down model forgetting during training. Our work motivates further exploration of different types of functions to leverage past gradients, beyond EMAs.

Published as a conference paper at ICLR 2025
THE ADEMAMIX OPTIMIZER: BETTER, FASTER, OLDER
Matteo Pagliardini EPFL
Pierre Ablin Apple
David Grangier Apple
ABSTRACT
Momentum based optimizers are central to a wide range of machine learning applications. These typically rely on an Exponential Moving Average (EMA) of gradients, which decays exponentially the present contribution of older gradients. This accounts for gradients being local linear approximations which lose their relevance the iterate moves along the loss landscape. This work questions the use of single EMA to accumulate past gradients and empirically demonstrates how this choice can be sub-optimal: single EMA cannot simultaneously give a high weight to the immediate past, and non-negligible weight to older gradients. Building on this observation, we propose AdEMA the Adam optimizer with mixture of two EMAs to better take advantage of past gradients. Our experiments on language modeling and image classification show- quite surprisingly-thh that gradients They help to converge faster, and often to lower minima: e.g., a 1.3B parameter AdEMAMix LLM trained on 101B tokens performs comparably model trained on 197B tokens (+95%). Moreover, our method significantly slows- down model forgetting during training. Our work motivates further exploration of different types of functions to leverage past gradients, beyond EMAs.
simple modification of
stay relevant for tens of thousands of steps.
Adam W
INTRODUCTION
With large neural networks, deep-learning has revolutionized numerous fields, such computer vision and natural language processing. At the heart of this paradigm lies the challenge of optimizing complex, non-convex loss functions using noisy gradient estimates. This optimization process is typically carried out using variants of Stochastic Gradient Descent (SGD) (Robbins & Monro, 1951) or adaptive methods such Adam and AdamW (Kingma & Ba, 2015; Loshchilov & Hutter, 2019), which have become ubiquitous in training state-of-the-an models (Devlin et al., 2019; Brown et al., 2020; Dosovitskiy al., 2021a; Radford al., 2021; Zhai et al., 2022; Dehghan al., 2023; Touvron et al., 2023; Dubey et al., 2024).
key component in many of these iterative optimization algorithms is momentum, which has long been shown to accelerate convergence (Nemirovskii & Nesterov, 1985) and often leads to solutions with superior generalization properties (Sutskever et al., 2013). By accumulating gradient vectors over successive optimization steps, momentum helps overcome small local variations of the loss landscape, potentially escaping shallow local minima, and accelerate in plateau regions (Qian, 1999; Ruder, 2016; Goh, 2017). Both SGD with momentum (SGD+M) and Adam incorporate momentum under the form of Exponential Moving Averages (EMAs) of past gradients GT
{g(0)
(
T
EMA(3, e EMA(3, (MA(3.g(T1)) + (1- (T
( Σ 8¹(1 - (T-)
(EMA)
Two considerations support the use of EMAs. From a practical standpoint, the recursive formula of EMA allows for efficient implementations, which do not require maintaining a buffer of past gradients. From a theoretical standpoint, gradient descent with momentum leads optimal convergence rates for quadratics (Polyak, 1964; Nesterov, 1983). However, those results do not guarantee any optimality for general non-quadratic cases (Goujaud et al., 2023).
Work done while interning Apple.

---

Published as a conference paper at ICLR 2025
3.1
3.4
Adamm creined on
Adaantx trained on
2.8
Adar ol teateed on
AMEM AMMi towred (179, 250, SJR3
3.0
AdtRM A Mie troived (20h,
2.7
NaFMaalie trained (245,, 10115, 1818) toarmo
3.2
2.9
2.6
Loss
2.8
2.5
3.0
2.7
2.4
2.6
2.3
2.8
2.5
2566
400k
500k
250k
1006
5006
256k Iterations
7706
Ieraaions
Iteratioons
(a) 110M parameters.
(h) 330M parameters.
(c) 1 1.3B parameters.
Comparing A and AdEMAMix on language modeling. In (a,b,c), we plot the loss obtained using AdamW and AdEMAMix (our optimizer) to train Transformer models of various sizes on the Redpajama dataset. In (a), we train multiple baselines for 256k, 400k, and 500%: iterations, resulting in processing from 17B to 33B tokens. Two AdamW runs with different number of iterations look very different as we use a cosine-decay for the learning rate. We compare those baselines to training AdEMAMix for 256k iterations. We observe that our method reaches a similar loss as an Adam model trained on nearly twice the number of tokens. Analogous comparisons can be derived from (b) and (c). Notably, in (c), 1.3B parameter AdEMAMix model trained on 101B tokens performs comparably to an Adam model trained on 197B tokens (95% more, blue horizontal line). See 4.1 1 and App. B.1 for a detailed description of our experimental setup, including hyperparameters.
Figure 1:
The use of momentum in optimization is grounded in the varying nature of gradients. As local linear approximations of the loss landscape, their information can quickly become outdated as the optimization process progresses (Pascanu et al., 2013). For this reason, practitioners typically employ moderate momentum values (i.e. 8 x 0.8 or 0.9), effectively creating a moving average of recent gradients while discarding older information. Selecting larger 3 values seems counter-intuitive, as it would suggest that older gradients maintain their relevance over extended periods of training. While it is tempting to see the use of small As as a confirmation of the limited temporal relevance of gradients, our work reveals instead that older gradients can efficiently be used. When we increase J3, we decrease the relative importance of recent gradients, and the iterate now fails to respond to local changes in the loss landscape. We observe that single EMA cannot both give a significant weight to recent gradients, and give non-negligible weight to older gradients (see Fig. 3a). However, a linear combination between a "fast-changing" (e.g. B = 0.9 or B = 0) and a "slow-changing" (e.g. = 0.9999) EMA allows the iterate to beneficiate from (i) the great speedup provided by the larger (slow-changing) momentum, while (ii) still being reactive to small changes in the loss landscape (fast-changing). More precisely, we find the following statement to convey an important intuition behind this approach:
While changing the direction of the slow momentum is difficult, any adjustment orthogonal to that direction is easy-which favors fast progress in sinuous valley-like landscapes: the fast EMA helping to stay in the valley while the slow EMA accelerates in the valley's 's direction.
A toy illustration of this can be seen in Fig. 2. Based on this idea, we propose AdEMAMix (Adaptive EMA Mixture), a novel Adam based optimizer which successfully leverages very old gradients to reach better solutions.
Contributions, Our contributions can be summarized as follows:
We propose AdEMAMix, a novel optimizer which better leverages past gradients by avoiding a common pitfall of EMA-based optimizers (sce $ 3).
We empirically demonstrate the superiority of our method over Adam by training ViT and language models (Transformers and Mamba) of up to 1.3B parameters (see $ 4). In addition, we show gains from switching mid-training from Adam AdEMAMix (see Fig. 5).
2

---

Published as a conference paper at ICLR 2025
Adam 8 = 0.9
AdEMAMix As = 00.99
Adaan
(6.0 (6.00090,0000 0.3.009990000... 0.000))
Adam
AdEMAMi
0.9999
8; -0.99
Distanee to solesion
102
AdRMAAbii 3. 0.0 0.19 (0.500 0.9m0))
Adam S, U. 9000
10
109
10"2
10"
2000
4000
Iterations
(a) [[x(( xe 1|2-
(b) Adam trajectorics.
(c) AdEMAMix trajectories.
Figure 2: Comparing Adam and AdEMAMix on the Rosenbrock function. Starting from
(x))) = [-3, -3,,5, 5). we minimize the Rosenbrock function f(11, 222 = (1 )1 + 100 - =²) The global minimum (+) is x [1, 1]. We use 32 0.999 for Adam and (B1, 82, a) (0.9, 0.999, 9) for AdEMAMix (see 3). We reduce the learning rate for AdEMAMix to compensate for the influence of a. We do sweep over 31 (resp. 33) for Adam (resp. for AdEMAMix). In (b), When Adam's 's 31 is small (e.g. 0.9), the iterates do not oscillate but convergence is slow. Increasing $1 makes the iterates move faster but with large oscillations. In contrast, for AdEMAMix in (c), we observe that despite 33 being large, the iterates moves fast and without oscillations. This results in reaching better solutions faster as can be seen in (a).
We show AdEMAMix forgets the training data slower when compared to Adam (see Fig. 4).
More broadly, our findings contribute to a deeper understanding of the optimal balance between using historical gradients and adapting to the rapidly changing loss landscape. Our work invites further research methods combining old and recent gradients, beyond EMAs
RELATED WORK
Work on understanding momentum. From the seminal work of Polyak (1964), many publications analyzed the effect of gradient descent + momentum (GD+M) in both convex and non-convex settings (Ghadimi al., 2015; Flammarion & Bach, 2015; Kidambi et al., 2018: Defazio, 2020; Sebbouh et al., 2021). While the acceleration in the noise-free setting has been long theorized for convex functions, several publications indicate this effect might not necessarily extend to stochastic settings (Yuan et al., 2016; Kidambi et al., 2018; Leclerc & Madry, 2020), emphasizing instead link between momentum and effective learning rate. Recent work have been seeking to understand the impact momentum on generalization through studying the implicit bias of momentum methods (Ghosh et al., 2023; Papazov et al., 2024), exposing preference of SGD+M for lower norm solutions. Those further exposed link between higher momentum and higher effective learning rate and higher variance reduction. Despite the volume of prior work on the subject, our understanding of momentum methods in non-convex stochastic settings is still incomplete (Yuan et al., 2016). Oscillatory behaviours, and the sometimes ambiguous effect of variance on optimization render the analysis tedious. From a theoretical standpoint, our work raises several questions. First, given that we gain from averaging very old gradients, what can it reveal of the loss landscape and the consistency of one batch's gradient during training? Second, would our approach not decrease the variance up a point that is harming generalization (Ghosh et al., 2023y? While no answer to those questions is given in this work, we provide a toy justification which indicates that large momentums can have positive impact in noise-free non-convex settings (see Fig. 2)-indicating the improvement of our approach is at least partially explainable without considering variance-rede reeuuccin effects. We moreover expose a link between momentum and forgetting the training data (see Fig. 4), which to knowledge is novel.
Work on deep-learning optimizers. Despite the popularity of Adam and AdamW (Kingma & Ba, 2015; Loshchilov & Hutter, 2019) in training deep neural networks, optimizer design is rich field of research and we focus on a few the works most relevant to this study. Adafactor (Shazeer & Stern, 2018) improves Adam's memory efficiency by factorizing the second moment estimate. Lamb (You
3

---

Published as a conference paper at ICLR 2025
et al., 2020) extends Adam by adding layerwise normalized updates. Chen et al. (2023a) use algorithm discovery to derive the Lion optimizer. Contrary to Adam, Lion uses a single momentum term and the sign function to produce updates with the same magnitude across dimensions. Interestingly, Chen et al. (2023a) also report better scores are obtained when using slightly larger momentum term (B = 0.99). In this work show how increasing the momentum well beyond this value can still be beneficial. See App. C.3.2 for a detailed comparison between AdEMAMix and Lion. Recently, Liu et al. (2023) introduced Sophia, a scalable second-order optimizer designed for LLM training. Sophia uses a Hessian- based pre-conditioner which better normalizes the step size, penalizing steps in high curvature direction and accelerating in low curvature directions. Understanding in which circumstances those novel optimizers bring improvements is still being investigated (Kaddour et al., 2023), and Adam's 's dominance remains mostly unchallenged.
Work incorporating an additional momentum term. uses pre-filtering filtering step on the gradient to amplify the low frequencies and help solve groking. When combined with Adam, it effectively applies the Adam's 's EMAs on top of another gradient averaging method (e.g. EMA for Grokfast-EMA). Somewhat similarly, Chen al. (2023b) refer to the Double EMA (DEMA) used in some finance applications (Mulloy, 1994) as motivation for their AdMeta optimizer. Our motivation behind AdEMAMix is to combine both a high sensitivity to the recent gradients as well as incorporating very distant gradient, in this respect, using nested EMAs is not the right candidate as it reduces the influence of recent gradients. A more detailed review of AdMeta and DEMA is in App. C.3.1. Most relevant to us, Lucas al. (2019, Agg also observe that using a combination of EMAS can enable the use of larger 3s, and incorporate sum of K momentum terms into GD. They show their approach reaches similar performances as baseline optimizers, with a faster convergence. In contrast, we modify Adam, and introduce schedulers that are critical to reaching good performances at larger scales. This allows us to outperform Adam by a significant margin. In App. C.3.3, we notice no further improvement by adding more momentum terms. Finally, Szegedy et al. (2024) propose a general framework to derive and study optimizers with linear combinations of memory soctors-wiiee encompasses EMA mixtures.
Lee et al. (2024) introduce Grokfast, which
Work on distributed optimization. Perhaps surprisingly, recent work on distributed optimization DiLoCo and SlowMo (Douillard et al., 2023; Wang et al., 2020) are relevant to our work. N workers (0)) ,...., (02 are trained independently for K steps (e.g. K = 500). Every K steps, the delta updates {A0, JMN, & (o!t+K) _ Ji=1 are averaged and applied to each worker using an outer optimizer with momentum 3: g(*+k) = 9!¹) r Opt( N Σ, A0, 33. The application of the outer momentum every 500 steps increases the importance of old gradients in the optimization trajectory. We believe this observation might in parts explain the strong results provided by those methods, and further motivated our work.
OUR METHOD: ADEMAMIX
Setup & notations. Let Lo X + RR be loss function parameterized by O, and mapping inputs x E X' to R. Given a sampled batch æ, let = VoLo(ar) be a stochastic gradient of the loss w.r.t. 0. To minimize the empirical loss, the Adam optimizer (Kingma & Ba, 2015) relies on first and second moments, resp. m and v, estimated via two EMAS parametrized by (81,8) E [0, 1[2. A weight-decay parameter À E R is often used as in Loshchilov & Hutter (2019):
m (t) = -m (1-) + (1 (1)
(t)) = m
-1-
(t) 13 + (1- (()
D(0) (0
(AdamW)
T-p3
0(0) 0(1-1) -
VDit+4 -n( 20((-11).
(1)
With t > 0 being the timestep, 17 being the learning rate, and e a small constant. Initially (t=0)) (c=0) O. To prevent the bias induced by the initial m!=0) and (!=0), the outputs of the two EMAs are corrected into m and (t) Those are used to compute the final weight update, scaled by the learning rate.
How far to look into the past? A typical value for B3 is 0.9. Fig. 3a shows the weights given to past gradients for different values of 3. The larger the B, the more uniform the average is. To put this in parspectve--bbeeyee that De (1 B) for 3 [0, 1 the number of successive previous
4

---

Published as a conference paper at ICLR 2025
x10-3
1.0
2.7
EMA 0.5
Ashun
AdaawW on 16B) whlets
FMAA2: 0.90
AdiiMAAli Py = 0.000
0.8
EMAA: 0.093
EMA:: 5.0950
2.6
AdEMMAbix dy
100k stops Uneor u doows
0,90000
AutSaAMi pained on (4B, 18B) toloza
3.5
0.6
2.5
3.4
0.4
2.1
3.2
2.3
0.0
3.1
0
3000
100000
50006 1o.5
1M
13M
0
61k
128k
Tterations
Iterations
(a) Limitation of EMAs.
(b) Constant 17scheduler.
(c) Mamba results (168M).
Figure 3: Limitation of EMAs, constant r-scheduler, & Mamba results. In (a), we plot the weights w, wo-for each past gradient g (*) given by different EMAS after 10k steps. For given B, EMA(8. g ()) .............. .g......9.. exponentially. A small 33 (e.g. 0.9) will give high weight to the immediate past and negligible weights to older timesteps. In contrast, a high 3 (e.g. 0.9999) is giving a relatively uniform, yet non-negligible weight to all gradients. No 3 value can simultaneously give high weight to the immediate past and a nonlneglie weight to very old timesteps. In (b), we train multiple 1.3B language models using 3k steps of warmup and then constant learning rate 17 10-4. This allows us to observe the gap between Adam W and AdEMAMix without the cosine decay as a confounder. We still use schedulers for a and B3 with Ta.63 500 5. Similar Zhai et al. (2022); Hu et al. (2024); Hägele et al. (2024), we decay the learning rate linearly at 1 = 1M and /. = 1.3M. The loss-gap between Adam and AdEMAMix increases at first, and then remains constant. AdEMAMix still outperforms Adam W after decaying the learning rate. See App.C.1.10 to see the impact of the linear decay duration. (c), we train 168M parameter Mamba models, showing how AdEMAMix's performances can generalize outside of the Transformer architecture.
(T) = Et (1 (T-)), which decays the contribution of past gradients
1. For B 0.9, thaif = 6, meaning that half of the weight is given to the previous six gradients. This observation can also be extended to SGD with e.g. polyak or nesterov momentums (Polyak, 1964; Nesterov, 1983), which typically relies on similar 3 values. The value of 3, is rarely increased beyond- ~ 0.9. In our experiments with Ad increasing 31 further degraded the performance (see App. C.1.7). Does this mean older gradients are outdated? We show that this is not the case, rather, increasing beta is reducing the sensitivity to recent gradients too much. We design AdEMAMix such that the sensitivity recent gradients is kept, while also incorporating information from much older gradients using an additional momentum term. This allows for the use of much larger 3 values e.g. 0.9999. To compare, for 3 = 0.9999, thalf & 6,930, spreading half of the mass over the previous 6,930 past gradients.
steps receiving a cumulative weight of 0.5, is thaly In(_))
(n.).
AdEMAMix. To keep a high sensitivity recent gradients, while also incorporating information from older gradients, we add a second EMA (changes compared to AdamW are in Blue):
(() = =B (1-1) (1-) + (1 - 0
(1)
m
(0) Bamm! (1-1) +1 - (1)
m2
(AdEMAMix)
D(0 = B2 + +1 (1-11800 (1)
2
I-
0 1((-1) - ( m tame
10)
vorl+
In our experiments, while the values of $r, 32 remain similar to those of equation AdamW, we often use 33- = 0.9999. We find a E 10] to work well in practice.
Tackling early training instabilities. Early training instabilities are commonplace when training deep models, and empirically motivated the introduction of methods such as learning rate warmup and gradient clipping. Gilmer et al. (2022) show how the use of learning rate warmup can be justified from a curvature perspective; allowing the iterates to move to parts of the optimization landscape where larger learning rates stable. While we use learning rate warmup in all our experiments,
5

---

Published as a conference paper at ICLR 2025
still noticed AdEMAMix models using a large S3 would diverge early. This, despite not using bias correction over m2, which lets the momentum buffer fill itself slowly. Those failed runs are characterized by updates of large magnitudes in the early phase of training (see App. C.1.10, Fig. 28). For this reason, we progressively increase the values of (33 and a using schedulers. For a we use linear scheduler. A linear scheduler for 33 would be ill-fitted as the same increment of 33 has different impact for different values of $3. For instance, observe that an increase of B of d3 0.0001 barely increases the fhalf for 3 0.9, while 0.999 + 0.999 + òg increases the thalf of 77. For this reason, we design the B3 scheduler to increase f linearly (see App. A. for derivation of that scheduler). The two schedulers are summarized below:
(0) laa(,, Ta) min( To (,))),
(Sa)
In( B start) In(3g) In (Sataar))) ),Bs).
060 ft, Bz, Sstare, To₂) min
exp
(fr.)
(3s) + To,
Taa
With To and To, resp. the warmup times for (+) and a(th reach their final and maximal values. In practice we always set those two to the same value: Ta = Ts, = -Ta,9g, and we typically use Ta,6. T, with T being the total number of iterations. Istart is always set 31 in our experiments. The use of those schedulers is not always required, especially, we found those have no impact when AdEMAMix is activated later during training (see Fig. 5). The full AdEMAMix optimizer, including the schedulers, is shown in Alg. 1.
Algorithm AdEMAMix optimizer. Differences with Adam are in blue.
1: Input: Data distribution D. Initial model parameters ((0), Number of iterations T. Learning rate 17. c a small constant. Adam W parameters: 31, B2 and x. AdEMAMix parameters 83, a.
Warmup parameter Ta.3y. note that we usually set it to T. Astort is usually set to 31.
2: Initialize timestep: + 0
3: Initialize EMAs: mi (0) <- 0 m2 (0)
+ 0, 0,2(0) +0
4: for E [T) do
5:
++++1
Optional: use schedulers (11 2841 fos Ba, Botart, Ts) and a + falt,a,To,o)
6:
7:
Sample batch: x ~
8:
Compute gradient: g (t) + VoLgit (x)
9:
Update the fast EMA mi: m
(1) f- 31 3m (1-1)
+ (1 (1)g(1)
10:
Update the slow EMA m2: m2 t ( m
(1)
(1-1) + (1 - )g0))
11:
Update the second moment estimate: (1) (2 (1 (0)
Apply bias corrections: m (t)
1-0p (1) + I-
ml (1)
m² (o
13:
Update parameters: 0(t) + 0(t (1)
m
+ x0(--11)
vDit+
14:
end for
15: Return optimized parameters (T)
Using B1 0 to save memory. By setting 31 O, we can save on memory by replacing m by (()), In this case, the memory cost of AdEMAMix is the same as for AdamW. We show in App. C. (Fig. 16b) and App. 1.8 that using 31 O often works, at the cost of potentially less stable training.
(0)
Computational overheads. Adding an additional EMA requires additional memory and compute. The added compute is negligible in comparison to what is required for the forward- backwrdd and has little impact over the total runtime as shown in Fig. 5a. Moreover -wen considering larger distributed setups-ir it is worth noting that AdEMAMix is not increasing communication (gradient reduction) over Adam. Therefore, we expect the runtime overhead of AdEMAMix to shrink in those settings, as data movements occupy a larger fraction of the total runtime. A more significant overhead is in terms of memory when 31 * 0, as AdEMAMix requires to allocate both m. and m2, which are of the same size as the model parameters O. We believe this issue is of lesser consequences as Fully- Sharded- -Datd---aa--aa Data-Paralelliist Paralellism (Zhao et al., 2023, FSDP) can always be used to distribute the optimizer states across compute nodes. As example, using the Megatron library (Shoeybi al., 2020) to

---

Published as a conference paper at ICLR 2025
3.2
teainion without B
Loss on botch
Adaut training using it 1
3.1
3.0
2.9
2.8
70k
170k
2564
70k
1904
256k
70k
230k
3.2
training willoul B
Loss 00 batch
3.1
using B as.
AdEMAMIx training
3.0
2.9
2.8
90k
2568
70k
170k
2564
70k
190k
256k
70k
230h
(a) 90k:
(b) Ls
170k.
(c)
190k.
(d) 230k.
Figure 4: Measuring forgetting using a held-out batch B. The top row is for AdamW, the bottom row is for AdEMAMix. We trained one AdamW and AdEMAMix model on a RedPajama dataset not containing the batch B, those runs are in blue. We then run multiple experiments where we inject B the training data specific timestep f Those runs are in orange. To inspect how much influence B had when it is injected at tg. we can observe the evolution of the gap between the blue and the orange curves, For both optimizers, we observe a rapid decrease of the loss on B right after training on B. The sharpness of this decrease in loss is more pronounced for AdamW compared to AdEMAMix. However, when using AdamW, W. the loss on B then increases faster, which interpret as the model forgetting B faster. In contrast, the curves for AdEMAMix are smoother, the loss on B goes back up slower, and ultimately B had a bigger impact on the training when using AdEMAMix-as
be seen by looking at the larger gap between the orange and blue curves for the last iteration. Finally, the forgetting behaviour evolves during training, with later training batches being remembered better. See App. C.1.3 for more experiments on forgetting.
train 8B model across 16 compute nodes, each having 8 GPUs. Assuming data- aata-parallelism (DP) of 64 (2 GPUs per replica). This would imply sharding the optimizer states over the 64 replicas, causing small Per-GPU offset of 0.5GB when compared with AdamW. In the cases where memory remains an issue, one mitigation strategy could be to apply factorization strategies as in Shazeer & Stern (2018); Zhao et al. (2024). reduce the size of the momentum buffers in Zhang et al. (2025).
Hyperparameter sensitivity. While we introduce up to four new hyperparameters: a, B3, Ta and Tag, in practice we always set T, = T3a- = =To.ss and use Ta,93 = T' in most cases. We show in App. C.1.5 that T, and T3a should only be large enough prevent instabilities early during training. While all of our experiments on language modeling use $3 = 0.9999, other values such as 0.999 or even 0.99999 still can outperform the AdamW bascline (see App. C.5. Fig. 13). On vision tasks, the scatter plots in Fig. 6 show all the AdEMAMix models trained for those experiments, highlighting how easy can be find good a and 33 values. Overall, we find the ranges of values of a, B3 and Ta,2, providing improvements over AdamW to be wide. See App. C.1.5 for hyperparameter sweeps.
Gradient elipping & B2. Adam can diverge in cases where Bg is not sufficiently large compared to 3i (Zhang et al., 2022). Intuitively, the denominator in the AdamW update can shrink faster than the numerator, causing a spike in update magnitude.. Given the proximity of our method AdamW, similar observation applies to AdEMAMix, and require 32 to be large enough. In practice, we observe that both 3g = 0.999 or 0.9999 work well. Moreover, we notice that gradient clipping mitigates the issue and we use small gradient clipping values in experiments (e.g. 0.5 or 0.1). We also believe clipping to help dealing with outlier gradients, which can have long-lasting and detrimental effect for AdEMAMix given the large 33.
Limitations, AdEMAMix consists in leveraging very old gradients. Therefore, the method is best suited to settings where the number of iterations is important. We report on this effect in App. C.1.6,
7

---

Published as a conference paper at ICLR 2025
60
AIEMAMEx
Adann
2.55
Adamm 10115 tobena
AIEMAAII 5000
Traluing line (Junzs)
3.1
AGEMAKiss
50
AGEEMAASEi Zrom 1008
3.0
AcciitAAS from
40-
2.9
2.35
30-
2.8
2:30
2.25
110M
13B
3004
300k:
500k
500k
600k
7706
Iterations
Itenations
(a) Train time for 256k iters.
(h) From Adam W (110M).
(c) From AdamW (1.3B).
Figure 5: Training time comparison starting AdEMAMix from Adam W. In (a), we compare the time required to train 110M and 1.3B parameter models for 256k iterations. The additional EMA renders AdEMAMix slightly slower than Ad However, we were to train Adam longer to compensate for this gap, would only train for additional 2379 and 5421 iterations for resp. 110M and 1.3B parameter models. Those additional iterations would not be sufficient to close the gap (see Fig. 1). In (b) and (c), we show-fo two different model sizes-the effect of switching from AdamW to AdEMAMix during training, AdEMAMix's 's additional parameter m2 is initialized O, no scheduler is used for a B3. For both model sizes, we observe the loss increases slightly at first before decreasing and outperforming the baseline. In both cases, the earlier AdEMAMix is used, the better the final loss. See App. C.1.10 for results using schedulers.
additionally showing how smaller values of A3 (e.g. 83 0.999) can be better for low iterations scenarios. Moreover, retaining gradient information over many thousands steps can pose a problem in domains requiring fast adaptation to sudden distribution shift.
RESULTS
In this section AdEMAMix train language models ($ 4.1 and 4.2) and vision transformers (& 4.3) ranging from 24M to 1.3B parameters.
4.1 TRANSFORMER LLM TRAINING
Experimental setup. We use transformer architecture (Vaswani al., 2017) with learnt positional encoding. All of our experiments are using sequences of 1,024 tokens. We experiment with three sizes of transformers: 110M, 335M and 1.3B. For the learning rate, we use 3k warmup steps followed by by-unless specified-
cosine decay to 10 We extensively tuned the hyperparameters for both AdamW and AdEMAMix models (see App. B.1). We use the RedPajama v2 (Computer, 2023) dataset for all of our experiments. We batch sizes of 64, 96 and 128 for respectively our 110M, 335M, and 1.3B parameter models. Depending on the model, we vary the number iterations from 256k to 1.5M. For AdEMAMix, we use 33 0.9999 and a E {5,8, 10) depending on the model. A full description of the architecture and hyperparameters used is in App. B. We train on up to 8 A100 NVIDIA GPUs using aata-paaallelism.
Why not simply increasing Adam W's 31? While our toy experiment in Fig. 2 already gives some intuition on why increasing Adam's 's 81 is likely not to have the same effect as having an additional EMA as in AdEMAMix, we verify this intuition by training multiple 110M models using Adam with large 31 E (0.99, 0.999, 0.9999, 0.99999}. When we use large B1 from the beginning of training, we observe instabilities for larger 31 values and no 31 > 0.9 improves upon the AdamW baseline. One can imagine this to be due to increasing 31 too early. Therefore, we also modify Adam and add the same scheduler on 31 training duration. While this mostly stabilizes the training, none of the experiments outperformed the baseline using 31 0.9. Moreover, to rule out any effect that could be due early training instabilities, we run the same experiments starting from a pre-trained Adam W checkpoint trained for 300k iterations. We simply resume training and either increase B1 suddenly or using scheduler.
use AdEMAMix's 83. 31 is now increased steadily over the entire
8

---

Published as a conference paper at ICLR 2025
Here again- Lunlke when using AdEMAMix none of those experiments outperform the baseline. The details of those experiments are available in App. C.1.7. Those experiments show that simply increasing the 31 value in Adam W is not enough, which justifies our design of AdEMAMix.
Better perplexity for the same number of steps. For all model sizes, all the number of iterations used ranging from 256k to 1M-, AdEMAMix always outperforms the AdamW baseline. In Fig. 1, we show the validation loss curves for Adam W and AdEMAMix models trained- for each size-on various numbers of tokens. For 110M parameter models, training for 256k iterations gives similar results as training an Adam W model for 500k iterations. The gap between baseline and our method seems to be increasing as we increase the number of iterations. For 1.3B parameter models, training using 770k steps is on par with training the baseline for 1.5M iterations reaching the same performance with 51% fewer tokens (an economy of 96M tokens). In Fig. 3b, we observe similar improvements when using a constant and then linear decay learning rate scheduler.
Training speed comparison. We measure the time per iteration for all of our experiments. In Fig. 5a we plot the time required to train our 110M and 1.3B parameter models for 256k iterations. We observe that the impact of using an additional EMA on the training speed is negligible. If we were to train new models with a fix time budget, the extra iterations of the baseline would not be sufficient to close the gap. For instance, to match a 110M parameter AdEMAMix model trained for 256k iterations, we need to train AdamW model for 500k iterations, and the time advantage of Adam would only allow us to do 2,379 additional iterations. Moreover, as mentioned in $ 3, we expect the time overhead decrease when molti-node training is used, as IOs would become important bottleneck, and AdEMAMix is not increasing IOs.
Consistency of the gain across loken-budget In Fig. 1, given that enough tokens have been used w.r.t. the model size, we observe consistent gains accross token budgets. It seems our method is able to bring a constant improvement over the baseline. This can be seen more clearly in Fig. 3b, showing results when using a constant learning rate scheduler-whie learning rate decay. We observe how, after an initial phase in which the gap grows, this gap becomes seemingly constant after sufficient number of iterations.
removes the confounder of the cosine
Resume from Adam vs. training AdEMAMix from scratch. So far we trained AdEMAMix models from scratch. We show it is also possible to switch from an Adam W to an AdEMAMix state in the middle of training. When switching to AdEMAMix at step fawitch. we initialize m (fswitch) and replace & by t - twwitch in the scheduler equations-if those are used. However, we find that schedulers are not required when resuming training and report results without them in the main paper (see App. 1.1.10 for more details). In Fig 5b and Fig.5c, we show that (i) it is possible to improve upon the baseline when activating AdEMAMix later during training, and (ii) the earlier the switch, the larger the gain, with diminishing returns. This indicates that the improvement of AdEMAMix cannot be attributed solely to early training dynamics, but rather, late training dynamics seem to play an important role. This is further corroborated by the reverse experiment- which switches from AdEMAMix to Adam W mid-training and show a performance degradation (see results in App. C.1.4).
AdEMAMix models forget the training data slower. As an attempt to understand the reason for AdEMAMix's improvements over A, we study how fast a training batch is forgotten after being used during training. We focus on following one specific batch B. We start by removing B from the RedPajama training data and train AdamW and AdEMAMix models. For those runs, B is therefore akin batch from the validation set. We measure the loss on B while training. Now we can resume training from various checkpoints, and inject B into the training data at various times tg. By comparing the two runs-one having trained on B. while the other never saw B-we can visualize how B is learned, and how it is forgotten. After training on B, we expect the loss on B to decrease suddenly and then increase as the model forgets the contribution of that batch. When comparing the forgetting curves for Adam and AdEMAMix in Fig. 4, we see striking differences. AdamW models forget much faster-the loss over B increases faster-than AdEMAMix models. Moreover, at the end of training, batches processed by AdEMAMiX see their loss being improved over many thousands of iterations. Additional experiments on forgetting can be found in App. C.1.3.
4.2 MAMBA LM TRAINING
Experimental setup & results. Our setup is similar to $ 4.1. except that we now train 168M parameter Mamba models (Gu & Dao, 2023) using the Web dataset (Penedo al., 2024). See
9

---

Published as a conference paper at ICLR 2025
App. B.2 for details. In Fig. 3c the improvement using AdEMAMix is consistent with our experiments on Transformer models, showing how AdEMAMix 's gains can extend beyond Transformers.
4.3
VIT TRAINING
2.60
2.45
Adamm
AdEMAMiis
AGEMAbiic ActIMAMI As... 3.8603
2.40
AGEMMANis 0.0993
2.55
1.2
Test loss
2.35
Test losss
2.50
1.1
2.25
2,20
2.40
0.9
2.1
2.3
0.1
Training loss
Training lose
Trainin loes
(a) 24M params, 11M images. (ImageNet-21k)
(b) 86M params, 11M images. (ImageNet-21kk)
(c) 86M params, 1.3M images. (ImageNet-Ik)
Figure 6: VIT results for different capacity/data ratio. As described in $ 4.3, we compare the train vs. test loss between AdEMAMix and AdamW, on multiple scenarios (a,b,c). Those have different capacity/data ratios. For each setting, we first tune an AdamW baseline by testing multiple hyperparameters, those are represented by blue dots. We then pick the hyperparameters of the best ceselinee-corrspooddiinn to the lowest test loss-and train multiple AdEMAMix models by simply testing various a E {1, 5, 10, 15, 20] and 33 E {0.999, 0.9999). Those are represented by orange squares and red triangles. See App. B.3 for detailed hyperparameters. The red area in the above plots represents the area in which we achieve both a better training and test loss compared to the best baseline. The diagonal black line represents the line would follow when we improve upon the baseline without adding any overfitting (i.e. an improvement in train loss yields the same improvement in test loss). (a) represents the most desirable setting, in which we have relatively lot of data compared to the model size. In this setting, AdEMAMix outperforms the baseline for all hyperparameters tested. If we increase the model size in (b). most AdEMAMix hyperparameters are improving upon the baseline. If we now decrease the dataset size, using 86M model with 1.3M images as in (c), doing total of 320 epochs, we do not observe an improvement from using AdEMAMix. See App. C.2 to see classification accuracies and loss curves for those experiments.
Experimental setup. In this section we consider a different setting in which the data is now a limited resource, and we do multiple epochs (e.g. 37 or 320). We use two subsets of the ImageNet (Russakovsky et al., 2015) dataset: (i) the widely used ImageNet- subset, consisting of 1.3M images and 1,000 possible classes; (ii) filtered and pre- processed version of the ImageNet-21k (Ridnik et al., 2021) containing 11M images corresponding to 10,450 classes. For each, we measure the test loss on held-out test set. We the ViT architecture (Dosovitskiy al., 2021b) two different scales: 24M and 86M parameters. Importantly, if it is common in the vision literature to pre-train large ViT models and finetune them on smaller datasets, this work focuses on pretraining optimization and we therefore train and test on the same distribution. The models' hyperparameters are detailed in App. B.3, we use batch size of 4096 for all our experiments. We used training hyperparameters from Dosovitskiy et al. (2021b) as a starting point and did some additional tuning of the learning-rate rate, dropout, and weight decay for our Adam W baselines. We then use the hyperparameters of the best AdamW baseline and train AdEMAMix models with various values of E {1, 5, 10, 15, 20} and (3g E (0.999, 0.9999). We train models for 320 and 37 epochs for resp. the ImageNet- 1k and ImageNet-2 datasets, corresponding in both cases to 100k iterations. Data-augmentation techniques have been shown to be central to the efficient training of ViTs Touvron et al. (2021; 2022). We simple aattt--ummettninns, which includes mixup (Zhang et al., 2018). We train 8 A100 NVIDIA GPUs using Pytorch Fully Sharded Paralleral (Zhao et al., 2023, FSDP).
AdEMAMix for different capacity/data ratios. We consider three scenarios that differ by their capacity/data ratios. First, we trained 24M parameter models on 11M images (ImageNet- for 37 epochs. In this setting, be seen in Fig. 6a, it is trivial to find AdEMAMix parameters
10

---

Published as a conference paper at ICLR 2025
outperforming the baseline in terms of both training and test accuracy. We now increase the model size to 86M parameters. Fig. 6b shows it is still easy to find parameters outperforming the baseline. We now keep the model size of 86M parameters and switch to the smaller ImageNet-Il Ik dataset (1.3M the number of epochs to 320. In this setting,
images), which-given our 100k iterations-incrent Fig. 6c shows that outprrorming the baseline becomes difficult. AdEMAMix seems perform best in scenarios with large volumes of data w.r.t. the capacity of the model.
5 CONCLUSION
In this work, we find that old gradients can be leveraged to efficiently train large language models and ViTs. Our proposed optimizer combines two momentum terms. A slow (large 3) momentum gathers information over many timestep, while fast (slow B) momentum can adapt the trajectory of the iterates to the rapidly changing loss landscape. We demonstrate the superiority of our optimizer over AdamW through a set of experiments text modeling and image classification. We moreover reveal how our optimizer forgets the training data at a slower pace.
REFERENCES
Aida Amini, Saadia Gabriel, Shanchuan Lin, Rik Koncel-Kedziorsk, Yejin Choi, and Hannaneh Hajishirzi. Mathqa: Towards interpretable math word problem solving with operation- based formalisms. In NAACL-HLT (1), pp. 2357-2367. Association for Computational Linguistics, 2019.
Sumithra Bhakthavatsalam, Daniel Khashabi. Tushar Khot, Bhavana Dalvi Mishra, Kyle Richardson, Ashish Sabharwal, Carissa Schoenick, Oyvind Tafjord, and Peter Clark. Think you have solved direct-answer question answering? try arc-da, the direct-answer A12 reasoning challenge. CoRR, abs/2102.03315, 2021.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In AAAI, pp. 7432-7439, AAAI Press, 2020.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao. Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. Gpt-neox-20b: An open-source autoregressive language model. CoRR, abs/2204.06745, 2022.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, and Skye Wanderman- Jax: composable transformations of python+numpy programs, 2018. URL https: htt//gtthuuuua thub. com/ http//ooole///eee/ jax.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In NewrIPS, 2020.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe V. Le. Symbolic discovery of optimization algorithms. In NeurIPS, 2023a.
Yineng Chen, Zuchao Li, Lefei Zhang, Bo Du, and Hai Zhao. Bidirectional looking with A novel double exponential moving average adaptive and non-adaptive momentum optimizers. In ICML, volume 202 of Proceedings of Machine Learning Research, pp. 4764-4803 PMLR, 2023b.
Christopher Clark, Kenton Lee, Ming- Weig Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolg: Exploring the surprising difficulty of natural yes/no questions. In NAACL-HLT (1), pp. 2924-2992 2936. Association for Computational Linguistics, 2019.
Together Computer. Redpajama: an open dataset for training large language models, 2023, URL https:/ / com wtt//ooo///t//eece tttetttttttttthteeeeeeeeee wwwptttwwwooooottrreeeeoo /Reee jama Da.

---

Published as a conference paper at ICLR 2025
Aaron Defazio. Understanding the role of momentum in non-convex optimization: Practical insights from a Iyapunov analysis. CoRR, abs/2010.00406, 2020. URL
otge abs/
1010.00406.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenat- ton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xíao Wang. Carlos Riquelme Ruiz, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Collier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Nader Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetic, Dustin Tran, Thomas Kipf, Mario Lucic, Xiaohua Zhai, Daniel Keysers, Jeremiah J. Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters. In ICML, volume 202 of Proceedings of Machine Learning Research, pp. 7480-7512. PMLR, 2023.
Jacob Devlin, Ming-Wairre Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT (I), pp. 4171 Association for Computational Linguistics, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR. OpenReview.net, 2021a.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale.
In ICLR. OpenReview.net, 2021b.
Arthur Douillard, Qixuang Feng, Andrei A. Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Mare' Aurelio Ranzato, Arthur Szlam, and Jiajun Shen. Diloco: Distributed low- communication training of language models. CoRR, 01//2211.08100 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan. Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Grégoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, and Kevin Stone. The llama 3 herd of models. CoRR, abs/2407.2178.. 2024.
Nicolas Flammarion and Francis R. Bach. From averaging to acceleration, there is only a step-size. In COLT, volume 40 of JMLR Workshop and Conference Proceedings, pp. 658-695. JMLR.org, 2015.
Leo Gao. Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang. and Andy Zou. framework for few-shot language model evaluation, 12 2023. URL https: //wp///ee.ooo////// org/ records /102568326666
12

---

Published as a conference paper at ICLR 2025
Euhanna Ghadimi, Hamid Reza Feyzmahdavian, and Mikael Johansson. Global convergence of the heavy-ball method for convex optimization. ECC, pp. 310-315. IEEE, 2015.
Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang. Implicit regularization in heavy-ball momentum accelerated stochastic gradient descent. In ICLR. OpenReview.nee net, 2023.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Edward Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Represen- lations, 2022. URL https: http:///ppeeeeeeeeeeeeeeeeo//eep net/ Sorrum?id=OKMMT
Gabriel Goh. Why momentum really works. Distill, 2017. doi: 10.23915/distil.000006. URL
pttppp//////m/mmmmttmmmmm ummment
Akhilesh Gotmare, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation. In ICLR (Poster). OpenReview.net, 2019.
Baptiste Goujaud, Adrien Taylor, and Aymeric Dieuleveut. Provable non-accelerations the heavy- ball method. arXiv preprint arXiv:2307.11291, 2023.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:2312.00752 2023.
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Werra, and Mar- tin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. CoRR, abs/2405.18392. 2024.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng. Yewei Fang, Yuxiang Huang. Weilin Zhao, Xinrong Zhang, Zhen Leng Thai, Kai Zhang. Chongyi Wang. Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. Minicpm: Unveiling the potential of small language models with scalable training strategies. CoRR, abs/2404.063955, 2024.
Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William W. Cohen, and Xinghua Lu. Pubmedqa: A dataset for biomedical research question answering. In EMNLP//JCNLP (!), pp. 2567-2577. Association for Computational Linguistics, 2019.
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt J. Kusner. No train no gain: Revisiting efficient training algorithms for transformer-baat language models. In NeurIPS, 2023,
Rahul Kidambi, Pranceth Netrapalli, Prateck Jain, and Sham M. Kakade. On the insufficiency of cxisting momentum schemes for stochastic optimization. In ICLR. OpenReview.net, 2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Eduardo Blanco and Wei Lu (eds.), Pro- ceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018. Association for Computational URL https////ccaaathooooo.ooo org/ 8- 2012..
Linguistics. doi: 10.18653/v1/DI
Nathan Lambert, Valentina Pyatkin, Jacob Morrison, IJJ Miranda, Bill Yuchen Lin, Khyathi Raghavi Chandu, Nouha Dziri, Sachin Kumar, Tom Zick, Yejin Choi, Noah A. Smith, and Hannanch Ha- jishirzi. Rewardbench: Evaluating reward models for language modeling. CoRR, ab3/2403.13787. 2024.
Guillaume Leclerc and Aleksander Madry. The two regimes of deep network training. CoRR, abs/2002. 10376, 2020.
Jaerin Lee, Bong Gyun Kang, Kihoon Kim, and Kyoung Mu Lee. Grokfast: Accelerated grokking by amplifying slow gradients. CoRR, abs/2405.20033. 2024.
13

---

Published as a conference paper at ICLR 2025
Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods. In ACL (1), pp. 3214-3252. Association for Computational Linguistics, 2022.
Hong Liu, Zhiyuan Li, David Hall, Perey Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. CoRR, abs/2305.14342 2023.
Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. In IJCAI, pp. 3622-3628. ijcai.org, 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ICLR (Poster). OpenReview.net, 2019.
James Lucas, Shengyang Sun, Richard S. Zemel, and Roger B. Grosse. Aggregated momentum: Stability through passive damping. In ICLR (Poster). OpenReview.net, 2019.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? A new dataset for open book question answering. In EMNLP, pp. 2381-2399 Association for Computational Linguistics, 2018.
Patrick G. Mulloy. Smoothing data with faster moving averages. Technical Analysis of Stocks and Commodities Magazine, 12, 1994.
Arkaddii S Nemirovskii and Yu E Nesterov. Optimal methods of smooth convex minimization. USSR Computationnl Mathematics and Mathematical Physics, 25(2):21-30, 1985.
Yurii Nesterov. A method for unconstrained convex minimization problem with the rate of convergence o(1/k²). In Doklady Akademii Nauk SSSR, 1983. URL https://app /apii
Semanticscholes www.:/llrrrrrr/aaoooooooo wttg//orpsss.rooooooooooo
Corpus ID: 200214442
Hristo Papazov, Scott Pesme, and Nicolas Flammarion. Leveraging continuous time to understand momentum when training diagonal linear networks. In AISTATS, volume 238 of Proceedings of Machine Learning Research, pp. 3556-3564. PMIR, 2024.
Razvan Pascanu, Tomás Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In (3)I. volume 28 JMLR Workshop and Conference Proceedings, pp. 1310-1318. JMLR.org, 2013.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z.. Yang, Zachary Veto,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In NeurIPS, pp. 8024-8035, 2019.
Guilherme Penedo, Hynek Kydlfček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. URL https: / aarxiv..iv.. org/ abs/2400000 22002446... 17557.
B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Com- putational Mathematics and Mathematical Physics, 3(4):864-878, 1964. ISSN 0041-5553. doi: https://doi.org/10.101600000 0000:/00000000001155533333333333 com/
URL https: /tt:://www. www.ecediieeettteccoo
Ning Qian. On the momentum term in gradient descent learning algorithms. Neu- ral Networks, 125151 1999. ISSN 0893-6080. doi: https://oi.orr/10 1016//008996608899990000000 URL http:://www.ccieeeeere..t... arricleepp1//009930000000
com/science/ httpc///eccctcccco//ccceecceee
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In ICML, volume 139 of Proceedings of Machine Learning Research, pp. 8748-8763. PMLR, 2021.
14

---

Published as a conference paper at ICLR 2025
Tal Ridnik, Emanuel Ben-Baruch Asaf Noy, and Lihi Zelnik-Manor. Imagenet-21k pretraining for the masses, 2021.
H. Robbins and S. Monro. A stochastic approximation method. Annals of Mathematical Statistics, 22:400-407, 1951.
Sebastian Ruder. An overview of gradient descent optimization algorithms. CoRR, wwwpaabs////// (etpl.1069099999 04747, 2016.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV). 115 (3):211-252) 2015. doi: 10070/ 1263-015-0816-y)
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In AAAI, pp. 7740- AAAI Press, 2020.
Othmane Sehbouh, Robert M. Gower, and Aaron Defazio. Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball. In COLT, volume 134 of Proceedings of Machine Learning Research, pp. 3935-3971. PMLR, 2021.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In ICML, volume 80 of Proceedings of Machine Learning Research, pp. 4603-4611. PMLR, 2018.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Im: Training multi-billion parameter language models using model parallelism, 2020. URL https: : / arxiv. org abs/ 909.9.99.00000 08053.
Guijin Son, Hanwool Lee, Sungdong Kim, Seungone Kim, Niklas Muennighoff, Tackyoon Choi, Cheonbok Park. Kang Min Yoo, and Stella Biderman. KMMLU: measuring massive multitask language understanding in korean. CoRR, abs/2402. 1 1548, 2024.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In Sanjoy Dasgupta and David McAllester (eds.), Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 1139-1147, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL https: http:////oceeedeeeeemm mlr press wtt//suttt/eeeeeeee 3 html.
Balázs Szegedy, Domonkos Czifra, and Péter Kórösi-Szabó. Dynamic memory based adaptive optimization, 2024. URL
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. ICML, volume 139 of Proceedings of Machine Learning Research, pp. 10347-10357. PMLR, 2021.
Hugo Touvron, Matthieu Cord, and Hervé Jégou. Deit III: revenge of the vit. In ECCV (24), volume 13684 of Lecture Notes in Computer Science, pp. 516-533. Springer, 2022.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian. Xiaoqing Ellen Tan, Binh Tang, Ross Tay Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodrigue7, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL https:://// / //aaiiivv......i.//aaaa//aaaavvt org/ abs/ 220033302222 09288.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, pp. 5998-6008, 2017.
15

---

Published as a conference paper at ICLR 2025
Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael G. Rabbat. Slowmo: Improving
communication-efficce distributed SGD with slow momentum. JCLR. OpenReview.net, 2020.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In NUT@ EMNLP, pp. 94-106. Association for Computational Linguistics, 2017.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. Huggingface's transformers: State-of-the-art natural language processing. CoRR, abs/1910.03771, 2019.
Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training BERT in 76 minutes. In ICLR. OpenReview.net, 2020.
Kun Yuan, Bicheng Ying, and Ali H. Sayed. On the influence of momentum acceleration on online learning. Journal of Machine Learning Research, 17(192):1-66, 2016. URL http:
//jmlr. warg/papers.. ppapers/v1/16--11
.htll.
Rowan Zellers, Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In ACL (1), pp. 4791-4800 Association for Computational Linguistics, 2019.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In CVPR, pp. 1204-1213. IEEE, 2022.
Hongyi Zhang. Moustapha Cissé, Yann N. Dauphin, and David Lopez-Paz, mixup: Beyond empirical risk minimization. In ICLR (Poster). OpenReview.net, 2018.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification on update rules. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.). Advances in Neural In- formation Processing Systems, volume 35, pp. 18386-28399) Curran Associates, Inc., 2022. URL
file/b https: htttpa://proceediii ngs neurips cc/paper_fiiiieeeeeeeeeeee wwtpapprrrill/ppppppeeeeeee wttfleespppeee/22222 -apeaaa-ee--rer---a-a------- Caper--Cnferencee .df..
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye. Zhi Zhi-Quar Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. In The Thirteenth International Conference on Learning Representations, 2025. URL https: // //openreviai ew net/ fast//oorii forum? ii Beah:=iEEEx
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham M. Kakade, Decon- structing what makes good optimizer for language models. CoRR, abs/2407.07972. 2024.
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao. Ajit Mathews, and Shen Li. Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023. URL https: / aaaxiiivv.v./... org abs 2304 11277.
16

---

Contents
Introduction
2 Related Work
3 Our method: AdEMAMix
Results
4.1 Transformer LLM Training
4.2 Mamba LM Training
4.3 VIT Training.
10
Conclusion
11
A Implementation details
18
A.1 Deriving the 83 scheduler
18
A.2 Pytorch implementation
19
A.3 Optax implementation
20
Experimental details
21
B.1 Transformer LLM experiments
21
B.2 Mamba LM experiments
23
B.3 VIT experiments
24
Additional results
25
C.I
LLM experiments
25
C.1.1
In-Context Learning (ICL) results
25
C.1.2
Results using 3B parameter models
26
C.1.3
More results on forgetting
27
C.1.4 Removing the second EMA mid- mid-training
................................
29
C.1.5 Hyperparameter sensitivity
30
C.1.6 Impact of training for fewer iterations
................................
32
C.1.7 Limitations of a single EMA
................................
33
C.1.8 Removing my by setting 81 = 0
35
C.1.9 Comparing my + am2 versus (1 mi + am
36
C.1.10 Miscellancous results
38
C.2 ViT experiments
40
C.3 Comparison with other methods
41
C.3.1 Comparison with AdMeta and DEMA
41
C.3.2 Comparison with Lion
42
C.3.3 Adding a third momentum term (~AggMo)
43
17

---

Published as a conference paper at ICLR 2025
A IMPLEMENTATION DETAILS
A.1 DERIVING THE 33 SCHEDULER
Let's consider Stt), the sum of the weights given to the last gradients by EMA parameterized by 3 E [0. 1/:
S()) - (1- Σ B
We want to know which timestep would correspond to a cumulative weight of 0.5:
In(0.5)
(1 B) 2 3r 0.5 g = 0.5 =
In(3)
k=0
Let f(3) = In(0.5) (n(
This function provides how many past consecutive gradients receive
cumulative weight of 0.5.
Its inverse is:
(t) = 0.5
want a scheduler which increases 3 from 3start to Bend such that f(3) increases linearly. Given an interpolating parameter μ E [0, 1], this scheduler can be written
B(x)) f f-1(111------- () uf (Bend)
By replacing f and f by their respective formula, one can arrive to:
In(Sstart) In (n(Ben))
((x) exp (11 11) (n(Bend) + + In In(Sstari) A
By setting Bend 33 and μ Toy arrive the 8g-scheduler introduced in $ 3. We show the shape of our scheduler and compare to a linear scheduler in Fig.
1.00
0.98
3(2)
0.94
Our s-scheduler: Astart = 0.9, Bend 0.9999
Lineur scheduler: Barart 0.9, Bend 0.9999
0.92
Our A-scheduler: Astart 0.9, Bend
0.999
Linear scheduler: Butrtt 0.9, 0.9, 0.999
0.0
0.2
0.4
0.6
0.8
1.0
Figure 7: AdEMAMix's By scheduler. We compare scheduler linear scheduler for Sstart 0.9 and Bend E (0.999, 0. 0.9999). While our scheduler looks more aggressive at first glance, increases fast for smaller values of and slowly for larger ones associated. This makes sense the same increase of 3 for larger 3 values has bigger impact than that same increase applied to smaller value of The linear schedulers look practically the same, despite values of Bend differing by one order of magnitude. This is not the case with our scheduler.
18

---

Published as a conference paper at ICLR 2025
A.2 PYTORCH IMPLEMENTATION
The following is a code skeleton for our AdEMAMix optimizer in Pytorch (Paszke et al., 2019).
The full implementation of AdEMAMix in Pytorch is provided in the following repository: https:
/ //github.com con/ haamn/pplleeee e ml ademamix
Listing 1: AdEMAMix code skeleton using Pytorch
Import math
Import torch
1rom torch. optim import Optimizer
class AdEXAREX (optimizer);
def
init
(self,
parana.
"phaa5.0,
13
"_beia3-D,
14
"_aipha=3,
15
eps-1e-8,
16
weigar_.eccam0000:
17
init the optimizer
18
19
#torch.no_grace
20
st fl:
group seLt.param_gEzops:
23
24
Ir - group["lr")
25
Inbda - groupp"seight..nn. ""egght____eaa""
26
group "eps"]
27
betal, betaz, betal_tinal
group("bstas")
28
T_bara group/"T ("r beta3"] ""rup("[_Betta")
29
I_slpha group "rrupp"[_aapaa aipha_tinai groupl"alppe
30
32
p group []
33
34
grad p.grad state selt.statelp)
35
36
37
State initialization it len (staze) == 3:
38
79
state["step")
step counter used for bias correct fon [m1"] secch..esos__iae
40
fast ENA
statel"mat [mm"] state [*au"] terech. like
EKA
n1, 82, nu state["ml"), state(""2"), stata("no*)
15
46
Blas correcticn: no correction for bata3's a ENA state "etee""tepp]] +-
47
48
ciass_correction - betsl " state("step"] piaa_correction2 -1 betal state("step")
49
50
51
Calling the schedulers for slpha and beca3
52
siphe Aipha_scheculerfee ("stepp"""""" start-0, end-aaphs_ T-I_slphs) beta3_acheduleg (state ("stte"" start-betal, and-betas_finn T-T_pata3)
Update (gra,, alpha-1 hetal] _mml [betall.add_
57
m2.mul_ (beta3) .add _add_ (grad, alpha-1 betx3] (beral nu.mul_ (imtal_ttttt]) (beta2) aadecmul_ igzac, grad, vaice=;, betaZ]
58
59
60
Compute step denon (quts / math seet (blas_ correettorr)))) aadd_ (epa) update (nl.div (wiavvviiissooeeee (sies_correctior]] alpha n2) denon.
61
62
63
Add weight alpha=lnbds)
67
Apply update scaled by the learning rate p.add_ update)
68
69
70
return 1os8
19

---

Published as a conference paper at ICLR 2025
A.3 OPTAX IMPLEMENTATION
The following is a code skeleton for our AdEMAMiX optimizer in Optax, an optimization library based Jax (Bradbury et al., 2018). The full implementation of AdEMAMix in Jax is provided in the following repository: https: //tp://gtttttt///////e wttphhu...ooppppppp/eppepp anl-addmmmix
Listing 2: AdEMAMix code skeleton using JAX and Optax
Fxom typing import NanedTuple
Import chex
1zom optax. inport transtorm, cempins, base, numerics From jax Import tree_ aree_att1 as jtu
import jax.numpy as Sap
class Scaleeyydemamikkkaat (Namee (NamedTuple):
count_n2: Arre.
Ussa-Updates
hase-Updatte
13 nu: Rese.Updater
14
15
16 def ademanix (Lr,
17
bi-0.9,
18
b2-0.999,
19
b390.9993,
20
phaaa.0,
bs__chedaler--eiee
d.phh__ccaediler-Nonn.
23
eps=le-8,
24
weighh__ecca=00001
25
return c
26
scsle scale_ by_annea (b1, b2, b3, alpha, schheeed,,,,, schhe_sch epsi,
27
(weigat__accay)
28
1er_learra1g (1r), tranaforn.aca
29
30
32 scale scale_by_ bydammaaxx
33
b2,
34
b3, aipaa,
35
36
b3 23__achdullr,, alpha_sciiie scheddeeee,, ops)
37
38
79
40
def fnit_ffilaaaan
(params)
EMA
cree_seros__ (parans)
nо txee_zeros_iie (params)
moment
44
return #RyAdemar ixState
15
count-jnp. zeros (eaut= (1), jnp.int32),
46
count_nuz= zeros jap.int22),
47
ni -i-n!. m2-n2, n0-hu
48
49
50
51
52
def update_fn (upddatee,,
params-Yone):
53
count_m2) if ss__scheellrr c_sipha sipha_schediier (state.count_
se__chhdulerr
aipha_scheduler noz Hone alpha (1-b1) updates
56
tree_ ppee_ppaata noment (updates, state.m), b1, nu tree_ updete__mmment nonent (updates, state.m2, c_bš, @aoe__udate__mmmee
37
m2 tree
38
poor_aion_nomm/ppoo count 'nc nuner ics.eefe (iffeinn int intee__nneeenntttt (state.
state.nu, b2, 21
59
60
count_n2_ine numerics. caarriccssaae... int32_ increeennnnnnnte (stats.count_r2. how_biaaa_ooretttleeeeee (m), 51, count_ine) tree
61
hat stre__iie_ccoreet cille_ccrreectin
62
hat
b2, count_ine)
63
updates jtu.tree_napp
(m) (ml_to_alphasm) (ttpa//s_////////r///r/ teps),
nithat, m2,
67
no_hat
68
)
69
return updates, ScalaäyademmmisStttee
70
count=count_it
71
Saun___22coun_____ia mi -nl, m2-n2,
72
73
74
77
return base-Grad www.-Grrddddntdd aeaeecaraddeenrrraaafttatttteo tn, update_tn)
20

---

Published as a conference paper at ICLR 2025
B EXPERIMENTAL DETAILS
B.1 TRANSFORMER LLM EXPERIMENTS
Architecture details. Our architecture is based on the transformer decoder of Vaswani al. (2017). We use learnt positional encoding. We use SentencePiece (Kudo & Richardson, 2018) tokenizer with a vocabulary of 32000 tokens. The model specific parameters for the different sizes are summarized in Table 1. Our implementation is using Jax (Bradbury et al., 2018), and we train using bfloat 6, except for normalization modules and softmax which use float32. The optimizer states and operations in float32.
Table 1: Model parameters for our LLM experiments.
Model parameters
110M
330M
1.3B
Hidden size
768
1024
2048
FFW expansion factor
Attention heads
12
16
16
Layers
12
24
24
How did tune the hyperparameters? Starting from smallest models (110M parameters), first tuned hyperparameters for our A models. We then use the best hyperparameters found as starting point for AdEMAMix and tuned 83 and a. When we use schedulers for a and $3, we To = T9, = T, with T being the total number of iterations. Table 2 summarizes the hyperparameters we tried for this model size. The impact of AdEMAMix's hyperparameters is discussed extensively in App. C.1.5. For our 330M parameter models, we mostly kept the best hyperparameters found for the 110M model and re-tuned the learning rate and gradient clipping. For AdEMAMix, we additionally tested multiple 33 and a values. We summarize this process in Table 3. Finally, for our 1.3B parameter models, we re-iterated the same process, re-tuning only the learning rate and gradient clipping parameters for our AdamW runs. When trying to transfer the best learning rate found to AdEMAMix, we found it to be too high, causing instabilities we couldn't 't fix using gradient clipping. For this reason, we also tuned the learning rate for AdEMAMix for this model size. This process is described in Table 4.
Table 2: Hyperparameter tuning for our 1 110M parameter LLM models. In this table we report the hyperparameters we tuned for our 110M parameter models. When multiple values are given, we bold the parameters found to give the best results.
Hyperparameter
Value
Learning rate
0.005, 0.002, 0.001, 0.0005, 0.0001 2000, 3000, 4000, 5000, 6000
Number of warmup steps
Sequence length
1024 0.1, 0.0 no-decay, cosine-decay 64
Weight decay
Learning rate decay scheduler Batch size
Gradient clipping
None, 5.0, 1.0, 0.5
Adam B1
0.9, 0.99, 0.999, 0.9999 0.95, 0.999
A B2
AdEMAMix 31
0.9
AdEMAMix 82
0.999
AdEMAMix 33
0.999. 0.9999. 0.99999 2, 4, 6, 8, 10, 15, 20
AdEMAMix a
Hyperparameters for experiments switching from Adam and AdEMAMix. For experiments in Fig. 5b and Fig. 5c, when we switch from AdamW to AdEMAMix during training, for our 110M parameter models (Fig. 5b), we
2, 83 0.9999 and T, Ta,83 0. For our 1.3B parameter
21

---

Published as a conference paper at ICLR 2025
Table Hyperparameter tuning for our 330M parameter LLM models. In this table we report the hyperparameters we tuned for our 330M parameter models. When multiple values are given, we bold the parameters found to give the best results.
Hyperparameter
Value
Learning rate ?)
0.001, 0.0005, 0.0001
Number of warmup steps
3000
Sequence length
1024
Weight decay À
0.1
Learning rate decay scheduler
cosine-decay 96
Batch size
Gradient clipping
1.0, 0.5, 0.1
A B1
0.9
AdamW B3
0.999
AdEMAMix 81
0.9
AdEMAMix 82
0.999
AdEMAMix 33
0.999, 0. 0.9999, 0.99999
AdEMAMix
5. 8, 10, 15
Table 4: Hyperparameter tuning for 1.3B parameter LLM models. In this table report the hyperparameters we tuned for our 1.3B parameter models. When multiple values are given, we bold the parameters we found give the best results.
Hyperparameter
Value
Number of warmup steps
3000
Sequence length
1024
Weight decay À
0.1
Learning rate decay scheduler
cosine-decay 128
Batch size
Gradient clipping
1.0, 0.5, 0.1
Learning rate 7) for AdamW
0.001, 0.0005, 0.0003, 0.0001, 0.00005 0.9
Adam W 31
AdamW W $2
0.999
Learning rate 77 for AdEMAMix
0.0005, 0.0003
AdEMAMix 81
0.9
AdEMAMix 82
0.999
AdEMAMix 33
0.999, 0. 0.9999, 0.99999
AdEMAMix a
1, 5, 8, 10, 15
models (Fig. 5c), from the AdamW model we are switching from.
1, 1, 0.9999 and T. Ta, 0. Other hyperparameters are inherited
Hyperparameters used our constant learning rate scheduler experiments. For Fig. 3b 7 = 10 4, the remaining of the hyperparameters are identical as those used for our other 1.3B experiments and provided in Table 4.
Hyperparameters used in our forgetting experiments. For the experiments in Fig. 4, we used 110M parameter models and hyperparameters from Table 2.
22

---

Published as a conference paper at ICLR 2025
B.2 MAMBA LM EXPERIMENTS
Architecture details. We use a Mamba architecture (Gu & Dao, 2023) with an embedding dimension of 768, expansion factor of 2, a state size of 16, and depth of 24. We batch size of 120 sequences 1024 tokens. We use the EllutherAA. http/gpteneoo-eeeeeeee 20b tokenizer from Black et al. (2022), using the HuggingFace library (Wolf et al., 2019).
Hyperparameter details. We used parameters mostly taken from Gu & Dao (2023):
Learning rate: 77 = 0.0006,
Warmup steps: 3k;
s--ccheullrr type: cosine decay to 10-5
Weight-decay: 0.1,
Gradient clipping value: 1,
Total training steps T E {64k, 128k).
For experiments with AdamW, we use B1 0.9 and B2 0.999. For AdEMAMix, we use 31 = 0.9, 32 = 0.999, 8g 0.9999, Ta,9g = T and o = 8.
23

---

Published as a conference paper at ICLR 2025
B.3 VIT EXPERIMENTS
Architecture details. We use a VIT architecture following Dosovitskiy et al. (2021b). The architecture details for both our 24M and 86M parameter models be seen in Table 5. We do of the iterates as our final model. Our implementation in Pytorch (Paszke et al., 2019), and use bfloat16.
EMA
Table 5: Model parameters for our ViT experiments.
Model parameters
24M
86M
Patch size
16
Number of patches
14
14
14 x 14
Image size
224
224
Embedding dim.
384
768
MLP dim.
1536
3072
Layers
12
12
Number of heads
12
Data augmentation and Mixup. While more sophisticated augmentation methods exist (Touvron et al., 2022), we simply use random-resized cropping, random horizontal flip (p = 0.5), and random color jitter (applied with probability 0.8).
How did we tune the hyperparameters? For each model size, we started by tuning the hyperparam- eters for the AdamW baseline. The hyperparameters three different settings are in Table 6. For our experiments on ImageNet- lk, given that all models overfit the dataset, select the best model according the minimum validation loss, akin to using early stopping. For AdEMAMix, for each setting, we use the hyperparameters of the best Adam model and only tune and B3. For each setting, train models using {1, 5, 10, 15, 20) and 33 E {0.999, 0.9999). All of the AdEMAMix models are shown in Fig. 6 and Fig. 31. Only the best AdamW W model is shown in Fig. 31.
tuned and the values retained for
Table 6: Hyperparameters tuning for our VIT Adam W experiments. We started our hyperpa- rameter search from the values given in (Dosovitskiy et al. (2021b), Table 3), which recommends using (learning-rate, weight decay, dropout)= (0.003, 0.3, 0.1) when training on ImageNet- Ik and (0.001, 0.03, 0.1) when training on ImageNet-21k: The various hyperparameters we experimented with are in the following table, values which gave us the lowest validation loss are in bold. For each setting, AdEMAMix experiments the same hyperparameters the best AdamW W basclines.
Setting
24M params, 11M images (ImageNet-21k)
86M params, 11M images (ImageNet-21kk
86M params, 1.3M images (ImageNet- 1k)
Learning rate
0.001, 0. 1..033,00 0.005
0.001, 0.003
0.001, 0.003, 0.005 0.03, 0.1, 0.3, 0.5, 0.7 0.1 0.9, 0.99. 0.999 4096
Weight decay
0.03, 0.1, 0.3 0.1 0.9, 0.99, 0.999 4096
0.03, 0.1, 0.3 0.1 0.9, 0.99, 0.999 4096
Dropout
AdamW 81
Batch size
24

---

Published as a conference paper at ICLR 2025
C ADDITIONAL RESULTS
C.I LLM EXPERIMENTS
C.1.1 IN-CONTEXT LEARNING (ICL) RESULTS
In-Context Learning (ICL) results. We evaluate our largest (1.3B) LLM models on in-context
learning tasks. We use the 1m-eval package (Gao et al., 2023). We consider the following tasks:
HellaSwag (Zellers et al., 2019), Winogrande (Sakaguchi et al., 2020), ARC (Bhakthavatsalamm et al., 2021), BoolQ (Clark et al., 2019), LogiQA (Liu al., 2020), MathQA (Amini et al., 2019), MMLU (Son et al., 2024), OpenbookQA QA (Mihaylov et al., 2018), PIQA (Bisk et al., 2020), PubmedQA (Jin et al., 2019), RewardBench (Lambert et al., 2024), Sciq (Welbl al., 2017). TruthfulQA (Lin et al., 2022).
The two AdEMAMis and AdamW models benchmark have been trained for 1M steps, with batch size of 128, corresponding to around 131B tokens. The results of the evaluation are Table 7. We observe how the model trained with AdEMAMix is outperforming the AdamW model on nearly all of the tasks, sometimes by large margin for e.g. PubmedQA. In addition to Table 7, we track the evolution of some of those scores during training for MMLU, RewardBench, ARC-Challenge, and ARC-Eas Easy, results can be seen in Fig. 8. For that figure, we modify the way MMLU is evaluated. Instead of appending all the multiple answers, we concatenate the question prompt and each answer separately, and pick the most likely answer. We found it allows much better comparison of smaller models, which otherwise fluctuates around random guessing.
Table 7: In-context learning results for our 1.3B parameter models. We compare Adam W and AdEMAMix models trained on 131B tokens. We use the Im-eval package.
Task
AdamW
AdEMAMix
Task
Adam
AdEMAMix
ARC-Challenge
0.262
0.274
OpenbookQA
0.240
0.238
ARC-Easy
0.612
0.619
PIQA
0.715
0.715
BoolQ
0.569
0.576
PubmedQA
0.556
0.632
HellaSwag
0.426
0,436
RewardBench
0.569
0.573
LogiQA
0.235
0.225
RewardBench (reasoning)
0.617
0.630
MathQA
0.226
0,236
Seiq
0.903
0.907
Winogrande
0.563
0.580
TruthfulQA
0.361
0.352
MMLU
0.244
0.248
0.633
0..02
Saliara Mine
6:26
11,00
Aceursey
0.20
Ansuaayy
Accuracy
Accussy
0.61
6:24
0.53
0.50
6:22
0.30
0.20
16B
70B
151B
15B
75B
1518
168
70B
131B
15B
70B
131B
1okmns
(a) MMLU.
(b) RewardBench.
(c) ARC-Challenge.
(d) ARC-Easy.
Figure 8: In-context learning results for our 1.3B models. We periodically measure the perfor- mances of AdamW and AdEMAMix models in-context learning tasks. Except for couple exceptions, AdEMAMix is performing better than AdamW.
25

---

Published as a conference paper at ICLR 2025
C.1.2 RESULTS USING 3B PARAMETER MODELS
Experimental setup. We train 3B parameter models on RedPajama v2 data. We use learning rate of = 0.0003, a weight decay À = 0.1, a gradient clipping value of 0.1, a batch size of 1024. The sequence length is 1024. We use cosine 17 decay and 3000 steps of warmup. For AdEMAMix we used B3 E {0.9999, (.99995} and 31 E {0.0,00 0.9}. For AdamW, we use 31 = 0.9. Both methods use B2 = 0.999. We train models for 477k iterations, on a total of 500B tokens.
Results. Results are shown in Fig. 9. We observe that AdEMAMix is outperforming the AdamW baseline. Moreover, using 81 0.0 AdEMA.ixx not deteriorating the results. More experiments required to estimate how many additional steps would be needed for the A baseline to reach the same performances as AdEMAMix.
AdbanW AdEMAMix 3, 0.9
Adami
2.4
24
AdEMAMIx
AdEMAMix
2.3
2.3
Loss
Loss
2.2
2.2
2.1
2.1
2.0
2.0
0
477k
645k
0
477k
Iterations
Itrraioon
(a) B3 0.9999, 31 E (0.0,0.9).
(b) 0/3 0.99995, 81 0.9.
Figure 9: Results using 3B parameter models. For AdamW models, we use 31 = 0.9, 3, = 0.999. In (a) we compare the AdamW baseline to AdEMAMix models trained with B3 0.9999 and 31 E (0.0, 0.9). Both AdEMAMix models reach a similar final loss and outperform the baseline. The convergence to the solution is slower at the beginning for 31 = 0.0. When 31 0.0, the memory footprint of AdEMAMix is the same as for AdamW, meaning that we can get the improvement in loss at no extra cost. In (b), we train an AdEMAMix model with a larger 33 = 0.99995. While we observe lower loss in the early training phase, the final loss is similar to the one obtained with By = 0.9999. AdEMAMix reaches the same loss as an Adam W model trained for 168k more steps.
26

---

Published as a conference paper at ICLR 2025
C.1.3 MORE RESULTS ON FORGETTING
Evolution of forgetting during training. In Fig. 4 in the main paper, follow the loss on one specific batch over the entire training. In the following experiment, we do a closeup on the forgetting of different batches as training progresses. The goal being to visualize how later batches are ultimately more remembered, and compare the forgetting profiles of Adam and AdEMAMix. For this, in Fig. 10a and Fig. 10b, we follow fixed batches at different stages of the training process. Every 10k: iterations, we measure the loss over a specific batch B before and after training on that batch.
0.0
0.0
2506
-0.5
Normalized loss
-0.2
2008
1.0
-0.4
1.5
150R
-0.6
2.0
of
100R
-0.8
2.5
1.0
3.0
5000
10000
5000
10000
lterations training on
Iseracions after training on
(a) AdamW.
(b) AdEMAMix.
W and AdEMAMix. In (a) and (b), we follow- over 10k iterations-the loss training batches B, for T E {30k, 40k; 250k]. We measure the loss right before training B,, and then monitor this loss over the next 10k steps. For instance, the most yellow curve represents the loss on batch B250k that we track for the last 6k iterations of training (we train for 256k iterations). The loss is normalized such that the loss for B, is 0 and the one at time T + 50 is -1 This allows us to overlap curves and observe the tendeney of the model to forget batches throughout training. We observe that AdEMAMix improves its loss on training batches B, over many timesteps. a striking difference compared to AdamW, which has its loss on B, increasing faster. Moreover-likely due the learning rate scheduler- the models forget significantly slower at the end of training. In this later training stage (T > 180k), the difference between the two optimizers is especially pronounced. While this visualization allows to efficiently compare forgetting at different stages of training, the normalization prevents us from comparing the drop loss between AdEMAMix and AdamW. To compare loss values, see Fig. 4.
Figure 10: Comparing forgetting between A
Why are training batches forgotten more the end of training? Given the previous observation about how both AdamW and AdEMAMix models are forgetting the training batches at different paces given different stages of training, one can wonder what is causing this phenomenon? To answer this question, we contrast the results of Fig. 4 with results obtained using different scheduler with a constant learning rate and linear decay. We train 110M parameter models for 300 iterations, using max learning rate of 0.001 and batch size of 64. Results for those experiments are in Fig. 11. Those results indicate that the decaying learning rate might be the main parameter controlling how much a batch is remembered during training. This has interesting implications when selecting which learning rate decay to use. A cosine decay, with a rather long period of decay, might remember more than constant learning rate scheduler with only small number of steps of linear decay at the end.
27

---

Published as a conference paper at ICLR 2025
3.1
AdbeaM truiniss leuining
3.0
on batch
2.9
2.8
2.7
70k
230k
300k
70k
2702
2502:
70k
2906
without training ul
3.0
batch
2.9
tg
2.6
70k
(a) LB 230k.
(b) LB 270k.
(c) lB 280k.
(d) (. 290k.
Figure 11: Measuring forgetting using held-out batch B, using a constant 1j-scheduler with linear decay. The top row is for AdamW, the bottom row for AdEMAMix. The setup similar to the from Fig. 4, except that now use different learning rate scheduler. We use 3000 steps of linear warmup. them 247k steps of constant learning rate 7 = 0.001, and finally 50k steps of linear decay (represented by the shaded green area in the above plots). In Fig. 4, observed that the batches used later during training are remembered more by both Adam W and AdEMAMix. Multiple hypotheses could explain this phenomenon: (i) it could be that the learning rate decay is the main parameter controlling forgetting, or (ii) it could be that a sufficient number of steps are required reach part of the optimization landscape where gradients remembered. The present experiment refutes that second hypothesis. Indeed, before the start of the n-decay (as seen in (a)), training batches are remembered the end of training. They only start to be remembered when the decay periods starts, as seen in (b,c,d). Moreover, we can make the same observation as in Fig. compared to Adam W, AdEMAMix models tend to keep the training data longer in memory.
more orthogonal and baches are
28

---

Published as a conference paper at ICLR 2025
C.1.4 REMOVING THE SECOND EMA MID-TRAINING
Removing the second EMA mid-training. In Fig. 5b and Fig. Sc, looked at what is happening when we switch from Adam W to AdEMAMix in the middle of training. In this section, we will look the opposite conversion: removing the m2 parameter of AdEMAMix during training, effectively switching back to AdamW. Results for this experiment can be seen in Fig. 12. Right after the switch, we observe drop of the loss, followed by an increase and finally back to convergence. Ultimately, the final loss value is in between the ones obtained training only using A and only using AdEMAMix.
3.4
AdEMAMix 10.0, Rs 0.9990, T. TTy
AdUMAMMx 10.0. 0.9089, V.
3.2
Loss
3.0
2.9
0
50000
100000 150000 200000 250000
Iterations
Figure 12: Switching from AdEMAMix to A at Lswitch = 128k. Using our AdEMAMix optimizer, train two 110M parameters models with different To TTg (90k; 256k]. For historical reasons, those experiments were done using a batch size of 32 instead of 64 (as we used for all the other 110M parameter experiments in this paper). We switch from AdEMAMix Adam swiitch = 128k. At first, the removal of the +am, contribution in the update is suddenly decreasing the effective learning rate, causing the loss to drop suddenly. The loss then increases slightly and ends up in between the AdamW W only (blue curve) and AdEMAMix only (grey curves) losses.
29

---

Published as a conference paper at ICLR 2025
C.1.5 HYPERPARAMETER SENSITIVITY
Hyperparameter sensitivity. Depending on whether the a and 83 schedulers are used, AdEMAMix adds up to 4 new hyperparameters: a, 83, To and T2g. In all our experiments we tied To = Tsp Ta. In this section we analyze the sensitivity of AdEMAMix to those hyperparameters. We study the impact of a and B3 in Fig. 13, revealing wide ranges of values for which hyperparameters are outperforming the Adam W baseline. We study the sensitivity of Ta,9, as well as justify the choice of using a scheduler on in Fig. 14. When training from scratch Ta.8% needs simply to be large enough to avoid early divergence. In Fig. 16 we study the sensitivity to the gradient clipping and AdEMAMix's 31 value. While gradient clipping can help stabilize training and smooth the loss curves, it has little impact over the final loss value. Reducing the value of 31, we observe some loss spikes, yet the final loss value is relatively unaffected.
3.4
3.4
3,7
AMESAMII E0050
AdEMANts -.566.
AGESSAMI E:206
--3
ActMANis de -55002 a
3.2
AMESSAMMK 5.906.
AdtDMANts 1558939
3.1
AGBAHAMK 3, E.956 5
AIBSAMM
6.955
3.2
3.0
2.9
288
0
50000
100.00
150000
200000
250000
0
50000
100000
150000
200000 250000
Eterations
Itcrations
(a) Sensitivity to a.
(b) Sensitivity to Sa.
Figure 13: Sensitivity of AdEMAMix to and 33. We test the sensitivity of 110M parameter
models to the values of a and 83. In (a), we vary a E {2, 6, 10, 15, 20, 30,30.40 403 while keeping all other parameters equal (33 = 0.9999, To, = 256k, clipping = 0,5, 0.001, = 0.1). We see that the larger the a, the more significant is the loss decrease early during training, yet this does not necessarily translates into a better final loss. The last iterate's s loss is larger for large values of a. For this model, a sweet spot seems to be reach for = 10. Moreover, the span of values resulting in good results is very large. In (b), we now vary 3/ E {0.99, 0.999, 0.9999, 0.99999, 0.999995). We first observe matt-eccept for 33- 0.999995-al those values are outperforming the AdamW baseline. Moreover, we observe that there a sweet spot, at B3 = 0.9999, after which the final loss increases.
AdaeW AdESSAMIC
Adoow
-.2--
-Ty
3.4
AMESSOHH Tuy
AEEAAMHI Tay
-T,
AARSRAMM To,
3.5
3.2
3.2
3.0
3.0
2.8
28
0
500000
100.00
150000
200000
250000
50000
100000
150000 300000
250000
[terations
Iterations
(a) Sensitivity to To T3p Tu.ps-
(b) Sensitivity to T...
Figure 14: Sensitivity of AdEMAMix to T. and To. For all the AdEMAMix experiments in this 10. In (a), test the sensitivity of 110M
figure, used 81 0.9, 32 0.999, B3g 0.9999, parameter models to the values of T, and Tzg- We tied the two values as in all the experiments in the main paper: T_ T TT We observe that when not enough warmup steps are used, the iterates either diverge or converge to slightly worse loss values. In (h) we set T = 256k and only vary T.. We observe that T. needs be large enough avoid divergence. This demonstrates the necessity of using a warmup on a. This being said, the necessity of increasing a linearly depends on the value of a. Using a smaller a value can work with To 0, but using an a-scheduler alleviates any instability issue and increases the range of optimal a values (see Fig. 15). Moreover, in other experiments, we observe how not using schedulers yields detrimental large updates in the early training phase, see App. 1.1.10, Fig. 28.
30

---

Published as a conference paper at ICLR 2025
3.2
AGENAASI 0.0
3.3
... Miaarati T.d 230t.
AdEarati 2ast mitoon
3.0
AGERGANES 0.0
3.2
Adtaratia Eudy 25st.00
3.1
3.0
2.0
2.1
Tterations
(a) No a-scheduler (110M).
(b) a-schcduler (1.3B).
Figure 15: The a-scheduler reduces the sensitivity to a. We compare AdEMAMix experiments with and without a-scheduler. The experiments with a-scheduler are more stable. (b), for 1.3B parameter models, only early), and reaches a final loss slightly worse than the AdEMAMix model trained using an a-scheduler over 256k steps. In (a), for smaller 110M parameter models, converge without a-scheduler for small a values, as well as (ii) that seems easier to reach good loss values with scheduler (the dotted lines have lower final loss average).
converges in stable fashion (the purple curve with
diverged
observe (i) that it is possible to
3.4
Albalt
3.3
AiRASAMis soad-dlm
3.4
3.2
1.80
AcRIAAAIi gra
Lone
3.0
3.0
2.0
2.8
100000
llunations
Beraaiios
(a) Sensitivity gradient clipping.
(b) Sensitivity $1.
Figure 16: Sensitivity of AdEMAMix to 31 and the gradient clipping value. Unless specified in the caption, all the AdEMAMix experiments in this figure used 31 = 0.9, 32 = 0.999, B3 0.9999, a 10, Ta,8, 256k. In (a) we vary the amount of gradient clipping used during training. We notice how gradient clipping can help to smooth the curves, yet bringing only minor improvements. In (b) we perform sweep over 31 {0, 0.1, 0.3, 0.5, 0.7, 0.9, 0.99). The value 3 0.0 is especially interesting as it allows us to remove entirely m. therefore saving memory. We observe that smaller 31 values yield noisier curves, with multiple loss spikes. At the 110M parameter scale, those spike do not significantly impact the convergence, we conjecture that they could become a problem at larger scales. We also notice that slightly better losses could be obtained using smaller 3 values (e.g. 3 = 0.3), we keep 31 = 0.9 in all of our experiments for the ease of comparison with AdamW.
31

---

Published as a conference paper at ICLR 2025
C.1.6 IMPACT OF TRAINING FOR FEWER ITERATIONS
Sensitivity to the number of iterations. As we rely on old gradients, on question that arises is whether AdEMAMix would still perform well when reducing the number of iterations. In this section we compare the loss obtained by 110M parameter models when halving the number of iterations and doubling the batch size, in such a way that the number of tokens consumed for training is always the same (17B). In Fig. 17, we observe that decreasing the number of iterations too much increases the final loss of the model. This effect is more pronounced for AdEMAMix. However, at both 32k and 64% iterations, AdEMAMix still outperforms Ad In Fig. 18 we observe that reducing {3 can mostly alleviate the problem.
3.4
AdlenWW: 2052 steps, ha 64
AdEMAMIix 2015
Adact 1311 AdIIMAMi 170)
Adhosti 1264 stope, 34 139
ASEMAMIS 1288 staps, be 139
Adaanii 646 stepe, ba i 250
AGEMAMII G12 stope, An 200
2.86
Adaan' 320 stape, da = 512
AGEMAMIIC 322 stope, de 512
3.2
3.2
Findl loss
Loss
2.84
3.0
3.0
2.82
2.8
2.8
2.80
5B
10B
17B
0
5B
10B
17B
32k
61k
128k
256k
Tolora
Tolaens
Number of training iterations
(a) AdamW with fewer steps.
(b) AdEMAMix with fewer steps.
(c) Comparing the final loss.
Figure 17: Impact of using fewer iterations on AdEMAMix vs. A We train multiple 110M parameter models. All models are trained on 17B tokens, but using different batch sizes (bs E {64, 128, 256, 512]), which results in different numbers of iterations: (32k, 64A, 128k, 256k]. For both AdamW and AdEMAMix, we use hyperparameters from Table 2 (33- = 0.9999 for AdEMAMix). In (a), we observe a drop in performance when reducing the number of steps, as the loss of the final iterates is higher for the model trained for 32k steps compared to the ones trained for more iterations. In (b) we observe a similar story for AdEMAMix models, with a slightly more pronounced performance degradation. Yet, as it can be seen in (c), the gap between AdEMAMix and AdamW still remains, even when training for 32k iterations. Importantly, this experiment does not suggest that AdEMAMix is performing worse when using larger batches. In our ViT experiments, we successfully use a batch size of 4096.
3.1
AREMAMii 2002 atepe. La 64
3.4
AGEMAMii 2004
04
Adamt L3R
AdTMMAMi 1964 strpe, 34 198
ARFMAMIi 1984 steps, 139
AdRMAMMis 178, ds 0.5000 AdEMAMiis 177, da 0.509
AdEMaMMi: 646 stope, bo 700
AHEMAMIl 616 stope, is 206
2.86
AdиMAMii 20u stepe, Aa - 513
ASEMAMII 302 stope, de 512
3.2
3.2
Final loss
Loss
2.84
3.0
3.0
2.82
2.8
2.8
2.80
5B
10B
17B
0
5B
1013
1713
32% 64k
128k
256k
Talons
Tolees
Number of training iterations
(a) AdEMAMix with fewer steps
(b) AdEMAMix with fewer steps
(c) Comparing the final loss.
33 0.9999 (same as Fig. 17h).
8g
0.999.
Figure 18: Reducing ,33 can work better with low numbers of iterations. We take the same experimental protocol Fig. 17 but this time compare training AdEMAMix models with 3 = 0.999 and 3 = 0.9999. In (a,b) we observe how using a smaller /33 = 0.999 solves-in a large part-the problem of decreasing performances when using a small number of iterations. In (c), while final loss for T' E {128k, 2566} is lower using 83 = 0.9999,, using 33 = 0.999 is better for T E {32k, 64 Using A3 0.999, the gap between AdEMAMix and Adam W increases as the number of iterations decreases.
32

---

Published as a conference paper at ICLR 2025
C.1.7 LIMITATIONS OF A SINGLE EMA
Results on a 2D toy example. A natural question that arises from our method is whether it is possible to obtain the same results without the additional EMA. In this section we aim to convince the reader that the two EMAs are required. First, we propose to study a small toy 2D optimization problem:
(x", (w","")= E arg min [(I, v); with f(x, = 8(x 1)2 x (1.3a² + 2x + 1) + 0.5(y 4)2
(2.0)
This function was introduced by Liu et al. (2023) as function with sharp curvature along the x-axis, and flatter curvature along the y-axis. Initializing the first iterate (0)) (0.3, 1.5), we start optimizing with (i) Adam with a large 31 = 0.999, and (ii) AdEMAMix with 31 = 0.9 and 33 = 0.999. In both cases 32 0.999. To make the experiment interesting, we initialize the EMA buffers for both methods to either ( (-3, 0) or (-0.8, -3). This has for effect to give an initial "speed" to the first iterate. As result of this speed, Adam with large 31 is unable to correct his trajectory. In contrast, using two FMAs, one with a small, and one with a large B-as in AdEMAMix AdEMAMixx-convrges
to the solution. The fast changing EMA can correct for the bias introduced by the slow changing EMA. See Fig. 19 for results.
Adam Fq C.508 AdEMAbii F1 G.D. I. 6,504
Adam Bt DURI AIESAASIS D. C.995
start
start
(a) m
m2 (0)
-13,,0 0j.
(0) (b) ml m2 |-0.8, -3).
(0)
Figure 19: Showing the necessity of a second momentum term on a toy example. In this experiment, we use Adam and AdEMAMix optimize simple function (same function as in Fig.2 in Liu et al. (2023),* * is a local optimum). We want to understand what is preventing the use of large 31 values in Adam, and how AdEMAMix is overcoming this problem. Starting from initialize m. (resp. m. and m2) for Adam (resp. AdEMAMix) to either [-3,0] or [-0.8, -3] shown with black arrows (not to scale). () is still initialized to O. This is similar to giving an initial speed to the iterate. In both cases (a) and (b), the iterates following the Adam trajectory fail to converge in the given number of iterations. This is due to the large 31 which requires many hundreds of iterations to significantly alter m. As a result, the iterates mostly follow the initial direction imposed by m (0) (v is still used scale m.). In contrast, AdEMAMix converges to the solution. The fast (smaller (3) EMA of AdEMAMix corrects the trajectory of the slow (high 3) EMA. Interestingly, AdEMAMix initially overshoots the optimum. It has to wait for
(x))) [0.3, -1.5),
the influence of m to fade away before finally converging to the solution. While this experiment does not explain why a larger momentum would be beneficial, it illustrates how using two EMAs can enable the use of large B values. A better intuition behind why a large momentum could be beneficial presented in Fig. 2.
Results training 110M parameters LMs. To further demonstrate that simply increasing 31 in Adam does not provide nearly the same gains as AdEMAMix, we run several additional experiments. In Fig. 20 we show what we obtain when training from scratch using a single EMA with a 31 E (0.99, 0.999, 0.9999, 0. 0.99999). Naively increasing 31 in Adam W does not work; adding a scheduler on 81 to smoothly increase its value during the entirety of the training also fails. In Fig. 21, we show results when increasing 3 in the middle of training, with and without scheduler on 31. This differs from the previous setting we bypass the initial training phase capable of causing instabilities (see
33

---

Published as a conference paper at ICLR 2025
3), well as bypass the initial iterations during which the bias correction done by AdamW can have an impact. Here again we observe increasing the 31 value does not provide any noticeable gain.
6.0
3.4
Adazsk 9,
Adamdh 83 & 1569
3.5
Adaziit
Adaat! 1.00.
4456)
Adaast bn 1.006
Adsalt
1.000, 75 06.29
5.0
Adamsk P) 1.0063
3.2
Adaatti
2.0066.22
4.5
3.1
4.0
29-
28-
0
500000
100,00 1500000
200000
250000
0
30000
10.0.0
130000
200000 250000
Iteraiion
Iterations
(a) Vanilla AdamW.
(b) AdamW with -cscchuuulrr.
Figure 20: Increasing 31 in AdamW. In (a) we perform a simple sweep over the 31 hyperparameter of Ad we observe how the optimizer is unable to cope with large values of 31. diverging for 31 > 0.999. In (b) we modify AdamW incorporate the tricks we used in AdEMAMix, i.e. we add scheduler on 31. While this prevents divergence, the final loss is getting worse as we increase B1. We had to increase 32 to stabilise the training when using 31 0.99999. Those experiments support the importance of the additional EMA in AdEMAMix.
3.2
Askanmol
3.2
Adaaed ELB Iuse woidd
Adus
Aduusi
0.5000
AdsedW 6.8933 2000
3.1
3.1
AdwaeW to0ours 7y 2000
1.000
2.9
2.9
2.8
200000 250000 300000 330000 4000DD 450000 500000
200000 250000 200000 350000 400000 450000 D00000
Barrstiun
(a) Increasing 31 without warmup.
(b) Increasing 81 with warmup.
Figure 21: Increasing 31 in Adam W later during training. In (a) we load an AdamW checkpoint and resume training using A with larger 31. We observe no gain over the baseline from increasing the value of 31. In (b), we add a scheduler over 31 (similar to the scheduler over 33 for AdEMAMix) to smooth the transition between 31 0.9 and larger values. We set T', 200k iterations. Despite the smoother transition, no gain over the baseline is observed.
34

---

Published as a conference paper at ICLR 2025
C.1.8 REMOVING my BY SETTING $1
Effect of the batch size when B1 0. In this section investigate the impact of setting 31 which allows to replace m by the current gradient g(), saving memory. In this scenario the memory complexity of AdEMAMix is the same the one of AdamW. We noticed in Fig. 16b how reducing 8 can make the training slightly more unstable, triggering some loss spikes. In Fig. 22 we study the impact of the batch size on those instability, revealing that 31 is less stable when using larger batch sizes. However, despite the spikes, the final loss for 31 = 0 and 31 = 0.9 are very similar 1 0.9 yielding slightly better results for smaller batch sizes.
3.8
4,8,
AGEMAMIC batch--ttze 6,8,
A-tIMMANIII 8,5
3.6
AJEMAMis batcchsiie 5,0,
0.0
AddSMAAbii batch-atsem 16. 0
AMRMMANII harcch-tion 16 0.0
3.4
AGEMAbix butchh-iie 52.8,
AdIMMAMIIc
AIEMANI hharch-sioon
3.2
AdEMAbli
3.0
AJEMAMix
128,
2.8
128k
256k
Itrations
Figure 22: Effect of batch size with 31 = O. For a 110M parameter model, we use 31 {0.0, 0.9) and vary the batch-size E {4, 8, 16, 32, 64, 128). Lighter curves use 31 darker curves are their 31 = 0.9 counterpart. For small batch sizes, slightly better performances. For larger batch sizes, instabilities are observed with using 31 = yet the final loss values are similar as the ones obtained with 31 0.9. This indicates that one can potentially save memory without compromising results.
instabilities are observed, and 33 = 0.9 reaches
Sefting 31 for 1.3B parameter models. In this section, using 1.3B parameter model, we re-do the same experiment as in Fig. 1, but we set 31 = O. In Fig. 23, we show the results of this experiment.
AdumW toleua
AdEMAMi 101B tolens 8, 0.9.A. 9.999 AdEMAMI 1011 takens 8, U.0.PPs 0.9999
2.7
2.5
2.
2.3
2004
500k
770k
Iterations
Figure 23: Using 31 with 1.3B parameter model. We use 82 0.999 for all those
experiments. We observe how it is possible to converge despite using 81 = 0. The final loss with 31 = 0 is slightly better than the one with 31 0.9.
35

---

Published as a conference paper at ICLR 2025
C.1.9 COMPARING m. am2 VERSUS (1 - a)m{ + amz
In the equation AdEMAMix, we combine my and m2 using m, + amng In this section we contrast this to using a convex combination (1 a))m; + am2, with a E 1].
Are those two somewhat equivalent? Given the combination s7(m) + am2), it possible rewrite it as convex combination as such:
a mm(m1 + am2) - fi((1 aimi + (m2)) with f - 17(a + 1) and à O +
Therefore, the two formulas equivalent up-to a reparametrization. However, in practice, use schedulers on Yj and a.. While at any timestep I, it is possible to reparametrize one formula into the other, there is not one single reparametrization that holds for all timesteps. Assuming that we do not change the nature of the schedulers for a and those two formulations are therefore not equivalent. In Fig. 24 we show the evolution of the weights given to my and m2 for the two approaches. We assume cosine scheduler for 1j and linear scheduler for a. We can see how the two formulations give very different weights to m. and m2.
x10
1.0
seight o01i 1
weigat for my: 8"
wright for my 20
weight for mms
0.5
0.0
0
100000
200000
100000
200000
(a) (0 10 (mi + a m2).
(b)) a ma m2).
Figure 24: Weights given to m, and my for the two formulations. In (a), we show the weights given to m, and my when we use ()()) + al)) We use a 8 (ax is the final value reached by the scheduler (")) In (b), we show the weights given to m2 and m2 when we use ( (+) ym + a m2). We use a 0.9. Importantly, both cases, we assume a cosine scheduler for 17()) and a linear scheduler for att), Given this choice of schedulers, there is no reparametrization which can make those two formulation equal. In Fig. 25 we show that using the convex combination formulation yields worse results.
Is the convex combination better? Now that we established that the two formulations are different, which one is better? To answer this we run two grid searches, varying E {0.4, 0.6, 0.8, 1, 1.2, 1.4, 1.6) and a E (2,4, 6, 8, 10, 12, 14) for the formulation in equation AdE- MAMix, and a E {0.3, 0.4, 0.5, 0.6, 0.7, 0.9} for the convex combination formulation. Results in Fig. 25 show how the formulation used AdEMAMix provides a better loss for all the puirs (n, a) tested, also outperforming AdamW for all of those. In contrast, for the convex combination, while most (17, a) pairs outperform Adam W. the improvement is smaller.
Getting the same results by changing the schedulers for the convex combination. We concluded before that-unless we tamper with the schedulers-it is not possible to find a reparametrization that would work for all timesteps and make the two formulations equivalent. What if we could change the schedulers? Which schedulers would we need to use with the convex formulation, to get the same results as with the other formulations. The answer is simple, let 1)()) and (t) denote the cosine and linear schedulers for resp. 7) and The convex combination would need to use the following schedulers:
(1
- (a +1),
((())
=
a +1
Therefore, to get equivalent results as for the original AdEMAMix formulation, the learning rate scheduler would need to be (1()), which is not a cosine scheduler. The a scheduler, (!t), would not be linear. We plot those schedulers in Fig. 26.
36

---

Published as a conference paper at ICLR 2025
3.07
AdaraW
3.07
0.7
3.06
3.06
8
0.6
3.05
3.05
Arlaaa
3.04
3.01
0.3
0.4 0.6
0.8
1.4
1.6
0.4
0.6
1.2
1.4 1.6
10
17
10
m2 + am2.
(b) a-m1 + am2-
Figure 25: 7 hyperparameter sweep for the two formulations. For pairs of (77 a), we train small (6-layers) transformer LMs on RedPajama v2 data. We plot the results the two formulations, the color represents the final validation loss, and is normalized between the two plots. We also do sweep over v to fine the value giving the best results for AdamW. We use green level curve to show the performance of the best AdamW model. Points in the interior of the green level curve represents configurations outperforming AdamW. (a), we color the entire frame in green as all the pairs of hyperparameters outperforming AdamW. W. It is clear from those two figures that using m. am2 yields better results, more consistently, compared to using (1 a m2 + am2.
n/0
4!
2.5
0.0
100000
200000
100000
200000
(a) schedulers.
(b) schedulers.
Figure 26: Using schedulers ()()) and (1) with the convex combination formulation is the same as using the original AdEMAMix. The original AdEMAMix forumulation uses cosine and linear schedulers for resp. and a. Those schedulers are in blue. For the weights given to m, and m2 to be the same when using the convex combination formulation, one need to use different schedulers ((() and (9) (in orange). Using the original AdEMAMix formulation allows to use simpler schedulers.
37

---

Published as a conference paper at ICLR 2025
C.1.10 MISCELLANEOUS RESULTS
Impact of B3 and schedulers when starting AdEMAMix from AdamW. As mentioned in 3, we do not necessarily need the a and B3 schedulers when resuming training from a sufficiently late checkpoint. Fig. 5b and Fig. 5c do not schedulers, unless start using AdEMAMix from scratch. In Fig. 27, we vary the warmup periods for a and B3 when starting AdEMAMix from an AdamW checkpoint fswitch 300k. We observe how increasing Tap Ta only slows down the convergence. This serves to illustrate that the schedulers for a and 33 are only required to stabilise the optimization during the early training phase.
3.2
AdEMAMix & 0.9160 2.0 7_ Tay
7, Gay
3.1
MEEMAMix 2a 0,9999 2.0. Ta, 5ON
AdRMAMix & 0,8000
2.0, 7, Vay 100%
Lave
3.0
2.8
200000 250000 300000
350000 400000 450000 500000
Ieerations
Figure 27: No need to use schedulers when switching from Adam W to AdEMAMix sufficiently late. For 110M parameter model, we switch from AdamW AdEMAMix after tawitth 300k iterations. We vary T3, = Ta E {0, 25k, 50k, 100k), the other prameters are fixed: B3- = 0.9999, a = 2, the learning rate inherited from Adam W and stays on his decaying course following cosine decay. AdEMAMix 's m2 is initialized to O. The larger T3, Ta. the slower the convergence is, the final loss value are very similar. This shows that the schedulers for a and 33 are not important when starting AdEMAMix from sufficiently late checkpoint.
AdEMAMix from scratch without schedulers. In this section we provide more justification for the use of schedulers and 83. We reveal that, without the use of schedulers, the norms of the updates increases significantly, even for small a values. Unstable and large weight updates are known to occur in the early phases of training when learning rate warmup is not used (Gotmare et al., 2019). Using large momentum values too early seem to increase this phenomenon. See Fig. 28.
6.0
AlaW
MIEMAMIS
Yo.r
MIEXANI Cd. 0,299
AJFMAAMS (33, Ax 6.00
AJEMANGS
Gd. 6.000
5.0
AdamWW
Ti.y Tu
Loxs
4.5
Noem
4.0
AEEIAABI 5 0.5000 Tu.dy
10a
3.5
3.0
200000
150000
25000
Iterations
Iterations
(a) Norm of 1000 consecutive updates.
(b) Loss.
Figure 28: The norm of the updates is exploding without schedulers. For historical reasons, those experiments were done with a 110M parameter model using a batch size of 32, and without weight decay (A 0). We use B2 0.999 and 31 0.9. We use 3000 steps of learning rate warmup for all the experiments. In (a), monitor the norm of 1000 consecutive updates: ||0(*)) ||2 [|0(++000)) - ||2- Even when using small values of the norm of (t()) is increasing significantly at the beginning of training. As result, in (b), we observe hiw-without the use of a scheduler on a and Bg-AdEMAMix models either diverging or fail to recover and improve upon AdamW.
38

---

Published as a conference paper at ICLR 2025
Impact of the linear decay duration when using constant 17-scheduler. In Fig. 3b show results using AdEMAMix with a linear warmup - constant - linear decay learning rate scheduler. We used 100k of linear decay. In this section we experiment with 200h steps of linear decay, the rest of the parameters are the same: we use a 1.3B parameter model with a max learning rate of 10-1 and remaining hyperparameters as in Table 4. Results can be seen in Fig. 29.
Adaw
2.7
AIEMAMii 200l anope Blnar deca?
2.3
241
2.3
2.2
5006 T Ta.05
9006
1.1M
Iterations
Figure 29: Testing different decay duration compared to Fig. 3b. Increasing the length of the learning rate decay improves the results for each methods. The gap between AdEMAMix using B3 = 0.99999 and A3 0.9999 is shrinking, but the gap between AdEMAMix models and AdamW remains.
Same figure as Fig. 1, including the Adam W trained on 197B tokens. In Fig. we represent the Adam experiment trained on 197B tokens by blue horizontal line for aesthetic reasons. In Fig. 30 we include this missing curve to the same plot.
2.8
Adamn trained (34B, 1018. 1918, 1978) tokeos
2.7
AdEMAMi trained (34B, 1018,1311) tokens
2.6
Loss
2.5
2.4
2.3
AdwaW. JOIB
256k
770k:
1M
1.5M
Iterations
Figure 30: Same as Fig. 1 adding Adam W 197B.
39

---

Published as a conference paper at ICLR 2025
C.2 VIT EXPERIMENTS
Top-k accuracies. In Fig. 6, we plot the test and train loss for the final iterates. In Fig. 31, we give more detailed view by reporting the evolution of the training loss, test loss, and top-1 accuracy. Looking at the first row, the training loss for AdEMAMix is systematically better than the AdamW baseline. The second row shows that in cases where the test loss correlates well with the train loss, AdEMAMix works well. The top-1 accuracy carries the same message.
3.5
3.5
AllanWW
Adumni
AGEMAAiix As D.90
AdEMAAkii d 0,500
AREMAMil As 0.995
3.0
AddEMAMiie A 0,5009
AdEMAMMi As 0.9000
Thin Loss
3.0
Teain Lass
Train Loss
2.5
2.5
2.0
1.5
.50000
100000
50000
100000
50000
100000
3.50
3.5
Tes Lxos
Test Loos
3.0
Tes Loos
2.75
2.5
2.50
2.25
2.0
50000
100000
50000
100000
50000
100000
45
50
80
Top-1 Accurecy
Top-1 Accurnay
Topp1 Accuracy
70
40
45-
355
40
30000
100000
50000 Ireestions
100000
50000
100000
Bosstions
Itorations
(a) 24M params, 11M images. (ImageNet-21k))
(b) 86M params, 11M images. (ImageNet-21k)
(c) 86M params, 1.3M images. (ImageNet-Ik)
Figure 31: AdEMAMix & capacity/data ratio. These figures complement Fig. 6 and provide the top-1 accuracy as well as the training and test losses for the three experimental setups used in our ViT experiments (see $ 4.3). From left to right, the ratio between data and model capacity worsen. In (a), we train a 24M parameter model on a dataset of 11M images (Imagenet-21k), doing a total of 37 epochs. In this setting the training and test loss curves very similar, signaling no blatant overfitting. We observe AdEMAMix trivially outperforms the AdamW V baseline. As the ratio between data and model capacity worsen, in (b) (more capacity) and (c) (more capacity and less data), it becomes more and more difficult to find hyperparameters outperforming the baseline. Yet, across all those settings, the training loss for AdEMAMix is always lower than that of the AdamW baseline. When decreasing training loss correlates well with decreasing test loss, AdEMAMiX can be expected work well.
40

---

Published as a conference paper at ICLR 2025
C.3 COMPARISON WITH OTHER METHODS
C.3.1 COMPARISON WITH ADMETA AND DEMA
Double Exponential Moving Average (DEMA). Originally introduced by Mulloy (1994), DEMA originally aimed to emphasize the weight of recent asset price fluctuations, making the DEMA indicator more reactive to changes compared to simple EMAs. Given notations introduced in the main paper, let EMA ( (T-N:T) & EMA(8, g (T~N), ..... .. (7), ) let N be the window size representing how many consecutive values are considered in the average, the formula for DEMA can be written as follows:
(T-N-T) - EMA(3, 3 ( g) (T-2N) 2N) ...)) = 2 EMA
(2 (T- ..... EMA
(T-N:T)
If a simple EMA tends to give significant weight to more recent observations, DEMA emphasizes this behaviour by removing some of the weight given to older observations. This is not what we suggest doing in this work, we want both high sensitivity to recent observations and non-negligible weights given to older observations.
AdMeta. Chen et al. (2023b) take inspiration over DEMA and use nested EMAs in their AdMeta-S optimizer:
(t)
3imm (t-1) 1)
(1)
mi
(t)
kggt)) pm
(t)
(AdMeta-S)
m. (1)
=82mL
(1
(1-1) (1 82) h)
((t) @(t-1) 1mm2 )
With j and parameterized by $1 [0, 1[ as such:
1
μ 25 10(31
+
B1
10
K
81
In their AdMeta-S experiments, they use 3, 0.9, corresponding to (p, k) (4.88, 2.11). B2 takes values ranging from 0.1 to 0.4. As a results, unlike AdEMAMix, AdMeta is not leveraging very old gradients. Analysing the AdMeta algorithm, we see that it consists in two nested EMAs. We show the shape of nested EMAs Fig. 32. In sharp contrast with our approach we observe that (i) it reduces the weights given to recent gradients and (ii) it gives small weight to old gradients.
Mianed Biike A.
Neatei Kidbi Doi
0.08
10-2
00-0
0.01
Wegght gaves
10:7
Weats
0.02
600
30
100
150
200
100
130
200
Brrnddepe 4
cliwetape
(a) Linear scale.
(b) Logarithmic scale.
Figure 32: Comparison between nested EMAS and a linear combination of EMAs. In (a) we observe how two nested EMAs actually decrease the influence of recent timesteps. This is the justification for the DEMA method, which aims to increase the sensitivity to recent gradients by subtracting nested EMAs. In (b), using log-scale for the y-axis, we see how, for the values used in AdMeta, the older gradients receive negligible weight. In our work, we show how very old gradients can be leveraged to get better results by keeping high sensitivity to recent gradients while giving non-negligible weights to older gradients.
41

---

Published as a conference paper at ICLR 2025
C.3.2 COMPARISON WITH LION
The Lion optimizer. Chen et al. (2023a) derived the following optimizer. We change notations to facilitate the comparison with AdEMAMix:
(0() 0 17 (sign(am (t 1t + (1 - (1) + 10( (1) (1 + 3)
(Lion)
The Lion optimizer uses a sign function and updates its EMA after updating the parameters. Moreover, it does not normalize the updates. While Lion and AdEMAMix are quite different from each others, we can draw one similarity. Indeed, the interpolation (m + (1 - ax (g)g()) is similar to combining (t-1)
two EMAs, one of them using 5 We also explore the possibility in 3, we show in App. C.1.8 and App. C.1.5 (Fig. 16b) that this often works despite a more unstable training. Interestingly. Chen al. (2023a) find that larger 3 0.99 values work best. Beside this similarity, the two optimizers behave very differently, the biggest difference being the use of the sign function, which Chen al. (2023a) claim can help regularize the training. We test the Lion optimizer on language modeling. To tune the hyperparameters, we took values from Chen et al. (2023a) as starting point, as well as recipes provided by the Optax Jax library. A summary of our hyperparameter tuning in Table 8:
setting AdEMAMix's $1 to
Table 8: Hyperparameter tuning for Lion 110M parameter models. When multiple values given, we bold the parameters we found to give the best results.
Hyperparameter
Value
Learning rate 7)
0.00005, 0. 0.0001, 1000000,000000 0.0002, 0.00004 0.0006, 0.0008, 0.001 3000
Number of warmup steps
Sequence length
1024
Weight decay À
0.01, 0.125, 0.1 0. 25, 0.5, 1.0
Learning rate decay scheduler
cosine-decay 64
Batch size
Gradient clipping
1.0, 0.5
Lion & Lion ,3
0.5, 0.7, 0.9 0.99, 0.9999
The training curve associated to the best hyperparameters is in Fig. 33. We observe that Lion is not outperforming our carefully tuned AdamW baseline. Moreover, no attempt to increase 3 beyond 0.99 successful, as those models mostly diverged. This emphasizes of our main contributions: the introduction of schedulers enabling the use of large momentum values.
AdamW
3.6
3.4
Lo08
3.2
3.0
0
50000
100000 150000
200000
250000
Iserations
Figure 33: Comparing AdamW with the best Lion model found. Hyperparameters used are in Table 8. We train many Lion models and select the one with the lowest final validation loss. Lion is outperforming our carefully tuned AdamW baseline.
42

---

Published as a conference paper at ICLR 2025
C.3.3 ADDING THIRD MOMENTUM TERM (~AGGMo)
Lucas al. (2019, AggMo) propose add an arbitrary number (K) of momentum terms to the gradient descent algorithm:
m (t)
= 3
(t-1) (t-1( + gg! for lissss K
(AggMo)
O(0) = -(it-l)- $ Eisisk m; (0)
In order to see whether more momentum terms can be beneficial, we modify our AdEMAMix optimizer to include a third momentum mgg. We name the resulting optimizer Ad3EMAMix:
m (t) }1m (t-1) + (1 (1)
(t)
m
mi
1-85 8
m. (1) Bym, (t-1) + (1 (g)g(0)
m (1) Bimg (t-1) (1 (1)
(Ad3EMAMix)
D(t) = (1 +1 82) (1)
D(1)
(( I-p5
((() ((--1) 11( m (mm 50
+ 10(
Ve
We Train a 110M parameter model with same hyperparameters as Table 2, but we use a instead of 8. We apply the same scheduler 84 as for $3. In Fig. 34 show the resulting validation loss curves for various (51, B3. 84) triplets. Our experiments show no advantage of adding an extra 31 EMM. Rather than carrying the message that should use more momentum terms, work shows how-in in order to use a large momentum a term should be added to stay sensitive to local fluctuations of the loss landscape. In that sense, believe terms should be enough.
Adaun
3.3
AdMEMAMIX 3.. 9. 0.9. 0.990,9999, 0,99999
AdREMAMIx 35. 0.9. 0.999, 0.99999
3.2
AdDEMAMIX 3..
0.999,,9.99990
AdREMAMix 35. 0.9. 0.99, 0,9999
Loss
3.1
AdSEMAMIx 3.. 0.99. 0.98999 0.990.99...999999
3.0
2.9
2.8
128k
256k
Iterations
Figure 34: Adding an extra momentum to AdEMAMix does not provide better results. We
train multiple 110M parameter Ad3EMAMix models with various (31, B3, 34) triplets (lighter curves). None of those models end up outperforming AdEMAMix, indicating that the additional my momentum term is scemingly not bringing any advantage.
43


Reviews:
[{"review_id": "n0FXVCyfdP", "reviewer": "Reviewer_zVt6", "summary": "The paper propose a new optimizer called AdEMAMix, which add an additional EMA sequence to Adam. The intuition is to to keep a high sensitivity to recent gradients (using m1), while also incorporating information from older gradients (using m2).", "questions": "1. The motivation on the Rosenbrock function is rather weak.  Why do you consider the Rosenbrock function? How does it relate to LLM training?  Why do you think the implication on the Rosenbrock function can be transferred to LLM?  To me, there is a huge gap between these two.\n\n2. Need more evidence that \"beta1 = 0.9 is indeed suboptimal for LLMs\". While I fully understand Figure 3, i still did not fully convinced that \"LLM training really needs to use more historical gradients, in one way or another\". More LLM-related evidence is needed to convince the that it is indeed a major bottleneck. Despite the impressive performance, the overall motivation still seems quite vague to me. I suggest the authors put more effort into designing more experiments for better motivation.\n\n3. It is difficult to understand the highlighted sentence in the introduction: \"While changing the direction of the slow momentum is difficult, any adjustment orthogonal to that direction is easy—which favors fast progress in sinuous canyon-like landscapes.\" I don't understand what you mean by \"change the direction of the slow momentum\" or \"any adjustment orthogonal to that direction\". Please explain more.  Also, the authors mentioned \"canyon-like landscapes (in Rosenbork function).\" but never connected it to LLM in the script. This also makes the motivation rather weak.\n\n4. I don't quite understand how to read Figure 2. For instance: \n\n   -- in Figure 2 (a), is AdaEMAMix considered better than Adam? It is not clear whether we can draw such conclusion based on the figure. In my opinion, at least 10 x more iterations are needed to draw valid conclusions.\n\n   -- in Figure 2 (c), did AdaEMAMix converge to the optimal solution? \n\n5. The proposed method boosts performance by using extra memory (to store an additional copy of momentum). Though many readers might regard it as a drawback,  I personally think such a trade-off is acceptable as long as the performance gain is significant.  Further, I think AdaEMAMix can be combined with some orthogonal methods to reduce memory. For instance, AdaEMAMix can be combined with the recent method Adam-mini [1] to reduce the memory for V. I suggest the authors try it out.\n\n\n\n6. Some missing related works: [2] proves that vanilla Adam can converge under a wide range of beta1  = 0, 0.5, 0.9, 0.99, etc., as opposed to the divergence result in Reddi et al. 2018. This result lays down a preliminary foundation of this work.  Without the theoretical guarantee, it would be dangerous to play with beta1 of Adam.\n\n[1] Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Ye, Y., ... & Sun, R. (2024). Adam-mini: Use fewer learning rates to gain more. *arXiv preprint arXiv:2406.16793*.\n\n[2] Zhang, Y., Chen, C., Shi, N., Sun, R., & Luo, Z. Q. (2022). Adam can converge without any modification on update rules. *Advances in neural information processing systems*, *35*, 28386-28399.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The prerformance of AdaEMAMix is quite impressive. The writing is excellent. The ablation studies are thorough.", "weaknesses": "The motivation is rather vague. See 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": "HZGoTxkYT6", "reviewer": "Reviewer_H5Kk", "summary": "This paper proposes AdEMAMix, a new optimizer which outperforms Adam on language model training and ViT training. Their empirical results show large benefits (~50% reduction) over Adam in the regime of noisy gradients i.e. small batch size or longer runs. The main idea behind the optimizer is to maintain two momentum terms and combine them to get the final movement direction. The coefficient of momentum and their combination are also dynamically adapted.", "questions": "The authors state “While no answer to those questions is given in this work, we\nprovide a toy justification which indicates that large momentums can have a positive impact in\nnoise-free non-convex settings (see Fig. 2)—indicating the improvement of our approach is at least partially explainable without considering variance-reduction effects.” Is there empirical support for this in the LLM experiments? looking at Figure 17 the benefit seems to drop with increasing batch size. Note that the maximum batch size used here (512k) is smaller than that used for LLMs like Llama (4m), though I agree that the model size is also small here. Could the authors provide an experiment with 2m batch size to see the trend?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "These are strong results on a very important problem. They also provide many optimizer ablations in the Appendix showing the robustness of their proposed optimizer.", "weaknesses": "Since many of the experiments are with small batch size it would have been interesting to explore the effect of weight averaging. For example, is it the case that weight averaging helps AdamW and AdEMAMix equally? Or not?", "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": "NEAO3UGAgL", "reviewer": "Reviewer_GCtm", "summary": "This paper proposes to use an additional momentum buffer for the Adam optimizer. The motivation behind is that a single momentum buffer may not be able to utilize both past and current gradient information efficiently. To this end, one buffer with large momentum parameter $\\beta$ is for slow-changing directions, and the other one with small $\\beta$ is for fast adaptation to current gradient. Then, the two momentum buffers are mixed according to some fixed ratio ($\\alpha$), and the rest of updates (i.e., pre-conditioning and weight-decay) follow those of AdamW. However, a large $\\beta$ (for slow-changing momentum buffer) can cause training instabilities in the initial stage. To this end, the paper proposes to gradually increase this parameter until the target value is reached.  Then, experiments are performed on various language modelling tasks to show that the proposed algorithm is faster than AdamW given a fixed computation budget.", "questions": "Overall, I think the idea introduced in this paper is interesting. The paper can be improved if the above  weaknesses can be addressed.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written and easy to follow. The experiments on the Rosenbrock function are convincing. There are some similar (loss landscape) models proposed recently to analyze learning rate schemes [1]. It maybe interesting to draw some theory/experimental connections in the case of momentum.   \n\n- The experiments seem to be comprehensive covering different settings and tasks. The improvement over baseline is shown.\n\n[1] Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective.", "weaknesses": "- There are some additional hyperparameters introduced. Noticeably, it seems that $\\alpha$ (that controls the mixing ratio)  is important for the algorithm. It would be important to study the sensitivity of the algorithm to this hyperparameter, given that it essentially controls the contribution of each momentum buffer to the current update. \n\n- It would be better if some convergence guarantees of the algorithm can be provided even in the convex setting. For example, what is the relationship between the two momentum parameters that would guarantee convergence?", "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": "lw3RysqbTt", "reviewer": "Reviewer_98Nr", "summary": "This paper introduces AdEMAMix, an optimizer modifying Adam by incorporating a mixture of two Exponential Moving Averages (EMAs) to better leverage past gradients. Traditional Adam or AdamW optimizers rely on a single EMA to smooth gradients, which prioritizes recent gradients over older ones. However, the authors argue that this approach is suboptimal because it cannot simultaneously emphasize both recent and very old gradients. They conduct experiments across language (LLMs and Mamba) and vision (ViTs), showing that AdEMAMix outperforms and tends to forget training data slower than Adam.", "questions": "1. Doesn’t the introduction of the additional beta_3 term potentially lead to an exploding step size if the incoming gradients are extremely small for many steps (the additional term in the numerator causes it to decay slower than the denominator in the update)? This isn’t unforeseeable in language model training, where the model could encounter multiple documents of rare tokens consecutively.\n2. Do the authors have any experiments adding an EMA to other optimizers like Adafactor or Signum? Does the additional EMA always provide a gain, even with factored gradients?\n3. Related to the previous question to address the memory overhead, does storing the slow-moving EMA in lower precision affect AdEMAMix performance?\n4. Why was the optimal Adam learning rate being too high for AdEMAMix only manifesting at higher scales? Does this perhaps suggest AdEMAMix being more difficult to tune at larger scales? I’m wondering if there was there further investigation into this eg. Was there a certain part of the network that was destabilizing?\n5. Did the authors use certain stabilization strategies, eg. QK-LayerNorm in your experiments? I’m wondering if this could be used for the author’s benefit, to stabilize training and offer greater ease in hyperparameter tuning. \n6. Why was alpha set to 2 and 1 for the switching Adam -> AdEMAMix experiments whereas the optimal values found in hyperparameter tuning were much higher?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is written and organized well. The optimizer is a simple modification to Adam and is conceptually easy to understand, with the algorithmic differences presented clearly.\n- The authors extensively benchmark their proposed optimizer against Adam across architectures, dataset domains and hyperparameters.\n- The spirit of AdEMAMix (having a ‘fast’ and ‘slow’ tracking of gradients) could be useful for mitigating forgetting in continual learning regimes, and is also shown to be relevant for current pretraining regimes where we often train much longer past eg. Chinchilla-optimal tokens.\n- The authors have experiments which address one’s immediate concerns of the proposed optimizer, like hyperparameter tuning and memory overhead.", "weaknesses": "- The proposed optimizer is not that novel, with several existing approaches in maintaining a longer horizon at the cost of memory which have been mentioned by authors; in particular, AggMo is essentially the same as AdEMAMix but applied to gradient descent, and the only difference in the setting of this work is applying the same principle to Adam.\n- Although the authors have performed experiments on the hyperparameter stability of AdEMAMix, they do introduce two new parameters beta_3 and alpha which require the use of schedulers for larger values to avoid divergence in early iterations, which may require more extensive tuning depending on the setting (eg. fast domain shifts).\n- There is an additional memory overhead in the order of the model size to incorporate the second EMA, which the authors propose can be mitigated by setting beta_1 to be 0; however, these runs exhibit instabilities and sometimes large spikes in training especially at larger batch sizes (Figure 22). These spikes are generally undesirable even though the loss seems to ‘recover’, it is unclear to me that such spikes wouldn’t have more detrimental effects at larger scales.", "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": "uHi5Cjthff", "reviewer": "Reviewer_A3u9", "summary": "This work proposes a new optimization algorithm named AdEMAMix, a variant of Adam(W). They add an aditional moving average of gradients to the algorithm, on top of the existing first and second moment moment averages. In contrast to the traditional  $\\beta_1 = 0.9 $ moving average typically used in Adam, the new moving average that updates much slower ( $\\beta_3 = 0.9999 $). This additional moving average is then used with the traditional moving average to update the parameters, in a similar fashion to the Adam update. The authors argue this assists in preventing the “forgetting” of data and better uses past information than the traditional Adam update, and show strong empiracle performance on some vision and language settings.", "questions": "- In figure 2, authors claim the proposed algorithm does not exhibit oscilations when compared to Adam. However, in both 2 (a) and 2 (c) it appears the AdEMAMix does in fact overshoot the minimizer, further than Adam at that. It looks to me that Adam with “default” hypterparamters appears to be the most stable of all the options. In general this is a little concerning as modern loss surfaces will be much more complex than the Rosenbrock function and these large oscilations/overshooting minimizers may be difficult to deal with or even diagnose. A seperate point on this figure but potentially contributing to my confusion is that the colour scheme in the three subplots is not consistent, if the authors could use a shared legend or something along those lines the results may be more interpretable.\n\n- Given this method requires an additional buffer for the second EMA, this authors discuss it potentially being slower and mention that it is also possible to set  $\\beta_1 $ to zero to compensate for this. I’m curious if any of the main experiments or the wall time comparison are using this  $\\beta_1= 0 $. The hyperparameter configuration for figure 1 uses  $\\beta_1 =0.9 $ so I would just  want ensure the wall time comparison is too to ensure figure 1 doesnt change substantially if the x axis is changed from steps to hours.\n\n- What does older gradients being outdated mean? Is there any way to quantify this? Any intuition? All the data is weighted equally as far as minimizing the loss is concerned despite being at different points in the loss surface when you evaluate it. In general this whole idea of older/newer gradients being less/more relevant seems vaugly okay, but it would be nice to see a more scientific defintion or explanation. This doesn’t need to be totally rigorus and have proof but the current explanations seem hand-wavey and leads to far more questions than answers.\n\n- Why not just use a normal non exponential moving average? You can have a uniformly weighted moving average that should not diverge, has this been tried?\n\n- Is there any intuition for why  $\\alpha $ is so large (in particlar  $\\ge1 $?) Does this lead to a smaller step size being necessary? Can you also set it to zero until some point to fight early instability?\n\n- Why is there no bias correction on  $m_2 $? I don’t really think it matters much much and most analysis ignores it but I’m curious why that decision was made in comparison to the typical Adam bias correction step.\n\n- Regarding figure 4, while it does appear that the proposed algorithm “forgets” slightly less quickly than AdamW, is the claim that this is why it works better? It appears to me that the biggest predictor of “forgetting” is the injection time  $t_B $ rather than the algorithm, although perhaps if this minor difference is compounded across batches this leads to the superior performance of the algorithm. \n\n- In many experiments, notably figure 5, the x axis is limited to not show the optimization performance in earlier parts of training. Given that instability is cited as a challenge for this algorithm, to a point where complex schedules need to be used, I’m curious how unstable the well tuned algorithm is. Are there instability issues in the 0-200k range in figure 5 (b) with AdEMAMix from 0 for example?\n\n- Recent analysis in [1] has shown Adam(W) should work well if the gradient and Hessian become correlated, and [2] provided a mechanism where that shows up in language models. Can the authors reconcile this with their modifications to the algorithm? It seems to me those works would suggest this new variant is not a great idea so I’m wondering if the authors have any thoughts for why that isn’t the case.\n\n- Any idea why there is less good performance on imagenet 1k? It seems AdamW is performing better there although all the plots are fairly close.\n\n\n[1]\nMichael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, Zhenxun Zhuang\n\nRobustness to Unbounded Smoothness of Generalized SignSGD\n\nhttps://openreview.net/forum?id=8oj_2Ypp0j\n\n[2]\nFrederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, Alberto Bietti\n\nHeavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models\n\nhttps://arxiv.org/abs/2402.19449", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The primary strength of this paper apperas to be an algorithm with strong empiracle performance. The authors conduct a thorough set of experiments and ablation in the setting of language modelling and vision using multiple architectures and configuarations. The idea is interesting and combines other ideas from the optimization literature (the Adam update, multiple momentum terms). I am particularly interested in the authors study of how the loss with respect to different batches changes over time depending on the order in which they where trained on. I think this direction is a promising area of research for understanding why some optimizers work better than others. The paper also contains many interesting questions that are of interest to the optimization community, although struggles to answer them (as does the rest of the field).", "weaknesses": "Some specific concerns are addressed in the questions section, but I think the fundamental challenge this paper faces is explaining why the algorithm is better. As is typical with deep learning optimizers, it would be too hard for any provable reason why the algorithm may be superior, but I do think some more effort in the paper could be focused on explaining why this is a good idea, especially with the increased memory requirments. As said in the strengths, I do like the amount of questions being asked but it would be nice to see if the answer to any of those questions is leading to a better optimizer. I think it’s possible that the “forgetting slower” direction may lead to an answer, but the majority of the experiments and argumentation for that is relegated to the appendix. While practically showing plots with better performance is good, from a scientific perspective understanding why that performance is better is important.\n\nI’m slightly concernerd by the number of additional hyperparameters being introduced, which appear to be somewhat sensitive given the use of schedules for most of them. This one can make the algorithm more difficult to use in practice, and two runs the risk of having some lucky configurations for a given problem that works great but in general may not be great. I also feel that the additional memory requirment is a non trivial drawback. While it does not increase communication costs, just running in FSDP is not an easy option for many people who do not have easy access to multi-gpu servers.\n\nA more cosmetic note, the figures in this work can lead to confusion at times. In multiple of the figures, there are three subplots where some and perhaps all are unrelated. In some cases, Some subplots are meant to be related, but use different colours for the same thing. This is also likely the reason many of the figure captions are  $\\approx $ half a page. I think it would be valuable for the authors to give the plots another pass before any camera ready version. \n\nOverall despite being a potentially effective algorithm, I do feel that this paper is somewhat lacking in the motivation and justification of that. I would be open to raising my score in the discussion period but in it’s current state I’m not sure this work is ready. I would be interested in why the authors think this is a good idea and if there are better ways to show that to be the case. This does not need to be large experiments or rigorous theory. Well designed small experiments or simple derivations can be just as effective give an intuition to explain the performance on minimizing the loss or solving some problem existing optimizers have.\n\nMore specific questions can be found 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": {}}]

Author rebuttal:

Decision: Accept (Poster)