Continuous-Time Analysis of Adaptive Optimization and Normalization

## Abstract

Abstract Adaptive optimization algorithms, particularly Adam and its variant AdamW, are fundamental components of modern deep learning. However, their training dynamics lack comprehensive theoretical understanding, with limited insight into why common practices—such as specific hyperparameter choices and normalization layers—contribute to successful generalization. This work presents a continuous-time formulation of Adam and AdamW, facilitating a tractable analysis of training dynamics that can shed light on such practical questions.
We theoretically derive a stable region for Adam’s hyperparameters ( β , γ ) 𝛽 𝛾 (\beta,\gamma) ( italic_β , italic_γ ) that ensures bounded updates, empirically verifying these predictions by observing unstable exponential growth of parameter updates outside this region. Furthermore, we theoretically justify the success of normalization layers by uncovering an implicit meta-adaptive effect of scale-invariant architectural components. This insight leads to an explicit optimizer, 2 2 2 2 -Adam, which we generalize to k 𝑘 k italic_k -Adam—an optimizer that applies an adaptive normalization procedure k 𝑘 k italic_k times, encompassing Adam (corresponding to k = 1 𝑘 1 k=1 italic_k = 1 ) and Adam with a normalization layer (corresponding to k = 2 𝑘 2 k=2 italic_k = 2 ). Overall, our continuous-time formulation of Adam facilitates a principled analysis, offering deeper understanding of optimal hyperparameter choices and architectural decisions in modern deep learning.

## 1 Continuous-time formulation of Adam

In this section we present a continuous-time formulation of the Adam (and AdamW) optimizer which forms the theoretical foundation for the rest of the paper. We derive a continuous-time expression for the Adam gradient update in Section 1.1 (utilized in Section 2 to derive a condition for bounded updates), and in Section 1.2 we formulate Adam as a second-order differential equation (used in Section 3 to interpret the implicit effect of scale invariance and motivate the $k$-Adam optimizer). We provide brief descriptions of relevant optimizers (momentum, RMSprop, Adam, AdamW) in Appendix A. Experimental details are left to Appendix F.

Notation. We write $||x||:=\sqrt{\sum_{i}x_{i}^{2}}$ and $||x||_{\infty}:=\max_{i}|x_{i}|$ for $x\in\mathbb{R}^{p}$, and denote the inner product $\braket{x,y}\equiv\sum_{i}x_{i}y_{i}$. We denote elementwise squaring by $x^{\odot 2}$. We consider a loss function $L:\Theta\to\mathbb{R}$ where $\Theta\subset\mathbb{R}^{p}$, and $\theta_{n}\in\Theta$ denotes a model’s parameters after $n$ discrete updates, with corresponding gradient $g_{n}:=\nabla_{\theta}L(\theta_{n})$, and initial parameter $\theta_{0}$. We will write $g_{0:n}\equiv(g_{0},\ldots,g_{n})$.

### 1.1 Parameter update in continuous-time

Neglecting weight decay for the moment (i.e. $\lambda=0$), Adam and AdamW possess the discrete-time update rule

$$ $\theta_{n+1}=\theta_{n}-\eta u_{n},\;\;\text{where}\;\;u_{n}:=\frac{\sqrt{1- \beta^{n+1}}}{1-\gamma^{n+1}}\frac{m_{n}}{\sqrt{v_{n}}}$ (1) $$

with learning rate $\eta$, and *adaptive hyperparameters* $(\beta,\gamma)\in[0,1]^{2}$ with corresponding moving-averages

$$ $m_{n}:=\gamma m_{n-1}+(1-\gamma)g_{n},\;\;\;\;\;\;v_{n}:=\beta v_{n-1}+(1- \beta)g_{n}^{\odot 2}$ $$

where $(m_{-1},v_{-1}):=(0,0)$. We say that $u_{n}$ is an *adaptive-normalization* of the gradients $g_{0:n}$ (see Appendix N for more details).

We will work in continuous-time, at a timescale of $\eta^{p}$, such that e.g. $m_{n}=m(t_{n})$, with $t_{n}:=n\eta^{p}$ for some $p\in\mathbb{R}$ (typically $p=1$).
We can then derive continuous-time expressions for the moving-averages $m_{n}$ and $v_{n}$ (details left to Appendix C), which, to leading order in $\eta^{p}$, take the form

$$ $m(t)=\int_{0}^{t}d\tau\;K_{\gamma}(\tau,t)g(\tau),\;\;\;\;\;\;v(t)=\int_{0}^{t }d\tau\;K_{\beta}(\tau,t)g(\tau)^{\odot 2},$ (2) $$

$$ $\text{where}\;\;\;K_{\alpha}(\tau,t):=\frac{1-\alpha}{\eta^{p}\alpha}\exp\left (-\frac{1-\alpha}{\eta^{p}\alpha}(t-\tau)\right)$ $$

with $g(\tau):=\nabla_{\theta}L(\theta(\tau))$ the gradient at continuous time $\tau$. This provides a continuous-time expression for Adam/AdamW’s update $u_{n}=u(t_{n})$,

$$ $u(t)=\frac{\sqrt{1-\beta^{1+t/\eta^{p}}}}{1-\gamma^{1+t/\eta^{p}}}\frac{\int_{ 0}^{t}d\tau\;K_{\gamma}(\tau,t)g(\tau)}{\sqrt{\int_{0}^{t}d\tau\;K_{\beta}( \tau,t)g(\tau)^{\odot 2}}}$ (3) $$

Including weight decay is simple and is described in Appendix C. As we will see, Equation 3 is a central result for Section 2, allowing us to derive theoretical guarantees regarding training stability with respect to Adam/AdamW’s adaptive hyperparameters $(\beta,\gamma)$.

### 1.2 Continuous parameter dynamics via differential equation

Figure: Figure 1: Continuous-time model closely agrees with discrete-time trajectories. We plot the discrete-time and continuous-time trajectories for 16 randomly chosen parameters from a transformer model (details in Appendix F).
Refer to caption: x1.png

From Equation 1 we can also determine $\theta_{n}=\theta(t_{n})$ in continuous-time via a governing differential equation (details left to Appendix C). To second-order, AdamW with weight decay $\lambda$ is governed by the differential equation

$$ $\lambda\eta\theta(t)+\eta^{p}\dot{\theta}(t)+\frac{\eta^{2p}}{2}\ddot{\theta}( t)=-\eta u(t)$ (4) $$

In combination with the continuous-time expression for $u(t)$ (Equation 3), we can numerically solve Equation 4 (method described in Appendix G) to obtain a continuous-time parameter trajectory $\theta(t)$. We demonstrate that the continuous-time trajectory (for timescale $p=1$) associated with Equation 4 closely agrees with the true discrete-time trajectory in Figure 1. In Section 3 we will use Equation 4 to interpret the implicit effect of scale invariance.

## 2 Theory of adaptive hyperparameters

We now begin to apply the continuous-time framework presented in the previous section to understanding aspects of adaptive optimization from a theoretical perspective. In this section we will present a theory-driven account of adaptive hyperparameter choice for Adam/AdamW, understanding the values of $(\beta,\gamma)$ that result in stable training and effective generalization.

First we will use the continuous-time expression for Adam’s update (Equation 3) to derive a theoretical region of hyperparameter space $\mathcal{B}_{+}$ for which updates are provably bounded (Section 2.1). We will then empirically verify that this region indeed exhibits stable training, with unstable training in the complementary region $\mathcal{B}_{-}$ (exhibiting a *predictable* exponential growth) (Section 2.2), and observe how generalization performance varies in the regions $\mathcal{B}_{+}$ and $\mathcal{B}_{-}$ (Section 2.3).

### 2.1 Deriving a condition for bounded updates

The continuous-time expression for Adam’s update (Equation 3) has an immediate consequence in regards to bounding the max-update $||u_{n}||_{\infty}\equiv||\theta_{n+1}-\theta_{n}||_{\infty}/\eta$. The max-update quantifies the maximal parameter change (across *all* parameters) at a given step, hence upper bounds on this quantity also hold identically for the parameter change of any arbitrary parameter. We have the result,

###### Max-update bound .

The max-update can be bounded as

$$ $||u_{n}||_{\infty}\leq\frac{\sqrt{1-\beta^{n+1}}}{1-\gamma^{n+1}}\frac{1- \gamma}{\gamma}\sqrt{\frac{\beta}{1-\beta}}B_{n}(\beta,\gamma)$ (5) $$

where

$$ $B_{n}(\beta,\gamma):=\begin{cases}1/\sqrt{C(\beta,\gamma)},&C(\beta,\gamma)>0, \\ \sqrt{n},&C(\beta,\gamma)=0,\\ \exp(n|C(\beta,\gamma)|/2)/\sqrt{|C(\beta,\gamma)|},&C(\beta,\gamma)<0\end{cases}$ $$

defining $C(\beta,\gamma):=(2\beta(1-\gamma)-\gamma(1-\beta))/\beta\gamma$.

Figure: Figure 2: Visualization of level curves $\mathcal{B}_{c}$ (solid lines) and normal curves $\mathcal{C}_{\beta,\gamma}$ (dashed red lines). Level curves are coloured based on their value of $C(\beta,\gamma)$, (i.e. purple has most positive value, red most negative). The bounded-update region $\mathcal{B}_{+}$ is highlighted in red.
Refer to caption: x2.png

We highlight that this bound is only possible because of the specific form of Adam’s update $u(t)$: a moving-average of $g$ divided by the square-root of a moving-average of $g^{\odot 2}$ (Equation 3), i.e. an adaptive-normalization of $g$ (see Appendix N), which allows us to apply the Cauchy-Schwarz inequality. We can therefore view this result as a theoretical justification for the form of Adam’s update. We leave the derivation of this bound to Appendix C.

We define the bounded-update region $\mathcal{B}_{+}:=\{(\beta,\gamma):C(\beta,\gamma)>0\}$, and the complementary region $\mathcal{B}_{-}:=\{(\beta,\gamma):C(\beta,\gamma)<0\}$. From Equation 5, we can bound the max-update by a constant independent of $n$ when $(\beta,\gamma)\in\mathcal{B}_{+}$. Outside of $\mathcal{B}_{+}$, the bound depends on $n$ and diverges over training as $n\to\infty$. This is suggestive that we may observe stable training when $(\beta,\gamma)\in\mathcal{B}_{+}$, and that the max-update may grow exponentially (at a rate/exponent proportional to $|C(\beta,\gamma)|$) when $(\beta,\gamma)\in\mathcal{B}_{-}$. Indeed, we will verify this phenomena empirically in Section 2.2. We comment on the case $(\beta,\gamma)\in\mathcal{B}_{0}$ in Appendix D. It is easy to show that $\beta>\gamma$ is a sufficient condition for $(\beta,\gamma)\in\mathcal{B}_{+}$, meaning that the choice of hyperparameters typically chosen in practice – e.g. the PyTorch default values $(\tilde{\beta},\tilde{\gamma}):=(0.999,0.9)$ – lie within $\mathcal{B}_{+}$.

For the following, we define the *level curves* $\mathcal{B}_{c}:=\{(\beta,\gamma):C(\beta,\gamma)=c\}$, and consider *normal curves* perpendicular to these level curves, visualized in Figure 2. We will denote the (unique) normal curve passing through the point $(\beta,\gamma)$ by $\mathcal{C}_{\beta,\gamma}$. In Section 2.2 and Section 2.3 we will analyse how $C(\beta,\gamma)$ correlates with training stability & generalization performance by observing how such properties change as we move along the normal curve $\mathcal{C}_{\tilde{\beta},\tilde{\gamma}}$ in hyperparameter space, since $C(\beta,\gamma)$ varies monotonically along a normal curve (by construction).

###### Max-update bound .

### 2.2 Empirical analysis of max-update

Figure: (a)
Refer to caption: x3.png

We now consider the empirical implications of Equation 5, assessing whether these bounds indeed hold in practice and whether $C(\beta,\gamma)$ is predictive of training stability. We train a decoder-only transformer model on a shakespeare text dataset (details left to Appendix F) and observe how the max-update $||u_{n}||_{\infty}$ evolves over training iterations $n$. Specifically, in Figure 3, we consider the normal curve $\mathcal{C}_{\tilde{\beta},\tilde{\gamma}}$ passing through the typical hyperparameter values $(\tilde{\beta},\tilde{\gamma})=(0.999,0.9)$, taking $64$ points $(\beta,\gamma)$ uniformly along this curve (which we visualize in Figure 3(c)). For each point $(\beta,\gamma)$, we plot the max-update trajectory (Figure 3(a)), finding that the theoretical bounds of Equation 5 are satisfied in both regions $\mathcal{B}_{+}$ and $\mathcal{B}_{-}$, observing well-behaved bounded growth in $\mathcal{B}_{+}$, whereas $\mathcal{B}_{-}$ exhibits exponential growth at a rate that appears correlated with $|C(\beta,\gamma)|$ as predicted by Equation 5. We more closely verify this phenomena in Figure 3(b), finding that in $\mathcal{B}_{-}$, the slope $d\log||u_{n}||_{\infty}/dn$ has a near-perfect agreement with the theoretically predicted growth rate of $|C(\beta,\gamma)|/2$. It is surprising that not only are the theoretical bounds (Equation 5) satisfied in practice, but the bounds are very accurate models of the true empirical dynamics, with exponential growth occurring almost immediately after entering $\mathcal{B}_{-}$. We look more closely at the results of Figure 3(a) in the case of $(\beta,\gamma)\in\mathcal{B}_{+}$ in Appendix J. We comment on correcting the exponential growth in $\mathcal{B}_{-}$ via learning rate annealing in Appendix L.

These results partly justify the success of typical values $(\tilde{\beta},\tilde{\gamma})=(0.999,0.9)$ in practice: these values lie in $\mathcal{B}_{+}$ and hence benefit from theoretical guarantees regarding training stability (which, as we have seen, are faithful to practice). Can we obtain a more fine-grained picture as to what choices of $(\beta,\gamma)$ within the region $\mathcal{B}_{+}$ will result in successful generalization? We will now explore this.

### 2.3 Properties of generalization in ℬ + subscript ℬ \mathcal{B}_{+} caligraphic_B start_POSTSUBSCRIPT + end_POSTSUBSCRIPT

We will now assess how generalization performance varies in $\mathcal{B}_{+}$, and particularly, how the value of $C(\beta,\gamma)$ correlates with generalization performance. In Figure 4 we observe how test loss varies along the normal curve $\mathcal{C}_{\tilde{\beta},\tilde{\gamma}}$. taking 128 points uniformly along the entire curve (visualized in Figure 4(c)). We see that larger values of $C(\beta,\gamma)$ display faster generalization, as seen in Figure 4(a) and highlighted at iteration 1000 in Figure 4(b). After sufficient training, i.e. at iteration 3000 in Figure 4(b), all points in $\mathcal{B}_{+}$ achieve roughly the same test loss. The test loss exhibits a rapid increase after entering the region $\mathcal{B}_{-}$, which is to be expected given the unstable exponential growth in the max-update shown in Section 2.2.

These results suggest that for a given normal curve, the point $(\beta,\gamma)$ with the largest value of $C(\beta,\gamma)$ generalizes the fastest along the curve. We provide supporting evidence for this claim, finding similar results for the normal curve that instead passes through the hyperparameter values $(0.95,0.9)$, in Appendix K. This would further justify why the values $(\tilde{\beta},\tilde{\gamma})$ succeed in practice (since $\tilde{\beta}\approx 1$, hence $C(\tilde{\beta},\tilde{\gamma})$ is large relative to other points along its normal curve), however unlike the results of Section 2.2, this claim lacks an explicit supporting theoretical result; it is only suggestive from Equation 5 that a larger value of $C(\beta,\gamma)$ may correlate with better training properties. We leave direct theoretical guarantees regarding the rate of generalization to future work.

Figure: (a)
Refer to caption: x6.png

## 3 Implicit effect of scale invariance

The presence of scale-invariant architectural components has become ubiquitous in modern machine learning. One particular instance of such components are normalization layers, such as layer-norm (Ba et al., ), batch-norm (Ioffe & Szegedy, ), and qk-norm (Dehghani et al., ; Gilmer et al., ). It has been observed in practice that such normalization provides an implicit beneficial effect to training, compared to e.g. explicitly fixing/constraining the weight norm after each update step (Tanaka & Kunin, ; Lubana et al., ; Santurkar et al., ). In this section we will apply the continuous-time framework of Section 1 to interpreting this implicit role of scale invariance, allowing us to gain an understanding as to how normalization layers contribute to successful generalization.

We will first briefly discuss scale-invariant maps (Section 3.1) and then use the second-order differential equation for Adam/AdamW (Equation 4) to solve for the norm dynamics and uncover an implicit *meta-adaptive* effect of scale invariance, which we then convert into an explicit optimizer, $2$-Adam (Section 3.2). Finally, we will introduce the $k$-Adam optimizer – a generalization of Adam and $2$-Adam – and compare its performance with Adam for training a CNN on CIFAR10 (Section 3.3).

### 3.1 Scale-invariant maps

A function $f:\Theta\to\mathbb{R}$ is *scale-invariant* with respect to a weight $W\in\mathcal{W}\subset\Theta$ if and only if $f(\theta)$ remains unchanged under the transformation $W\mapsto\alpha W$ for all $\alpha\in\mathbb{R}$. As a concrete example, qk-norm (Dehghani et al., ; Gilmer et al., ) applies a layer-norm to the query and key projections in a transformer, with the attention logits $z_{ij}$ of a particular attention head taking the form,

$$ $z_{ij}:=\text{LN}(W_{Q}x_{i})^{T}\text{LN}(W_{K}x_{j})$ $$

for embeddings $x_{i}\in\mathbb{R}^{d}$, query matrix $W_{Q}$, key matrix $W_{K}$, and layer-norm $\text{LN}(\cdot)$. As a result, the attention logits $z_{ij}$ are scale-invariant with respect to $W_{Q}$ and $W_{K}$, and hence the loss function $L$ is equivalently scale-invariant. One can show (see Appendix I) that if $L$ is scale-invariant with respect to a weight $W$, then

$$ $\braket{W(t),g_{W}(t)}=0$ (6) $$

at all times $t$, where $g_{W}(t):=\nabla_{W}L(\theta(t))$ is the gradient associated with $W$. Equation 6 will be utilized in the following to show that scale-invariant maps possess an implicit *meta-adaptive* effect.

### 3.2 Uncovering an implicit meta-adaptive effect

To theoretically analyse the implicit effect of scale invariance, we will use the second-order differential equation for AdamW described by Equation 4 and choose a timescale of $p=1$ (we verified the accuracy of this setup in Section 1.2). For an arbitrary scale invariant weight $W$ (e.g. a key/query matrix under qk-norm), the continuous-time dynamics of $W$ are governed by

$$ $\frac{1}{2}\eta\ddot{W}(t)+\dot{W}(t)+\lambda W(t)=-u_{W}(t)$ (7) $$

to second-order, with $u_{W}$ defined analogously to $u$ (replacing $g$ with $g_{W}$ in Equation 3). From here we can derive (see Appendix H) the following expression for the weight norm $||W(t)||$,

$$ $||W(t)||^{2}\approx||W(0)||^{2}e^{-2\lambda t}+\int_{0}^{t}d\tau\;e^{-2\lambda (t-\tau)}(\eta||u_{W}(\tau)||^{2}-2\braket{W(\tau),u_{W}(\tau)})$ (8) $$

We empirically verify that Equation 8 is accurate to the true discrete trajectory for $||W(t)||^{2}$ in Figure 5. In order to derive Equation 8 we did not need to use the property of scale invariance. We will now make use of this property, in combination with the following two simplifying assumptions:

Figure: Figure 5: Norm approximation agrees with true norm. The true trajectory of $||W||^{2}$ compared to the approximation (dashed line) of Equation 8 for 16 randomly chosen query/key matrices.
Refer to caption: x9.png

###### Assumption 1 .

We will assume that

$$ $\braket{W(t),g_{W}(\tau)}\approx 0\;\;\forall\;\tau\leq t$ $$

###### Assumption 2 .

We will use a coarse-graining approximation for $v_{W}(t)$ (using Equation 2):

$$ $v_{W}(t)\equiv\int_{0}^{t}d\tau\;K_{\beta}(\tau,t)g_{W}(\tau)^{\odot 2}\approx \int_{0}^{t}d\tau\;\tilde{K}_{\beta}(\tau,t)||g_{W}(\tau)||^{2}=:\tilde{v}_{W} (t)$ $$

with $\tilde{K}_{\beta}(\tau,t):=K_{\beta}(\tau,t)/N$ where $N$ is the total number of entries of $W$, i.e. we replace the entries of $g_{W}(\tau)^{\odot 2}$ with the average entry across the tensor.

Assumption 1 is intuitive given the properties of high-dimensional spaces(^1^11For example the Johnson–Lindenstrauss lemma implies that in high-dimensional spaces, randomly chosen vectors are likely to be almost orthogonal.), but also Equation 6 is suggestive of such a result. We provide empirical support for this assumption in Appendix E. We also show (in Appendix E) that replacing $v_{W}$ with the coarse-grained version $\tilde{v}_{W}$ has little effect on training dynamics. A consequence of these assumptions is that $\braket{W(t),u_{W}(t)}\approx 0$, allowing us to approximate,

$$ $||W(t)||^{2}\approx||W(0)||^{2}e^{-2\lambda t}+\eta\int_{0}^{t}d\tau\;e^{-2 \lambda(t-\tau)}||u_{W}(\tau)||^{2}$ (9) $$

Interestingly, Equation 9 has an identical form to an exponential moving-average of $||u_{W}||^{2}$ in continuous-time (see Appendix H). Further (also see Appendix H), the evolution of the unit direction $\hat{W}$ – which is most relevant to consider, given that $W$ is scale invariant – is governed by

$$ $\frac{1}{2}\eta\ddot{\hat{W}}(t)+\dot{\hat{W}}(t)+\Lambda\hat{W}(t)=-\frac{1}{ ||W(t)||}u_{W}(t)$ (10) $$

where we define $\Lambda:=\frac{1}{2}\eta r||\dot{\hat{W}}||^{2}$. Equation 10 has the same form as Equation 7, with weight decay $\lambda=\Lambda$ and updating via $u_{W}/||W||$. With the aforementioned interpretation of $||W||$ as the square-root of a moving-average of $||u_{W}||^{2}$, we can "undo" the coarse-graining assumption (Assumption 2) and view $\hat{W}$ as updating by an adaptive-normalization of $u_{W}$ without additional momentum.

Converting implicit effect into explicit optimizer. This concept of *meta-adaptive normalization*, i.e. adaptive-normalization of an already adaptively-normalized gradient $u_{W}$, can be converted into an explicit optimizer which we call $2$-Adam, with hyperparameters $(\beta_{1:2},\gamma_{1:2})$ and update rule

$$ $2\text{-Adam:}\;\;\;\theta_{n+1}=\theta_{n}-\eta u_{n}^{(2)},\;\;\text{where} \;\;u_{n}^{(2)}:=\frac{\sqrt{1-\beta_{2}^{n+1}}}{1-\gamma_{2}^{n+1}}\frac{m_{n }^{(2)}}{\sqrt{v_{n}^{(2)}}},$ $$

where

$$ $m_{n}^{(2)}:=\gamma_{2}m_{n-1}^{(2)}+(1-\gamma_{2})u_{n},\;\;\;v_{n}^{(2)}:= \beta_{2}v_{n-1}^{(2)}+(1-\beta_{2})u_{n}^{\odot 2}$ $$

are moving-averages – with hyperparameters $(\beta_{2},\gamma_{2})$ – of Adam/AdamW’s update $u_{n}$, which is as defined in Equation 1 but instead with respect to hyperparameters $(\beta_{1},\gamma_{1})$. That is, $2$-Adam applies an adaptive-normalization procedure *twice* to the gradients $g_{0:n}$. We can view a scale-invariant weight trained under Adam as analogous to $2$-Adam with $\gamma_{2}=0$.

###### Assumption 1 .

###### Assumption 2 .

### 3.3 The k 𝑘 k italic_k -Adam optimizer

The $2$-Adam optimizer applies an adaptive-normalization procedure twice with respect to hyperparameters $(\beta_{1:2},\gamma_{1:2})$. The $k$-Adam optimizer extends this concept, applying an adaptive-normalization procedure $k$ times with respect to hyperparameters $(\beta_{1:k},\gamma_{1:k})$, defined by the update rule

$$ $k\text{-Adam:}\;\;\;\theta_{n+1}:=\theta_{n}-\eta u_{n}^{(k)}$ (11) $$

where for $i=1,\ldots,k$,

$$ $u_{n}^{(i)}:=\frac{\tilde{m}_{n}^{(i)}}{\sqrt{\tilde{v}_{n}^{(i)}}},\;\;\; \tilde{m}_{n}^{(i)}:=\frac{1}{1-\gamma_{i}^{n+1}}m_{n}^{(i)},\;\;\;\tilde{v}_{ n}^{(i)}:=\frac{1}{1-\beta_{i}^{n+1}}v_{n}^{(i)}$ $$

$$ $m_{n}^{(i)}:=\gamma_{i}m_{n-1}^{(i)}+(1-\gamma_{i})u_{n}^{(i-1)},\;\;\;v_{n}^{ (i)}:=\beta_{i}v_{n-1}^{(i)}+(1-\beta_{i})(u_{n}^{(i-1)})^{\odot 2}$ $$

with $u_{n}^{(0)}:=g_{n}\equiv\nabla_{\theta}L(\theta_{n})$ and $(m_{-1}^{(i)},v_{-1}^{(i)})=(0,0)$ for all $i=1,\ldots,k$. We present an explicit algorithm for implementing $k$-Adam in Algorithm 1. A more succinct (but equivalent) description of $k$-Adam is presented in Appendix N. Note that $1$-Adam is equivalent to Adam.

Figure: Algorithm 1 $k$-Adam update rule

Update bounds. Theoretical guarantees in the form of bounds – as shown for Adam/AdamW in Section 2 – hold analogously for $k$-Adam. Specifically, since $u_{n}^{(i)}$ is an adaptive-normalization of $u_{0:n}^{(i-1)}$ for each $i=1,\ldots,k$, we have the bound

$$ $||u_{n}^{(i)}||_{\infty}\leq\frac{\sqrt{1-\beta_{i}^{n+1}}}{1-\gamma_{i}^{n+1} }\frac{1-\gamma_{i}}{\gamma_{i}}\sqrt{\frac{\beta_{i}}{1-\beta_{i}}}B_{n}( \beta_{i},\gamma_{i})$ (12) $$

with $B_{n}$ defined as in Equation 5.

Hyperparameter choice. For evaluating $k$-Adam, we will consider 4 different strategies for choosing hyperparameters $(\beta_{1:k},\gamma_{1:k})$,

$$ Inverse exp: $\displaystyle\beta_{i}:=1-(1-\tilde{\beta})^{1/k},$ $\displaystyle\gamma_{i}:=1-(1-\tilde{\gamma})^{1/k}$ Exp: $\displaystyle\beta_{i}:=\tilde{\beta}^{k},$ $\displaystyle\gamma_{i}:=\tilde{\gamma}^{k}$ Scaled: $\displaystyle\beta_{i}:=\tilde{\beta}/k,$ $\displaystyle\gamma_{i}:=\tilde{\gamma}/k$ Naive: $\displaystyle\beta_{i}:=\tilde{\beta},$ $\displaystyle\gamma_{i}:=\tilde{\gamma}$ $$

Figure: Figure 6: Optimization with $k$-Adam optimizer. Plot of the best test accuracy against $k$ after training a CNN for $100$ epochs on CIFAR10 using the $k$-Adam optimizer. We highlight the $k=1$ case in red, corresponding to Adam/AdamW.
Refer to caption: x10.png

for all $i=1,\ldots,k$, where $(\tilde{\beta},\tilde{\gamma}):=(0.999,0.9)$ are the typical Adam/AdamW hyperparameter values. In Appendix M we motivate the inverse exp strategy; the other strategies have been chosen heuristically. We note that for these strategies, $(\beta_{i},\gamma_{i})\in\mathcal{B}_{+}$ for all $i=1,\ldots,k$ (since $\beta_{i}>\gamma_{i}$) hence we expect stable training as a result of Equation 12.

Evaluating $k$-Adam. We motivated interpreting the implicit effect of scale invariance based on its beneficial influence on training (Tanaka & Kunin, ; Lubana et al., ; Santurkar et al., ). We would therefore expect $2$-Adam, and perhaps $k$-Adam for $k>2$, to outperform Adam; is this indeed the case?

We train a CNN (architecture details in Appendix F) on the CIFAR10 dataset for 100 epochs using $k$-Adam, for values $k=1,\ldots,10$ and using each hyperparameter strategy defined above (under both decoupled and coupled weight decay). We display the results in Figure 6, finding that $k$-Adam outperforms Adam/AdamW at various $k>1$ under the first three strategies, however the naive strategy performs particularly badly for $k>2$; we show this in more detail in Appendix M. We find that $2$-Adam with the inverse exp strategy and decoupled weight decay performs best.

Though we find $k=2$ to perform best in this case, it is possible that a larger $k$ may be beneficial in certain contexts, e.g. at small batch sizes (to possibly alleviate noisy gradients). There may also be a more natural, theoretically-motivated hyperparameter strategy compared to the strategies we consider here. We leave a study of these aspects to future work.

## 4 Related work

Previous work has studied adaptive optimization theoretically (Chen et al., ; Li & Orabona, ; Zhou et al., ; Barakat & Bianchi, ; da Silva & Gazeau, ), but these works focus on asymptotic convergence rates and hence do not cover the results presented here. These works perform a discrete-time analysis, except for (da Silva & Gazeau, ) which uses the same differential equation representation of Adam to the one we consider in Section 1, and (Barakat & Bianchi, ) which considers a different differential equation for Adam, though these works do not explicitly solve for Adam’s moving-averages $m(t)$, $v(t)$ or update $u(t)$ in continuous-time as we do in this paper.

There has been previous work on interpreting the effect of scale invariance on training dynamics from a theoretical perspective (Tanaka & Kunin, ; Kunin et al., ; Zhao et al., ; Lyu et al., ; Lubana et al., ). Most relevant to our work is (Tanaka & Kunin, ), which – using a Lagrangian-based approach – shows that scale-invariant weights trained under gradient descent with momentum exhibits an adaptive effect analogous to RMSprop. Since this work does not consider an adaptive learning rate, the theoretical analysis is greatly simplified, however the results are less faithful to practical machine learning where optimizers with adaptive learning rates (i.e. Adam/AdamW) are widely used; our analysis in Section 3 extends this argument to the case of Adam/AdamW, where we observe a meta-adaptive effect. We further use our theoretical insights to motivate a novel optimizer, $k$-Adam, which captures the implicit role of scale invariance.

Various extensions to Adam have been proposed – such as the AMSGrad (Reddi et al., ), LAMB (You et al., ), and Lion (Chen et al., ) optimizers – however they differ from our proposed optimizer $k$-Adam: AMSGrad extends Adam by scaling the learning rate by an adaptive factor $1/\sqrt{\hat{v}_{n}}$ where $\hat{v}_{n}$ is a maximum of all past moving-averages $v_{n}$; LAMB uses layer-wise adaptive learning rates and a trust ratio mechanism; Lion (Chen et al., ) uses the *sign* of a moving-average of the gradient in order to update the parameters. $k$-Adam extends Adam uniquely, applying an adaptive normalization procedure $k$ times in succession, and is motivated by theoretical insight.

## 5 Discussion

In this work we have presented a continuous-time framework for Adam, facilitating a mathematically tractable analysis of training dynamics. We have demonstrated the utilty of this framework for answering questions of practical interest regarding hyperparameter choice and the role of normalization.

We have seen that the derived bound on the max-update (Equation 5) is surprisingly accurate to empirical training dynamics, however our bound does not rigorously justify (only suggests) why we observe a faster rate of generalization for larger $C(\beta,\gamma)$ along a given normal curve, hence there is still progress to be made for a complete understanding of adaptive hyperparameter choice.

In this paper we do not perform a rigorous analysis of $k$-Adam’s performance – i.e. evaluating its performance across various datasets and architectures – as our focus in this paper is a theory-driven analysis of training dynamics. We note that the meta-adaptive effect we describe in this paper is not a complete account of the benefits of normalization layers. Normalization layers possess additional benefits, such as avoiding rank collapse (Daneshmand et al., ; Dong et al., ; Noci et al., ) and contributing towards the smoothness of the loss landscape (Santurkar et al., ; Kohler et al., ; Karakida et al., ; Lyu et al., ); it is unclear whether such phenomena are related to the meta-adaptive effect we describe in this paper.

## 6 Acknowledgements

We thank the CBS-NTT Physics of Intelligence program for their support. We would also like to thank Bo Zhao for helpful discussions on early drafts of this work and Ekdeep Singh Lubana for his valuable comments.

## Appendix A Descriptions of optimizers

Momentum (Polyak, ) features an adaptive gradient direction – updating by a moving-average of the gradient – as described by the update rule,

$$ $\text{Momentum:}\;\;\;\theta_{n+1}:=\theta_{n}-\eta m_{n},$ $$

$$ $m_{n}:=\gamma m_{n-1}+(1-\gamma)g_{n}$ $$

for $n=0,1,2,\ldots$, with moving-average hyperparameter $\gamma$ and $m_{-1}:=0$.

RMSprop (Tieleman & Hinton, ) features an adaptive learning rate – normalizing the update gradient $g_{n}$ by a moving-average of the squared gradient $g_{n}^{\odot 2}$ – as described by the update rule,

$$ $\text{RMSprop:}\;\;\;\theta_{n+1}:=\theta_{n}-\eta\frac{g_{n}}{\sqrt{v_{n}}},$ $$

$$ $v_{n}:=\beta v_{n-1}+(1-\beta)g_{n}^{\odot 2}$ $$

with $v_{-1}:=0$. Note that $g_{n}^{\odot 2}$ describes an element-wise squaring, and the division $g_{n}/\sqrt{v_{n}}$ is also element-wise. Here we have neglected weight decay; we will discuss weight decay strategies (coupled vs decoupled) below.

Adam (Kingma & Ba, ) and its variant AdamW (Loshchilov & Hutter, ) are a combination of momentum and RMSprop, and also use a bias correction factor, which we justify in Appendix B. Neglecting weight decay, Adam & AdamW are equivalent, with update rule

$$ $\text{Adam/AdamW:}\;\;\;\theta_{n+1}:=\theta_{n}-\eta\frac{\tilde{m}_{n}}{ \sqrt{\tilde{v}_{n}}},$ $$

$$ $\tilde{m}_{n}:=\frac{1}{1-\gamma^{n+1}}m_{n},\;\;\;\tilde{v}_{n}:=\frac{1}{1- \beta^{n+1}}v_{n},$ $$

$$ $m_{n}:=\gamma m_{n-1}+(1-\gamma)g_{n},\;\;\;v_{n}:=\beta v_{n-1}+(1-\beta)g_{n }^{\odot 2}$ $$

where a tilde denotes bias correction. The difference between Adam and AdamW comes from how they apply weight decay: Adam uses *coupled* weight decay, equivalent to transforming the loss $L(\theta)\mapsto L(\theta)+\frac{\lambda}{2}||\theta||^{2}$ such that $g_{n}\mapsto g_{n}+\lambda\theta_{n}$, and AdamW uses *decoupled* weight decay, which involves subtracting $\lambda\eta\theta_{n}$ from the RHS of the update rule, i.e.

$$ $\text{AdamW:}\;\;\;\theta_{n+1}:=\theta_{n}-\lambda\eta\theta_{n}-\eta\frac{ \tilde{m}_{n}}{\sqrt{\tilde{v}_{n}}}$ $$

We note that the PyTorch implementation of RMSprop uses coupled weight decay, i.e. RMSprop is Adam with $\gamma=0$.

## Appendix B Motivating bias correction in Adam

Consider an exponential moving average of a sequence $\{x_{0},x_{1},x_{2},\ldots\}$ of tensors,

$$ $\displaystyle y_{n}$ $\displaystyle:=\beta y_{n-1}+(1-\beta)x_{n}$ $\displaystyle=\cdots$ $\displaystyle=(1-\beta)(x_{n}+\beta x_{n-1}+\cdots+\beta^{n}x_{0})$ $$

Consider the stationary case with $\mathbb{E}[x_{n}]$ independent of $n$, then

$$ $\mathbb{E}[y_{n}]=\mathbb{E}[x_{n}](1-\beta)(1+\beta+\cdots+\beta^{n})=\mathbb {E}[x_{n}](1-\beta^{n+1})$ $$

hence if we want an unbiased exponential moving average, we should instead consider a *bias-corrected* form of $y_{n}$:

$$ $\tilde{y}_{n}:=\frac{1}{1-\beta^{n+1}}y_{n}$ $$

such that

$$ $\mathbb{E}[\tilde{y}_{n}]=\mathbb{E}[x_{n}]$ $$

## Appendix C Deriving the parameter update bound

First we will derive a continuous-time expression for $m_{n}$ by Taylor expanding its definition,

$$ $\displaystyle m(t_{n})$ $\displaystyle=\gamma m(t_{n}-\eta^{p})+(1-\gamma)g(t_{n})$ $\displaystyle=\gamma m(t_{n})-\eta^{p}\gamma\dot{m}(t_{n})+(1-\gamma)g(t_{n})+ O(\eta^{2p})$ $$

and so to leading order,

$$ $\dot{m}(t)+\frac{1-\gamma}{\eta^{p}\gamma}m(t)=\frac{1-\gamma}{\eta^{p}\gamma} g(t)$ $$

$$ $\frac{d}{dt}\left[\exp\left(\frac{1-\gamma}{\eta^{p}\gamma}t\right)m(t)\right] =\frac{1-\gamma}{\eta^{p}\gamma}\exp\left(\frac{1-\gamma}{\eta^{p}\gamma}t \right)g(t)$ $$

$$ $\displaystyle\implies m(t)$ $\displaystyle=\frac{1-\gamma}{\eta^{p}\gamma}\int_{0}^{t}d\tau\;\exp\left(- \frac{1-\gamma}{\eta^{p}\gamma}(t-\tau)\right)g(\tau)$ $\displaystyle\equiv\int_{0}^{t}d\tau\;K_{\gamma}(\tau,t)g(\tau)$ $$

Similarly for $v_{n}$,

$$ $\dot{v}(t)+\frac{1-\beta}{\eta^{p}\beta}v(t)=\frac{1-\beta}{\eta^{p}\beta}g(t) ^{\odot 2}$ $$

$$ $\displaystyle\implies v(t)$ $\displaystyle=\frac{1-\beta}{\eta^{p}\beta}\int_{0}^{t}d\tau\;\exp\left(-\frac {1-\beta}{\eta^{p}\beta}(t-\tau)\right)g(\tau)^{\odot 2}$ $\displaystyle\equiv\int_{0}^{t}d\tau\;K_{\beta}(\tau,t)g(\tau)^{\odot 2}$ $$

As a result, the ratio of moving-averages present in Adam’s update $u(t)$ takes the form,

$$ $\frac{m(t)}{\sqrt{v(t)}}=\eta^{-p/2}\frac{1-\gamma}{\gamma}\sqrt{\frac{\beta}{ 1-\beta}}\frac{\int_{0}^{t}d\tau\;\exp\left(-\frac{1-\gamma}{\eta^{p}\gamma}(t -\tau)\right)g(\tau)}{\sqrt{\int_{0}^{t}d\tau\;\exp\left(-\frac{1-\beta}{\eta^ {p}\beta}(t-\tau)\right)g(\tau)^{\odot 2}}}$ (13) $$

Note that this describes an element-wise division of tensors. This expression is particularly amenable to the Cauchy-Schwarz inequality, which says that for (square-integrable) functions $f,g:\mathbb{R}\to\mathbb{R}$,

$$ $\bigg{|}\int dx\;f(x)g(x)\bigg{|}\leq\sqrt{\int dx\;f(x)^{2}\int dx\;g(x)^{2}}$ $$

We can apply the Cauchy-Schwarz inequality (element-wise) to the integral in the numerator of Equation 13 to provide an upper bound,

$$ $\displaystyle\bigg{|}\int_{0}^{t}d\tau\;\exp$ $\displaystyle\left(-\frac{1-\gamma}{\eta^{p}\gamma}(t-\tau)\right)g(\tau)\bigg {|}$ $\displaystyle\equiv\bigg{|}\int_{0}^{t}d\tau\;\left[\exp\left(-\frac{C(t-\tau) }{2\eta^{p}}\right)\right]\left[\exp\left(-\frac{1-\beta}{2\eta^{p}\beta}(t- \tau)\right)g(\tau)\right]\bigg{|}$ $\displaystyle\leq\sqrt{\int_{0}^{t}d\tau\;\exp\left(-\frac{C(t-\tau)}{\eta^{p} }\right)}\sqrt{\int_{0}^{t}d\tau\;\exp\left(-\frac{1-\beta}{\eta^{p}\beta}(t- \tau)\right)g(\tau)^{\odot 2}}$ $\displaystyle=\sqrt{\frac{\eta^{p}}{C}(1-\exp(-Ct/\eta^{p}))}\sqrt{\int_{0}^{t }d\tau\;\exp\left(-\frac{1-\beta}{\eta^{p}\beta}(t-\tau)\right)g(\tau)^{\odot 2}}$ $$

where we have defined

$$ $C:=\frac{2\beta(1-\gamma)-\gamma(1-\beta)}{\beta\gamma}$ $$

Note that the second factor is exactly equal to the denominator in the expression for $m(t)/\sqrt{v(t)}$, and so they cancel to give the following bound:

$$ $\displaystyle\bigg{|}\frac{m(t)}{\sqrt{v(t)}}\bigg{|}$ $\displaystyle\leq\frac{1-\gamma}{\gamma}\sqrt{\frac{\beta}{1-\beta}}\sqrt{ \frac{1-\exp(-Ct/\eta^{p})}{C}}$ $\displaystyle\leq\frac{1-\gamma}{\gamma}\sqrt{\frac{\beta}{1-\beta}}\begin{ cases}1/\sqrt{C},&C>0,\\ \sqrt{t/\eta^{p}},&C=0,\\ \sqrt{\exp(|C|t/\eta^{p})/|C|},&C<0\end{cases}$ $$

We can then obtain an expression for $u_{n}=u(t_{n})$ using $n=t_{n}/\eta^{p}$.

Weight decay. Note that extending Equation 3 to include weight decay is simple: if we want coupled weight decay (as in Adam), we can transform $g(\tau)\mapsto g(\tau)+k\theta(\tau)$; if we want decoupled weight decay (as in AdamW), we can subtract $k\eta\theta_{n}$ from the RHS of Equation 1 and proceed identically.

Differential equation form. We can obtain a differential equation for Adam/AdamW by noting that,

$$ $\theta_{n+1}-\theta_{n}\equiv\theta(t_{n}+\eta^{p})-\theta(t_{n})=\eta^{p}\dot {\theta}(t_{n})+\frac{\eta^{2p}}{2}\ddot{\theta}(t_{n})+O(\eta^{3p})$ $$

by Taylor expansion, which we can then substitute into the update rule $\theta_{n+1}=\theta_{n}-\eta u_{n}$.

Normal curves. The level curve $\mathcal{B}_{c}$ is determined by

$$ $f(\beta,\gamma)=c,\;\;\;\;\;f(\beta,\gamma):=\frac{2\beta(1-\gamma)-\gamma(1- \beta)}{\beta\gamma}\equiv C(\beta,\gamma)$ $$

and normal curves $\gamma=\gamma(\beta)$ satisfy

$$ $\frac{d\gamma(\beta)}{d\beta}=\frac{\partial f/\partial\gamma}{\partial f/ \partial\beta}$ $$

$$ $\implies\gamma(\beta)=(k-2\beta^{3})^{1/3}$ $$

for an integration constant $k$. The normal curve passing through $(\beta_{0},\gamma_{0})$ has integration constant $k=2\beta_{0}^{3}+\gamma_{0}^{3}$.

## Appendix D Empirical analysis of ℬ 0 subscript ℬ 0 \mathcal{B}_{0} caligraphic_B start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT

Equation 5 suggests that $||u_{n}||_{\infty}$ may scale like $\sqrt{n}$ when $(\beta,\gamma)\in\mathcal{B}_{0}$. Do we observe this in practice? In Figure 7 we instead see bounded behaviour analogous to that of the region $\mathcal{B}_{+}$ (as seen in Figure 3), i.e. there is bounded growth rather than a $\sqrt{n}$ scaling in the max-update. We note that the bound Equation 5 is not violated; the predicted upper bound is indeed met, though the prediction is just not *tight* relative to the true dynamics empirically. We suspect that the condition $C(\beta,\gamma)=0$ may be sensitive to numerical precision issues in practice, which is perhaps the reason why we observe behaviour analogous to $C(\beta,\gamma)>0$ instead.

Figure: (a)
Refer to caption: x11.png

## Appendix E Assumptions of Section 3.2

Assumption 1. To verify Assumption 1 we consider the transformer model setup described in Appendix F (same setup Figure 5) with qk-norm (Dehghani et al., ; Gilmer et al., ) enabled, such that the loss is scale invariant under each query and key matrix. In Figure 8 we plot the cosine similarity $\text{sim}(W_{100},g_{n}):=\frac{\braket{W_{100},g_{n}}}{||W_{100}||||g_{n}||}
\in[-1,1]$ for iterations $n=1,\ldots,100$, for each query and key matrix $W$ in the model (a total of 32 matrices). We observe a cosine similarity of $\approx 0$ at $n=100$ (supporting Equation 6) and an increase in cosine similarity as $n$ decreases, however the cosine similarity still remains negligible. We also demonstrate that without qk-norm, the cosine similarity is no longer well behaved.

Figure: (a) With qk-norm
Refer to caption: x13.png

Assumption 2. To verify Assumption 2 we choose 16 random parameter values from random query/key matrices $W$, and plot their trajectory in regular training vs. their trajectory when using coarse-graining on $W$. We display the results in Figure 9, finding a strong agreement.

Figure: Figure 9: Plot of the trajectories, with and without coarse-graining, for 16 randomly selected parameters from query/key matrices.
Refer to caption: x15.png

A consequence of these assumptions is that $\braket{W(t),m_{W}(t)}\approx 0$ (from Assumption 1 and Equation 2), and hence (in combination with Assumption 2) we have $\braket{W(t),u_{W}(t)}\approx 0$. We use this in Section 3.2 in order to obtain an adaptive-normalization interpretation of training dynamics. Note that though the cosine similarity is larger at earlier iterations in Figure 8, the kernel $K_{\gamma}$ present in $m_{W}$ places a lower weight on earlier iterations such that a larger cosine similarity maintains a negligible effect.

## Appendix F Experimental setups

Architecture details. Throughout the paper we consider a nanoGPT model (Karpathy, ) – which is a decoder-only transformer model (Vaswani et al., ) – and a CNN model (Lecun et al., ). The nanoGPT model architecture we use has $6$ layers, with $6$ attention heads per layer. The embedding dimension is $384$. We use a dropout of $0.2$. When training we use the shakespeare dataset associated with the nanoGPT repository (Karpathy, ), using character-level tokenization, for a max input length of $256$ tokens. The CNN model architecture we use is shown in Figure 10 which we use for Figure 6.

Continuous-time trajectory. For experiments that require continuous-time trajectories, we use the method described in Appendix G.

Figure: Figure 10: CNN model architecture used in Figure 6. All convolutional layers use a kernel size of $3\times 3$ and a padding of $1$.

Figure 1. We use the method described in Appendix G – with $K=100$, $p=1$ and $\eta=10^{-3}$ – to numerically solve Equation 4, obtaining a continuous-time trajectory in parameter space. We randomly select 16 parameters from a nanoGPT model (trained using setup described above).

Figure 2. We describe the equation for normal curves in Appendix C.

Figure 3. We train the nanoGPT model using Adam at a learning rate of $10^{-3}$ (and batch size 64) for a range of adaptive hyperparameters $(\beta,\gamma)$ for 1500 iterations, saving the max-update $||u_{n}||_{\infty}$ at every iteration $n$. We compute the slope of $\log||u_{n}||_{\infty}$ at $n=500$ using the values at $n=495$ and $n=505$.

Figure 4. We use the same setup as Figure 3, and additionally use a learning rate warmup (linear, for the first 100 iterations) and after warmup, a cosine decay down to $10^{-4}$. We train for 3000 iterations in order for training to converge such that we can assess generalization performance. We evaluate on a test split of the shakespeare dataset (on 200 iterations worth of test data) every 100 iterations to obtain the test losses.

Figure 5. We use the same setup as Figure 1, with $(\beta,\gamma)=(0.999,0.9)$, training a nanoGPT model for 100 iterations and comparing the predicted value of $||W(t)||^{2}$ (Equation 8), computed via the method described in Appendix G, with the true value. We do this for 16 randomly selected query/key matrices.

Figure 6. We train a CNN (using architecture in Figure 10) on CIFAR10 for 100 epochs and for each $k\in\{1,\ldots,10\}$, we sweep over learning rates $\eta\in\{\text{6e-5, 1e-4, 3e-4, 6e-4, 1e-3, 3e-3}\}$ and weight decays $\lambda\in\{\text{1e-6, 1e-5, 1e-4, 1e-3, 1e-2}\}$ and plot the best test accuracy across this sweep. We use a linear warmup for the first 2 epochs, and a cosine decay to $\eta/10$ for the remaining epochs. We evaluate on the test split every 5 epochs. We use $(\tilde{\beta},\tilde{\gamma})=(0.999,0.9)$ and a batch size of 512.

## Appendix G Numerically solving Equation 4

Numerically solving for continuous dynamics. Defining $\psi(t):=\dot{\theta}(t)$, we can write Equation 4 as two coupled first-order DEs,

$$ $\dot{\theta}(t)=\psi(t)$ $$

$$ $\dot{\psi}(t)=-\frac{2}{\eta^{2p}}\left(\lambda\eta\theta(t)+\eta^{p}\psi(t)+ \eta\frac{\sqrt{1-\beta^{1+t/\eta^{p}}}}{1-\gamma^{1+t/\eta^{p}}}\frac{m(t)}{ \sqrt{v(t)}}\right)$ $$

and numerically solve via Euler’s method,

$$ $\theta((n+1)\Delta t)\approx\theta(n\Delta t)+\Delta t\psi(n\Delta t)$ $$

$$ $\text{with}\;\;\;\psi(n\Delta t)\approx\psi((n-1)\Delta t)+\Delta t\dot{\psi}( (n-1)\Delta t)$ $$

defining $\Delta t:=\eta^{p}/K$. In our numerical validation experiments, we use the timescale $p=1$ and $K=100$.

To compute $\dot{\psi}(n\Delta t)$ via the above DE we must compute

$$ $\displaystyle m(t)$ $\displaystyle=\int_{0}^{t}d\tau\;M(\tau,t)g(\tau)=\int_{0}^{t-\Delta t}d\tau\; M(\tau,t)g(\tau)+\int_{t-\Delta t}^{t}d\tau\;M(\tau,t)g(\tau)$ $\displaystyle=\exp\left(-\frac{1-\gamma}{\eta^{p}\gamma}\Delta t\right)\left[ \underbrace{\int_{0}^{t-\Delta t}d\tau\;M(\tau,t-\Delta t)g(\tau)}_{\text{from $n-1$th update step}}+\underbrace{\int_{t-\Delta t}^{t}d\tau\;M(\tau,t-\Delta t )g(\tau)}_{\approx\Delta t\frac{1-\gamma}{\eta^{p}\gamma}g(t-\Delta t)}\right]$ $$

Empirically validating $\varphi(t)$. We note that we can write Equation 8 approximately as

$$ $\varphi(n\Delta t)\approx\varphi_{0}e^{-2\lambda n\Delta t}+\Delta t\sum_{m=0} ^{n-1}e^{-2\lambda(n-m)\Delta t}f(m\Delta t)$ $$

where

$$ $f(t):=\eta||u_{W}(t)||^{2}-2\braket{W(t),u_{W}(t)}$ $$

We then compare this expression to the true value of $||W(n\Delta t)||^{2}$ to validate Equation 8 in Figure 5.

## Appendix H Implicit meta-adaptive effect

Deriving an expression for $||W(t)||^{2}$. Applying the inner product $\braket{W(t),\cdot}$ to both sides of Equation 7, we arrive at the differential equation

$$ $\frac{1}{2}\eta\ddot{\varphi}+\dot{\varphi}+2\lambda\varphi\approx\eta||\dot{W }||^{2}-2\braket{W,u_{W}}$ $$

where we have defined the squared norm $\varphi(t):=||W(t)||^{2}$. One can show (see further below) that to leading order, $||\dot{W}||^{2}\approx\left\|u_{W}\right\|^{2}$ and $\eta\ddot{\varphi}=O(\eta^{2})+O(\lambda\eta)$, hence we neglect the $\ddot{\varphi}$ term, with the resulting first-order differential equation,

$$ $\dot{\varphi}+2\lambda\varphi\approx\eta||u_{W}||^{2}-2\braket{W,u_{W}}$ $$

which when integrated, gives

$$ $||W(t)||^{2}\approx||W(0)||^{2}e^{-2\lambda t}+\int_{0}^{t}d\tau\;e^{-2\lambda (t-\tau)}(\eta||u_{W}(\tau)||^{2}-2\braket{W(\tau),u_{W}(\tau)})$ $$

$||W(t)||^{2}$ as a moving average. Restating Equation 9,

$$ $||W(t)||^{2}\approx||W(0)||^{2}e^{-2\lambda t}+\eta\int_{0}^{t}d\tau\;e^{-2 \lambda(t-\tau)}||u_{W}(\tau)||^{2}$ $$

which takes an analogous form to a moving-average of $||u_{W}||^{2}$ in continuous-time. Specifically, consider a moving-average of $(x_{0},x_{1},x_{2},\ldots)$,

$$ $w_{n}=(1-\alpha)w_{n-1}+\alpha x_{n}$ $$

Then in continuous time, with $w_{n}=w(t_{n})$ and $x_{n}=x(t_{n})$, we have shown in Appendix C that

$$ $w(t)=w(0)\exp\left(-\frac{1-\alpha}{\eta^{p}\alpha}t\right)+\frac{1-\alpha}{ \eta^{p}\alpha}\int_{0}^{t}d\tau\;\exp\left(-\frac{1-\alpha}{\eta^{p}\alpha}(t -\tau)\right)x(\tau)$ $$

Comparing this expression to Equation 9 allows us to view $||W||^{2}$ as a moving-average of $||u_{W}||^{2}$.

Weight direction dynamics. To obtain Equation 10 from Equation 7, we write $W=r\hat{W}$ with $r:=||W||$ (and $||\hat{W}||=1$), which after substituting into Equation 7 becomes

$$ $(\frac{1}{2}\eta\ddot{r}+\dot{r}+\lambda r)\hat{W}+(\eta\dot{r}+r)\dot{\hat{W} }+\frac{1}{2}\eta r\ddot{\hat{W}}=-u_{W}$ (14) $$

Note that taking the derivative of $\braket{\hat{W},\hat{W}}=1$ implies $\braket{\dot{\hat{W}},\hat{W}}=0$, and further $\braket{\ddot{\hat{W}},\hat{W}}=-||\dot{\hat{W}}||^{2}$. As a result, applying $\braket{\cdot,\hat{W}}$ to Equation 14 gives

$$ $\frac{1}{2}\eta\ddot{r}+\dot{r}+\lambda r=\frac{1}{2}\eta r||\dot{\hat{W}}||^{2}$ $$

Substituting this expression back into Equation 14 and neglecting higher-order terms (noting that from above, $\eta\dot{r}=O(\eta^{2})$), we arrive at

$$ $\frac{1}{2}\eta\ddot{\hat{W}}+\dot{\hat{W}}+\Lambda\hat{W}=-\frac{1}{r}u_{W}$ $$

where we define $\Lambda:=\frac{1}{2}\eta r||\dot{\hat{W}}||^{2}$.

Leading order approximations. We can write

$$ $\dot{W}=-\lambda W-\frac{1}{2}\eta\ddot{W}-u_{W}$ $$

Squaring both sides and using $\braket{W,u_{W}}\approx 0$ (a consequence of Assumption 1 and Assumption 2) results in

$$ $||\dot{W}||^{2}=||u_{W}||^{2}+\lambda\eta\braket{W,\ddot{W}}+\eta\braket{\ddot {W},u_{W}}+\frac{\eta^{2}}{4}||\ddot{W}||^{2}+\lambda^{2}||W||^{2}$ $$

hence to leading order in $(\eta,\lambda)$, $||\dot{W}||^{2}\approx||u_{W}||^{2}$.

## Appendix I Symmetries of the transformer architecture

Say $\Phi_{s}:\Theta\to\Theta$ describes some continuous invariance parameterised by some continuous parameter $s\in\mathbb{R}^{S}$, i.e.

$$ $L(\Phi_{s}(\theta))=L(\theta)\;\;\;\forall\;s\in\mathbb{R}^{S}$ $$

with $\Phi_{0}(\theta)=\theta\;\forall\;\theta\in\Theta$. Differentiating the above with respect to $s_{i}$ and evaluating at $s=0$, we obtain

$$ $\braket{\frac{\partial\Phi_{s}}{\partial s_{i}}\bigg{missing}}{{}_{s=0},\nabla _{\theta}L(\theta)}=0\;\;\;\;\;\forall\;i=1,\ldots,S$ (15) $$

Scale invariance. In the special case of scale invariance with respect to a weight $W$, if we write $\theta=(W,\phi)$ then the relevant transformation is $\Phi_{s}(\theta)=(sW,\phi)$, and Equation 15 becomes,

$$ $\braket{W,\nabla_{W}L(\theta)}=0$ $$

## Appendix J Max-update trajectory in ℬ + subscript ℬ \mathcal{B}_{+} caligraphic_B start_POSTSUBSCRIPT + end_POSTSUBSCRIPT

To visualize the theoretical bounds in the case of $(\beta,\gamma)\in\mathcal{B}_{+}$ in more detail, we plot the results of Figure 3(a) but restricted to $(\beta,\gamma)\in\mathcal{B}_{+}$ in Figure 11. We indeed see that the theoretically predicted bounds (dotted lines) are satisfied.

Figure: (a)
Refer to caption: x16.png

## Appendix K Supporting data for Section 2.3

In Section 2.2 we found that a larger value of $C(\beta,\gamma)$ along a normal curve correlated with a faster generalization. We further verify this in Figure 12 by instead taking points along the normal curve through $(0.95,0.9)$, but otherwise an identical experimental setup to the one used to produce Figure 4, finding similar results.

Figure: (a)
Refer to caption: x17.png

We also consider the level curve $\mathcal{B}_{C(\tilde{\beta},\tilde{\gamma})}$, the results of which we plot in Figure 13. We see a slower rate of generalization for points closer to the origin, as expected, given that such points will have a smaller value of $C(\beta,\gamma)$ relative to other points along their normal curve.

Figure: (a)
Refer to caption: x20.png

## Appendix L Attempting to correct exponential growth with learning rate annealing

Figure: (a)
Refer to caption: x23.png

Since the parameter update is $\theta_{n+1}-\theta_{n}\equiv-\eta u_{n}$, it appears that a well chosen step-dependent learning rate $\eta_{n}$ may correct the exponential growth in the case of $C(\beta,\gamma)<0$ described by Equation 5, which occurs via the factor $\exp(n|C(\beta,\gamma)|/2)$. Particularly, given the tightness of the bound seen in Figure 3(b), we could consider the learning rate

$$ $\eta_{n}=\eta_{0}\exp(-n|C(\beta,\gamma)|/2)$ (16) $$

at step $n$, which results in a bound on $||\theta_{n+1}-\theta_{n}||_{\infty}\equiv\eta_{n}||u_{n}||_{\infty}$ that is equivalent to the case of $C(\beta,\gamma)>0$ with learning rate $\eta_{0}$. However when we run this learning rate annealing strategy (Equation 16) in practice, we find that it is too strong, as shown in Figure 14; the max-update becomes very small, and we do not observe good generalization performance. Hence, even though this annealment strategy for $\mathcal{B}_{-}$ results in the regions $\mathcal{B}_{+}$ and $\mathcal{B}_{-}$ possessing the same bounds on $||\theta_{n+1}-\theta_{n}||_{\infty}$, their training properties are still quite different. Analysis of weaker annealing methods are out of scope for this work. It is unclear whether correcting such exponential growth in $\mathcal{B}_{-}$ by choosing an appropriate annealing method would provide any benefits to simply choosing $(\beta,\gamma)\in\mathcal{B}_{+}$.

## Appendix M Choice of k 𝑘 k italic_k -Adam hyperparameters

Figure: Figure 15: Zoomed out version of Figure 6 to show performance of the naive strategy.
Refer to caption: x27.png

The moving-average $m_{n}^{(k)}$ will have a prefactor of $(1-\beta_{1})\cdots(1-\beta_{k})$, and we can consider a uniform strategy $\beta_{1}=\cdots=\beta_{k}$ and enforce that this prefactor is equal to the typical prefactor for Adam/AdamW, i.e.

$$ $(1-\beta_{i})^{k}=1-\tilde{\beta}$ $$

$$ $\implies\beta_{i}=1-(1-\tilde{\beta})^{1/k}$ $$

Naive strategy. In Figure 15 we show a zoomed-out version of Figure 6 to show the performance of the naive strategy across all $k$, emphasizing the importance of well-selected $k$-Adam hyperparameters.

We also found the strategy $\beta_{i}=\tilde{\beta}^{1/k}$, $\gamma_{i}=\tilde{\gamma}^{1/k}$ to perform badly.

## Appendix N Adaptive normalization and k 𝑘 k italic_k -Adam

For a sequence of tensors $x_{0:n}\equiv(x_{0},\ldots,x_{n})$, we define the bias-corrected moving-average operator $M_{\gamma}$,

$$ $M_{\gamma}(x_{0:n}):=\frac{1-\gamma}{1-\gamma^{n+1}}(x_{n}+\gamma x_{n-1}+ \cdots+\gamma^{n}x_{0})$ $$

We justify bias correction in Appendix B. We then define the *adaptive-normalization* operator $A_{\gamma,\beta}$,

$$ $A_{\gamma,\beta}(x_{0:n}):=\frac{M_{\gamma}(x_{0:n})}{\sqrt{M_{\beta}(x_{0:n}^ {\odot 2})}}$ $$

Adaptive optimization. We can describe Adam/AdamW with hyperparameters $(\beta,\gamma)$ simply as

$$ $\text{Adam/AdamW:}\;\;\;\theta_{n+1}=\theta_{n}-\eta A_{\gamma,\beta}(g_{0:n})$ $$

i.e. updating by the adaptive-normalization of gradients $g_{0:n}$, where $g_{i}\equiv\nabla_{\theta}L(\theta_{i})$.

$k$-Adam extends Adam/AdamW by applying the adaptive normalization operator $k$-times with hyperparameters $(\beta_{1:k},\gamma_{1:k})$,

$$ $k\text{-Adam:}\;\;\;\theta_{n+1}=\theta_{n}-\eta(A_{\gamma_{k},\beta_{k}}\circ \cdots\circ A_{\gamma_{1},\beta_{1}})(g_{0:n})$ $$

which is equivalent to Equation 11. Specifically, the intermediate updates take the form $u_{n}^{(i)}\equiv(A_{\gamma_{i},\beta_{i}}\circ\cdots\circ A_{\gamma_{1},\beta
_{1}})(g_{0:n})$ (meaning that $u_{n}^{(i+1)}$ is an adaptive-normalization of $u_{n}^{(i)}$).

Continuous-time properties. We move into continuous-time, with $x_{n}\equiv x(t_{n})$ where $t_{n}=n\eta^{p}$. Appendix C has shown that

$$ $M_{\gamma}(x_{0:n})=\int_{0}^{t_{n}}d\tau\;K_{\gamma}(\tau,t)x(\tau),\;\;\;\; \;K_{\gamma}(\tau,t):=\frac{1-\gamma}{\eta^{p}\gamma}\exp\left(-\frac{1-\gamma }{\eta^{p}\gamma}(t-\tau)\right)$ $$

to first-order. Then using Cauchy-Schwarz on $A_{\gamma,\beta}$, Appendix C has also shown that

$$ $||A_{\gamma,\beta}(x_{0:n})||_{\infty}\leq\frac{\sqrt{1-\beta^{n+1}}}{1-\gamma ^{n+1}}\frac{1-\gamma}{\gamma}\sqrt{\frac{\beta}{1-\beta}}B_{n}(\beta,\gamma)$ (17) $$

where $B_{n}$ is defined as in Section 2.1. Note that this bound is in general independent of the elements $x_{0:n}$. Since each $k$-Adam sub-update $u_{n}^{(i)}=A_{\gamma_{i},\beta_{i}}((A_{\gamma_{i-1},\beta_{i-1}}\circ\cdots
\circ A_{\gamma_{1},\beta_{1}})(g_{0:n}))$, this bound holds analogously for $u_{n}^{(i)}$.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: - It is unclear how to get equation 7. 

-  can you clarify what does META-ADAPTIVE EFFECT mean? 

- can you clarify why the exponential moving average ($\|W\|$ vs $\|u_W\|$) is significant and surprise ? 

- I don't understand this "η  ̈φ = O(η2) + O(λη)," , and why you can drop this term, which hasn't mentioned in the main text.
soundness: 2
strengths: * The paper proposes a framework to approximate the parameter dynamics of Adam(W), which could be potentially very useful.
* This framework can help estimate the stable region of hyperparameters $(\beta, \gamma)$, which can be helpful in practice for hyperparameter selection.
* The framework is relatively easy to understand.
summary: This paper presents a continuous-time formulation of Adam(W). The authors derive the continuous-time ODE to approximate the parameter dynamics of Adam(W) via a Taylor expansion. The main application of the theory is to predict the stable region of hyperparameters $(\beta, \gamma)$ in which the Adam update won't blow up. The authors also use a similar method to study the dynamics of scale-invariant networks and identify the so-called "meta-adaptive effect," which needs further clarification.

However, I am a bit suspicious regarding the theoretical results of the paper. Most of the derivations are non-rigorous and hand-wavy.  Furthermore, the experimental results lack diversity and scale.

Overall, I think the research direction is interesting and the results can be quite useful. However, the theoretical derivations lack rigor and clarity. I am afraid I don't feel comfortable to support acceptance of the paper in its current form.


------- update -----------
thanks for addressing some of the concerns and updating the paper + experiments; 
i raised my score to 5; but i echoed Reviewer GyMg's point, the paper would benefit from further refinement and resubmission.
weaknesses: - The whole theoretical framework is non-rigorous and based on flawed derivations/assumptions. 
    - I don't understand the crucial deviation in Section C, and it should be put in the main text. 
    - First, why is $m$ differentiable, why can you drop the higher order (second derivative is bounded?) of it?
   -  Second, why is it $g(t_n)$ not $g(t_n - \eta^p)$? In the latter setting, you would need to have an extra term $g'(t_n)$ and an error term $g''(t_n)$?
   - I think these are crucial questions, on which the whole paper is based. I also don't know why we should expect a continuous formation of the dynamics of $m$ in the first place.

* The experimental results are not convincing enough. I would expect larger scale + more diversity experiments, given the theoretical results are not rigorous. 
   - The training steps ($n=100$ or $1500$) are too few. should consider $n= O(10k)$
   - The scale and diversity of the experiments are limited (a small transformer on Shakespeare). I would expect the authors to run experiments on a more realistic dataset, like a 20M-100M parameter model on a subset of C4 or similar.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - How tight is the bound of eq 5? Is there a corresponding lower bound?
- Can you add more experimental verification of Adam-k?
soundness: 3
strengths: - The paper is well written and crisp.
- Continuous time analysis of Adam is interesting and novel (AFAIK, but I’m not an expert). Since Adam is an important algorithm, any analysis of it can be impactful.
summary: The author presents a continuous time analysis of Adam. They first present the continuous version of Adam and a corresponding differential equation (eq 4). With this formulation they present a bound on the parameter updates (eq 5) which depends on the adam hyperparamters. They find that this bound holds reasonably well in practice, which provides some justification for why certain Adam hyperparameters work well in practice. The authors also motivate why Adam+normalization behaves like adam applied multiple times, a method they call Adam-k. The authors provide some small-scale experiments with Adam-k, showing benefits for training on CIFAR10.
weaknesses: - Assumptions 1 and 2 are not very intuitive, at least not to me.
- CIFAR10 results are not very convincing, it’s small scale.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See Weaknesses.
soundness: 2
strengths: - **Originality:** The authors present a novel and quantitatively precise stability region for Adam's betas in the non-stochastic, continuous-time regime, filling a gap in the literature. While several papers discuss Adam’s performance under stochastic gradients, this work focuses on a deterministic model and provides insights under simplifying assumptions.
  
- **Quality:** The theoretical contributions are sound and grounded in ODE analysis. The derivation of a stability region for Adam is supported by both mathematical rigor and experimental validation. The authors maintain an interesting balance between theory and practical insight, though certain aspects (such as stochastic gradients) are left for future exploration.

- **Clarity:** The paper communicates its main theoretical contribution—the stability region for the betas—relatively clearly, though some sections (like the explanation of k-Adam and the role of normalization) could be refined for better comprehension.

- **Significance:** The work provides potentially valuable insights into the dynamics of Adam in a deterministic setting. If this stability region has practical relevance beyond full-batch training, it could offer a useful guideline for tuning $\beta_1$ and $\beta_2$. Additionally, the introduction of k-Adam as a new optimizer inspired by the theoretical results opens avenues for further exploration.
summary: This paper studies Adam through the lens of Ordinary Differential Equations (ODEs). The authors utilize this continuous-time model to reveal aspects of Adam's dynamics in the **full-batch setup, i.e. no stochasticity in the gradients**. In particular, they provide theoretical arguments to identify a stability region for the hyperparameters $\beta_1$ and $\beta_2$, which characterizes the maximum size of the Adam updates. Importantly, they provide some experimental validation of their claims and offer a conjecture that correlates generalization to this region. Finally, they study the beneficial impact of scale-invariant architectural components and leverage their insights to introduce a new optimizer, dubbed k-Adam.
weaknesses: Since there is no predefined space to provide a "General Comment", I provide it here.

**Overall Comment:**  
The paper addresses interesting and mathematically sound questions, but the manuscript quality is lacking in several areas. It needs substantial improvement, and I recommend **rejecting** the paper, with strong encouragement to resubmit once the following issues have been addressed. In particular, I highlight: i) Missing discussion of literature; ii) Missing discussion on contributions: What is the final product? How can this help practitioners? iii) Better visualizations are needed; iv) Unclear whether the continuous-time model was NECESSARY over the discrete-time setup; Unclear discussion regarding normalization; k-Adam does not seem to be properly justified nor properly tested.

I will now provide a "Detailed Feedback" that covers some weak points and incorporates some questions:

**Detailed Feedback:**

The following points are presented in the same order as they appear in the paper. Unfortunately, due to format alterations in the PDF, I cannot reference line numbers, but I will do my best to highlight clear landmarks:

1. **Literature Review:**  
   A review of the literature on continuous-time models for optimizers is crucial, with a particular focus on those that cover Adam. Notably, there is no citation of Malladi et al., who derive an SDE model for Adam, thus addressing the stochastic aspect, which is lacking in this paper.

2. **Limitations Section:**  
   A brief section outlining the limitations of the approach would be useful. Additionally, a subsection comparing the method to existing literature could provide more clarity.

3. **Merging Contributions:** In "1. Continuous-time formulation (Section 2)", why not merge these two points into a simpler one? It seems to me that you simply derived an ODE for Adam, which also entails the derivation of Adam's update.

4. **Better explanation and context:** Your main contribution seems to be the stability region for the betas. Parameter sets inside that region imply that the maximum increment of Adam is controlled, while outside of it, an exponential explosion is "predicted". i) 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"? ii) Is it beneficial to avoid such explosions? Adam is often successfully used even when loss spikes are observed. iii) RMSprop and SignSGD are subcases of Adam: Are they covered by this framework? It seems that RMSprop is always in the stable region, so adding momentum to RMSprop might make it unstable. This is a peculiar observation, as momentum is usually thought of as a stability enhancer.

5. Is it possible that betas in the instability region could benefit from better tuning of the learning rate? I conducted a small experiment on a simple convex quadratic and identified your regions experimentally (via visual inspection, not by plotting your lines). I tested this with learning rates spanning three orders of magnitude and 100 values for each beta: 30,000 small runs. I observed that the stability region expands as the learning rate drops. Could you comment on this? Including a graphical representation for a cheap setup like the one described above would be helpful. If you need the details of my experiment, feel free to ask. Expanding this toy example to the stochastic setting while varying noise levels is also crucial: Could noise shrink the stability region? How does weight decay alter this region?

6. When discussing contribution (f), as well as later in the text, the exact contribution is unclear. Since this is a main point, it must be crystal clear.

7. Regarding contribution (g), it is unclear what advantages k-Adam has. While a high-level description is fine, more details are needed about k-Adam’s advantages, both theoretically and empirically.

8. **Eq. 2:** These formulas for $m$ and $v$ are simply the continuous-time versions of their discrete counterparts, which are already known. Are these surprising in any way? Do they reveal something that the discrete-time version does not?

9. **Eq. 4:** Is your second-order ODE equivalent to the formulation in terms of three first-order ODEs (see Eq. 3.3 of Barakat)? Is there an intrinsic advantage to this point of view?

10. **Eq. 4:** It seems that even if the RHS (e.g., $u(t)$) is set to 0, Adam would still follow a non-trivial second-order ODE that can likely be solved explicitly. This suggests some oscillatory dynamics. Could you elaborate? It seems odd that parameters would evolve even if Adam’s increments are 0.

11. **Figure 1:** I suggest using a color-blind-friendly palette and adding line markers. Rather than comparing the trajectories of some weight, you could calculate the average squared error between the entire trajectory of the real weights and that of the ODE and plot that instead. This would give a more global intuition. I also suggest plotting for longer time and different values of $p$.

12. **Max-update bound:** As highlighted earlier, elaborating more on RMSprop and SignSGD is key.

13. **Eq. 5:** Is this bound independent of the learning rate and $p$? My experiments suggest that the stability region enlarges with a smaller learning rate. Reorganizing this as a well-stated proposition might help.

14. **"We highlight that this bound is only possible because...":** Specify that the reason why you can do this is actually because you do not use stochastic gradients. Also, was it absolutely necessary to use the continuous-time model for this bound, or could it have been derived from the discrete-time version expressions of $m$ and $v$?

15. **Figure 3:** The plots could benefit from reordering. Panel (c) is described before (a) in the caption. Clear legends and a color-blind palette would be helpful.

16. **Section 3.3:** Of course, if you leave the stability region, generalization will likely worsen. However, inside the region, is there a "monotonicity" effect? Is it the case that the further inside the region, the better the generalization? If so, can one always get some guidance in selecting the betas? This is a good point for future research.

17. **Section 4.2:** I’ve reread this multiple times, and it is unclear what the message of this section is. If there is a clear takeaway, I suggest highlighting it. Frame technical results as Lemmas, minimize references to the appendix by grouping them and provide a clear paragraph interpreting the results. Experimental validation is already present, which is good.

18. **k-Adam:** I struggle to see the takeaway from this discussion. First, I suggest removing the 2-Adam discussion currently found before Section 4.3. Otherwise, simply incorporate it into the following section where you properly describe k-Adam. Then, I question the purpose of i) introducing a new optimizer based on loosely justified intuition from Section 4.2, ii) not highlighting clearly its advantages over Adam, and iii) not evaluating it on state-of-the-art experiments. I wonder whether this should be in the main paper as a major contribution.

19.  Finally, this analysis is conducted in a deterministic setting: How can it be generalized to cover stochastic gradients? I expect noise to interact in a non-trivial way with all the moving elements of these analyses.

**Of course**, I reserve the right to alter my score based on the discussion we will have during the rebuttal period.

Malladi et al: On the SDEs and Scaling Rules for Adaptive Gradient Algorithms.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: I suggest that the authors adjust their template since I do not find the indexes of lines in their manuscript.
soundness: 2
strengths: This paper contains some interesting empirical results.
summary: In this paper, the authors propose continuous ode approximations for the first and second-order moments $m_t$ and $v_t$. Based on these continuous ode approximations, the authors provide a theoretical guarantee for the stability of the update of Adam. Besides, the authors also demonstrate that the iterates of Adam is an exponential averaging of $m_t/\sqrt{v_t}$ when applied to a scale-invariant architecture. Based on this property, the authors propose a new variant of Adam with a normalization layer.
weaknesses: 1. This paper lacks sufficient novelty. Specifically, the conclusion that the update of Adam is bounded given specific exponential moving average rates—one of the major technical contributions of this work—has already been discussed in several prior studies [1, 2, 3, 4], utilizing the Cauchy-Schwarz inequality. Additionally, reference [2] provides a tight bound for the accumulated updates of Adam. Moreover, similar techniques involving continuous approximation for the momentum terms $m_t$ and $v_t$ have also been explored in [5, 6].

2. This paper lacks a thorough discussion of previous works. [6] provides a comprehensive analysis of Lion-K algorithms by constructing an ODE approximation of the dynamic system and establishing a Lyapunov function for this ODE system. By the specific pattern of  the Lyapunov function, they provide an interesting and novel analysis of each hyperparameter in the optimization problem. I believe this contribution is worth highlighting.

3. The assumptions presented in the paper are relatively problematic. Regarding assumption 4.1, while I agree that high-dimensional vectors tend to be nearly orthogonal, the statement that they are "approximately 0" does not imply they are exactly 0. For a rigorous mathematical derivation, the authors should provide a bound for this term and explain why it can be considered negligible in relation to the other terms, rather than simply omitting it as they have done in the current proof. As for assumption 4.2, it is quite unusual. If this assumption is valid, it implies that each coordinate of $v_t$ is equal, which results in the absence of a coordinate-wise adaptive learning rate effect. Under such an assumption, I fail to see any significant distinction between gradient descent with momentum and 'Adam'. 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.

4. The scaling of Euler's approximation is inconsistent in this paper. In the derivation of ode approximation of $m_t$ and $v_t$, the authors choose the $\eta^p$ as the stepsize and omit all terms containing high order terms, while in the derivation of formula (4), the authors remain a term with coefficient $\eta^{2p}$.

[1] Zou, D., Cao, Y., Li, Y. and Gu, Q. (2023). Understanding the generalization of adam in learning neural networks with proper regularization. In The Eleventh International Conference on Learning Representations, ICLR 2023.

[2] Xie, S. and Li, Z. (2024). Implicit bias of adamw: ℓ∞ norm constrained optimization.

[3] Zhang, C., Zou, D. and Cao, Y. (2024). The Implicit Bias of Adam on Separable Data. arXiv preprint arXiv:2406.10650.

[4] Hong, Y. and Lin, J. (2024). On Convergence of Adam for Stochastic Optimization under Relaxed Assumptions. arXiv preprint arXiv:2402.03982.

[5] Wang, B., Meng, Q., Chen, W. and Liu, T.-Y. (2021). The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning. PMLR.

[6] Chen, L., Liu, B., Liang, K. et al. (2023). Lion secretly solves a constrained optimization: As lyapunov predicts. In The Twelfth International Conference on Learning Representations.

[7] Balles, L. and Hennig, P. (2018). Dissecting adam: The sign, magnitude and variance of stochastic gradients. In Proceedings of the 35th International Conference on Machine Learning,

Meta-review: This paper introduces a continuous-time formulation of AdamW, deriving a stability region for its hyperparameters and exploring the role of normalization layers, leading to the proposal of k-Adam. While the approach offers a novel perspective on Adam's dynamics and suggests practical guidelines for hyperparameter selection, the theoretical framework relies on non-rigorous derivations and lacks strong experimental validation, particularly regarding the proposed k-Adam optimizer.  Given these weaknesses, the paper's contributions are not sufficiently robust to warrant acceptance in its current form.

Decision: Reject