Weight decay (L2 regularization) in PyTorch

Introduction

Weight decay helps prevent a model from memorizing training data by keeping its weights small and simple.

When your model performs well on training data but poorly on new data (overfitting).

When training deep neural networks that tend to have many parameters.

When you want to improve your model's ability to generalize to unseen examples.

When you want to add a penalty to large weights during training to keep the model simpler.

Syntax

PyTorch

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, weight_decay=decay_rate)

weight_decay is the L2 regularization strength; typical values are small like 0.0001 or 0.01.

It is added as a penalty to the loss to discourage large weights.

Examples

SGD optimizer with learning rate 0.1 and L2 regularization strength 0.01.

PyTorch

optimizer = torch.optim.SGD(model.parameters(), lr=0.1, weight_decay=0.01)

Adam optimizer with a small weight decay to reduce overfitting.

PyTorch

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0001)

Sample Model

This code trains a simple linear model to fit y=2x+1 using SGD optimizer with weight decay to keep weights small. It prints loss each epoch and final weights.

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

# Simple linear model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
    def forward(self, x):
        return self.linear(x)

# Create model
model = SimpleModel()

# Dummy data: y = 2x + 1
x_train = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y_train = torch.tensor([[3.0], [5.0], [7.0], [9.0]])

# Loss function
criterion = nn.MSELoss()

# Optimizer with weight decay (L2 regularization)
optimizer = optim.SGD(model.parameters(), lr=0.1, weight_decay=0.01)

# Training loop
for epoch in range(10):
    optimizer.zero_grad()
    outputs = model(x_train)
    loss = criterion(outputs, y_train)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

# Print final weights
for name, param in model.named_parameters():
    print(f"{name}: {param.data.flatten().tolist()}")

Important Notes

Weight decay is applied automatically by the optimizer during the update step.

It is different from manually adding L2 penalty to the loss but mathematically equivalent.

Too large weight decay can cause underfitting by making weights too small.

Summary

Weight decay adds a penalty to large weights to reduce overfitting.

In PyTorch, set weight_decay in the optimizer to enable L2 regularization.

It helps models generalize better by keeping weights small and simple.

Practice

1. What is the main purpose of weight decay (L2 regularization) in training a PyTorch model?

easy

A. To reduce overfitting by penalizing large weights

B. To increase the learning rate automatically

C. To add more layers to the model

D. To speed up the training process

Solution

1. Step 1: Understand weight decay concept

   Weight decay adds a penalty to large weights during training to prevent the model from fitting noise in the data.

2. Step 2: Connect to overfitting reduction

   By keeping weights small, the model generalizes better and avoids overfitting.

3. Final Answer:

   To reduce overfitting by penalizing large weights -> Option A

4. Quick Check:

   Weight decay = reduces overfitting [OK]

Hint: Weight decay shrinks weights to avoid overfitting [OK]

Common Mistakes:

• Confusing weight decay with learning rate changes
• Thinking weight decay adds layers
• Assuming weight decay speeds training

2. Which of the following is the correct way to apply weight decay in a PyTorch optimizer?

easy

A. optimizer = torch.optim.SGD(model.parameters(), lr=0.01, wd=0.001)

B. optimizer = torch.optim.SGD(model.parameters(), lr=0.01, decay_weight=0.001)

C. optimizer = torch.optim.SGD(model.parameters(), lr=0.01, weightDecay=0.001)

D. optimizer = torch.optim.SGD(model.parameters(), lr=0.01, weight_decay=0.001)

Solution

1. Step 1: Recall PyTorch optimizer syntax

   PyTorch optimizers accept a parameter named weight_decay to apply L2 regularization.

2. Step 2: Identify correct parameter name

   Only optimizer = torch.optim.SGD(model.parameters(), lr=0.01, weight_decay=0.001) uses the exact parameter weight_decay correctly.

3. Final Answer:

   optimizer = torch.optim.SGD(model.parameters(), lr=0.01, weight_decay=0.001) -> Option D

4. Quick Check:

   Correct parameter name is weight_decay [OK]

Hint: Use exact parameter name 'weight_decay' in optimizer [OK]

Common Mistakes:

• Using wrong parameter names like decay_weight or wd
• Capitalizing parameter names incorrectly
• Confusing weight decay with learning rate

3. Consider this PyTorch code snippet:

import torch
model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, weight_decay=0.01)
initial_weight = model.weight.data.clone()
optimizer.zero_grad()
output = model(torch.tensor([[1.0, 2.0]]))
loss = output.sum()
loss.backward()
optimizer.step()
updated_weight = model.weight.data
print((initial_weight - updated_weight).abs().sum().item())

What does the printed value represent?

medium

A. The total change in weights after one optimization step including weight decay

B. The learning rate value

C. The loss value before backward pass

D. The sum of model outputs

Solution

1. Step 1: Understand code flow

   The code runs one optimizer step with weight decay, then measures how much weights changed.

2. Step 2: Interpret printed value

   The printed value is the sum of absolute differences between initial and updated weights, showing total weight change including weight decay effect.

3. Final Answer:

   The total change in weights after one optimization step including weight decay -> Option A

4. Quick Check:

   Weight change sum = printed value [OK]

Hint: Weight decay affects weight updates, so weight change includes it [OK]

Common Mistakes:

• Thinking printed value is loss or learning rate
• Ignoring weight decay effect on weights
• Confusing output sum with weight change

4. You wrote this PyTorch optimizer code:

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.1)

But your model is overfitting badly. What is a likely mistake?

medium

A. Weight decay value is too high, causing poor training

B. Weight decay should be set to zero to reduce overfitting

C. Weight decay is applied to biases by default, so overfitting remains

D. Learning rate is too low to affect weight decay

Solution

1. Step 1: Recall weight decay behavior in PyTorch

   By default, weight decay is applied to all parameters, including biases and batch norm weights, unless explicitly excluded.

2. Step 2: Understand overfitting cause

   If weight decay is applied to all parameters including biases, it may not reduce overfitting effectively because biases are not regularized properly.

3. Final Answer:

   Weight decay is applied to biases by default, so overfitting remains -> Option C

4. Quick Check:

   Biases often excluded from weight decay for better regularization [OK]

Hint: Check if weight decay excludes biases to reduce overfitting [OK]

Common Mistakes:

• Assuming weight decay does not apply to biases
• Setting weight decay to zero to fix overfitting
• Blaming learning rate for weight decay issues

5. You want to apply weight decay only to the weights of a PyTorch model's linear layers but not to biases. Which code snippet correctly sets this up?

hard

A. optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.01)

B. params = [
{'params': [p for n, p in model.named_parameters() if 'weight' in n], 'weight_decay': 0.01},
{'params': [p for n, p in model.named_parameters() if 'bias' in n], 'weight_decay': 0.0}
]
optimizer = torch.optim.Adam(params, lr=0.001)

C. optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0)

D. params = [
{'params': model.parameters(), 'weight_decay': 0.01}
]
optimizer = torch.optim.Adam(params, lr=0.001)

Solution

1. Step 1: Understand selective weight decay

   To apply weight decay only to weights, separate parameters into groups with and without weight decay.

2. Step 2: Check code correctness

   params = [ {'params': [p for n, p in model.named_parameters() if 'weight' in n], 'weight_decay': 0.01}, {'params': [p for n, p in model.named_parameters() if 'bias' in n], 'weight_decay': 0.0} ] optimizer = torch.optim.Adam(params, lr=0.001) creates two groups: weights with weight_decay=0.01 and biases with weight_decay=0.0, correctly excluding biases.

3. Final Answer:

   Code snippet that separates weights and biases with different weight_decay values -> Option B

4. Quick Check:

   Separate params for weight decay control [OK]

Hint: Group parameters by name to apply weight decay selectively [OK]

Common Mistakes:

• Applying weight decay to all parameters blindly
• Not separating biases from weights
• Using wrong parameter names in filtering