Exercises — Adam and AdamW optimizers

• First moment $m_t$ = EMA of the gradient: $m_t={\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t$. It smooths the direction — this is the momentum idea.
• Second moment $v_t$ = EMA of the squared gradient: $v_t={\beta }_2{v}_{t-1}+(1-{\beta }_2)g_t^2$. It smooths the magnitude — this is the RMSProp idea. One tracker says where to go, the other says how big a step is safe there.

• (P) is AdamW. The decay term $-\eta \lambda {\theta }_{t-1}$ sits outside the $\sqrt{{\hat{v}}_t}$ fraction, so it is never divided by the magnitude tracker. See Weight Decay and L2 Regularization.
• (Q) is Adam+L2. The $\lambda {\theta }_{t-1}$ is folded into $g_t$, so it flows through $m_t,v_t$ and gets divided by $\sqrt{{\hat{v}}_t}$ — the very bug AdamW removes.

It is a signal-to-noise ratio. ${\hat{m}}_t$ is the mean gradient, $\sqrt{{\hat{v}}_t}$ is its typical size.

• Ratio $\approx 1$: the gradient points consistently the same way (mean $\approx$ magnitude) ⇒ confident, take a full step of size $\approx \eta$.
• Ratio $\approx 0$: the gradient jitters around zero (mean $\ll$ magnitude) ⇒ unsure, take a tiny step. This is a form of variance-aware stepping.

• $m_1=0.9(0)+0.1(0.5)=0.05$.
• $v_1=0.999(0)+0.001(0.25)=2.5\times 10^{-4}$.
• ${\hat{m}}_1=0.05/(1-0.9)=0.5$. (At $t=1$ this fully un-shrinks $m$: recovers $g_1$.)
• ${\hat{v}}_1=2.5\times 10^{-4}/(1-0.999)=0.25$. (Recovers $g_1^2$.)
• step $=\eta {\hat{m}}_1/\sqrt{{\hat{v}}_1}=0.01\cdot 0.5/0.5=0.01$.
• ${\theta }_1=1.0-0.01=0.99$. Notice: the step magnitude is exactly $\eta$, independent of $|g_1|$ — the "compass step".

• $m_2=0.9(0.05)+0.1(0.5)=0.045+0.05=0.095$.
• $v_2=0.999(2.5\times 10^{-4})+0.001(0.25)=2.4975\times 10^{-4}+2.5\times 10^{-4}=4.9975\times 10^{-4}$.
• ${\hat{m}}_2=0.095/(1-0.9^2)=0.095/0.19=0.5$.
• ${\hat{v}}_2=4.9975\times 10^{-4}/(1-0.999^2)=4.9975\times 10^{-4}/0.001999=0.25$.
• step $=0.01\cdot 0.5/\sqrt{0.25}=0.01\cdot 0.5/0.5=0.01$, so ${\theta }_2=0.99-0.01=0.98$. On a constant gradient the corrected estimates stay pinned at $\hat{m}=g$, $\hat{v}=g^2$, so every step keeps size $\eta$. This is the steady-state behaviour previewed in the parent note.

The decay part is $-\eta \lambda {\theta }_{t-1}$, so $\theta \leftarrow \theta (1-\eta \lambda )$. Surviving fraction $=1-\eta \lambda =1-0.001\cdot 0.1=1-0.0001=0.9999$. Each step shrinks every weight by $0.01\%$ — the same for all weights, which is exactly what classical weight decay wants.

At $t=1$: $m_1=(1-{\beta }_1)g_1$ and $v_1=(1-{\beta }_2)g_1^2$. Bias correction: ${\hat{m}}_1=\frac{(1-{\beta }_1)g_1}{1-{\beta }_1}=g_1$, and ${\hat{v}}_1=\frac{(1-{\beta }_2)g_1^2}{1-{\beta }_2}=g_1^2$. Step $=\eta \frac{{\hat{m}}_1}{\sqrt{{\hat{v}}_1}}=\eta \frac{g_1}{\sqrt{g_1^2}}=\eta \frac{g_1}{|g_1|}=\eta \text{sign}(g_1)$. The magnitude is $|\eta \text{sign}(g_1)|=\eta$, free of $|g_1|$. Look at figure s01: the SGD step (yellow) grows with the gradient, while the Adam step (blue) is clamped to length $\eta$ in both cases.

Replace $g_i\to c{g}_i$ for all $i$. Both EMAs are linear in their inputs:

• $m_t\to c{m}_t$ (since $m$ is a weighted sum of gradients), so ${\hat{m}}_t\to c{\hat{m}}_t$.
• $v_t\to c^2{v}_t$ (inputs are $g_i^2\to c^2{g}_i^2$), so ${\hat{v}}_t\to c^2{\hat{v}}_t$, hence $\sqrt{{\hat{v}}_t}\to c\sqrt{{\hat{v}}_t}$. Step $=\eta \frac{c{\hat{m}}_t}{c\sqrt{{\hat{v}}_t}}=\eta \frac{{\hat{m}}_t}{\sqrt{{\hat{v}}_t}}$ — the $c$ cancels. Consequence: Adam is immune to global gradient rescaling (e.g. a change of loss units). Plain SGD is not — its step $\eta cg$ scales with $c$.

• With $ϵ=0$: step $=\eta \frac{10^{-4}}{\sqrt{10^{-10}}}=\eta \frac{10^{-4}}{10^{-5}}=10\eta$ — a 10× blow-up.
• With $ϵ=10^{-8}$: $\sqrt{{\hat{v}}_t}+ϵ=10^{-5}+10^{-8}\approx 1.001\times 10^{-5}$, step $\approx \eta \frac{10^{-4}}{1.001\times 10^{-5}}\approx 9.99\eta$. Barely different here.
• Now push ${\hat{v}}_t\to 0$: with $ϵ=0$ the step $\to \infty$; with $ϵ$ the step is capped at $\eta {\hat{m}}_t/ϵ$. So $ϵ$ is not just anti-divide-by-zero — it bounds the maximum effective learning rate in near-flat regions. It is a real hyperparameter.

Decay contribution = the part of the step coming from the $\lambda \theta$ term. (i) Adam+L2: the $\lambda \theta$ is folded into the gradient, so it is divided by $\sqrt{\hat{v}}$: $\text{decay}=\eta \frac{\lambda \theta }{\sqrt{\hat{v}}}$.

• $A$: $1\cdot \frac{0.05\cdot 2}{0.02}=\frac{0.1}{0.02}=5.0$.
• $B$: $1\cdot \frac{0.05\cdot 2}{2.0}=\frac{0.1}{2.0}=0.05$. $A$ is decayed 100× harder than $B$ — purely because of gradient history. Backwards. (ii) AdamW: $\text{decay}=\eta \lambda \theta =1\cdot 0.05\cdot 2=0.1$ for both $A$ and $B$. Uniform, gradient-independent — the correct behaviour. Figure s02 shows the two decay bars.

  

  Figure s02 caption: grouped bar chart. Left group = weight $A$, right group = weight $B$. Red bars = Adam+L2 decay contribution ($5.0$ for $A$, $0.05$ for $B$ — a 100× spread). Green bars = AdamW decay contribution ($0.1$ for both — equal height). The visual punch: red bars are wildly uneven, green bars are identical, showing AdamW's uniform shrinkage.

Unroll the EMA from $m_0=0$: $m_t=(1-{\beta }_1){\sum }_{i=1}^t{\beta }_1^{t-i}{g}_i.$ Take expectations, using $\mathbb{E}[g_i]=\bar{g}$: $\mathbb{E}[m_t]=\bar{g}(1-{\beta }_1){\sum }_{i=1}^t{\beta }_1^{t-i}.$ The sum is a geometric series ${\sum }_{k=0}^{t-1}{\beta }_1^k=\frac{1-{\beta }_1^t}{1-{\beta }_1}$. So $\mathbb{E}[m_t]=\bar{g}(1-{\beta }_1)\cdot \frac{1-{\beta }_1^t}{1-{\beta }_1}=\bar{g}(1-{\beta }_1^t).$ This is smaller than the target $\bar{g}$ by the factor $(1-{\beta }_1^t)$. Dividing $m_t$ by that factor gives $\mathbb{E}[{\hat{m}}_t]=\bar{g}$ — unbiased. As $t\to \infty$, ${\beta }_1^t\to 0$ and the correction $\to 1$ (vanishes). The identical argument on the squared stream gives ${\hat{v}}_t=v_t/(1-{\beta }_2^t)$.

flowchart TD
  G["raw gradient g_t"] --> M["update m_t EMA of g"]
  G --> V["update v_t EMA of g squared"]
  M --> MH["bias correct m hat"]
  V --> VH["bias correct v hat"]
  MH --> R["ratio m hat over root v hat"]
  VH --> R
  R --> U["subtract eta times ratio"]
  TH["current weight theta"] --> D["decay subtract eta lambda theta"]
  U --> NEW["new weight"]
  D --> NEW

Node TH is the current weight ${\theta }_{t-1}$; node D is the decay step $-\eta \lambda {\theta }_{t-1}$ applied straight to it. The decay arrow (TH --> D) feeds the weight in directly, never touching the ratio node R. Adam+L2's mistake is to instead route decay through node G (fold $\lambda \theta$ into the gradient), so it passes through V/VH and gets divided by $\sqrt{\hat{v}}$. AdamW keeps the two paths separate.

Observable: the norm of the very first update, $\Vert \Delta {\theta }_1\Vert$, relative to $\eta$.

• Correct Adam: at $t=1$, per-coordinate $|\Delta \theta |=\eta$ exactly (from L3-Q1). So $\Vert \Delta {\theta }_1\Vert /\eta =\sqrt{d}$ for $d$ parameters — a clean, gradient-independent value.
• Bug (both corrections off): use raw $m_1=(1-{\beta }_1)g_1$ and raw $v_1=(1-{\beta }_2)g_1^2$. Per-coordinate step $=\eta \frac{(1-{\beta }_1)|g_1|}{\sqrt{(1-{\beta }_2)}|g_1|}=\eta \frac{1-{\beta }_1}{\sqrt{1-{\beta }_2}}$. The $|g_1|$ cancels, giving a gradient-independent constant: $\frac{1-{\beta }_1}{\sqrt{1-{\beta }_2}}=\frac{0.1}{\sqrt{0.001}}=\frac{0.1}{0.0316}=\sqrt{10}\approx 3.162$. Signature: the buggy first-step norm is $\Vert \Delta {\theta }_1\Vert /\eta \approx 3.162\sqrt{d}$ instead of the correct $\sqrt{d}$. Concretely, dividing the measured first-step norm by $\eta \sqrt{d}$ should give $\approx 1$ if correct, $\approx 3.16$ if both corrections are missing — a clean, model-size-independent tell.

Pick AdamW. Mechanistic argument:

1. Transformer parameters have wildly different gradient scales across layers/embeddings, so a per-parameter adaptive rate (RMSProp-style $\sqrt{\hat{v}}$) is worth more than a single SGD rate.
2. Heavy overfitting means regularization strength must be predictable. Adam+L2 makes decay gradient-dependent (L4-Q1: 100× spread between weights), so you cannot reason about how hard any weight is being pulled to zero.
3. AdamW's decoupled $\eta \lambda \theta$ shrinks every weight by the identical factor $(1-\eta \lambda )$ each step, giving clean, tunable regularization on top of the adaptive step. Net: you get the adaptivity of Adam and honest weight decay — exactly why AdamW is the default for transformers.

The AdamW decay term is $-\eta \lambda \theta$. If $\eta \approx 0$ during warmup, then decay $\approx 0$ too — regularization is effectively off while $\eta$ is small, because decay is scaled by the same $\eta$ as the gradient step. Fixes:

• Use a decay strength that is applied on a schedule decoupled from $\eta$ (some implementations let you scale decay by a fixed rate independent of the warmup $\eta$).
• Or accept that regularization ramps in with the warmup, and set the schedule/$\lambda$ so that over the full run the accumulated shrinkage ${\prod }_t(1-{\eta }_t\lambda )$ reaches the intended target. The lesson: in AdamW, $\eta$ multiplies both the step and the decay, so any $\eta$-schedule silently reshapes regularization over time.

First-step Adam magnitude (any $g\ne 0$) ::: exactly $\eta$ Steady-state Adam step on constant $g$ ::: $\eta \text{sign}(g)$ AdamW surviving fraction per step from decay ::: $1-\eta \lambda$ What multiplying all gradients by $c$ does to the Adam step ::: nothing (scale-invariant) The factor $m_t$ is biased low by at step $t$ ::: $(1-{\beta }_1^t)$ The one thing decay in AdamW is NOT decoupled from ::: the learning rate $\eta$