Decoupled Weight Decay and Adaptive Scaling

Weight decay in a first-order optimizer can be expressed two ways that agree exactly for plain stochastic gradient descent with momentum and diverge for adaptive methods. In the first, the regulariser’s gradient $\lambda \theta$ is added to the loss gradient before the optimizer’s update rule is applied. In the second, the parameter is multiplied by $(1 - \eta\lambda)$ at each step, independently of whatever direction the loss gradient produced. For optimizers that scale each coordinate by an estimate of its gradient magnitude, the first form produces a decay whose strength varies across parameters by the ratio of their gradient statistics, and the second does not.

Two updates that look the same

Consider the $L_2$-regularised objective $f(\theta) + \frac{\lambda}{2}|\theta|^2$. Its gradient is $g + \lambda\theta$, and SGD with learning rate $\eta$ performs $\theta \leftarrow \theta - \eta g - \eta\lambda \theta = (1-\eta\lambda)\theta - \eta g$. The algebra is identical to applying $f$‘s gradient and then shrinking. For SGD the two formulations are the same function of the parameter, so a framework’s choice of which to implement is invisible.

Insert Adam between the gradient and the parameter and the equivalence breaks. Adam’s update is $\theta \leftarrow \theta - \eta , \hat{m}/(\sqrt{\hat{v}} + \epsilon)$, and if the regulariser’s gradient was folded into $g$ then $\hat{m}$ and $\hat{v}$ both contain contributions from $\lambda\theta$. The decay’s effect on the update is then scaled by $1/(\sqrt{\hat{v}}+\epsilon)$ per coordinate. A parameter with a large gradient history has a large $\hat{v}$ and receives proportionally less shrink; a parameter with a small gradient history, typically one that has not been doing useful work, receives proportionally more. The regulariser intended to penalise large weights has become a penalty on large weight-to-gradient ratio.

The magnitude of the difference is bounded by the range of the denominator. If two coordinates have second-moment estimates differing by a factor of a hundred, their decay strengths differ by a factor of ten, since the denominator enters under a square root.

What the coupled form does to the second moment

There is a second, less obvious effect. Because the regulariser’s contribution enters $v_t$, the decay itself inflates the denominator for large weights. A weight that grows large contributes $\lambda^2\theta^2$ to $g^2$ at every step, which stabilises $v$ at a value that includes the decay’s own magnitude. The decay thus attenuates its own effective strength in a way that depends on the weight’s size and on $\beta_2$: it is a fixed point, and where it settles is not the value the hyperparameter suggests. This is why $L_2$ regularisation with Adam and the equivalent coefficient for SGD do not produce comparable models at the same $\lambda$.

The decoupled update

The decoupled form applies the shrink multiplicatively to the parameter, after or before the adaptive step, without touching $m$ or $v$:

$$\theta \leftarrow \theta - \eta\left(\frac{\hat{m}}{\sqrt{\hat{v}}+\epsilon}\right) - \eta\lambda,\theta$$

Each coordinate now shrinks by the same relative amount per step, so a parameter decays toward zero with a time constant of $1/(\eta\lambda)$ steps regardless of its gradient history. The decay is a property of the optimizer, not of the regularised loss, and the gradient statistics remain a pure estimate of the loss’s geometry.

def step_decoupled_adam(p, m, v, grad, t, lr, b1, b2, eps, wd):
    m = b1 * m + (1 - b1) * grad
    v = b2 * v + (1 - b2) * grad * grad
    m_hat = m / (1 - b1 ** t)
    v_hat = v / (1 - b2 ** t)
    # decay is applied to the parameter itself, not through the gradient
    p = p * (1 - lr * wd) - lr * m_hat / (v_hat ** 0.5 + eps)
    return p, m, v

The order of the multiply and the adaptive step matters at second order: $(1-\eta\lambda)(\theta - \eta\Delta)$ differs from $\theta - \eta\Delta - \eta\lambda\theta$ by $\eta^2\lambda\Delta$, negligible for small $\eta\lambda$ but not zero. Frameworks differ in which order they implement, and the difference is only meaningful when the product approaches one.

The time constant is the useful way to tune $\lambda$ in this form. With $\eta = 10^{-3}$ and $\lambda = 10^{-2}$, the product is $10^{-5}$ per step, and a weight far from zero halves in about $7\times10^4$ steps. Expressed that way, the decay can be checked against the total number of steps the run will take: a decay time constant much longer than the run does nothing, and one much shorter erases weights before the loss can move them.

Where the decoupled form is not equivalent to a prior

The multiplicative decay is not the gradient of any fixed penalty in parameter space when $\eta$ or the adaptive scaling varies. For constant $\eta$ the iteration $\theta \leftarrow (1-\eta\lambda)\theta$ is a contraction toward zero whose fixed point is zero and whose continuous-time limit is $\dot\theta = -\lambda\theta$, which is the gradient flow of $\frac{\lambda}{2}|\theta|^2$ with a rescaled time. With a schedule in which $\eta$ changes, the accumulated decay depends on the integral of $\eta_t \lambda$ over the run, not on the number of steps, so a decay coefficient tuned for a schedule with warmup and cosine decay is tied to that schedule’s shape. Changing the peak rate or the decay shape changes the total shrinkage.

This is easiest to see by taking the log: the accumulated multiplicative factor is $\prod_t (1-\eta_t\lambda) \approx \exp(-\lambda\sum_t \eta_t)$. The sum of learning rates, not the step count, determines shrinkage. A run that spends a long time at a small rate shrinks less than the same number of steps at a large rate.

Interaction with normalization layers

The scale invariance of normalisation layers changes what weight decay does to the parameters preceding them. If a sublayer’s output is followed by a normalisation that divides by the output’s norm, then multiplying the sublayer’s weights by any positive constant leaves the normalized output unchanged and leaves the loss unchanged. Gradient descent has no component that changes that constant, so the radial direction is unconstrained by the data. Weight decay is the only term that acts on it, and it pulls the norm toward zero until the normalised output’s direction begins to change — that is, until the weights become small enough that the additive $\epsilon$ or the bias term dominates. The effective decay is therefore not a penalty on model complexity but a slow motion along a flat direction, which does not affect the loss until it does. The variance preservation argument that fixes an appropriate initial scale also fixes the scale at which this motion starts, which is one reason decay interacts with initializer choice in ways that a fixed $\lambda$ does not capture.

Failure modes

Three are common. Setting the decay coefficient to the value used for SGD, whose coupled form has a different effective strength, produces under-regularisation for the coupled Adam form and roughly the intended effect for the decoupled one, which is why decoupled implementations often describe their coefficient as “the same $\lambda$ as SGD”. Applying decay to bias and normalisation gains, which are not weights whose magnitude encodes complexity, tilts the model toward zero offsets and unit gains; the usual practice of excluding them from the decay list is a model of which parameters the regulariser’s prior applies to, not an optimisation detail.

The third is the interaction with gradient clipping, since the clip is applied to the loss gradient before the optimizer sees it, and the decay is applied after. A clip that binds reduces the loss’s contribution to the step but not the decay’s, so under sustained clipping the parameter shrinks while the loss gradient cannot counteract it. The bound described in gradient clipping and norm scaling is thus not a neutral safety measure when a multiplicative decay is in force.

A fourth is diagnostic rather than functional: because the shrink is a fixed multiplicative factor per step, the parameter’s norm under pure decay decays geometrically, so a logged norm that declines along a straight line on a log scale is decaying and one that plateaus is being held by the loss gradient. Reading that difference off a training log distinguishes a decay that is doing work from one that is overpowered by the loss within a few hundred steps, and it costs nothing to compute.