AdamC

Why gradient norms climb at the end of training, and the one-line change to weight decay that stops it.

In long language model runs, people log the size of the gradient as a health check. Near the end of training that number often turns around and climbs. In the run that motivated AdamC it more than doubled. This is strange, because the learning rate is shrinking toward zero at that point and the loss looks fine.

The explanation turns out to be bookkeeping, not instability. Three ordinary ingredients interact. They are normalization layers, weight decay, and a learning rate schedule. AdamW ties the strength of weight decay to the learning rate in a way that makes the weights shrink as the schedule decays, and shrinking weights make gradients look bigger.

AdamC is the fix proposed by Aaron Defazio in a 2025 arXiv note. It multiplies the weight decay coefficient by the same factor as the learning rate schedule. The rest of this page builds up exactly why that works, starting from things you already know.

Roadmap

AdamW in one line, and where weight decay meets the learning rate.
A normalized layer cannot see the length of its weights.
Two consequences for the gradient. It is perpendicular to the weights, and it scales like one over their length.
Perpendicular steps always make the weights longer.
Weight decay pulls back, and the two forces settle at an equilibrium.
A learning rate schedule moves that equilibrium. This is the rising gradient norm.
The same story for Adam, where it is stronger.
The fix. Scale weight decay with the schedule.
A lab where you run AdamW and AdamC side by side.
What you actually buy. The schedule means what it says.

Notation

Symbol	Meaning
$\w_t$	The weights of one layer at step $t$ , flattened into a vector with $d$ entries. Blue everywhere on this page.
$\g_t$	The minibatch gradient of the loss with respect to $\w_t$ . Crimson.
$\lr_t$ , $\lr_{\max}$	Learning rate at step $t$ , and its peak value over the whole schedule. Amber.
$\wd$	Weight decay coefficient. Teal.
$\\|\cdot\\|$	Euclidean length of a vector. $\langle a,b\rangle$ is the dot product.
$u_t$	Adam's normalized update direction (introduced in Step 1).

I follow the usual AdamW write-up and use $\lr$ for the learning rate and $\wd$ for weight decay. Defazio's paper writes $\gamma$ for the learning rate and $x$ for the weights. Same objects, different letters.

Step 1

AdamW, and where weight decay meets the learning rate

You already know Adam. It keeps a running average of the gradient $m_t$ and of the squared gradient $v_t$ , and it steps along their ratio. Call that ratio the update direction.

u_t=\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}

Here $\hat m_t$ and $\hat v_t$ are the bias-corrected averages and $\epsilon$ is a tiny constant that avoids division by zero. AdamW then does two things to the weights in each step. It moves along $u_t$ , and it shrinks the weights slightly toward zero.

\w_{t+1}=\w_t-\lr_t\,u_t-\lr_t\,\wd\,\w_t

Look at the last term. The shrink per step is $\lr_t\wd$ , the product of learning rate and weight decay. This is how the AdamW paper writes it and how torch.optim.AdamW implements it.

Tiny example

A typical recipe for a small language model uses peak $\lr=0.003$ and $\wd=0.1$ . At the peak, each step multiplies the decayed weights by $1-0.003\times0.1=0.9997$ . Late in the schedule, when $\lr_t=0.0003$ , each step multiplies them by $0.99997$ . So the schedule turns down weight decay along with the learning rate, by the same factor.

That sounds harmless, even sensible. The next few steps show that "the same factor" is the wrong amount.

Takeaway. In AdamW the per-step shrink is

\lr_t\wd

, so a learning rate schedule is silently also a weight decay schedule.

Step 2

A normalized layer cannot see the length of its weights

A normalization layer such as RMSNorm or LayerNorm rescales its input vector to a fixed size. Suppose a weight matrix $W$ produces $h=Wx$ , and the very next operation normalizes $h$ . If you double every entry of $W$ , then $h$ doubles, and the normalization divides the doubling right back out.

Tiny example

Take $h=(3,4)$ . Its length is $5$ , so dividing by the length gives $(0.6,\,0.8)$ . Now double the weights so that $h=(6,8)$ . Its length is $10$ , and dividing gives $(0.6,\,0.8)$ again. Everything downstream sees the identical vector, so the loss is identical.

This property has a name. A loss $L$ is scale invariant in $\w$ if stretching $\w$ by any positive factor $c$ leaves the loss unchanged.

L(c\,\w)=L(\w)\quad\text{for every } c>0

Geometrically, the loss depends only on the direction of $\w$ . Picture rays coming out of the origin. The loss is constant along each ray and only changes when you rotate from one ray to another.

A caveat to carry with you. This is exact when a linear layer is immediately followed by normalization, as in a ResNet with BatchNorm. In a pre-norm transformer the matrices feed a residual stream that is normalized later, so the property holds only approximately. The paper handles this with a heuristic that I describe in Step 8.

Takeaway. For a normalized layer the loss depends on the direction of the weights and not on their length.

Step 3

Two consequences for the gradient

If moving along a ray never changes the loss, then the gradient has no reason to point along the ray. It must point sideways. And if you are far from the origin, you have to travel further sideways to rotate by the same angle, so the loss changes more slowly per unit of distance. Both facts drop out of one line of calculus.

Consequence A. The gradient is perpendicular to the weights.

Start from scale invariance and differentiate both sides with respect to the stretch factor $c$ .

\frac{d}{dc}L(c\,\w)=\frac{d}{dc}L(\w)

The right side does not contain

c

, so its derivative is

0

. The left side needs the chain rule. The input

c\,\w

changes at rate

\w

c

changes.

\big\langle \nabla L(c\,\w),\ \w\big\rangle=0

Now set

c=1

and write

\g=\nabla L(\w)

\langle \g,\w\rangle=0

Consequence B. Longer weights mean smaller gradients.

This time differentiate both sides of $L(c\,\w)=L(\w)$ with respect to $\w$ , holding $c$ fixed.

c\,\nabla L(c\,\w)=\nabla L(\w)

Chain rule again. Each entry of the input is multiplied by

c

, so a factor of

c

comes out.

\nabla L(c\,\w)=\frac{1}{c}\,\nabla L(\w)

Double the weights and the gradient halves. So the product $\|\g\|\cdot\|\w\|$ stays the same along a ray. I will call that product $\kappa$ . It measures how steep the loss is in direction space, with the arbitrary length of $\w$ divided out.

Tiny example

Use the toy loss $L(\w)=1-\cos(\text{angle between }\w\text{ and }(1,0))$ , which only depends on direction. At $\w=(3,4)$ its gradient is $\g=(-0.128,\ 0.096)$ .

Check A. $\langle\g,\w\rangle=-0.128\times3+0.096\times4=-0.384+0.384=0$ .

Check B. $\|\g\|=0.16$ . At the doubled point $(6,8)$ the gradient is $(-0.064,\ 0.048)$ with length $0.08$ , exactly half. The product $\kappa=\|\g\|\|\w\|$ is $0.8$ both times.

Try this. Drag the blue point outward along a ray and watch the crimson gradient arrow shrink while staying at a right angle. Then drag it around a circle and watch the gradient size change with direction but the right angle stay.

Real computation of the toy loss above. Grey rays are lines of equal loss. The dashed ray is the best direction. The gradient arrow is drawn 3 times its true length so it stays visible.

Takeaway. For a scale invariant layer,

\g\perp\w

and

\|\g\|=\kappa/\|\w\|

. The raw gradient norm is partly a statement about how long the weights are.

Step 4

Perpendicular steps always make the weights longer

Take one plain gradient step with no weight decay. You stand at $\w$ and move at a right angle to it. By Pythagoras, the new point is further from the origin than the old one. It does not matter which sideways direction you chose.

\|\w-\lr\g\|^2=\|\w\|^2-2\lr\langle\g,\w\rangle+\lr^2\|\g\|^2

This is just expanding the square

\|a-b\|^2=\|a\|^2-2\langle a,b\rangle+\|b\|^2

. By Consequence A the middle term is zero.

\|\w-\lr\g\|^2=\|\w\|^2+\lr^2\|\g\|^2

Tiny example

From $\w=(3,4)$ with $\g=(-0.128,0.096)$ , take a deliberately huge step $\lr=10$ so the numbers are visible. The new point is $(4.28,\ 3.04)$ . Its squared length is $18.32+9.24=27.56$ . The formula predicts $25+100\times0.0256=27.56$ . The length grew from $5$ to $5.25$ .

So without any opposing force the weights of a normalized layer grow forever. That growth is not free. By Consequence B a longer $\w$ means a smaller gradient, and the same step now rotates $\w$ by a smaller angle. Van Laarhoven (2017) calls $\lr/\|\w\|^2$ the effective step size, and here it decays in a way nobody asked for.

Takeaway. Every gradient step on a normalized layer adds

\lr^2\|\g\|^2

to the squared length of the weights.

Step 5

Weight decay pulls back, and the two settle at an equilibrium

Weight decay is the opposing force. Each step it multiplies the weights by a number slightly below one. Now there is a tug of war. Gradient steps push the length up, and decay pulls it down.

The new idea in this step is equilibrium, also called a steady state. Think of a bathtub with the tap running and the drain open. The drain lets out more water when the tub is fuller. So the level rises until outflow matches inflow, and then it stays there. If you start too full, it drains down to that same level. An equilibrium is a value that the process keeps returning to because the two opposing effects cancel exactly there.

Our tug of war has the same shape. Decay removes more length when the weights are long. Gradient steps add more length when the weights are short, because short weights have large gradients. So there is one length where the two cancel. To keep the algebra clean I use plain SGD with weight decay first and return to Adam in Step 7.

Derivation

\w_{t+1}=\w_t-\lr\g_t-\lr\wd\w_t=(1-\lr\wd)\,\w_t-\lr\g_t

The SGD update with decay, with the two

\w_t

terms grouped.

\|\w_{t+1}\|^2=(1-\lr\wd)^2\|\w_t\|^2-2(1-\lr\wd)\lr\langle\g_t,\w_t\rangle+\lr^2\|\g_t\|^2

Expand the square exactly as in Step 4.

\|\w_{t+1}\|^2=(1-\lr\wd)^2\|\w_t\|^2+\lr^2\|\g_t\|^2

The middle term vanishes because

\g_t\perp\w_t

At equilibrium the length stops changing, so set $\|\w_{t+1}\|=\|\w_t\|=\|\w\|$ .

\|\w\|^2=(1-2\lr\wd+\lr^2\wd^2)\|\w\|^2+\lr^2\|\g\|^2

Expanded

(1-\lr\wd)^2

(2\lr\wd-\lr^2\wd^2)\,\|\w\|^2=\lr^2\|\g\|^2

Moved the

\|\w\|^2

terms to the left.

2\lr\wd\,\|\w\|^2\approx\lr^2\|\g\|^2

Dropped

\lr^2\wd^2

because it is tiny next to

2\lr\wd

. With the example values from Step 1,

2\lr\wd=0.0006

and

\lr^2\wd^2=0.00000009

\frac{\|\g\|}{\|\w\|}=\sqrt{\frac{2\wd}{\lr}}

Divided both sides by

\lr^2\|\w\|^2

and took the square root.

This is the central formula of the paper. At equilibrium the ratio of gradient length to weight length is pinned by the two hyperparameters and by nothing else. Since $\g_t$ is a noisy minibatch gradient, read it as holding on average.

We can go one step further using Consequence B, $\|\g\|=\kappa/\|\w\|$ . Substituting gives each length separately.

\|\w\|=\Big(\frac{\lr\,\kappa^2}{2\wd}\Big)^{1/4}\qquad\qquad \|\g\|=\sqrt{\kappa}\,\Big(\frac{2\wd}{\lr}\Big)^{1/4}

This second line needs an extra assumption, that $\kappa$ stays roughly constant over time. That is reasonable in single-pass training where minibatch noise keeps the gradient from vanishing. It is the route taken by van Laarhoven in 2017. Defazio's main argument avoids it by only talking about the ratio.

Tiny example

Take $\lr=0.1$ , $\wd=0.05$ , $\kappa=8$ . Then $\|\w\|^4=0.1\times64/0.1=64$ , so $\|\w\|=2.83$ and $\|\g\|=8/2.83=2.83$ . The ratio is $1$ , and the formula says $\sqrt{2\times0.05/0.1}=1$ .

Pause and predict.

The plot below starts the same recurrence from two lengths, one far too short and one far too long. Before you look, where do you expect each curve to end up? And if you double $\wd$ , does the final length go up or down, and by a lot or a little?

Try this. Move the sliders and check that both starting points always land on the dashed prediction. Notice how slowly the level responds to weight decay. That is the fourth root at work.

learning rate weight decay

Real iteration of

\|\w_{t+1}\|^2=(1-\lr\wd)^2\|\w_t\|^2+\lr^2\kappa^2/\|\w_t\|^2

with

\kappa=8

. No noise and no neural network, just the recurrence.

One more fact from this picture matters later. Reaching equilibrium takes time, on the order of $1/(\lr\wd)$ steps. A smaller learning rate means a slower approach.

Takeaway. With a fixed learning rate, a normalized layer settles where

\|\g\|/\|\w\|=\sqrt{2\wd/\lr}

Step 6

A schedule moves the equilibrium

Everything so far used a constant learning rate. Real runs use a schedule, usually warmup and then a decay toward zero. Put $\lr_t$ into the formula and the equilibrium becomes a moving target.

\text{target ratio at step } t=\sqrt{\frac{2\wd}{\lr_t}}

The learning rate sits in the denominator. As the schedule decays, the target ratio rises, and it rises without bound if the schedule goes all the way to zero.

Tiny example

With $\wd=0.1$ , a learning rate of $0.003$ gives a target of $\sqrt{66.7}=8.2$ . A learning rate of $0.0003$ gives $\sqrt{667}=25.8$ . Ten times less learning rate, about $3.2$ times more ratio.

How does the ratio actually rise? From the two formulas at the end of Step 5, the weights shrink like $\lr_t^{1/4}$ and the gradient grows like $\lr_t^{-1/4}$ . Mechanically, a lower learning rate makes the outward push $\lr^2\|\g\|^2$ collapse faster than the inward pull $2\lr\wd\|\w\|^2$ . Decay wins, the weights shrink, and by Consequence B shorter weights have bigger gradients.

That is the whole phenomenon. The gradient norm rises at the end of training because the schedule is dragging the equilibrium, and the weights are following it inward.

Pause and predict.

Suppose the learning rate has decayed to one sixteenth of its peak and the layer has kept up with its equilibrium. By what factor has the gradient norm grown under SGD with weight decay?

Reveal

The gradient scales like $\lr_t^{-1/4}$ , and $16^{1/4}=2$ . So it has doubled. Step 7 shows that for AdamW the exponent is $1/2$ , which gives a factor of $4$ .

Try this. Switch between cosine and linear decay. Look at how late the blow-up happens under cosine, and compare the three curves in the bottom panel.

Computed from the equilibrium formulas, assuming the layer tracks its target perfectly. Real training lags behind near the very end, as the next paragraph explains. The vertical axis is cut off at 10.

Two honest caveats. First, the equilibrium takes about $1/(\lr_t\wd)$ steps to reach, and that time grows as the learning rate falls. So near the very end the layer cannot keep up, and the real ratio rises less than the target. Defazio's ResNet plot shows exactly this. The prediction tracks the measurement closely until the final stretch. Second, short runs may never reach equilibrium at all. The paper notes that a 10 billion token run was too short to show the uptick that a 200 billion token run showed clearly.

Takeaway. A decaying learning rate raises the target

\sqrt{2\wd/\lr_t}

, which shrinks the weights and inflates the gradient norm.

Step 7

The same story for Adam, where it is stronger

Adam changes one thing in the argument. Its step does not get smaller when the gradient gets smaller. Adam divides the gradient by its own running size, so each entry of the update $u_t$ is roughly $+1$ or $-1$ . The paper calls this the sign-descent regime. With $d$ entries, $\|u_t\|^2\approx d$ no matter how long the weights are.

Rerun the tug of war with this constant-size push. I treat $u_t$ as uncorrelated with $\w_t$ on average, which plays the role that exact perpendicularity played for SGD.

\|\w_{t+1}\|^2\approx(1-\lr\wd)^2\|\w_t\|^2+\lr^2 d

Same expansion as Step 5, with

\|u_t\|^2\approx d

in place of

\|\g_t\|^2

2\lr\wd\,\|\w\|^2\approx\lr^2 d

Set the two lengths equal and dropped the tiny

\lr^2\wd^2

term, as before.

\|\w\|\approx\sqrt{\frac{\lr\,d}{2\wd}}

Under SGD the push weakened as the weights shrank, which softened the effect to a fourth root. Under Adam the push has fixed size, so the weights respond with a full square root. Then Consequence B gives the gradient.

\|\g\|=\frac{\kappa}{\|\w\|}\approx\kappa\sqrt{\frac{2\wd}{\lr\,d}}

Tiny example

Take $d=64$ , $\lr=0.01$ , $\wd=0.5$ . The predicted length is $\sqrt{0.01\times64/1}=0.8$ . The lab in Step 9 uses exactly these values and measures about $0.78$ . If the learning rate falls by $4$ , the predicted length halves to $0.4$ and the gradient norm doubles.

How this relates to the sources. The formula $\sqrt{\lr d/(2\wd)}$ matches the AdamW equilibrium norm in Kosson, Messmer and Jaggi's rotational equilibrium paper. Defazio states his Adam result differently. He shows the ratio $\sqrt{2\wd/\lr}$ holds when both lengths are measured in norms reweighted by Adam's preconditioner, and then simplifies it to say the largest weight entry is about $\sqrt{\lr/(2\wd)}$ . That is the per-entry version of the formula above. He labels that simplification as non-rigorous, and the derivation ignores momentum. Adam's update is also not exactly perpendicular to the weights. So treat this step as a good approximation and not a theorem.

Takeaway. Under AdamW the equilibrium weight length scales like

\sqrt{\lr_t}

, so the gradient norm scales like

1/\sqrt{\lr_t}

Step 8

The fix

The target ratio is $\sqrt{2\wd/\lr_t}$ . It moves because the denominator moves and the numerator does not. So let the numerator move too. Replace the constant $\wd$ with a scheduled value that follows the learning rate.

\wd_t=\wd\cdot\frac{\lr_t}{\lr_{\max}}

Plug it into the target.

\sqrt{\frac{2\wd_t}{\lr_t}}=\sqrt{\frac{2\wd}{\lr_t}\cdot\frac{\lr_t}{\lr_{\max}}}=\sqrt{\frac{2\wd}{\lr_{\max}}}

The

\lr_t

cancels. The target is now a constant for the whole run.

Written as an update rule, AdamC changes only the decay term, and only for normalized layers.

\w_{t+1}=\w_t-\lr_t\,u_t-\frac{\lr_t^{\,2}}{\lr_{\max}}\,\wd\,\w_t

Compare with Step 1. The per-step shrink was $\lr_t\wd$ . Now it is $\lr_t^2\wd/\lr_{\max}$ . When the schedule halves the learning rate, AdamW halves the shrink and AdamC quarters it.

Tiny example

Same values as Step 1, $\lr_{\max}=0.003$ and $\wd=0.1$ . At the peak both methods shrink by $0.0003$ per step. When $\lr_t=0.0003$ , AdamW shrinks by $0.00003$ and AdamC by $0.0003^2\times0.1/0.003=0.000003$ . Ten times gentler, which is what it takes to stop the weights from collapsing inward.

Three design details

Why divide by the peak. Dividing by $\lr_{\max}$ makes AdamC identical to AdamW at the peak learning rate. So $\wd$ keeps its old meaning and a value tuned for AdamW is a sensible starting point.

Why not make decay independent of the learning rate. You might guess the cure is to stop multiplying decay by $\lr_t$ at all. Redo Step 5 with a shrink of $\wd$ per step and you get a ratio of $\sqrt{2\wd}/\lr_t$ . That depends on the schedule even more strongly. The paper makes this point explicitly.

Which parameters. The argument only applies to scale invariant weights. Algorithm 1 in the paper applies the corrected decay to normalized layers and ordinary AdamW decay to everything else. For the transformer experiments the paper treats every linear layer as normalized except the output layer. That is a heuristic, since transformer layers are not literally followed by a norm.

In PyTorch you do not need a new optimizer class. AdamW already multiplies weight_decay by lr, so scheduling weight_decay with the same factor as the learning rate gives the AdamC update. This is my own sketch and not code from the paper.

# groups[0]: weight matrices treated as normalized
# groups[1]: output layer and anything else that gets plain AdamW decay
opt = torch.optim.AdamW(groups, lr=lr_max, weight_decay=wd)

for step in range(total_steps):
    f = schedule(step)                 # lr_t / lr_max, between 0 and 1
    for group in opt.param_groups:
        group["lr"] = lr_max * f
    opt.param_groups[0]["weight_decay"] = wd * f   # the AdamC correction
    ...
    opt.step()

Takeaway. AdamC sets

\wd_t=\wd\,\lr_t/\lr_{\max}

on normalized layers, which freezes the equilibrium at

\sqrt{2\wd/\lr_{\max}}

for the whole schedule.

Step 9

Lab. Run both and watch

This lab trains a real, tiny, scale invariant problem in your browser. The weight vector has $d=64$ entries. The loss is half the squared distance between the normalized weights and a noisy target direction, so only the direction of $\w$ matters. Each step draws fresh noise, like a new minibatch, so the gradient never dies out. Both optimizers see identical noise. The error panel reports $1-\cos$ of the angle between $\w$ and the true direction, so zero means perfectly aligned.

Pause and predict.

With a cosine schedule, sketch in your head the weight length and gradient norm for AdamW and for AdamC. Which one keeps a flat weight length? Then guess which one rotates the weights more per step in the middle of training.

Try this. Start with the defaults. Then switch the schedule to constant and confirm the two methods coincide. Then drag weight decay to its minimum and see the effect fade, because the layer no longer has time to reach equilibrium.

peak learning rate weight decay

━ uncorrected ━ corrected ┄ predicted equilibrium

Real training run of 4000 steps, recomputed on every change. Curves are averaged in blocks of 20 steps. "Uncorrected" is AdamW or SGD with standard decay. "Corrected" is AdamC or SGDC. Adam uses

\beta_1=0.9

\beta_2=0.95

. This is a toy with one layer, so it shows the mechanism and says nothing about how large the benefit is in a real model.

With the defaults you should see the uncorrected gradient norm more than double while the corrected one stays flat. The uncorrected weight length follows its dashed prediction inward. The error panel usually ends a little lower for the corrected run, though that gap moves around when you draw new noise.

Takeaway. In a problem that satisfies the assumptions, the predicted equilibria match the simulation closely, and the correction flattens both the weight length and the gradient norm.

Step 10

What you actually buy. The schedule means what it says

A flat gradient norm plot is pleasant, but it is not the point. The point is in the "rotation per step" panel of the lab.

For a scale invariant layer, the only thing that changes the network's function is the direction of $\w$ . So the honest measure of learning speed is how far the direction rotates per step. For a step of size $\lr\|u\|$ taken sideways from a point at distance $\|\w\|$ , the rotation angle is about their ratio.

\text{rotation per step}\approx\frac{\lr_t\,\|u_t\|}{\|\w_t\|}

Under AdamW, substitute $\|u\|\approx\sqrt d$ and the equilibrium length from Step 7.

\frac{\lr_t\sqrt d}{\sqrt{\lr_t d/(2\wd)}}=\sqrt{2\,\lr_t\,\wd}

The

\sqrt d

cancels and

\lr_t/\sqrt{\lr_t}=\sqrt{\lr_t}

So under AdamW the rotation per step follows the square root of your schedule. If you ask the learning rate to fall by a factor of 100, the layer only slows down by a factor of 10. The shrinking weights quietly undo half of your schedule, in log terms.

Under AdamC the length is frozen at $\sqrt{\lr_{\max}d/(2\wd)}$ . Substitute that instead.

\frac{\lr_t\sqrt d}{\sqrt{\lr_{\max} d/(2\wd)}}=\lr_t\sqrt{\frac{2\wd}{\lr_{\max}}}

Now the rotation is proportional to $\lr_t$ . A hundredfold decay in the schedule is a hundredfold decay in learning speed.

Provenance. The AdamW expression $\sqrt{2\lr\wd}$ appears in Kosson, Messmer and Jaggi's table of equilibrium quantities, with an extra factor for momentum that I have left out. The AdamC line is my own algebra using the same recipe. Defazio's note does not phrase its result in terms of rotation. It argues from the gradient to weight ratio and from the observed loss curves.

Is the square-root behaviour of AdamW actually bad? That is an empirical question and not something the derivation settles. One could tune a schedule under AdamW that already compensates for it. What AdamC offers is that the learning rate schedule you write down is the schedule the normalized layers experience.

Takeaway. Under AdamW the effective learning rate of a normalized layer decays like

\sqrt{\lr_t}

. Under AdamC it decays like

\lr_t

What the evidence shows, and its limits

Defazio's experiments. The language model run is a 120 million parameter model with the Llama 3 architecture, trained on FineWeb-Edu for 200 billion tokens with a cosine schedule, a batch of 256 thousand tokens and weight decay 0.05. The learning rate was swept from 0.001 to 0.02 for both optimizers. AdamC largely removes the gradient norm increase in the second half of training, keeps the weight norm far more stable, and reaches lower loss. For ResNet-50 on ImageNet with SGD, the corrected version removes the sharp rise at the end, but a slower upward drift across the whole run remains and the theory does not explain it.

Limits. This is a short single-author arXiv note with two test problems. The transformer application relies on the heuristic from Step 8. The Adam analysis is approximate. The paper itself says that many standard benchmarks do not show the uptick at all, because they lack normalized layers, are too short, or overfit.

Follow-up work. Jason Chou's "Correction of Decoupled Weight Decay" (arXiv, December 2025) revisits the question. From its abstract, it agrees that the per-step decay should scale like $\lr^2$ to keep weight norms stable, but it disputes the perpendicularity argument as the reason. It derives the same scaling from a different assumption, that updates become independent of the weights at steady state. I have only verified that paper at the level of its abstract.

Summary of the whole chain

AdamW shrinks weights by $\lr_t\wd$ per step, so the schedule also schedules the decay.
Normalization makes the loss blind to the length of the weights feeding it.
That forces $\g\perp\w$ and $\|\g\|=\kappa/\|\w\|$ .
Perpendicular steps lengthen $\w$ by Pythagoras.
Decay shortens it. The two balance where $\|\g\|/\|\w\|=\sqrt{2\wd/\lr}$ .
A decaying $\lr_t$ raises that target. Weights shrink and the gradient norm climbs.
Under Adam the weights scale like $\sqrt{\lr_t}$ , so the effect is stronger.
Setting $\wd_t=\wd\,\lr_t/\lr_{\max}$ cancels the schedule out of the target.
The result is a constant weight length, a flat gradient norm, and an effective learning rate that follows the schedule you wrote.

Common misconceptions

"A rising gradient norm means training is going unstable." Not here. For normalized layers the gradient norm is $\kappa/\|\w\|$ . It can double purely because the weights halved, with the loss surface in direction space unchanged.
"Decoupled weight decay means decay is independent of the learning rate." Decoupled means the decay is applied directly to the weights instead of being added to the gradient before Adam's normalization. In the standard implementation it is still multiplied by $\lr_t$ .
"AdamC just turns weight decay off at the end." It reduces decay in proportion to the learning rate, which is exactly the amount needed to hold the equilibrium still. Removing decay entirely would let the weights grow.
"The correction applies to every parameter." The argument only covers scale invariant weights. Other parameters keep ordinary AdamW decay.
"A flat gradient norm is the goal." It is a symptom. The substantive change is in the effective learning rate of Step 10.
"AdamC is a new optimizer." It is AdamW with a schedule on one hyperparameter for some parameter groups.

Blind spots worth knowing about

Effective learning rate as the real knob. Zhang and coauthors' "Three Mechanisms of Weight Decay" argues that in normalized networks weight decay mainly works by controlling the effective learning rate, not by regularizing in the classical sense. This reframes $\wd$ as an optimization hyperparameter. It also explains why the best learning rate and the best weight decay tend to move together when you tune them.
Rotational equilibrium. Kosson, Messmer and Jaggi develop the equilibrium picture for AdamW, SGD with momentum and Lion, per neuron and per layer. They also connect the transient phase before equilibrium to the need for learning rate warmup. That is the natural next paper to read.
Equilibrium has a time scale. It takes on the order of $1/(\lr\wd)$ steps to reach. With the Step 1 values that is about $1/(0.003\times0.1)\approx3300$ steps at the peak. A run of only ten thousand steps or so has time to approach equilibrium but not much to spare, so whether it shows the late uptick is worth checking in its gradient norm logs rather than assuming.
Removing the degree of freedom entirely. If the weight length carries no information, you can project the weights back to a fixed length after every step. Then there is no equilibrium to manage. That is the idea behind the rotational optimizer variants of Kosson and coauthors. AdamC is the minimal fix and projection is the structural one.
Adam's update is not exactly perpendicular. While building the lab I first used a target aligned with a single coordinate axis. Adam's per-coordinate rescaling then amplified a tiny radial gradient component, and the corrected run's weight length drifted upward instead of staying flat. A random target direction removed the artifact. The lesson is that the clean theory depends on Adam's preconditioner not singling out the radial direction.
Gradient clipping lives in absolute units. Many language model recipes clip the global gradient norm at 1.0. If weight lengths drift, the same threshold bites more or less often over time. This is my own inference and not a claim from the paper.
Which parameters are decayed at all. Most recipes already exclude embeddings, normalization parameters and biases from decay. So the decayed set is mostly the weight matrices, which is also where AdamC applies.
Earlier work on the same symptom. Xie and coauthors (2023) documented gradient norms rising because of weight decay and proposed scheduling the decay. Defazio's contribution is a derivation that says exactly which schedule cancels the effect.

Sources and further reading

Aaron Defazio, Why Gradients Rapidly Increase Near the End of Training, arXiv 2506.02285, June 2025. The AdamC paper.
Twan van Laarhoven, L2 Regularization versus Batch and Weight Normalization, 2017. The original equilibrium argument.
Guodong Zhang, Chaoqi Wang, Bowen Xu, Roger Grosse, Three Mechanisms of Weight Decay Regularization, ICLR 2019.
Atli Kosson, Bettina Messmer, Martin Jaggi, Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks, ICML 2024.
Jason Chuan-Chih Chou, Correction of Decoupled Weight Decay, arXiv 2512.08217.
Kaiyue Wen and coauthors, The Hitchhiker's Guide to the Weight Norm Theory. A longer blog post on how weight norms evolve under weight decay.
Ilya Loshchilov, Frank Hutter, Decoupled Weight Decay Regularization, ICLR 2019. The AdamW paper.