In Search of Adam’s Secret Sauce

We thank the reviewer for their time, their insightful comments, and for finding our work comprehensive, practical, and novel!

We answer your questions below:

1) Other domains

The reviewer is absolutely correct: this question is of extreme interest. We limited ourselves to the LM setting, as this discussion already encompasses multiple scenarios (batch sizes, regularization, data, and scale), is compute-intensive and often used to showcase the Adam-SGD gap. We, however, note that a concurrent paper, "Matias D. Cattaneo and Boris Shigida. Tuning Adam(w): Default beta2 may be too large, 2025", which we cite in line 332, provides ablations on small-scale vision domains confirming the default $\beta_2=0.999$ in Adam is too high.

2) Variational Interpretation for $\beta_1\ne\beta_2$

Our variational interpretation stems directly from our empirical finding that $\beta_1=\beta_2$ often leads to optimal results (e.g., Figure 2 and Figure 4).

1. Let us elaborate better what we mean: while for a fixed arbitrary $\beta_1$ it might indeed be that $\beta_2>\beta_1$ leads to slightly better results (maybe due to noise, e.g. Figure 2, $\beta_1=0.9875$), at the optimal $\beta_1$ we have both in Figure 2 and Figure 4 that $\beta_2=\beta_1$ leads to best perplexity. This implies that if one tunes only $\beta_1$ and sets $\beta_2=\beta_1$, then the optimal result can be achieved.
2. Our point is that Adam with $\beta_1=\beta_2$ is not only a good method but also a statistically principled method. Our novelty in this regard is to demonstrate that only in this setting does Adam admit a strong statistical viewpoint, as we show in Appendix C.2. We will reference this appendix and provide further commentary on it in our manuscript. Thanks for the question!
3. To conclude, $\beta_1=\beta_2$ leads to a simplified and statistically principled method, which performs well (= much better than Signum) and often optimally in practice. This is not only interesting but also gives clear guidelines for tuning.

3) Bias Correction

We thank the reviewer for this careful question. Our analysis doesn't include bias correction because of the way we initialize the momentum buffers in Eq (1). Bias correction was introduced because the momentum (or EMA) estimate is biased for stationary signals when you initialize momentum at zero. We take care of this issue by instead initializing our momentum estimates at the first sample. This already guarantees an unbiased estimate of the momentum estimate, see Section 2.3 here for a discussion.

In any case, we also show in Section B.5. (referenced at line 221) that bias correction does not affect qualitatively our results. We agree that some works devote much focus on this, but also point out that bias correction can be seen as simply modifying learning rate scheduling (e.g. see equation 5 in https://arxiv.org/pdf/2003.02395), and our statistical interpretation is independent of scheduling.

4) Other methods, e.g., Shampoo, Lion

In this paper, we address the question of understanding Adam with simpler methods that claim theoretical similarity with this optimizer. Both Shampoo and Lion are designed to surpass Adam, not to simplify it or to explain its update rule. For this reason, we believe comparing our methods with these would distract from the message in our paper, which is not devoted to finding the best method, but rather to understanding the Adam optimizer and how to simplify its tuning process.

[1] MoMo: Momentum Models for Adaptive Learning Rates, Fabian Schaipp, Ruben Ohana, Michael Eickenberg, Aaron Defazio, Robert M. Gower, ICLR 2024

Thank you for the time and effort you have put into reading our paper.

Below we answer all of your questions, including your main question on verifying the proof of Theorem 4.1 (which resulted from a small type-O). See below. We hope, given our response, that the reviewer will reconsider their position. We thank the reviewer for this process, which helps us improve clarity in our work.

Q1: “Reproduce the given objective function when substituting Eqs. 5 and 6 into Eq. 4”

A1: We very much appreciate the reviewer going through our proofs. The reason you could not re-deduce the objective is because of a type-O in our appendix: The appendix uses $1/\lambda$ instead of $\lambda$ multiplying the KL divergence. That is, in the appendix we have considered the objective $- \log p(g_{k+1} \mid m, \sigma^2) + \tfrac{1}{\lambda }\mathrm{KL}\left(\mathcal{N}(m_k, \sigma_k^2)\,\|\,\mathcal{N}(m, \sigma^2)\right)$

This is the only reason your deduced objective does not match the one in the objective in the proof. We apologize for this type-O. In our revision, we will fix this, and take your advice and give all the explicit steps of the proof, starting by stating “Substituting (5) and (6) into (4) gives the objective …”' This will result in a clearer proof.

Q2: "I do not understand how the authors went from $\beta \delta\_k + \beta(1-\beta)(m\_k-g\_{k+1})^2$ to \beta EMA\{\beta}(m\_k - g\_{k+1})^2?.. provide derivation, … may increase my score."

A2: Certainly, this derivation relies on applying the recursive definition of EMA given in Eq (1). The short answer, is that the equivalence holds because the EMA as an operator is linear, and consequently $\beta EMA_{\beta}[(m_{k-1}-g_k)^2] =EMA_{\beta}[\beta(m_{k-1}-g_k)^2].$ In case it helps, we also give a more detailed proof below using induction.

Our objective is to prove that, given $\delta_k$ defined recursively via:

$\delta\_{k+1} = \beta \delta\_k + \beta(1-\beta)(m\_k-g\_{k+1})^2$

it follows that $\delta_k = \beta EMA_{\beta}[(m_{k-1}-g_k)^2]$, where by definition the EMA is also defined recursively via

$EMA_{\beta}[(m_{k}-g_{k+1})^2] = \beta EMA_{\beta}[(m_{k-1}-g_k)^2] +(1-\beta) (m\_k-g\_{k+1})^2.$

Base case. First consider the base case $k=0,$ where we initialize $m_{-1}=m_0 = g_0,$ and thus

$\delta\_0 = 0 = \beta (m_{-1}-g_0)^2 = \beta EMA\_{\beta}[(m_{-1}-g_0)^2]$,

since we also initialize EMA at the first element of the sequence. Thus the equivalence holds for $k=0.$

Induction hypothesis. For $k>0$, assume we have that $\delta_k = \beta EMA_{\beta}[(m_{k-1}-g_k)^2],$ and let us now prove the equivalence for $k+1$ using this induction hypothesis. By the recursive definition of $\delta_{k+1}$ we have that

$\delta_{k+1} = \beta \delta_k + \beta (1-\beta) (m_k - g_{k+1})^2$

$= \beta^2 EMA_{\beta}[(m_{k-1}-g_k)^2] + \beta (1-\beta) (m_k - g_{k+1})^2$

$= \beta( \beta EMA_{\beta}[(m_{k-1}-g_k)^2] + (1-\beta) (m_k - g_{k+1})^2)$

$= \beta EMA_{\beta}[(m_{k}-g_{k+1})^2]$

Where in the second equality we used the induction hypothesis, and in the last equality we used the recursive definition of EMA.

Q3: interpretation of the $m_k^2$ factor in the denominator of Eq. 3?

A3: The interpretation we give is that, because of the denominator which includes the $m_k^2$ term, we can re-write Adam (for $\beta_1=\beta_2$) as Eq (9), where it is clear that Adam is a mollified variant of Signum.

Q4: There are several small typos in the maths

A4: Thank you very much for these, and reading our paper and appendix so carefully. We corrected all of them in our latex file.

Q5: Could we see a comparison of Adam and the update proposed in Eq.9?

A5: The method in Eq. 9 is equivalent to Adam with $\beta_1 =\beta_2$ and $\epsilon=0$, which is the method we most experimented with.

Q6: What regard their claim that Adam is an online estimate of the mean and variance of the gradient is novel compared to what is stated in the initial paper?

A5: We apologize that we may not have understood your question here. When you ask “novel compared to what is stated in the initial paper”, here we assume you mean the origin Adam paper [1]. Please let us know if this is not the case, and we will gladly interact and answer.

This viewpoint we give of Adam in Theorem 4.1 as a formal online estimator of mean and variance is entirely new. Despite years of research into Adam, to the best of our knowledge, we have never seen this viewpoint of Adam before. As compared to the original Adam paper [1], the authors introduce Adam as a “ ... the method computes individual adaptive learning rates for different parameters from estimates of first and second moments of the gradients”. That is, they refer to the $v_k$ buffer as a type of second moment estimator of the gradient. But there is no formal proof, or variational viewpoint that shows this. Compare this to our Theorem 4.1 that formally shows how Adam arise from finding the mean and variance that maximizes the likelihood of observing a gradient drawn from a Gaussian, subject to KL regularization term that encodes how the distribution over gradients changes gradually over iterations.

The only formal proof in the original Adam paper [1] is Theorem 4.1, which attempts to establish the online regret of the iterates of Adam, and which is a notoriously incorrect proof (a counter example emerged in [2])

To conclude, Let us also emphasize a crucial aspect of our theory: it is derived from the $\beta_1=\beta_2$ insight. Proposition 2 in the appendix shows that under no other choice we can get an explicit variance-like term in the reformulation. This is important, since the connection to SignSGD with momentum is only possible only if one can properly "split" the second order term $v$ into a squared mean and variance term. To the best of our knowledge, we are also the first to analyze this setting.

[1] Adam: A Method for Stochastic Optimization, Diederik P. Kingma, Jimmy Ba, ICLR 2014

[2] Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. "On the Convergence of Adam and Beyond." International Conference on Learning Representations, ICLR 2018.

We thank the reviewer for their work and their pleasant comments on our paper! We also thank them for their questions and suggestions, which will help further improve the quality of our work.

We provide separate answers for your main points.

1) Disconnected Appendix

The reviewer is absolutely right, the main reason for this is that we did not want to mention precise places in the appendix due to the separate main paper and appendix deadlines (e.g., adding a new figure would compromise the ordering stated in the main paper). We will make sure everything in the appendix is linked semantically to the main paper and properly connected. Another reason for this disconnect was the tight space in the main paper, insufficient to properly comment all our side results. This can be though of course done in the potential camera-ready.

2) Improved Figures and Conclusion

We are committed to delivering the best possible graphical results, and therefore will devote efforts to addressing all reviewers' comments, such as those regarding Figure 4. We will also expand our conclusion section to address the theoretical aspects of the paper better.

3) Other points

1. We will define Gclip in the equation. This helps clarity, thanks!
2. With "landscape-based arguments" we mean arguments that attribute the Adam-SGD gap not to noise in the gradient but on the Hessian structure of Transformers. An example is the paper "Why Transformers Need Adam: A Hessian Perspective" by Zhang et al. We will specify in our updated version.
3. Typos and other suggestions: thanks so much, we already updated our LaTeX file with your comments!

We thank the reviewer for their reading and appreciation of our work. It truly means a lot to us to hear that you also believe our empirical evidence to be extensive and of practical interest — especially for future research.

We also like to answer in detail your questions, as we believe those are very interesting: we will highlight these points in our revised manuscript and also include the new results we present next.

1) Optimality

While $\beta_1=\beta_2$ performs in our experiments better than the usual choice $0.9, 0.95$ we do not claim that $\beta_1=\beta_2$ is always the best choice -- but instead that

a. it is quite convenient for tuning, while leading to near-optimal results.

b. brings about theoretical insights.

2) Training on more tokens.

This is a setting that definitely helps deliver our point, and we thank the reviewer for bringing this up. We were able to run a quite comprehensive ablation. The results below are in the setting of Figure 2, but we train for 6.4B tokens instead of 3.2B. We performed a total of 168 full training runs, each taking approximately 10 hours on a A100 with 80GB. For each choice of $\beta_1,\beta_2\in [0.9, 0.95, 0.975, 0.9875]\times[0.8 ,0.9, 0.95, 0.975, 0.9875, 0.99375]$ we report the best learning rate in the grid $[0.005, 0.001, 0.002, 0.004, 0.008, 0.016, 0.032]$. We highlight the choice $\beta_1=\beta_2$ below. We report validation perplexity.

As the reviewer can notice, $\beta_1=\beta_2$ is always near-optimal (gap of 0.01 ppl), with a slight imperfection at $\beta_1=0.9$. Note that this imperfection is also quite apparent in Figure 2. This does not quite matter because $\beta_1=0.9$ leads anyways to suboptimal results. Notice how, as reported in Figure 3 in the paper, we notice a strong correlation between $\beta_1$ and $\beta_2$: the higher $\beta_1$, the higher $\beta_2$ has to be for optimality.

Note that crucially, if one only tries values on the diagonal $\beta_1=\beta_2$, one gets to a perplexity which is only suboptimal by 0.01 ppl. This type of optimality also holds in Fig. 2 and 4.

3) Question: Signum vs Adam + Batch Size effect

This is a crucial point, which we again thank the reviewer for bringing up, this is worth a paragraph in our potential camera-ready. We believe our results are in perfect agreement with the 2 papers you mention. We would, however, like to note that we have already partially addressed this point in our paper (lines 196 in the main text and 491 in the appendix).

Let us expand: we first summarize what the two papers you mention say and then compare them with our claims.

Summary: Zhao et al. report that Signum and Adam perform quite similarly. Srećković et al. show that the gap between Adam and SGD shrinks to almost zero at very low batch-sizes. Srećković et al. was published after the Neurips deadline, so that is why we could not compare with this work.

Comparison: The finding of Srećković et al. was additionally confirmed by [A], which goes one step further and test a few optimizers other than Adam and SGD. In particular, they compare Adam with Adafactor — which (similarly to Signum) can be seen as a simplification of Adam. Note how in Figure 1(a) in [A] they show how the gap between Adam and Adafactor shrinks as the batch size decreases. We also notice a similar fact and already commented in line 196 of our submission, pointing to our Figure 12 in the appendix: at a shorter sequence length of 512 — i.e. a smaller number of tokens/iteration — the gap between Signum and Adam shrinks. We did this experiment to compare indeed with Zhao et al., that uses a fixed sequence length of 512. Hence, we believe our results are in perfect agreement with the 2 papers you mention.

There is a final important point to stress regarding shrinkage of the gap at low batch sizes: this setting is not practical for parallel computing, where batch sizes are often a standard. Our result in Figure 1 illustrates the gap between Signum and Adam in a practical scenario, with commonly used values of batchsize $\times$ sequence length.

[A] Marek, Martin, et al. "Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful." arXiv preprint arXiv:2507.07101 (2025).

Thank you for the clarification and additional experiments! They have helped me gain a better understanding of the paper.

I have two remaining concerns. I understand they may not be easily addressed in a short time, and this doesn't affect my view of the paper as a valuable contribution (I'm not sure if I will raise the score further, as I believe a strong accept should be reserved for truly exceptional cases.)

1. The theoretical and empirical results on $\beta_1=\beta_2$ are indeed important, as they offer a new perspective and simplify hyperparameter tuning a lot. Still, the fact that $\beta_1=\beta_2$ is near-optimal rather than truly optimal suggests that there are other components in the "secret sauce" of AdamW beyond what is identified in this paper.
2. As pointed out by the authors, this paper leads to similar conclusions to prior work that AdamW works better than Signum when batch size is large. However, according to eq. (9), when $\sigma$ decreases (which should occur when batch size increases), AdamW should be more similar to Signum, and this doesn't seem to align with the empirical results.