Weight Decay

Dive into Deep Learning · §2.7

Taming overfitting by shrinking the weights
the \ell_2 penalty · the geometry · the spectral why · the Bayesian reading.

When more data is not an option

Motivation

  • Overfitting fades with more data, but data is often costly or fixed.
  • Dropping features is too blunt: with k inputs there are \binom{k-1+d}{k-1} degree-d monomials.
  • Instead, restrict the values the weights may take.

Among all f, the constant f=0 is the simplest. Measure complexity by how far the weights sit from zero.

\|\mathbf{w}\|_2^2 = \sum_i w_i^2

A single number that grows as the weights stretch away from the origin.

Add the size of the weights to the loss

The idea

Penalize a large weight vector by adding its squared norm to the objective, scaled by a knob \lambda \ge 0:

L(\mathbf{w}, b) \;+\; \frac{\lambda}{2}\,\|\mathbf{w}\|_2^2.

  • \lambda = 0 recovers the plain loss; larger \lambda pulls \mathbf{w} harder toward zero.
  • The \tfrac{1}{2} is cosmetic: it cancels the 2 from differentiating the square.

\ell_2-regularized linear regression is the classic ridge regression; the \ell_1 version is lasso.

Ridge shrinks, lasso selects

The geometry

A budget \|\mathbf{w}\| \le t turns the penalty into a constraint: the answer is where a loss contour first touches the constraint region.

Loss contours centred on the unconstrained optimum \hat{\mathbf{w}} grow until they meet the constraint at \mathbf{w}^\star. Left (\ell_2 ball): contact is tangential, so both coordinates shrink. Right (\ell_1 diamond): contact is at a corner, forcing w_2 to exactly zero.

A round ball generally produces continuous shrinkage without exact zeros; a pointed diamond can touch at a corner, producing sparsity. Thus lasso can perform feature selection whereas ridge generally does not.

Why it is called weight decay

The update

The penalty adds \lambda\mathbf{w} to the gradient, so each SGD step gains a shrink factor:

\mathbf{w} \leftarrow (1 - \eta\lambda)\,\mathbf{w} \;-\; \frac{\eta}{|\mathcal{B}|}\!\sum_{i \in \mathcal{B}} \mathbf{x}^{(i)}\bigl(\mathbf{w}^\top\mathbf{x}^{(i)} + b - y^{(i)}\bigr).

Before fitting the data at all, every weight is decayed toward zero by the factor 1 - \eta\lambda.

\lambda controls shrinkage continuously, whereas deleting a feature imposes a discrete constraint.

Usually the bias is left undecayed.

01

From Scratch

an underdetermined example with and without regularization

An underdetermined regression problem

From Scratch

A tiny linear signal in 200 inputs plus faint noise, y = 0.05 + \sum_i 0.01\,x_i + \epsilon:

class Data(d2l.DataModule):
    def __init__(self, num_train, num_val, num_inputs, batch_size):
        self.save_hyperparameters()                
        n = num_train + num_val 
        self.X = d2l.randn(n, num_inputs)
        noise = d2l.randn(n, 1) * 0.01
        w, b = d2l.ones((num_inputs, 1)) * 0.01, 0.05
        self.y = d2l.matmul(self.X, w) + b + noise

    def get_dataloader(self, train):
        i = slice(0, self.num_train) if train else slice(self.num_train, None)
        return self.get_tensorloader([self.X, self.y], train, i)

20 examples for 200 parameters leaves many parameter directions weakly constrained and creates substantial potential for overfitting.

The penalty, then the model

From Scratch

The penalty is a single line:

def l2_penalty(w):
    return d2l.reduce_sum(w**2) / 2

Subclass the scratch regressor and fold it into the loss:

class WeightDecayScratch(d2l.LinearRegressionScratch):
    def __init__(self, num_inputs, lambd, lr, sigma=0.01):
        super().__init__(num_inputs, lr, sigma)
        self.save_hyperparameters()
        
    def loss(self, y_hat, y):
        return (super().loss(y_hat, y) +
                self.lambd * l2_penalty(self.w))

Nothing else changes: same linear model, same squared loss. The loss includes the penalty scaled by lambd.

\lambda = 0: training and validation diverge

From Scratch · the overfit

train_scratch(0)

L2 norm of w: 0.009527594782412052

Training loss decreases sharply while validation loss remains high. The printed \|\mathbf{w}\|^2 shows that the unregularized fit uses large weights.

\lambda = 3: regularization reduces the gap

From Scratch · the rescue

train_scratch(3)

L2 norm of w: 0.00048313659499399364

Training loss is higher, but validation loss decreases and \|\mathbf{w}\|^2 is an order of magnitude smaller.

Why shrinkage helps: damp directions weakly constrained by the data

From Scratch · the why

Ridge keeps a closed form, and the SVD \mathbf{X} = \mathbf{U}\mathbf{D}\mathbf{V}^\top shows exactly what shrinks: the response along the j-th principal direction is damped by

\frac{d_j^2}{d_j^2 + \tilde{\lambda}} \qquad\Rightarrow\qquad \textrm{df}(\tilde{\lambda}) = \sum_j \frac{d_j^2}{d_j^2 + \tilde{\lambda}}.

Strong directions (d_j large) pass through nearly untouched; the weakly constrained ones are suppressed hardest.

Twenty examples pin down at most 20 of 200 directions (\textrm{df}(0) = 20); \lambda = 3 prices the model at \textrm{df} \approx 15 effective parameters.

02

Framework Implementations

where the framework keeps the decay

Decay lives in the optimizer

Concise · MXNet

Gluon’s Trainer takes wd directly; set wd_mult = 0 on the bias so only the weights decay:

class WeightDecay(d2l.LinearRegression):
    def __init__(self, wd, lr):
        super().__init__(lr)
        self.save_hyperparameters()
        self.wd = wd
        
    def configure_optimizers(self):
        for p in self.collect_params('.*bias').values():
            p.wd_mult = 0
        return gluon.Trainer(self.collect_params(),
                             'sgd', 
                             {'learning_rate': self.lr, 'wd': self.wd})

One Trainer argument replaces the hand-written penalty.

Same effect, less code

Concise · result

Fit with wd = 3: the validation curve matches the from-scratch run.

L2 norm of w: 0.0006210104911588132

A framework’s weight_decay adds \lambda\mathbf{w} to the gradient; the scratch penalty added \tfrac{\lambda}{2}\|\mathbf{w}\|^2 to the loss. Converged norms need not match exactly, only the effect.

The adaptive-optimizer reading: AdamW

Beyond linear models

Inside an Adam-style update each coordinate gets its own step size, so folding the penalty into the gradient rescales it per coordinate: it stops being uniform shrinkage.

Decoupling the decay from the adaptive step restores the intent of plain 1-\eta\lambda shrinkage. This is AdamW, a default for training large models.

For deep networks, weight decay is commonly applied to every layer’s weights, with exceptions such as biases configured separately.

The Bayesian reading: a prior on the weights

Beyond linear models

Put a zero-mean Gaussian prior on \mathbf{w}:

\mathbf{w}\sim\mathcal{N}(\mathbf{0},\lambda^{-1}\mathbf{I}) \;\Rightarrow\; -\log p(\mathbf{w}) = \tfrac{\lambda}{2}\|\mathbf{w}\|^2 + \textrm{const}.

Add it to the Gaussian-noise NLL from the linear-regression section:

\underbrace{-\log p(\mathbf{y}\mid\mathbf{X},\mathbf{w})}_{\textrm{MLE: }\,\frac{1}{2\sigma^2}\sum(\hat{y}-y)^2} \;\; \underbrace{-\log p(\mathbf{w})}_{=\,\frac{\lambda}{2}\|\mathbf{w}\|^2} \;\Rightarrow\; \textrm{MAP} = \textrm{ridge}.

MAP = MLE + a prior. The linear-regression section obtained squared loss from Gaussian noise; a Gaussian prior on \mathbf{w} adds weight decay, with \lambda=\sigma^2/(n\tau^2) under the averaged-loss convention.

A Gaussian prior centred at zero pulls the maximum-likelihood estimate back toward the origin.

Summary

Wrap-up

  • Weight decay = original loss +\ \tfrac{\lambda}{2}\|\mathbf{w}\|_2^2; per step it shrinks the weights by 1 - \eta\lambda before the data update.
  • Geometry: ridge shrinks (round ball), lasso selects (pointed diamond).
  • Spectral view: each direction damped by d_j^2/(d_j^2+\tilde\lambda); the 200-knob model ran at \textrm{df} \approx 15 effective parameters, a continuous dial tuned on a validation set.
  • The 20{\times}200 rig: \lambda=0 memorizes; \lambda=3 trades training error for a falling validation loss.
  • Frameworks expose decay in the optimizer (or layer / gradient transform).
  • Same idea scales up: AdamW for big models, a Gaussian prior in disguise.