AdamW's Weight RMS Asymptotic Estimate (Part 1)

AdamW的Weight RMS的渐近估计（上） 苏剑林 · 2025-10-01 · Mathematics estimation gradient optimizer mean field

In 《Why Is Adam's Update RMS 0.2?》, we used a mean-field approximation to estimate Adam's Update RMS. Not long afterward, reader @EIFY pointed out that the same result had already appeared in the paper 《Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks》. After reading it, I found that it contains not only an estimate of the Update RMS, but also an estimate of the Weight RMS.

In other words, for a model trained with AdamW, the RMS of its weights can be estimated asymptotically ahead of time. Does anyone else find this conclusion a bit surprising? I certainly did the first time I saw it — intuitively, the magnitude of the weights is something the model learns from the training data on its own, yet it turns out this magnitude is already implicitly determined by the optimizer's hyperparameters, which feels quite counterintuitive.

In this post, we'll again use the mean-field approximation method to reproduce the asymptotic estimate for the Weight RMS. more

A Shift in Perspective

Let's first recall AdamW's update rule:

\begin{equation}\text{Adam}\color{skyblue}{\text{W}}:=\left\{\begin{aligned} &\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\ &\boldsymbol{v}_t = \beta_2 \boldsymbol{v}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t^2\\ &\hat{\boldsymbol{m}}_t = \boldsymbol{m}_t\left/\left(1 - \beta_1^t\right)\right.\\ &\hat{\boldsymbol{v}}_t = \boldsymbol{v}_t\left/\left(1 - \beta_2^t\right)\right.\\ &\boldsymbol{u}_t =\hat{\boldsymbol{m}}_t\left/\left(\sqrt{\hat{\boldsymbol{v}}_t} + \epsilon\right)\right.\\ &\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}) \end{aligned}\right.\end{equation}

Again, note that all bold symbols here are by default vectors of dimension $\mathbb{R}^d$, and multiplication/division of vectors (including squaring and square roots) is by default element-wise Hadamard product/quotient.

As in 《Why Is Adam's Update RMS 0.2?》, we consider $t\to\infty$ (for $\beta_1,\beta_2$) and $\epsilon\to 0$, so that $\boldsymbol{u}_t=\boldsymbol{m}_t/\sqrt{\boldsymbol{v}_t}$. Let's first consider the case where $\eta_t,\lambda_t$ are all constants, so their subscripts can be omitted, and let us write $\beta_3 = 1-\eta\lambda$. We then have

\begin{equation}\boldsymbol{\theta}_t = \beta_3\boldsymbol{\theta}_{t-1} + (1-\beta_3)(-\boldsymbol{u}_t/\lambda)\label{eq:ema-wd}\end{equation}

This formula shows that we can understand Weight Decay from the perspective of an Exponential Moving Average (EMA) of the update. This is a very meaningful change of perspective, and it underlies works such as 《How to set AdamW's weight decay as you scale model and dataset size》 and 《Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training》.

Weighted Averages

Based on Equation $\eqref{eq:ema-wd}$, we can expand $\boldsymbol{\theta}_t$ into a weighted-average form:

\begin{equation}\boldsymbol{\theta}_t = \beta_3^t\boldsymbol{\theta}_0 + (1-\beta_3)\sum_{i=1}^t \beta_3^{t-i} (-\boldsymbol{u}_i/\lambda)\label{eq:theta-t}\end{equation}

Similarly, $\boldsymbol{m}_t$ and $\boldsymbol{v}_t$ can also be expanded as

\begin{equation}\boldsymbol{m}_t = (1 - \beta_1)\sum_{i=1}^t \beta_1^{t-i}\boldsymbol{g}_i,\qquad \boldsymbol{v}_t = (1 - \beta_2)\sum_{i=1}^t \beta_2^{t-i}\boldsymbol{g}_i^2\label{eq:mv-roll}\end{equation}

There's a small detail here: in the expression for $\boldsymbol{\theta}_t$ we kept the term $\boldsymbol{\theta}_0$, but in the expressions for $\boldsymbol{m}_t$ and $\boldsymbol{v}_t$ we did not keep $\boldsymbol{m}_0$ and $\boldsymbol{v}_0$. There are two reasons for this: (1) the initializations of $\boldsymbol{m}$ and $\boldsymbol{v}$ are usually zero; (2) even if their initializations were not zero, the corresponding $\beta_1^t$ and $\beta_2^t$ would still be close enough to zero, so the effect of the initialization can be neglected.

However, $\boldsymbol{\theta}$ represents the model weights, and its initialization is generally not zero. Moreover, $\beta_3$ tends to be very close to 1, so over the course of an entire training run, $\beta_3^t$ may not become sufficiently close to zero. For this reason we explicitly keep $\beta_3^t$ and $\boldsymbol{\theta}_0$, to be dropped or retained as needed.

A Quick Estimate

Our task is to estimate the Weight RMS, i.e. $\Vert\boldsymbol{\theta}_t\Vert_{RMS}$, which as the name suggests is the Root Mean Square over the components:

\begin{equation}\Vert\boldsymbol{\theta}\Vert_{RMS} = \sqrt{\frac{1}{d}\sum_{i=1}^d \theta_i^2},\qquad\qquad \text{where}\boldsymbol{\theta} = (\theta_1,\theta_2,\cdots,\theta_d)\end{equation}

It differs from the norm only by an extra division by $\sqrt{d}$, so most properties of the norm carry over directly to the RMS. For $\Vert\boldsymbol{\theta}_t\Vert_{RMS}$, there is a quick — though not entirely accurate — derivation: simply taking $\Vert\cdot\Vert_{RMS}^2$ of both sides of Equation $\eqref{eq:ema-wd}$, we get

\begin{equation}\begin{aligned} \Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 =&\, \Vert\beta_3\boldsymbol{\theta}_{t-1} + (1-\beta_3)(-\boldsymbol{u}_t/\lambda)\Vert_{RMS}^2 \\[5pt] =&\, \beta_3^2\Vert\boldsymbol{\theta}_{t-1}\Vert_{RMS}^2 + (1-\beta_3)^2\Vert\boldsymbol{u}_t\Vert_{RMS}^2/\lambda^2 - 2\beta_3(1-\beta_3)\boldsymbol{\theta}_{t-1}\cdot\boldsymbol{u}_t/(\lambda d) \end{aligned}\end{equation}

Assume that $\boldsymbol{\theta}_{t-1},\boldsymbol{u}_t$ are nearly orthogonal, so that $\boldsymbol{\theta}_{t-1}\cdot\boldsymbol{u}_t\approx 0$ — which is typically a good approximation in high-dimensional spaces (see 《The Angle Distribution Between Two Random Vectors in n-Dimensional Space》). We've already computed $\Vert\boldsymbol{u}_t\Vert_{RMS}$ before, and the answer is approximately $\sqrt{\frac{1-\beta_1}{1+\beta_1}}$. Finally, since we're interested in the steady-state result, we have $\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2=\Vert\boldsymbol{\theta}_{t-1}\Vert_{RMS}^2$, giving

\begin{equation}(1-\beta_3^2)\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx (1-\beta_3)^2 \frac{1-\beta_1}{1+\beta_1} /\lambda^2\qquad\Rightarrow\qquad \Vert\boldsymbol{\theta}_t\Vert_{RMS} \approx \sqrt{\frac{1-\beta_1}{1+\beta_1}\frac{\eta}{2\lambda}}\end{equation}

Going from the left-hand side to the right-hand side also uses the approximation $\beta_3\approx 1$. The final result carries some error, since $\boldsymbol{\theta}_t\cdot\boldsymbol{u}_t\approx 0$ doesn't actually hold that strictly, but the conclusion $\Vert\boldsymbol{\theta}_t\Vert_{RMS}\propto \sqrt{\eta/\lambda}$ is correct. A similar derivation also appears in 《Why Gradients Rapidly Increase Near the End of Training》.

A Better Approximation

In many cases it suffices to know $\Vert\boldsymbol{\theta}_t\Vert_{RMS}\propto \sqrt{\eta/\lambda}$, and this is a fairly general conclusion. For readers who want a more precise result, we can obtain a better approximation using the mean-field method — at the cost of a considerably more complex derivation, but with the benefit of gaining a clearer, deeper understanding.

Step One

We start from Equation $\eqref{eq:theta-t}$. The summation term itself already has the form of a weighted average, so let's apply the first mean-field approximation:

\begin{equation}\underbrace{\frac{1-\beta_3}{1-\beta_3^t}\sum_{i=1}^t \beta_3^{t-i} \boldsymbol{u}_i}_{\text{denote}\bar{\boldsymbol{u}}_t} = \frac{1-\beta_3}{1-\beta_3^t}\sum_{i=1}^t \beta_3^{t-i} \frac{\boldsymbol{m}_i}{\sqrt{\boldsymbol{v}_i}}\approx \frac{\bar{\boldsymbol{m}}_t \,\,\triangleq\,\, \frac{1-\beta_3}{1-\beta_3^t}\sum_{i=1}^t \beta_3^{t-i}\boldsymbol{m}_i}{\sqrt{\bar{\boldsymbol{v}}_t \,\,\triangleq\,\, \frac{1-\beta_3}{1-\beta_3^t}\sum_{i=1}^t \beta_3^{t-i}\boldsymbol{v}_i}}\label{eq:u-bar}\end{equation}

Now let's return to Equation $\eqref{eq:theta-t}$. Since $\boldsymbol{\theta}_0$ is a random initialization vector, we can assume that $\boldsymbol{\theta}_0$ is orthogonal to $\bar{\boldsymbol{u}}_t$, so we have

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \beta_3^{2t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1-\beta_3^t)^2 \lambda^{-2}\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2\end{equation}

Now we need $\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2$. Based on past experience, we need to assume that $\boldsymbol{g}_j$ is i.i.d. and follows $\mathcal{N}(\boldsymbol{\mu},\boldsymbol{\sigma}^2)$, and then compute

\begin{equation}\mathbb{E}[\bar{\boldsymbol{u}}_t^2] \approx \mathbb{E}\left[\frac{\bar{\boldsymbol{m}}_t^2}{\bar{\boldsymbol{v}}_t}\right] \approx \frac{\mathbb{E}[\bar{\boldsymbol{m}}_t^2]}{\mathbb{E}[\bar{\boldsymbol{v}}_t]}\end{equation}

Finally, averaging over the individual components of $\mathbb{E}[\bar{\boldsymbol{u}}_t^2]$ gives us an approximation for $\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2$.

Step Two

Combining with Equation $\eqref{eq:mv-roll}$, we get

\begin{gather} \sum_{i=1}^t \beta_3^{t-i}\boldsymbol{m}_i = (1 - \beta_1)\sum_{i=1}^t \beta_3^{t-i} \sum_{j=1}^i \beta_1^{i-j}\boldsymbol{g}_j = (1 - \beta_1)\sum_{j=1}^t \frac{\beta_3^{t-j+1} - \beta_1^{t-j+1}}{\beta_3 - \beta_1}\boldsymbol{g}_j\\ \sum_{i=1}^t \beta_3^{t-i}\boldsymbol{v}_i = (1 - \beta_2)\sum_{i=1}^t \beta_3^{t-i} \sum_{j=1}^i \beta_2^{i-j}\boldsymbol{g}_j^2 = (1 - \beta_2)\sum_{j=1}^t \frac{\beta_3^{t-j+1} - \beta_2^{t-j+1}}{\beta_3 - \beta_2}\boldsymbol{g}_j^2\\ \end{gather}

The final double sum can be simplified; if you don't see how, feel free to hand it to Kimi (see this link). From the above equation we can see that $\bar{\boldsymbol{m}}_t,\bar{\boldsymbol{v}}_t$ are, respectively, weighted averages of the gradient and the squared gradient, so computing $\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2$ is essentially the same as computing $\Vert \boldsymbol{u}_t\Vert_{RMS}^2$ in 《Why Is Adam's Update RMS 0.2?》 — the only difference being the weighting coefficients.

Step Three

Let's first compute the denominator:

\begin{equation}\begin{aligned} \mathbb{E}[\bar{\boldsymbol{v}}_t] =&\, \frac{(1 - \beta_3)(1 - \beta_2)}{1 - \beta_3^t}\sum_{j=1}^t \frac{\beta_3^{t-j+1} - \beta_2^{t-j+1}}{\beta_3 - \beta_2}\mathbb{E}[\boldsymbol{g}_j^2] \\ =&\, \frac{(1 - \beta_3)(1 - \beta_2)}{1 - \beta_3^t}\sum_{j=1}^t \frac{\beta_3^{t-j+1} - \beta_2^{t-j+1}}{\beta_3 - \beta_2}(\boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2) \\ =&\, \frac{(1 - \beta_3)(1 - \beta_2)}{(1 - \beta_3^t)(\beta_3 - \beta_2)}\left(\frac{\beta_3 - \beta_3^{t+1}}{1 - \beta_3} - \frac{\beta_2 - \beta_2^{t+1}}{1 - \beta_2}\right)(\boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2) \\[5pt] \approx &\, \boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2 \end{aligned}\end{equation}

The approximate-equality in the last step holds because, in practice, $\beta_3$ will be sufficiently close to 1 and $\beta_2^{t+1}$ will be sufficiently close to 0, but $\beta_3^{t+1}$ is not necessarily so. Hence we replace $\beta_2^{t+1}$ with zero, and after simplification replace the independent $\beta_3$ with $1$, and finally add the approximation $\beta_3^{t+1}\approx \beta_3^t$.

Step Four

Next comes $\mathbb{E}[\bar{\boldsymbol{m}}_t^2] = \mathbb{E}[\bar{\boldsymbol{m}}_t]^2 + \mathbb{V}ar[\bar{\boldsymbol{m}}_t]$. The computation of $\mathbb{E}[\bar{\boldsymbol{m}}_t]$ is similar to that of $\mathbb{E}[\bar{\boldsymbol{v}}_t]$, giving the result $\boldsymbol{\mu}$. For $\mathbb{V}ar[\bar{\boldsymbol{m}}_t]$, we make use of the additivity of variance for squares:

\begin{equation}\begin{aligned} \mathbb{V}ar[\bar{\boldsymbol{m}}_t] =&\, \frac{(1 - \beta_3)^2(1 - \beta_1)^2}{(1-\beta_3^t)^2}\sum_{j=1}^t \left(\frac{\beta_3^{t-j+1} - \beta_1^{t-j+1}}{\beta_3 - \beta_1}\right)^2\mathbb{V}ar[\boldsymbol{g}_j] \\ =&\, \frac{(1 - \beta_3)^2(1 - \beta_1)^2}{(1-\beta_3^t)^2}\sum_{j=1}^t \left(\frac{\beta_3^{t-j+1} - \beta_1^{t-j+1}}{\beta_3 - \beta_1}\right)^2 \boldsymbol{\sigma}^2 \\ =&\, \frac{(1 - \beta_3)^2(1 - \beta_1)^2}{(1-\beta_3^t)^2(\beta_3 - \beta_1)^2}\left(\frac{\beta_3^2 - \beta_3^{2(t+1)}}{1 - \beta_3^2} + \frac{\beta_1^2 - \beta_1^{2(t+1)}}{1 - \beta_1^2} - 2\frac{\beta_1\beta_3 - \beta_1^{t+1}\beta_3^{t+1}}{1 - \beta_1\beta_3}\right) \boldsymbol{\sigma}^2 \\[5pt] \approx &\, (1 - \beta_3)(1 + \beta_3^t)\boldsymbol{\sigma}^2/2(1 - \beta_3^t) \end{aligned}\end{equation}

The approximate-equality holds for the same reason as before.

Step Five

Substituting the results from the previous two sections, we have

\begin{equation}\mathbb{E}[\bar{\boldsymbol{u}}_t^2] \approx \frac{\boldsymbol{\mu}^2 + (1 - \beta_3)(1 + \beta_3^t)\boldsymbol{\sigma}^2/2(1 - \beta_3^t)}{\boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2}\end{equation}

then

\begin{equation}\Vert\bar{\boldsymbol{u}}_t\Vert_{RMS}^2 \approx \frac{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + (1 - \beta_3)(1 + \beta_3^t)/2(1 - \beta_3^t)}{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + 1} \end{equation}

which finally gives us

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \beta_3^{2t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1-\beta_3^t)^2 \frac{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + (1 - \beta_3)(1 + \beta_3^t)/2(1 - \beta_3^t)}{\lambda^2(\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + 1)}\label{eq:theta-rms}\end{equation}

A Brief Look at the Result

Equation $\eqref{eq:theta-rms}$ looks fairly complicated, so let's examine a few special cases. First consider the case $\boldsymbol{\mu}=\boldsymbol{0}$, in which

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \beta_3^{2t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1-\beta_3^{2t}) (1 - \beta_3)/2\lambda^2 = \beta_3^{2t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1-\beta_3^{2t}) \eta/2\lambda\label{eq:theta-rms-mu0}\end{equation}

In particular, if we further consider $t\to\infty$, or if $\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2$ is initialized as $\eta/2\lambda$, then we get

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS} \approx \sqrt{\frac{\eta}{2\lambda}}\label{eq:theta-rms-simple}\end{equation}

This is exactly the result given by the paper 《Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks》, consistent with the original paper's assumptions — it is the steady-state result of a zero-mean random walk. If instead of $t\to\infty$ we consider the limit $\lambda\to 0$, then from Equation $\eqref{eq:theta-rms-mu0}$ we obtain

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + \eta^2 t\end{equation}

This shows that, without weight decay, $\Vert\boldsymbol{\theta}_t\Vert_{RMS}$ grows roughly at the rate of $\eta\sqrt{t}$, which also implies that without weight decay we could stabilize the Weight RMS by choosing a specific learning-rate schedule. On the other hand, if the batch size is large enough that the signal-to-noise term $\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2$ dominates, then from Equation $\eqref{eq:theta-rms}$ we get

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \beta_3^{2t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1-\beta_3^t)^2 \frac{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2}{\lambda^2(\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + 1)}\end{equation}

This might apply to special cases where a model needs to actively increase its Weight RMS. Empirically, though, this situation tends to arise fairly rarely.

Simulation Experiment

We can use the following simulation script to give a quick sanity check on the accuracy of the above:

import numpy as np

N, T = 10000, 100000
beta1, beta2 = 0.9, 0.95
m, v = 0, 0
w = np.random.randn(N) * 0.1
for i in range(T):
    g = np.random.randn(N)
    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * g**2
    w = w - 0.001 * (m / v**0.5 + 0.1 * w)

weight_rms = (w**2).mean()**0.5
print(weight_rms)

Feel free to change the weight initialization or the mean/variance of the gradients yourself, and see how closely the final result matches Equation $\eqref{eq:theta-rms}$. I tried this out a number of times myself, and overall it holds up quite well.

The Sign Version

Only a small adjustment to the above derivation is needed to make it apply to the "SignSGDM + Weight Decay" combination:

\begin{equation}\text{SignSGDM}\color{skyblue}{\text{W}}:=\left\{\begin{aligned} &\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\ &\boldsymbol{u}_t = \newcommand{sign}{\mathop{\text{sign}}}\sign(\boldsymbol{m}_t)\\ &\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}) \end{aligned}\right.\end{equation}

The change is needed because $\sign(\boldsymbol{m}_t)=\boldsymbol{m}_t/\sqrt{\boldsymbol{m}_t^2}$, so we need to redefine $\bar{\boldsymbol{v}}_t$ as

\begin{equation}\bar{\boldsymbol{v}}_t \triangleq \frac{1-\beta_3}{1-\beta_3^t}\sum_{i=1}^t \beta_3^{t-i}\boldsymbol{m}_i^2\end{equation}

then

\begin{equation}\mathbb{E}[\bar{\boldsymbol{v}}_t] = \frac{1-\beta_3}{1-\beta_3^t}\sum_{i=1}^t \beta_3^{t-i}\mathbb{E}[\boldsymbol{m}_i^2] \approx \frac{1-\beta_3}{1-\beta_3^t}\sum_{i=1}^t \beta_3^{t-i}\mathbb{E}\left(\boldsymbol{\mu}^2 + \frac{1-\beta_1}{1 + \beta_1}\boldsymbol{\sigma}^2\right) = \boldsymbol{\mu}^2 + \frac{1-\beta_1}{1 + \beta_1}\boldsymbol{\sigma}^2\end{equation}

where for computing $\mathbb{E}[\boldsymbol{m}_i^2]$, we can refer to either 《Why Is Adam's Update RMS 0.2?》 or 《Rethinking Learning Rate and Batch Size (Part 4): EMA》. Using this result, we obtain

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \beta_3^{2t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1-\beta_3^t)^2 \frac{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + (1 - \beta_3)(1 + \beta_3^t)/2(1 - \beta_3^t)}{\lambda^2\left(\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + \frac{1-\beta_1}{1 + \beta_1}\right)}\end{equation}

In particular, considering the limit $\boldsymbol{\mu}=0,t\to\infty$, we have

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \sqrt{\frac{\eta}{2\lambda}\frac{1+\beta_1}{1 - \beta_1}}\end{equation}

This result also makes sense: since SignSGDMW's Update RMS is $\sqrt{\frac{1+\beta_1}{1 - \beta_1}}$ times that of AdamW, for the same $\eta,\lambda$ its Weight RMS is likewise $\sqrt{\frac{1+\beta_1}{1 - \beta_1}}$ times as large.

Related Analysis

As mentioned earlier, the result $\eqref{eq:theta-rms-simple}$ agrees with the paper 《Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks》, but our derivation method is completely different, and it allows us to obtain a more general expression for $\eqref{eq:theta-rms}$. That said, the original paper does contain some interesting ideas worth appreciating, such as its concept of Total Update Contribution (TUC).

The idea behind TUC is as follows: due to the momentum mechanism, the current gradient $\boldsymbol{g}_t$ doesn't just affect the current step — it also influences future steps (albeit with a "discount"). So, assuming the number of training steps tends to infinity, we can consider the current gradient's $\boldsymbol{g}_t$ total contribution to the entire training process. Specifically, for Adam we have $\boldsymbol{u}_t=\boldsymbol{m}_t/\sqrt{\boldsymbol{v}_t}$: the current $\boldsymbol{g}_t$'s contribution to $\boldsymbol{u}_t$ is $(1-\beta_1)\boldsymbol{g}_t/\sqrt{\boldsymbol{v}_t}$, and at the next step $\boldsymbol{g}_t$ this will be discounted (multiplied by $\beta_1$), with the denominator also changed to $\boldsymbol{v}_{t+1}$, and so on. We can therefore define the total contribution as

\begin{equation}\tilde{\boldsymbol{u}}_t = \sum_{k=t}^{\infty} (1-\beta_1)\beta_1^{k-t}\frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{v}_k}}\end{equation}

In this way we decompose the update $\boldsymbol{u}_1,\boldsymbol{u}_2,\boldsymbol{u}_3,\cdots$ into a sequence of updates $\tilde{\boldsymbol{u}}_1,\tilde{\boldsymbol{u}}_2,\tilde{\boldsymbol{u}}_3,\cdots$, with the benefit that each $\tilde{\boldsymbol{u}}$ contains only a single-step gradient, allowing us to repeat the derivation from the Quick Estimate section:

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 = \Vert\beta_3\boldsymbol{\theta}_{t-1} + (1-\beta_3)(-\tilde{\boldsymbol{u}}_t/\lambda)\Vert_{RMS}^2 \approx \beta_3^2\Vert\boldsymbol{\theta}_{t-1}\Vert_{RMS}^2 + (1-\beta_3)^2\Vert\tilde{\boldsymbol{u}}_t\Vert_{RMS}^2/\lambda^2 \label{eq:tilde-u-rms}\end{equation}

The final approximation relies on $\boldsymbol{\theta}_{t-1}\cdot\tilde{\boldsymbol{u}}_t\approx 0$: we claim that $\boldsymbol{\theta}_{t-1}\cdot\tilde{\boldsymbol{u}}_t$ is closer to zero than $\boldsymbol{\theta}_{t-1}\cdot\boldsymbol{u}_t$, because $\tilde{\boldsymbol{u}}_t$ depends only on the current gradient $\boldsymbol{g}_t$, whereas $\boldsymbol{\theta}_{t-1}$ has not yet come into contact with $\boldsymbol{g}_t$, so they are independent variables — assuming $\boldsymbol{g}_t$ has zero mean, $\boldsymbol{\theta}_{t-1}\cdot\tilde{\boldsymbol{u}}_t\approx 0$ then tends to hold quite naturally. In order to estimate $\Vert\tilde{\boldsymbol{u}}_t\Vert_{RMS}^2$, the original paper directly assumes that $\boldsymbol{g}_t/\sqrt{\boldsymbol{v}_k}$ share the same direction and have unit RMS, so that

\begin{equation}\Vert\tilde{\boldsymbol{u}}_t\Vert_{RMS} = \sum_{k=t}^{\infty} (1-\beta_1)\beta_1^{k-t}\left\Vert\frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{v}_k}}\right\Vert_{RMS} = \sum_{k=t}^{\infty} (1-\beta_1)\beta_1^{k-t} = 1\end{equation}

Substituting into Equation $\eqref{eq:tilde-u-rms}$, and combining with the same approximation used in the Quick Estimate section, we solve to obtain

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS} \approx \sqrt{\frac{\eta}{2\lambda}}\end{equation}

However, if we restrict ourselves to what's stated in the original paper, we find that many of the approximations feel somewhat arbitrary — for instance, $\boldsymbol{v}_t$ also contains $\boldsymbol{g}_t$, so the claim that $\tilde{\boldsymbol{u}}_t$ only includes the influence of the current $\boldsymbol{g}_t$ isn't entirely accurate; likewise, the assertion about $\Vert\boldsymbol{g}_t/\sqrt{\boldsymbol{v}_k}\Vert_{RMS}=1$ feels a bit forced. But viewed through the lens of this post, we can see that under the mean-field approximation, the various operations in the original paper turn out to be quite reasonable — so in fact, the original paper had already implicitly made use of the mean-field method.

Summary

In this post, we used the mean-field approximation to derive an interesting and perhaps surprising conclusion: for a model trained with AdamW, the RMS of its weights can also be estimated asymptotically, and in general it depends only on the learning rate and the weight decay.