🎓 A1Learner
🔍

PyTorch Weight Decay

Weight decay is a penalty that pulls weights toward zero during training. Smaller weights make a simpler model, and a simpler model often generalises better.

Weights that fit noise

Plain training puts no cost on large weights. A weight can grow just to fit one noisy row.

loss = loss_fn(model(x), y)
loss.backward()
optimizer.step()  # every weight is free to grow

Decay in the optimizer

The optimizer takes a weight_decay value at creation. The stored value reads back from the param groups.

import torch

torch.manual_seed(0)
model = torch.nn.Linear(3, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05, weight_decay=0.1)
print(optimizer.param_groups[0]["weight_decay"])

Decay as a loss penalty

A penalty adds a term built from the squared weights to the loss. Training then pays for large weights on every step.

import torch

torch.manual_seed(0)
model = torch.nn.Linear(3, 1)
penalty = 0.1 * sum((p ** 2).sum() for p in model.parameters())
print(f"Penalty added to the loss: {penalty.item():.4f}")

Example

import torch

torch.manual_seed(0)
x = torch.randn(32, 20)
true = torch.zeros(20)
true[0] = 2.0
true[1] = -1.0
true[2] = 0.5
y = x @ true + torch.randn(32) * 0.5
loss_fn = torch.nn.MSELoss()


def train(decay):
    torch.manual_seed(1)
    model = torch.nn.Linear(20, 1)
    optimizer = torch.optim.SGD(
        model.parameters(), lr=0.05, weight_decay=decay)
    for step in range(300):
        optimizer.zero_grad()
        loss = loss_fn(model(x).squeeze(1), y)
        loss.backward()
        optimizer.step()
    norm = sum(p.norm().item() ** 2 for p in model.parameters()) ** 0.5
    return loss.item(), norm


for decay in [0.0, 0.1]:
    loss, norm = train(decay)
    print(f"Decay {decay}: final loss {loss:.4f}, weight norm {norm:.4f}")

Key points

  • Weight norm: the length of all weights treated as one list, found with norm()
  • Use the weight_decay argument in an optimizer to turn decay on
  • Decay shrinks every weight a little on every step, before the gradient update
  • A penalty term in the loss is the other route, and it adds the weight sizes to the loss
  • Generalise: to work well on new data, not only on the training data
  • Weight decay trades a little training fit for smaller weights, as the closing example shows

In this example, the data has 20 input columns and only 3 true weights. The seed fixes the data, and the other 17 columns are pure noise. A train function builds the same Linear model for each decay value. Each run trains for 300 steps with SGD. Decay 0 ends with the lower loss, 0.0702, and the larger weight norm, 2.6155. Decay 0.1 ends with loss 0.1131 and a smaller norm, 2.2646. Decay gave up a little fit and kept the weights smaller.

Note

Weight decay is not a smaller learning rate, because decay pulls toward zero while the rate sets step size.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does weight decay do on every training step?