The AdEMAMix Optimizer: Better, Faster, Older

We thank the reviewer for the time taken to review our work. On the intuition behind our method. We provide an intuition in our Rosenbrock experiments. In those experiments, the fast EMA (m1) helps the algorithm stay in the valley, and the slow EMA (m2) helps the algorithm crawl along the valley, remembering the valley's direction and accelerating. As pointed out by reviewer GCtm, a recent line of work [1] suggests a “river-valley” structure of the loss landscape of LLMs, which does share some similarities with the Rosenbrock function. We provide another intuition through studying the forgetting of the training data during training (Fig.4). This shows how the larger momentum can help the algorithm to be more data-efficient.

Now answering the reviewer’s questions:

1. In Fig.2.a, we indeed observe the distance to the solution going down and then up again. This is due to the momentum values causing the iterates to overshoot the solution. When we mention the lack of oscillations for AdEMAMix, we refer to the trajectory not oscillating around the bottom of the valley. The inability of Adam to follow a similar trajectory when using a large $\beta_1$ is what we want to focus on. It shows that—for larger momentum values—the iterate fails to respond to local variations of the loss landscape. This motivates us to also include a momentum term with a small $\beta$ as in AdEMAMix. The oscillations of momentum methods around the solution on deterministic objectives are well documented, see e.g. https://francisbach.com/continuized-acceleration/, https://distill.pub/2017/momentum/, or [2, fig5]. In practice, to the best of our knowledge, oscillations due to momentum have never been an issue in deep learning training (i.e. the phenomenon of overshooting the solution seems to be only observed in toy settings, see Fig.20 and Fig.21 to see the lack of oscillations even with very large $\beta_1$ in a real world setting using Adam). This is why we focus on the left part of the curves, corresponding to the first time required to arrive in the vicinity of the solution.
2. The times in Fig.5.a use $\beta_1=0.9$.
3. Given a past timestep $t$, a sample $x$, the gradient $\nabla_\theta \ell(x, \theta_{t})$ is outdated at time $t+K$ if its inner product with $\nabla_\theta \ell(x, \theta_{t+K})$ is non-positive. Such a gradient would be outdated as it pushes the iterate in a direction that is no longer relevant to decrease the loss on $x$. A simple illustration could be overshooting the minimum of a quadratic function. Gradients are local approximations, we mention outdated gradients as it is surprising for us that one can take advantage of very old gradients. Our intuition was that those would no longer be relevant given the local loss landscape. We believe that the fact this raises many questions is a good thing.
4. In a limited setting, we also tried with uniform averages, we found those to also improve over the baseline but much less. We therefore decided to focus on EMAs.
5. We do use a linear scheduler on alpha. Starting from $0$, alpha grows up to its final value over the same number of iterations as $\beta_3$. Setting it to $0$ during the early iterations can work, this corresponds to our experiments on switching from AdamW to AdEMAMix (Fig.5.b and Fig.5.c). Depending on when the switch occurs, the final loss—while better than the baseline—might not be as good as training AdEMAMix from scratch. In Fig.5.b and Fig.5.c, we do not use any schedulers on beta3 nor alpha.
6. The bias correction is there to compensate for the initialization of the buffers to $0$. In our case, it is desirable to not do the correction as it implies the values of m2 are initially very small, and increase little by little. Intuitively, this means that, for the early steps, AdEMAMix behaves similarly to AdamW. Moreover, when switching from AdamW to AdEMAMix, the effective learning rate increases, which we believe explains the small increase in loss that can be observed right after the switch in Fig.5.b and Fig.5.c. Not doing any bias correction might help to smooth the transition.
7. We were curious on what was explaining the correlation between forgetting and iterations and ran additional experiments in App.C.1.3 (see Fig.11). In that experiment we use a WSD learning rate scheduler instead of a cosine decay. This allowed us to show that the forgetting is tied not only to the optimizer, but also to the learning rate. Setting this dependency aside, it is clear from Fig.4 that AdEMAMix forgets training batches slower than AdamW.

I thank authors for their responses. In regards to their question about why the work I cited may not agree with their approach, I preface  by saying this is my  inutition and not rigorous. In [1], it is shown that Adam will benefit when the norm of the gradients and the trace of the Hessians become correlated. Then, [2] shows a situation where this occurs naturally where Adam  is superior to SGD. My intuition is that, since now we have two EMAs (with the slower changing one being more dominant due to $\alpha\ge1$) for the gradients but only one for the squared  gradient, this correlation will be weakened. In future work, it may be a good experiment to run in order to figue out if this algorithm is effective for similar  reasons to Adam. If the authors  have any thoughts on this I would be interested.

Overall I do still feel that there could more investigation of why this algorithm works better even with the response probvided by the authors. I will  raise my score to a weak accept based on the authors responses and hopw that that question will be addressed in future work. I would again like to point out that some of the figures a bit confusing as is stated  in my original review,  so if this work ends up getting accepted I reccommend the authors take  a  second  pass on  them.

We thank the reviewer for the time taken to review our work. First, responding to the weaknesses identified by the reviewer:

1. On the novelty of our approach compared to AggMo. While the AggMo method introduced by Lucas et al. also adds additional EMAs with larger beta values, we strongly disagree that it implies our work lacks novelty. As discussed in the related work, the claims from Lucas et al. are that additional momentum terms speed up convergence and hyperparameter stability. In our work we show our optimizer is not only converging faster, but reaches better solutions (lower loss). Moreover, Lucas et al., relies on small-scale settings where SGD works well, using MNIST, CIFAR-10 and CIFAR-100, and training LSTM LMs on the Penn Treebank dataset. In contrast, we focus on significantly larger real-world settings. Those larger settings heavily favor Adam over SGD. Getting our optimizer to work in those settings was not as easy as adding an additional momentum buffer, it required tackling training instabilities through deriving schedulers, without which the method would not work. Finally, AggMo is suggesting using many (e.g. 4) momentum terms. We show in App.C.3.3 that using more than two momentum terms in AdEMAMix is not providing any benefit.
2. On using AdEMAMix in fast domain-shift settings. This is a valid concern for that setting. We did not study non-stationary training distributions and leave that question for future work. Concerning schedulers, we observed that our design-choice of always setting $T_{\alpha,\beta_3}=T$ yielded stable training runs without additional validation, even when switching across model size, to state space models and vision models.
3. We will change the text to say that using beta_1=0 needs to be confirmed at larger scale. We want to emphasize that, in many cases, the memory overhead is not a critical issue. Especially, in the case of distributed training, the memory overhead can be mitigated by sharding the optimizer state across devices.

We thank the reviewer for the positive appreciation of our work.

In Fig.17 we keep the total number of training tokens constant, and vary the number of steps (in {$32k, 64k, 128k, 256k$}). To keep the number of training tokens constant, this forces us to increase the batch size as we decrease the number of steps. Looking at Fig.17.a and Fig.17.b, we observe that both methods suffer when we trade a large number of iterations for a larger batch size and fewer steps. AdEMAMix—while still outperforming Adam—is more affected. To understand this phenomenon, we notice that when we do fewer steps, we have fewer gradients to accumulate in $m_2$. Given our very large $\beta_3$ values (e.g. $0.9999$), this can become a problem. We show in Fig.18 that the problem is mitigated by reducing the $\beta_3$ value to e.g. $0.999$. Interestingly, comparing AdamW in Fig.17.a and AdEMAMix with $\beta_3=0.999$ in Fig.18.b, we can see that AdEMAMix is now less affected than AdamW when we trade a large number of iterations for a larger batch size and fewer steps (the increase in final loss is smaller for AdEMAMix when increasing the batch size). In general, from our experiments on both images and text, we did not notice any disadvantage of ADEMAMix over AdamW when increasing the batch size. Our largest experiments, in App.C.1.2, train 3B parameter models using a batch size of $1024^2$ tokens. We still observe improvements.

We thank again the reviewer for appreciating our work and stay at their disposal to answer any further questions.

We thank the reviewer for the time taken to review our work.

1. App.C.1.5 studies extensively the sensitivity to hyperparameters. The sensitivity to the mixing ratio alpha is shown in Fig.13.a. It is shown that the range of alpha values outperforming the Adam baseline is very wide.
2. While a theoretical proof of convergence of our method would be welcome, deriving such proof poses significant challenges. We can take as an example Adam, which convergence has been challenged by Reddi et al. [1], which still failed to explain convergence for hyperparameter values used in practice as discussed in [2]. Despite the enduring gap between theory and practice, Adam remained the workhorse of deep learning optimization. Understanding it from a theoretical standpoint is a line of research on its own. As such, we believe providing convergence guarantees—albeit desirable—lies outside the scope of our work. This being said, convergence bounds in convex settings for a simpler method (AggMo) combining GD with a linear combination of EMAs have been shown in [3].

We hope we have addressed the reviewer’s concerns. Sensitivity to hyperparameters is studied in App.C.1.5, and providing theoretical backing seems tedious given a theoretical understanding of even Adam is still an active field of research. We hope the reviewer will consider raising their score or provide further details justifying their rejection of our work.

[1] Reddi et al., 2019: On the Convergence of Adam and Beyond [2] Zhang et al., 2022: Adam can converge without any modification on update rules. [3] Lucas et al., 2019: Aggregated momentum: Stability through passive damping

We thank the reviewer for taking the time to review our work.

1. We agree that the gap between Rosenbrock and LLM training is important. We do not claim that this toy setting captures the complexity of LLM training dynamics. Instead, it serves as an illustration, showing that—quite remarkably—our method can be motivated even in a noise-free toy setting. Moreover, while understanding the training dynamics of LLMs is still an active line of research, a recent line of work [1] (also pointed out by reviewer GCtm) suggests a “river-valley” structure of the loss landscape of LLMs, which does share some similarities with the Rosenbrock function. We will cite this work in our next revision.
2. We are not exactly sure where our "effort into designing more experiments" should concentrate. Could the reviewer specify the type of experiments that would be needed? We provide two key intuitions that help understand the reason behind the superiority of our approach. First, while limited, our Rosenbrock example shows that using an EMA with a large beta can be beneficial yet needs adjustment of the optimizer (it doesn’t work to simply increase beta1 in Adam). Saying that “$\beta_1=0.9$ is suboptimal for LLMs” is ambiguous, and this is not what we aim to convey. It seems $\beta_1=0.9$ is optimal for Adam (see Fig.20). Using Adam with $\beta_1=0.9999$ does not work, which makes sense as $m_1$ is no longer responsive to local variations of loss landscape, therefore Adam with $\beta_1=0.9999$ fails to optimize the underlying function (Fig.3.a). Our core message is that this can be solved by combining it with a term that stays responsive to the local loss landscape (e.g. an EMA with a small beta). Then, it is possible to gain from using much older gradients. In the Appendix, section C.1.7 is entirely devoted to limitations of using a single EMA in Adam. In that section, we show increasing $\beta_1$ in Adam does not work, even if we bypass the initial training instabilities and start from a pretrained AdamW checkpoint (Fig.21), even when adding schedulers on $\beta_1$ as for AdEMAMix (Fig.20.b). Even when increasing $\beta_2$ from $0.999$ to $0.9999$ to stabilize training (again Fig.20.b). The second intuition we provide relates to our analysis of forgetting (Fig.4). We show that larger beta values as in AdEMAMix can be more data-efficient, improving the final loss on training samples when compared to Adam. While we concede that our understanding of the phenomenon and motivations are not exhaustive, we nonetheless prove empirically that more historical gradients can be used efficiently, while providing several possible justifications.
3. The larger the beta, the smaller the contribution of each gradient, and therefore many gradients are needed to change the direction of $m_2$. In contrast, changing $m_1$ only requires a few iterations. Therefore, updating the weights in a direction going against $m_2$ requires pushing against $m_2$ for many iterations. In contrast, updating the weights in a direction orthogonal to $m_2$ is easy. In essence, what we are trying to convey is: the fast EMA ($m_1$) helps the algorithm stay in the valley, and the slow EMA ($m_2$) helps the algorithm crawl along the valley, remembering the valley's direction and accelerating. We will clarify in our next revision.
4. Fig.2 aims to show that AdEMAMix can use larger momentum values to reach good solutions faster, with less oscillations. In Fig.2.a, we observe the distance to the solution going down and then up again. This is due to the momentum causing the iterates to overshoot the solution. The oscillations of momentum methods on deterministic objectives are well documented, see e.g. https://francisbach.com/continuized-acceleration/, https://distill.pub/2017/momentum/, or [2, Fig.5]. In practice, to the best of our knowledge, oscillations due to momentum have never been an issue in deep learning training (i.e. the phenomenon of overshooting the solution seems to be only observed in toy settings, see Fig.20 and Fig.21 to see the lack of oscillations even with very large $\beta_1$ in a real world setting). This is why we focus on the left part of the curves, corresponding to the first time required to arrive in the vicinity of the solution. Ultimately, all the methods tested converge to the solution given enough iterations, we only claim that AdEMAMix can find good solutions relatively fast.

Thanks for the rebuttal and I sincerely apologize for the late reply. I have carefully read the rebuttal. I still think this is a professionally-written paper with strong empirical evidence, yet the motivation is still a bit weak. I vote for acceptance and I will keep my score.

Here are some follow-up comments. I think they would be helpful to improve the paper quality.

1. Regarding my Q4: I am clearly aware of "oscillations of momentum methods". This is not my comment. My comment is "it is not clear whether we can draw such a conclusion based on the figure. In my opinion, at least 10 x more iterations are needed to draw valid conclusions.". Anyhow, the authors did not provide new experiments.

2. Regarding paper presentation. Some notations are unclear and are not consistent. For instance, in line 249 "This allows for the use of much larger beta values e.g. 0.9999" what is beta here? Is it beta1 or beta3?

3. Regarding motivation. I still think it is too weak to motivate using Rosenbrock function. Here are my suggestions: please try linear model + 2-class cross-entropy loss classification, which also have "river-valley" landscape. Show that AdEMAmix has an advantage on this task (perhaps a bit more theory would be cool), then generalize to 1-hidden-layer-NN + cross-entropy or 1-layer-Transformer + cross-entropy. These experiments and discussions will provide much stronger insight and motivation. At least much better than Rosenbrock.

4. Regarding your rebuttal. In rebuttal, the authors mentioned that "we propose a simpler solution which consists in setting beta1=0. ". This makes me quite confused. When beta1= 0, how is AdEMAmix different from AdamW with beta1 = 0.9999? They seem to be the same up to a constant alpha, right? If so, then does it mean that AdamW (beta1 = 0.9999) works better than AdamW (beta1 = 0.9)? However, the authors also claim that AdamW (beta1 = 0.9999) does not work well. So I am confused what is going on here.

   Further, if AdEMAmix (beta1 = 0) works as well as beta1 >0, then why bother introducing the idea of "balancing fast and slow changing signals"? It seems that we do not need to balance anything at all. This further makes me confused about the motivation of the paper.

5. Missing important discussions. It is good to have the discussion on line 407 "Why not simply increase AdamW's beta1?" But this discussion is rather numerically guided. I suggest adding a rigorous math clarification that "increasing beta1" is NOT equivalent to "linear combination over an additional EMA copy". (I assume they are indeed different, right? I didn't have time to carefully check)