Asymptotic Estimation of AdamW's Weight RMS (Part 2)

AdamW的Weight RMS的渐近估计（下） 苏剑林 · 2025-11-17 · Mathematics differential equation estimation gradient optimizer mean field

In the blog post "Asymptotic Estimation of AdamW's Weight RMS (Part 1)", we derived an asymptotic expression for the RMS of weights trained with AdamW. However, at that time we assumed that weight decay and the learning rate stay fixed throughout training, which doesn't quite match real-world training. So in this post we'll generalize the previous conclusions to a dynamic version.

By "dynamic version" we mean allowing both weight decay and the learning rate to vary with the training step, as in classic schedules like cosine decay or WSD (Warmup Stable Decay), thereby making the result more broadly applicable.

Step One

We still start from the definition of AdamW:

\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}more

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

\begin{equation}\boldsymbol{\theta}_t = (1 - \eta_t\lambda_t)\boldsymbol{\theta}_{t-1} -\eta_t\boldsymbol{u}_t \approx e^{- \eta_t\lambda_t}\boldsymbol{\theta}_{t-1} -\eta_t\boldsymbol{u}_t\end{equation}

Denoting $\kappa_t = \sum_{i=1}^t \eta_i\lambda_i$, direct expansion gives

\begin{equation}\boldsymbol{\theta}_t \approx e^{-\kappa_t}\boldsymbol{\theta}_0 - \sum_{i=1}^t e^{-(\kappa_t - \kappa_i)}\eta_i\boldsymbol{u}_i = e^{-\kappa_t}\left(\boldsymbol{\theta}_0 - \sum_{i=1}^t e^{\kappa_i}\eta_i\boldsymbol{u}_i\right)\end{equation}

Then, setting $z_t = \sum_{i=1}^t e^{\kappa_i}\eta_i$, the mean-field approximation gives

\begin{equation}\bar{\boldsymbol{u}}_t\triangleq\frac{1}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i \boldsymbol{u}_i = \frac{1}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i \frac{\boldsymbol{m}_i}{\sqrt{\boldsymbol{v}_i}}\approx \frac{\bar{\boldsymbol{m}}_t \,\,\triangleq\,\, \frac{1}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i\boldsymbol{m}_i}{\sqrt{\bar{\boldsymbol{v}}_t \,\,\triangleq\,\, \frac{1}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i\boldsymbol{v}_i}}\end{equation}

which leads to

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0 - z_t \bar{\boldsymbol{u}}_t\Vert_{RMS}^2 \approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{-2\kappa_t}z_t^2\Vert\bar{\boldsymbol{u}}_t\Vert_{RMS}^2\end{equation}

Step Two

Following our earlier approach, in order to estimate $\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2$, we need to assume that $\boldsymbol{g}_j$ are i.i.d. and follow $\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 components of $\mathbb{E}[\bar{\boldsymbol{u}}_t^2]$, the result can serve as an approximation for $\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2$.

Expanding $\boldsymbol{m}_t,\boldsymbol{v}_t$ gives

\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}

We also have the identity

\begin{equation}\sum_{i=1}^t \sum_{j=1}^i a_i b_j = \sum_{j=1}^t \sum_{i=j}^t a_i b_j\end{equation}

Using these two results, we can write

\begin{gather} \bar{\boldsymbol{m}}_t = \frac{1}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i\boldsymbol{m}_i = \frac{1 - \beta_1}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i\sum_{j=1}^i \beta_1^{i-j}\boldsymbol{g}_j = \sum_{j=1}^t\boldsymbol{g}_j\underbrace{\frac{1 - \beta_1}{z_t}\sum_{i=j}^t e^{\kappa_i}\beta_1^{i-j}\eta_i}_{\text{denote}\bar{\beta}_1(j,t)} \\ \bar{\boldsymbol{v}}_t = \frac{1}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i\boldsymbol{v}_i = \frac{1 - \beta_2}{z_t}\sum_{i=1}^t e^{\kappa_i}\eta_i\sum_{j=1}^i \beta_2^{i-j}\boldsymbol{g}_j^2 = \sum_{j=1}^t\boldsymbol{g}_j^2\underbrace{\frac{1 - \beta_2}{z_t}\sum_{i=j}^t e^{\kappa_i}\beta_2^{i-j}\eta_i}_{\text{denote}\bar{\beta}_2(j,t)} \\ \end{gather}

Step Three

Let's first deal with the denominator. When $t$ is large enough (i.e. $\beta_1^t, \beta_2^t$ small enough), $\sum_{j=1}^t \bar{\beta}_1(j,t)$ and $\sum_{j=1}^t \bar{\beta}_2(j,t)$ get close enough to 1 (since these are essentially doubly-weighted averages, just with the summation order swapped), so we have

\begin{equation}\mathbb{E}[\bar{\boldsymbol{v}}_t] = \sum_{j=1}^t\bar{\beta}_2(j,t) \mathbb{E}[\boldsymbol{g}_j^2] = \sum_{j=1}^t\bar{\beta}_2(j,t) (\boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2) \approx \boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2 \end{equation}

The same reasoning applies to $\mathbb{E}[\bar{\boldsymbol{m}}_t]$, giving $\boldsymbol{\mu}$, and with $\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 variances of squares we get

\begin{equation}\mathbb{V}ar[\bar{\boldsymbol{m}}_t] = \sum_{j=1}^t\bar{\beta}_1(j,t)^2 \mathbb{V}ar[\boldsymbol{g}_j] = \sum_{j=1}^t\bar{\beta}_1(j,t)^2 \boldsymbol{\sigma}^2\end{equation}

so

\begin{equation}\mathbb{E}[\bar{\boldsymbol{u}}_t^2] \approx \frac{\boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2\sum_{j=1}^t\bar{\beta}_1(j,t)^2}{\boldsymbol{\mu}^2 + \boldsymbol{\sigma}^2}\end{equation}

and

\begin{equation}\Vert\bar{\boldsymbol{u}}_t\Vert_{RMS}^2 \approx \frac{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + \sum_{j=1}^t\bar{\beta}_1(j,t)^2}{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + 1} \end{equation}

which finally gives

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{-2\kappa_t}z_t^2\frac{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + \sum_{j=1}^t\bar{\beta}_1(j,t)^2}{\Vert\boldsymbol{\mu}\Vert^2/\Vert\boldsymbol{\sigma}\Vert^2 + 1}\end{equation}

If you're reading this post without having gone through Part 1, some of these steps may feel a bit rushed. In that case, it's worth revisiting "Asymptotic Estimation of AdamW's Weight RMS (Part 1)" to get familiar with the reasoning behind each approximation.

Example One

Let's first consider $\boldsymbol{\mu}=\boldsymbol{0}$. Substituting the expression for $\bar{\beta}_1(j,t)$ into the above gives

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{-2\kappa_t}(1-\beta_1)^2\sum_{j=1}^t\left(\sum_{i=j}^t e^{\kappa_i}\beta_1^{i-j}\eta_i\right)^2\label{eq:w-rms-mu0}\end{equation}

Now consider the simple special case $\lambda_t=0$, i.e. no weight decay, in which case

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1-\beta_1)^2\sum_{j=1}^t\left(\sum_{i=j}^t \beta_1^{i-j}\eta_i\right)^2\end{equation}

If $\beta_1\to 0$, then immediately $\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + \sum_{j=1}^t\eta_j^2$, which tells us that in the absence of weight decay, as the number of training steps $t\to\infty$ grows, for the weight RMS not to blow up, the sum of squares of the learning-rate sequence must converge — this is exactly one of the classical conditions in traditional optimization theory. In fact, even when $0 < \beta_1 < 1$, this condition is both necessary and sufficient, i.e.

\begin{equation}\sum_{j=1}^{\infty}\left(\sum_{i=j}^{\infty} \beta_1^{i-j}\eta_i\right)^2 < \infty \qquad\Leftrightarrow\qquad \sum_{j=1}^{\infty}\eta_j^2 < \infty\end{equation}

The proof isn't too hard. Let's transform the left-hand side:

\begin{equation}\begin{aligned} \sum_{j=1}^{\infty}\left(\sum_{i=j}^{\infty} \beta_1^{i-j}\eta_i\right)^2 = \sum_{j=1}^{\infty}\left(\sum_{i=0}^{\infty} \beta_1^i\eta_{i+j}\right)^2 =&\, \sum_{j=1}^{\infty}\left(\sum_{i_1=0}^{\infty} \beta_1^{i_1}\eta_{i_1+j}\right)\left(\sum_{i_2=0}^{\infty} \beta_1^{i_2}\eta_{i_2+j}\right) \\ =&\, \sum_{i_1=0}^{\infty}\sum_{i_2=0}^{\infty} \beta_1^{i_1 + i_2}\sum_{j=1}^{\infty}\eta_{i_1+j}\eta_{i_2+j} \end{aligned}\end{equation}

This shows that if the left-hand side converges, then for $\forall i_1, i_2$, the sum $\sum_{j=1}^{\infty}\eta_{i_1+j}\eta_{i_2+j}$ also converges, which in turn means $\sum_{j=1}^{\infty}\eta_j^2$ converges — establishing necessity. As for sufficiency, we can start from the above expression and use the Cauchy-Schwarz inequality:

\begin{equation}\begin{aligned} \sum_{i_1=0}^{\infty}\sum_{i_2=0}^{\infty} \beta_1^{i_1 + i_2}\sum_{j=1}^{\infty}\eta_{i_1+j}\eta_{i_2+j} \leq&\, \sum_{i_1=0}^{\infty}\sum_{i_2=0}^{\infty} \beta_1^{i_1 + i_2}\sqrt{\left(\sum_{j=1}^{\infty}\eta_{i_1+j}^2\right)\left(\sum_{j=1}^{\infty}\eta_{i_2+j}^2\right)} \\ \leq&\, \sum_{i_1=0}^{\infty}\sum_{i_2=0}^{\infty} \beta_1^{i_1 + i_2}\sqrt{\left(\sum_{j=1}^{\infty}\eta_j^2\right)\left(\sum_{j=1}^{\infty}\eta_j^2\right)} \\ =&\, \frac{1}{(1-\beta_1)^2} \sum_{j=1}^{\infty}\eta_j^2 \end{aligned}\end{equation}

so convergence of $\sum_{j=1}^{\infty}\eta_j^2$ implies convergence of the left-hand side, establishing sufficiency.

Example Two

Next, let's consider the case where weight decay is constant and the learning rate varies, i.e. $\kappa_t = \lambda\sum_{i=1}^t \eta_i$. If we want to keep training indefinitely and get as close as possible to the theoretically optimal solution, 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 theoretically optimal solution should be independent of the initialization). Interestingly, this too is one of the classical conditions in traditional optimization theory.

In the general case, computing $\eqref{eq:w-rms-mu0}$ is fairly difficult, but we can consider further approximations based on practical conditions. In actual training, we typically have $\lambda_t \eta_t \ll 1$, so $e^{\kappa_i}$ grows much more slowly than $\beta_1^i$ decays, and meanwhile the learning rate $\eta_i$ is usually slowly varying compared to $\beta_1^i$. So we can consider the approximation

\begin{equation}\sum_{i=j}^t e^{\kappa_i}\beta_1^{i-j}\eta_i \approx \sum_{i=j}^t e^{\kappa_j}\beta_1^{i-j}\eta_j = e^{\kappa_j}\eta_j\sum_{i=j}^t\beta_1^{i-j}\approx e^{\kappa_j}\eta_j\sum_{i=j}^{\infty}\beta_1^{i-j} = \frac{e^{\kappa_j}\eta_j}{1-\beta_1}\end{equation}

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

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{- 2\kappa_t}\sum_{j=1}^t e^{2\kappa_j}\eta_j^2\label{eq:w-rms-simp}\end{equation}

From here, we can only compute case by case for specific $\eta_j$. For instance, when $\lambda_j,\eta_j$ are constants, we can compute $\kappa_t = \lambda\eta t$, along with

\begin{equation}\begin{aligned} \Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx&\, e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{- 2\kappa_t}\sum_{j=1}^t e^{2\kappa_j}\eta_j^2 \\ =&\, e^{-2\lambda\eta t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{-2\lambda\eta t}\sum_{j=1}^t e^{2\lambda\eta j}\eta^2 \\ =&\, e^{-2\lambda\eta t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + \frac{e^{2\lambda\eta}(1 - e^{-2\lambda\eta t})}{e^{2\lambda\eta} - 1}\eta^2 \\ \approx&\, e^{-2\lambda\eta t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1 - e^{-2\lambda\eta t} )\frac{\eta}{2\lambda} \end{aligned}\end{equation}

which agrees with the result from the previous post.

Differential Equation

For numerical computation purposes, $\eqref{eq:w-rms-simp}$ is already fairly concise, but if we want a general analytical result for arbitrary $\lambda_t,\eta_t$, this is usually still quite difficult, so we need to look for yet another computational tool.

Since integrals are generally easier to compute than sums, let's try approximating the sum with an integral:

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{- 2\kappa_t}\sum_{j=1}^t e^{2\kappa_j}\eta_j^2\approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + e^{- 2\kappa_t}\int_0^t e^{2\kappa_s}\eta_s^2 ds \label{eq:w-rms-int}\end{equation}

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 differentiating gives $\frac{d}{dt}(e^{2\kappa_t}\rho_t) \approx e^{2\kappa_t}\eta_t^2$, which rearranges to

\begin{equation}\frac{d}{dt}\rho_t \approx -2\lambda_t\eta_t\rho_t + \eta_t^2\end{equation}

This is the differential equation satisfied by the squared RMS, and it's not too complicated. If, as $t\to\infty$, $\rho_t$ converges to a constant, then the left-hand side equals 0, giving

\begin{equation}\lim_{t\to\infty} \rho_t \approx \lim_{t\to\infty} \frac{\eta_t}{2\lambda_t}\end{equation}

This tells us that for a decaying learning-rate schedule, the final learning rate should not be set to 0, otherwise there is a risk of weight collapse under long-term training. Alternatively, we could choose to set $\lambda_t\propto \eta_t$ (as in AdamC) to avoid weight collapse.

The Field of Averaging

Scenarios that can be treated as $t\to\infty$ typically correspond to multi-epoch supervised training, whereas in pretraining, training is usually single-epoch, and in this case $\kappa_t$ is often $\mathcal{\Theta}(1)$, because the sample weight at the beginning of training is comparable to the weight of the $\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2$ term, $e^{-2\kappa_t}$, and an overly large $\kappa_t$ risks "forgetting" the early training samples.

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

\begin{equation}e^{- 2\kappa_t}\int_0^t \eta_s^2 ds \leq e^{- 2\kappa_t}\int_0^t e^{2\kappa_s}\eta_s^2 ds = \int_0^t e^{2\kappa_s- 2\kappa_t}\eta_s^2 ds \leq \int_0^t \eta_s^2 ds\end{equation}

That is, the target integral itself is sandwiched between $e^{- 2\kappa_t} \nu_t$ and $\nu_t$, where $\nu_t = \int_0^t \eta_s^2 ds$. And when $\kappa_t=\mathcal{\Theta}(1)$, $e^{-2\kappa_t}$ won't be much smaller than 1, which means $\nu_t$ itself may already be a good approximation. Of course, we can also be a bit more careful and estimate a reasonable multiplicative factor for $\nu_t$:

\begin{equation}e^{- 2\kappa_t}\int_0^t e^{2\kappa_s}\eta_s^2 ds \approx e^{- 2\kappa_t}\int_0^t e^{2\kappa_s} (\nu_t / t) ds = \frac{\nu_t e^{- 2\kappa_t}}{t}\int_0^t e^{2\kappa_s} ds\end{equation}

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

\begin{equation}e^{- 2\kappa_t}\int_0^t e^{2\kappa_s}\eta_s^2 ds \approx \frac{\nu_t e^{- 2\kappa_t}}{t}\int_0^t e^{2\kappa_s} ds \approx \frac{\nu_t e^{- 2\kappa_t}}{t}\int_0^t e^{2(\kappa_t/t)s} ds = \frac{\nu_t}{2\kappa_t}(1 - e^{- 2\kappa_t})\end{equation}

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

\begin{equation}\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx e^{-2\kappa_t}\Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + (1 - e^{- 2\kappa_t})\frac{\nu_t}{2\kappa_t}\end{equation}

Example Three

Let's return once more to the common setting we care most about — "fixed weight decay, variable learning rate" — and compute $\kappa_t,\nu_t$ for a few concrete examples. First, the linear learning rate:

\begin{equation}\eta_s = \eta_a + (\eta_b - \eta_a) s / t\end{equation}

Here, $\eta_a,\eta_b$ are the initial and final learning rates respectively; we might have $\eta_b > \eta_a$ (as in warmup) or $\eta_b < \eta_a$ (linear decay), and $t$ is the expected total number of training steps. Integrating gives

\begin{gather} \kappa_t = \int_0^t \lambda\eta_s ds= \lambda (\eta_a + \eta_b) t / 2 \\ \nu_t = \int_0^t \eta_s^2 ds = (\eta_a^2 + \eta_a \eta_b + \eta_b^2) t / 3 \end{gather}

Next, cosine decay:

\begin{equation}\eta_s = \eta_{\min} + (\eta_{\max} - \eta_{\min})\left(\frac{1}{2} + \frac{1}{2}\cos \frac{s\pi}{t}\right)\end{equation}

Integrating gives

\begin{gather} \kappa_t = \int_0^t \lambda\eta_s ds= \lambda (\eta_{\min} + \eta_{\max}) t / 2 \\ \nu_t = \int_0^t \eta_s^2 ds = (3\eta_{\min}^2 + 2\eta_{\min} \eta_{\max} + 3\eta_{\max}^2 ) t / 8 \end{gather}

Finally, WSD (Warmup Stable Decay):

\begin{equation}\eta_s = \left\{\begin{aligned} \frac{s}{t_1}\eta_{\max}, \quad s \in [0, t_1] \\[5pt] \eta_{\max} , \quad s \in [t_1, t_2] \\[5pt] \frac{t-s}{t-t_2}\eta_{\max}, \quad j \in [t_2, t] \end{aligned}\right.\end{equation}

giving

\begin{gather} \kappa_t = \int_0^t \lambda\eta_s ds= \lambda \eta_{\max} (t + t_2 - t_1) / 2 \\ \nu_t = \int_0^t \eta_s^2 ds = \eta_{\max}^2 (t + 2t_2 - 2t_1) / 3 \end{gather}

Numerical Verification

We can also verify all the above approximations through numerical simulation:

import numpy as np

N, T = 10000, 10000
beta1, beta2 = 0.9, 0.95
m, v = 0, 0
w = np.random.randn(N) * (init_std := 0.1)
lr_max, lr_min, wd = 0.001, 0.0001, 0.1
lr = lr_min + (lr_max - lr_min) * (1 + np.cos(np.arange(T) / T * np.pi)) / 2
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 - lr[i] * (m / v**0.5 + wd * w)

# 直接计算 ≈ 0.0744
weight_rms = (w**2).mean()**0.5

# 级数近似 ≈ 0.0742
kappa = wd * lr.cumsum()
approx1 = ((np.exp(kappa * 2) * lr**2).sum() + init_std**2)**0.5 * np.exp(-kappa[-1])

# 平均场近似 ≈ 0.0760
kappa = wd * (lr_max + lr_min) / 2 * T
nu = (3 * lr_max**2 + 2 * lr_max * lr_min + 3 * lr_min**2) / 8 * T
approx2 = ((np.exp(kappa * 2) - 1) * nu / kappa / 2 + init_std**2)**0.5 * np.exp(-kappa)

print(weight_rms)
print(approx1)
print(approx2)

Summary

This post generalized the results of the previous one to a dynamic version, allowing us to estimate the weight RMS of AdamW-trained models under time-varying learning rates and weight decay.