Asymptotic Estimation of Weight RMS of AdamW (Part 2)

By 苏剑林 | November 17, 2025

In the blog post "Asymptotic Estimation of Weight RMS of AdamW (Part 1)", we derived the asymptotic expression for the RMS of model weights trained with AdamW. However, at that time, we assumed that Weight Decay and the learning rate were fixed throughout the training process, which doesn't perfectly match actual training. Therefore, in this article, we will generalize the previous conclusions into a dynamic version.

The so-called dynamic version means allowing both Weight Decay and the learning rate to change as the number of training steps increases, such as the classic Cosine Decay, WSD (Warmup Stable Decay), etc., making the conclusions more general.

Step 1

Our starting point is still the definition of AdamW:

Since $\eta_t\lambda_t\ll 1$, we can write:

Let $\kappa_t = \sum_{i=1}^t \eta_i\lambda_i$, then expanding directly gives:

Then let $z_t = \sum_{i=1}^t e^{\kappa_i}\eta_i$, then by the mean-field approximation we get:

In this way:

Step 2

Following the previous strategy, to estimate $\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2$, we need to assume that $\boldsymbol{g}_j$ are independently and identically distributed following $\mathcal{N}(\boldsymbol{\mu},\boldsymbol{\sigma}^2)$, and then solve for:

Finally, we average the various components of $\mathbb{E}[\bar{\boldsymbol{u}}_t^2]$, and the result can serve as an approximation of $\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2$.

Expanding $\boldsymbol{m}_t$ and $\boldsymbol{v}_t$ gives:

Meanwhile, we also have the identity:

Utilizing these results, we can write:

Step 3

First calculate the denominator. When $t$ is sufficiently large (so $\beta_1^t, \beta_2^t$ are sufficiently small), $\sum_{j=1}^t \bar{\beta}_1(j,t)$ and $\sum_{j=1}^t \bar{\beta}_2(j,t)$ will be sufficiently close to 1 (because it is essentially a form of double-weighted average, just with the summation symbols swapped), so we have:

Similarly for $\mathbb{E}[\bar{\boldsymbol{m}}_t]$, the result is $\boldsymbol{\mu}$, and since $\mathbb{E}[\bar{\boldsymbol{m}}_t^2] = \mathbb{E}[\bar{\boldsymbol{m}}_t]^2 + \mathbb{V}ar[\bar{\boldsymbol{m}}_t]$, using the additivity of variance for independent terms, we get:

Therefore:

And:

Finally:

If readers are looking at this article directly, they might feel that some steps are a bit of a leap; in that case, one might as well revisit "Asymptotic Estimation of Weight RMS of AdamW (Part 1)" first to become familiar with the ideas behind each approximation step.

Example 1

First consider $\boldsymbol{\mu}=\boldsymbol{0}$, then substituting the expression for $\bar{\beta}_1(j,t)$ into the above formula gives:

Then consider the simple case where $\lambda_t=0$, meaning there is no Weight Decay. At this point:

If $\beta_1\to 0$, then we immediately have $\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + \sum_{j=1}^t\eta_j^2$. This shows that in the absence of Weight Decay as the training steps $t\to\infty$, if we want the Weight RMS not to explode, the sum of the squares of the learning rate sequence must converge, which is one of the classic conditions in traditional optimization theory. In fact, even with $0 < \beta_1 < 1$, this condition is necessary and sufficient, namely:

The proof is not difficult; we perform a transformation on the left side:

This shows that if the left side converges, then for $\forall i_1, i_2$, the sum $\sum_{j=1}^{\infty}\eta_{i_1+j}\eta_{i_2+j}$ must converge, which naturally implies $\sum_{j=1}^{\infty}\eta_j^2$ converges—Necessity proven. As for Sufficiency, we can start from the above equation and use Cauchy's inequality:

Thus $\sum_{j=1}^{\infty}\eta_j^2$ converging implies that the left side converges—Sufficiency proven.

Example 2

Next, we consider the case where Weight Decay is constant and the learning rate is variable. Here $\kappa_t = \lambda\sum_{i=1}^t \eta_i$. If we want to train infinitely and obtain a solution as close as possible to the theoretical optimum, the learning rate should satisfy $\sum_{i=1}^{\infty} \eta_i \to \infty$, so that the first term $e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2$ can completely "forget" the initialization (the theoretical optimal solution should be independent of initialization). Interestingly, this is also one of the classic conditions in traditional optimization theory.

For the general case, calculating the expression $\eqref{eq:w-rms-mu0}$ is relatively difficult, but we can consider further approximations based on actual conditions. In actual training, typically $\lambda_t \eta_t \ll 1$, so the growth rate of $e^{\kappa_i}$ is far slower than the decay rate of $\beta_1^i$. Meanwhile, compared to $\beta_1^i$, the learning rate $\eta_i$ is usually slowly-varying. Therefore, we can consider the approximation:

Substituting this approximation back into formula $\eqref{eq:w-rms-mu0}$ gives:

From here, one must calculate specifically for each $\eta_j$. For example, when both $\lambda_j$ and $\eta_j$ are constants, we can calculate $\kappa_t = \lambda\eta t$, and:

This is consistent with the results of the previous article.

Differential Equations

For numerical calculation, formula $\eqref{eq:w-rms-simp}$ is already quite concise. However, if we want to obtain general analytical results for $\lambda_t, \eta_t$, it is usually still difficult, so we need to look for further computational tools.

Considering that integrals are usually easier to compute than sums, we can try to approximate the summation with an integral:

Where $\kappa_t = \int_0^t \lambda_s\eta_s ds$. Letting $\rho_t = \Vert\boldsymbol{\theta}_t\Vert_{RMS}^2$, multiplying both sides by $e^{2\kappa_t}$ and then taking the derivative gives $\frac{d}{dt}(e^{2\kappa_t}\rho_t) \approx e^{2\kappa_t}\eta_t^2$. Rearranging yields:

This is the differential equation satisfied by the square of the RMS, which is arguably not very complex. If $\rho_t$ converges to a constant as $t\to\infty$, then the left side equals 0, thus:

This tells us that for Decay-type learning rate sequences, the final learning rate should not be set to 0; otherwise, there is a risk of model weight collapse under long-term training. Of course, we can also choose to set $\lambda_t \propto \eta_t$ (for example, AdamC) to avoid weight collapse.

Mean Field

Situations that can be viewed as $t\to\infty$ are generally Multi-Epoch supervised training. In pre-training scenarios, training is usually Single-Epoch. In these cases, $\kappa_t$ is often $\mathcal{\Theta}(1)$ because the weight of early training samples is comparable to the weight $e^{-2\kappa_t}$ of the $\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2$ term. Too large of a $\kappa_t$ might "forget" early training samples.

Under the assumption that $\kappa_t=\mathcal{\Theta}(1)$, we can consider mean-field approximation. Starting again from the integral form $\eqref{eq:w-rms-int}$, by definition, within $[0,t]$, $\kappa_s$ is a monotonically increasing function starting from 0 and ending at $\kappa_t$. Since $e^{\kappa_s} \geq 1$ and $e^{\kappa_s - \kappa_t} \leq 1$, we have:

That is, the target integral itself is bounded between $e^{- 2\kappa_t} \nu_t$ and $\nu_t$, where $\nu_t = \int_0^t \eta_s^2 ds$. When $\kappa_t=\mathcal{\Theta}(1)$, $e^{-2\kappa_t}$ will not be much smaller than 1, meaning that $\nu_t$ itself could be a good approximation. Of course, we can be more careful and estimate a reasonable multiple for $\nu_t$:

Considering that $\kappa_s$ is a monotonically increasing function from 0 to $\kappa_t$, we approximate it with $(\kappa_t/t)s$:

Substituting into formula $\eqref{eq:w-rms-int}$ yields:

Example 3

Returning once more to the common setting we care about—"Weight Decay fixed, learning rate variable"—let's calculate $\kappa_t, \nu_t$ for several specific examples. First is the linear learning rate:

Here $\eta_a, \eta_b$ are the initial and final learning rates, respectively. It could be that $\eta_b > \eta_a$ (like Warmup) or $\eta_b < \eta_a$ (Linear Decay), and $t$ is the total expected training steps. Integrating gives:

Next is Cosine Decay:

Integrating gives:

Finally, WSD (Warmup Stable Decay):

Correspondingly:

Simulation Verification

We can also verify the various approximations through numerical simulation:

Article Summary

This article generalizes the results from the previous post into a dynamic version, allowing us to estimate the Weight RMS of AdamW under a learning rate and Weight Decay that change over time.