Weight Decay › Code Implementation
Language AI HandbookAll chapters Weight Decay › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
This section walks through a complete implementation demonstrating weight decay in practice: comparing the effect of different values on a small neural network and showing how AdamW differs from Adam with L2 regularization.
Setting Up the Experiment
We will train a small multi-layer perceptron on a regression task with limited data, where overfitting is likely without regularization.
import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset
torch.manual_seed(42)
np.random.seed(42)
# Generate a small dataset where overfitting is likely
# True function: y = sin(2*pi*x) + cos(4*pi*x)
n_train = 80
n_val = 200
def true_function(x):
return np.sin(2 * np.pi * x) + np.cos(4 * np.pi * x)
x_train = np.random.uniform(-1, 1, n_train).astype(np.float32)
y_train = true_function(x_train) + np.random.normal(0, 0.3, n_train).astype(
np.float32
)
x_val = np.linspace(-1, 1, n_val).astype(np.float32)
y_val = true_function(x_val).astype(np.float32)
# Convert to tensors
X_train_t = torch.tensor(x_train).unsqueeze(1)
y_train_t = torch.tensor(y_train).unsqueeze(1)
X_val_t = torch.tensor(x_val).unsqueeze(1)
y_val_t = torch.tensor(y_val).unsqueeze(1)
train_loader = DataLoader(
TensorDataset(X_train_t, y_train_t), batch_size=16, shuffle=True
)Training samples: 80 Validation samples: 200 Noise std: 0.3 Function: sin(2*pi*x) + cos(4*pi*x)
The training set is intentionally small (80 samples) to make overfitting visible. The true function combines sine and cosine components, which a neural network can approximate, but with only 80 noisy observations there is ample room to overfit. We add Gaussian noise with standard deviation 0.3 to simulate the measurement uncertainty common in real-world data. A four-layer MLP with 128 hidden units has far more parameters than the 80 training points, making it the ideal candidate to demonstrate overfitting.
Defining the Model and Training Loop
import torch.nn as nn
import torch.optim as optim
class MLP(nn.Module):
def __init__(self, hidden_size=128, n_layers=4):
super().__init__()
layers = [nn.Linear(1, hidden_size), nn.Tanh()]
for _ in range(n_layers - 1):
layers += [nn.Linear(hidden_size, hidden_size), nn.Tanh()]
layers.append(nn.Linear(hidden_size, 1))
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
def train_model(weight_decay, use_adamw=True, n_epochs=300, lr=1e-3):
"""Train MLP with specified weight decay and return train/val loss histories."""
model = MLP()
if use_adamw:
optimizer = optim.AdamW(
model.parameters(), lr=lr, weight_decay=weight_decay
)
else:
# Adam with L2 regularization (adds wd to gradients)
optimizer = optim.Adam(
model.parameters(), lr=lr, weight_decay=weight_decay
)
criterion = nn.MSELoss()
train_losses = []
val_losses = []
for epoch in range(n_epochs):
model.train()
epoch_loss = 0.0
for xb, yb in train_loader:
optimizer.zero_grad()
loss = criterion(model(xb), yb)
loss.backward()
optimizer.step()
epoch_loss += loss.item() * len(xb)
train_losses.append(epoch_loss / n_train)
model.eval()
with torch.no_grad():
val_loss = criterion(model(X_val_t), y_val_t).item()
val_losses.append(val_loss)
return model, train_losses, val_lossesComparing Weight Decay Values
# Train with different weight decay values
weight_decays = [0.0, 1e-3, 1e-2, 1e-1]
results = {}
for wd in weight_decays:
model_wd, train_losses, val_losses = train_model(weight_decay=wd)
results[wd] = {
"model": model_wd,
"train_losses": train_losses,
"val_losses": val_losses,
"final_train": train_losses[-1],
"final_val": val_losses[-1],
}Weight Decay Comparison (Final MSE): lambda Train MSE Val MSE Overfit Gap ------------------------------------------------------- 0.0000 0.0819 0.0297 -0.0523 0.0010 0.1756 0.0803 -0.0953 0.0100 0.1326 0.0830 -0.0497 0.1000 0.0930 0.0314 -0.0616
This seeded run does not produce a monotonic ranking by . The training loss is measured against noisy observations, while the validation loss is measured against the noiseless generating function. Training MSE can therefore exceed validation MSE, and the reported validation-minus-training difference is negative for every setting. The useful comparison here is the absolute validation MSE: both zero decay and finish below the two intermediate settings.
On datasets where training and validation losses are computed against targets with the same noise characteristics, their difference can be a useful generalization diagnostic. That interpretation does not apply directly to this experiment because the targets differ. This is a reminder to compare like with like when using a train-validation gap and to prioritize held-out loss when tuning weight decay.
The training and validation loss curves over epochs make this progression concrete:
Comparing AdamW vs. Adam with L2
# Compare AdamW vs Adam+L2 at the same weight_decay value
wd_compare = 1e-2
_, adamw_train, adamw_val = train_model(weight_decay=wd_compare, use_adamw=True)
_, adam_l2_train, adam_l2_val = train_model(
weight_decay=wd_compare, use_adamw=False
)
adamw_final_val = adamw_val[-1]
adam_l2_final_val = adam_l2_val[-1]Weight decay lambda = 0.01 AdamW final val MSE: 0.0786 Adam+L2 final val MSE: 0.9398 AdamW improvement: 0.8611
AdamW typically achieves a lower validation loss than Adam with L2 regularization at the same nominal value. The improvement comes from the uniform application of weight decay across all parameters, rather than the uneven regularization that emerges when L2 is folded into Adam's adaptive gradient scaling.
In this particular seeded run the difference is pronounced: AdamW continues learning, while Adam+L2 remains near its initial loss. The magnitude should not be generalized from one small stochastic experiment, but the qualitative distinction remains useful: decoupling decay avoids scaling the regularization term through Adam's adaptive gradient machinery.
The learning curve comparison makes the difference visible over training:
Inspecting Weight Magnitudes
# Compare weight magnitudes for no-decay vs. moderate decay
model_no_wd = results[0.0]["model"]
model_with_wd = results[1e-2]["model"]
def get_weight_norms(model):
"""Return list of per-layer weight norms (excluding biases)."""
norms = []
for name, param in model.named_parameters():
if "weight" in name:
norms.append((name, param.data.norm().item()))
return norms
norms_no_wd = get_weight_norms(model_no_wd)
norms_wd = get_weight_norms(model_with_wd)Weight norms by layer: Layer No WD WD=0.01 Reduction % ----------------------------------------------------------------- net.0.weight 8.4290 7.7771 7.7% net.2.weight 8.3294 8.0106 3.8% net.4.weight 6.6294 6.4784 2.3% net.6.weight 7.8198 7.6035 2.8% net.8.weight 0.7927 0.8328 -5.1%
The weight norms show that the effect is strong but not uniform in this single run. Weight decay reduces the norm of each hidden-layer weight matrix, while the small output matrix ends slightly larger. The broader mechanism remains the same: decay continually penalizes magnitude, but the equilibrium for each layer also depends on its task gradients.
A pattern you might notice is that weight norms tend to be larger in middle layers than in the first or last layer. This reflects the model's information processing: middle layers encode intermediate representations that must support many different output patterns, while early layers process raw inputs and late layers project to the output space. Weight decay applies the same pressure at all depths, but the equilibrium norm at each layer reflects how much gradient signal flows through it. Middle layers, which are influenced by gradients from many output patterns via backpropagation, may sustain larger weights even under regularization.
Visualizing Weight Decay Trajectory
One more visualization helps build intuition about how weight decay affects individual weights over the course of training. We can track a single weight's value over training epochs to see the interplay between gradient updates and the decay pressure.
# Track a specific weight's value throughout training for two conditions
def train_and_track_weight(
weight_decay, use_adamw=True, n_epochs=300, lr=1e-3, layer_idx=0
):
"""Train and track the first weight of the first linear layer."""
model_track = MLP()
if use_adamw:
optimizer_track = optim.AdamW(
model_track.parameters(), lr=lr, weight_decay=weight_decay
)
else:
optimizer_track = optim.Adam(
model_track.parameters(), lr=lr, weight_decay=weight_decay
)
criterion_track = nn.MSELoss()
weight_history = []
for epoch in range(n_epochs):
model_track.train()
for xb, yb in train_loader:
optimizer_track.zero_grad()
loss = criterion_track(model_track(xb), yb)
loss.backward()
optimizer_track.step()
# Track the L2 norm of the first weight matrix
first_weight_norm = model_track.net[0].weight.data.norm().item()
weight_history.append(first_weight_norm)
return weight_history
weight_traj_no_wd = train_and_track_weight(0.0)
weight_traj_low_wd = train_and_track_weight(1e-3)
weight_traj_med_wd = train_and_track_weight(1e-2)
weight_traj_high_wd = train_and_track_weight(1e-1)
The trajectory plot shows how decay pressure interacts with, rather than simply overrides, the task gradient. The independently trained models follow different paths, so their norms are not monotonically ordered by at every epoch. The strongest setting, , is the only curve that turns downward substantially after reaching its peak, while the lower-decay runs continue growing or level off.
Notice that none of the equilibrium norms is zero: each reflects a balance between the expected task gradient and the weight decay term. The flattening portions show the models approaching that balance, while continued growth or decline shows that the two forces have not yet settled.