Continuous-Time Analysis of Adaptive Optimization and Normalization

We thank the reviewer for their comprehensive feedback on our paper, and for providing us with many helpful references relevant to our work. We appreciate your recognition of our “interesting empirical results”, and in this response we aim to address your central concerns. We have made the following changes towards addressing these concerns:

1. Provided a more comprehensive review of relevant literature and how our results relate to these works in Section 5, including the references that you have kindly provided.
2. Clarified our assumptions more precisely and framed our technical results as lemmas and theorems throughout the paper to aid readability.
3. Provided stronger empirical validation of Assumption 4.1 and 4.2 in Appendix G, verifying these assumptions up to 1,000 steps (equivalent to 100,000 continuous-time iterations).
4. Refined our presentation of Section 4.2 and 4.3 to make clearer our contributions towards understanding the influence of scale invariance on training dynamics.

We will now address the specific concerns detailed in your review.

This paper lacks sufficient novelty (...) the conclusion that the update of Adam is bounded given specific exponential moving average rates (...) has already been discussed in several prior studies [1, 2, 3, 4]

Thank you for the references. We agree that the boundedness of Adam’s update under certain conditions has been explored before, however, we believe that the observation of predictable exponential growth outside of the bounded region, and an empirical study of the bound’s implications – particularly its tightness to empirical discrete dynamics – is a novel contribution. Furthermore, our subsequent analysis of how scale invariance influences the training dynamics of Adam, as well as the construction of 2-Adam to capture its effect, is a further novel contribution. We also highlight that our analysis is generally applicable, independent of any specific architecture, loss function or data distribution, which contrasts with the analyses in [1, 3]. We have now incorporated these works in Section 5 to provide a more comprehensive literature review and how their contributions relate to our own.

Moreover, similar techniques involving continuous approximation for the momentum terms $m_t$ and $v_t$ have also been explored in [5, 6]. (...) This paper lacks a thorough discussion of previous works.

Thank you for making us aware of these works. We have now discussed these works in Section 5 of the updated manuscript.

The assumptions presented in the paper are relatively problematic (...) the statement that they are "approximately 0" does not imply they are exactly 0

We agree that this statement could be more precise, and we have now modified the manuscript to be more explicit in what we mean by this assumption: that the inner product is second-order in $\eta$ and $\lambda$. We have further provided stronger empirical validation of this assumption in Appendix G, validating for up to 1,000 steps (equivalent to simulating 100,000 continuous-time iterations). Further, we note that Figure 5 provides implicit verification of this assumption since the proof of Theorem 2 is dependent upon this assumption. We hope that this point is now clearer, and please let us know if you have any more concerns.

As for assumption 4.2, (...) I fail to see any significant distinction between gradient descent with momentum and 'Adam'.

We highlight that under assumption 4.2, $\tilde{v}_W(t)$ still depends on time and hence still maintains an adaptive learning rate effect that momentum lacks. As a result, we claim that this assumption still captures the important adaptive aspects of Adam relating to a dynamic learning rate over training. We also draw your attention to the now stronger empirical verification of this assumption in Appendix G, finding that training dynamics under assumption 4.2 remain close to Adam’s dynamics.

Finally, the conclusion that $||dot W_t||_2 \approx ||u_t||_2$ based on these two assumptions might also be insuitable. As [2, 3, 7] demonstrated, Adam aligns more with $\ell_{\infty}$ norm instead of $\ell_2$ norm.

Thank you for the references. It appears that these works describe how Adam/AdamW provably converges to the KKT minimizer of an $\ell_{\infty}$ constrained optimization problem, however we fail to see precisely how this finding is related to the norm of the continuous-time derivative $||\dot{W}_t||_2 \approx ||u_t||_2$. Could you provide more clarification on exactly how these two concepts are related?

I thank the authors for their efforts, while their responses hardly address my concerns.

• I believe this paper has poor theoretical novelty and contribution. After reviewing the revised manuscript and comparing it to the related works I mentioned in my original review, I feel even more confident in this assessment. I list the weaknesses in the following:
  • As I indicated in my original review, both the conclusion and proof techniques regarding the continuous form of $v_t$ have been proposed in [5], and this can be directly extended to $m_t$ without any modifications. However, in the revised manuscript, the authors still claim that "continuous-time formulation" is their major contribution in their introduction section without even acknowledging [5]. I also did not find any discussion of [5] surrounding Lemma 1 and its proof, which I believe is essential. While I acknowledge the discussion in Section 5, it is not sufficient. Moreover, if the authors believe the "continuous-time formulation" is their technical contribution, I hope they can provide a detailed illustration of the essential distinctions between the conclusion and proof technique of Lemma 1 and those presented in section A.3 of [5] (I know you are using different discretization scale and they consider updating rules with stability constant $\epsilon$, however, neither of these points constitutes an essential difference).
  • As I mentioned in my original review, the conclusion that the update of Adam is bounded given specific exponential moving average rates, has been widely utilized in theoretical studies regarding Adam [1, 2, 3, 4]. Specifically, I would like to emphasize that the conclusion regarding the rates approaching infinity outside stable regions can also be directly obtained without additional calculations. Take the proof of Lemma 6.5 in [3] as an example, they derived that $\frac{m_t[k]}{\sqrt{v_t[k]}} \leq \Big(\sum_{\tau=0}^t \frac{\beta_1^{2\tau}(1-\beta_1)^2}{\beta_2^\tau(1-\beta_2)}\Big)^\frac{1}{2}.$ It is straightforward that when $\beta_1^2\leq \beta_2$, the RHS will be bounded by a constant. When $\beta_1^2=\beta_2$, the RHS will be bounded at the scale $\sqrt t$, matching your results $\sqrt n$. And $\beta_1^2>\beta_2$, the RHS will be bounded at the scale of $a^t$ with $a>1$, matching your exponential growth results. Although [1, 2, 3, 4] do not directly present results concerning the unbounded growth rate, as their studies do not focus on this aspect, it does not mean that their findings cannot imply such conclusions. This is evident because they all employ the same proof techniques as you as I mentioned in my original review. Specifically, these works consider practical discretized updates, whereas you are examining an idealized continuous update. This difference suggests that their results are more practical and generalizable. Moreover, I would like to point out that [1, 3] specify the architecture because their focus differs from that of this paper; however, their derivation for the conclusions mentioned above does not depend on any specific architecture or loss function.
  • I do not think any theoretical conclusion in Section 4 is convincing to me, as they are derived based on two problematical assumptions, which will be discussed in detail.

We thank the reviewer for their in-depth engagement with our work and for providing us with detailed comments and suggestions. We are happy to hear that you found our work to be “sound and grounded” and to “maintain an interesting balance between theory and practical insight”. To address the concerns detailed in your review, we have made the following substantial changes to the paper:

1. We have provided a more comprehensive review of related literature in Section 5, including how our results relate to the discrete-time approaches of previous works.
2. We have discussed our contributions and their practical implications, the limitations of our framework, and future research directions in Section 6.
3. Revised our presentation of adaptive normalization and k-Adam in Section 4.2 and 4.3, now more clearly presenting our technical results as lemmas and theorems and highlighting the implications of our findings.
4. Evaluated k-Adam on a larger scale language modeling task, including training a 30M parameter transformer model on the C4 dataset, where we again find 2-Adam to outperform Adam in Appendix P.
5. Provided a discussion of how our approach relates to Adam’s special cases, RMSprop and signSGD, in Appendix D.
6. Discussed the non-trivial dynamics in the case of $u(t) \equiv 0$ in Appendix C.

Below we provide responses to your specific comments and hope that we have addressed your central concerns with our work.

A review of the literature on continuous-time models for optimizers is crucial, with a particular focus on those that cover Adam. (...) there is no citation of Malladi et al., who derive an SDE model for Adam (...) Additionally, a subsection comparing the method to existing literature could provide more clarity.

Thank you for your helpful feedback and reference to Malladi et al. To address these concerns, we have now expanded our literature review in Section 5 to include Malladi et al. and its relation to our work, as well as additional works that make use of both discrete-time and continuous-time models in the context of adaptive optimization.

A brief section outlining the limitations of the approach would be useful.

We have now provided a more complete discussion of our approach’s limitations in Section 6, now emphasizing that our analysis neglects stochastic gradients, and that we lack a complete theoretical understanding (it is only suggested from our theoretical analysis) of why larger values of $C(\beta, \gamma)$ correspond with a faster rate of generalization (along a given normal curve). We hope this has addressed your concern, and please let us know if you have any further concerns.

In "1. Continuous-time formulation (Section 2)", why not merge these two points into a simpler one?

This is a good point, and we have now combined these two sections into Section 2.2. We have also formatted the associated technical results more clearly as lemmas and propositions, which we hope aids readability.

While this is interesting, it is unclear if it is practically relevant. Does this give practitioners new insights on selecting the betas? Or are they always in the "safe region"?

Thank you for the question. The motivation behind Section 3 was to justify why the hyperparameter values chosen by practitioners result in successful generalization, where we indeed found that commonly chosen hyperparameters lie within the safe region. Furthermore, in pursuing this, we discovered a natural quantity $C(\beta, \gamma)$ that relates to training stability, finding empirically (in Section 3.3) that this quantity correlates strongly with the rate of generalization in the following sense: along a given normal curve, choices of $(\beta, \gamma)$ with a larger value of $C(\beta, \gamma)$ possess a faster rate of generalization. As a result, our framework has provided novel insight into optimal hyperparameter selection through the measure of stability $C(\beta, \gamma)$, which we believe can indeed aid practitioners in hyperparameter choice.

Is it beneficial to avoid such explosions? Adam is often successfully used even when loss spikes are observed.

This is an interesting question, and indeed was a motivating factor for our study of generalization performance in Section 3.3. In Section 3.3 (particularly Figure 4) we found that hyperparameters outside of the stable region experienced worsened generalization, monotonically in their distance from the stability boundary. There may be some specific contexts for which effective generalization can occur outside of the stability region, perhaps under sufficient learning rate annealing, which is an interesting direction for future research.

We thank the reviewer for their helpful feedback and positive assessment of our work. We are very happy to hear that our findings were “well written and crisp”, and that our “continuous time analysis of Adam is interesting and novel”. To address your concerns, we have made the following additions to the paper:

1. Improved experimental verification of k-Adam, now evaluating k-Adam in a language modeling setting, training a transformer model with 30M parameters on the C4 dataset and similarly finding 2-Adam to outperform Adam, with results in Appendix P.
2. Provided a more detailed discussion of Assumptions 1 and 2 in Section 4.2, as well as including a stronger empirical validation of these assumptions in Appendix G.

We have also addressed your specific feedback below, which we hope fully addresses your concerns.

Assumptions 1 and 2 are not very intuitive, at least not to me.

Thank you for your feedback. To aid intuition regarding these assumptions, we have now extended our discussion of both assumptions in Section 4.2. Specifically, we more clearly highlight the orthogonality property of scale invariance (described by Lemma 2) and how it supports Assumption 1, and the interpretation that Assumption 2 corresponds to assuming that every entry in the weight $W$ experiences the same amount of adaptive scaling. We also draw your attention to the stronger empirical verification that we have now provided for these assumptions, which can be found in Appendix G. We hope this addresses your concern. Please let us know if you have any questions regarding these assumptions.

How tight is the bound of eq 5?

In Figure 3(b) we showcase that, in the case of $C(\beta, \gamma) < 0$, the bound is extremely tight to the true empirical dynamics, with the empirical rate of exponential growth in close agreement with the theoretically predicted rate of $|C(\beta, \gamma)|/2$ (denoted by a red dashed line). In the case of $C(\beta, \gamma) > 0$, we look more closely at the bound in Appendix L, finding it to remain satisfied as expected, and to be reasonably tight.

Is there a corresponding lower bound?

Thank you for the interesting question. In our analysis we did not determine a lower bound on the max-update $||u_n||_{\infty}$, though we do find empirically that dynamics stay quite tight to the upper bound (as displayed in Figure 3(b) and Appendix L). Exploring this aspect further from a theoretical perspective is an interesting future research direction.

CIFAR10 results are not very convincing, it’s small scale (...) Can you add more experimental verification of Adam-k?

Thank you for the feedback. We agree that the experimental verification of k-Adam could be stronger. To address this point, we have now evaluated k-Adam in a language modeling setting on the C4 dataset, considering a transformer model of 30M parameters. We include the results in Appendix P, where we again find 2-Adam to outperform Adam. We hope this addition has addressed this concern, and please let us know if you have any further concerns.

Summary: We would like to thank the reviewer for providing very helpful suggestions and comments, and for their positive assessment of our findings. On the basis of your feedback, we have provided a stronger assessment of k-Adam’s performance in a larger scale language modeling setting, and have better clarified and verified the assumptions that facilitate our theoretical analysis. We are readily available to address any further concerns.

We would like to gently remind the reviewer that the discussion period ends within the next four days and that we have responded to their review. We hope you will let us know if we have addressed your concerns, and if you have any further concerns, we would be happy to address them.

Thank you for thoroughly reading our paper and providing us with constructive technical feedback! We are delighted that you found our research on establishing the continuous-time framework of Adam(W) to be "potentially very useful" and our findings on stable hyperparameter regions to be "helpful in practice." We also greatly appreciate your recommendations to clarify and verify the crucial assumptions underlying our theory and your thoughtful suggestions for experiments to fully validate our framework's predictions.

To address your concerns, we have taken the following steps:

1. We have made the set of theoretical assumptions more explicit and expanded the technical results with greater detail in the main text. We have also updated our draft after incorporating your suggestions.
2. We have conducted the experiments you have recommended, including training a transformer model with 30M parameters on the C4 dataset, to validate these assumptions comprehensively.

Below, we detail how we have addressed your feedback, and we hope this fully addresses your concerns.

I don't understand the crucial deviation in Section C (...) First, why is $m$ differentiable, why can you drop the higher order (second derivative is bounded?) of it?

Thank you for carefully reviewing our derivation in Appendix C. We sincerely appreciate your insightful feedback. First, we would like to clarify that the goal of this section is to approximate the trajectories of the discrete-time update rule using a continuous-time differential equation. Accordingly, the differentiability of $m$ is an assumption we make, rather than a result derived from our analysis, an important point that we should have emphasized and discussed more thoroughly in the main text. Regarding the dropping of the higher order terms, we now explicitly address that we assume $p \geq 1$ throughout the paper and clarify that we consider $\eta$ to be small, consistent with practical implementations (see comment at line number 125). These are indeed crucial assumptions that we agree should have been more explicitly stated.

In response to your feedback, we have made the following updates:

1. We now motivate the continuous-time formulation by referencing prior influential works in Section 5 that have derived practically insightful claims under similar assumptions.
2. We provide a comment in Section 2.2 (line number 125) explaining why dropping the second-order term is reasonable, particularly given the small learning rate $\eta$ typically used in practical implementations of Adam (e.g., $\eta = 10^{-3}$ by default in its PyTorch implementation, making higher-order terms practically negligible).
3. We have substantially increased the step counts in Figures 1 and 5, extending the validation of our theory to 1,000 steps (corresponding to 100,000 continuous-time iterations)—significantly more than the 100 steps in the previous version—to address your suggestions more comprehensively.

Overall, we believe these updates better motivate, clarify, and validate the assumptions in our work, aligning more closely with your suggestions.

Second, why is it $g(t_n)$ not $g(t_n-\eta^p)$?

This is an interesting point. We take this formulation to align with the standard implementation of Adam (cf. PyTorch). If we instead defined $m_n = \gamma m_{n-1} + (1-\gamma) g_{n-1}$, such that the continuous-time expression includes $g(t_n-\eta^p)$, we note that updates $u_n$ based on $m_n$ would use old gradients $g_{n-1}$ from one-step behind, with the most recent gradient $g_n$ being neglected. It seems it would be best to update using the most up-to-date parameter value $\theta_n$, and indeed this is what the official PyTorch implementation of Adam does (https://pytorch.org/docs/stable/generated/torch.optim.Adam.html), updating via the gradient associated with the most recent parameter value. We therefore define $m_n = \gamma m_{n-1} + (1-\gamma) g_n$ to accurately match practice, which results in a $g(t_n)$ term rather than $g(t_n-\eta^p)$.

We thank the reviewer for increasing their score to 5 (though, as Reviewer GyMg comments, I don't think the official rating in your review has been updated?). We emphasize that the assumption that higher-order terms are small relative to leading-order terms is a standard assumption in the literature which has been used to uncover meaningful properties of practical optimization [1, 2, 3, 4]. For example, see Equation 3 of [4] (which was published as a spotlight paper at ICLR 2024). We argue that the overall contribution due to higher-order terms is still negligible in comparison to the overall contribution of leading-order terms. Importantly, the role of our empirical validation in Figure 1 and Figure 5 was to indeed verify that this assumption is reasonable and provides an accurate model of Adam.

[1] Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms

[2] André Belotto da Silva and Maxime Gazeau. A general system of differential equations to model first order adaptive algorithms

[3] Anas Barakat and Pascal Bianchi. Convergence and dynamical behavior of the adam algorithm for non-convex stochastic optimization

[4] Lizhang Chen, Bo Liu, Kaizhao Liang, and Qiang Liu. Lion Secretly Solves a Constrained Optimization: As Lyapunov Predicts