AdamW › Using AdamW in PyTorch
Language AI HandbookAll chapters AdamW › Using AdamW in PyTorchOne section · Sections in this chapter9 / 15 ↓
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
In practice, you will use PyTorch's implementation rather than writing your own. PyTorch provides torch.optim.AdamW with a clean API that correctly implements the decoupled weight decay.
import torch
import torch.nn as nn
import torch.optim as optim
# Set random seed for reproducibility
torch.manual_seed(42)
# Create a simple regression model
X_tensor = torch.FloatTensor(X_norm).squeeze()
y_tensor = torch.FloatTensor(y_norm)
model = nn.Sequential(
nn.Linear(1, 64), nn.GELU(), nn.Linear(64, 32), nn.GELU(), nn.Linear(32, 1)
)
# AdamW with standard hyperparameters for small models
optimizer = optim.AdamW(
model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01
)
criterion = nn.MSELoss()# Training loop
n_epochs = 300
losses_pt = []
for epoch in range(n_epochs):
optimizer.zero_grad()
y_pred = model(X_tensor.unsqueeze(-1)).squeeze()
loss = criterion(y_pred, y_tensor)
loss.backward()
optimizer.step()
losses_pt.append(loss.item())Final loss: 0.2440 Best loss: 0.2440 (epoch 300)
This shows a standard PyTorch training loop with AdamW. The weight_decay parameter directly corresponds to in the AdamW formulation. Unlike PyTorch's torch.optim.Adam optimizer where weight_decay adds L2 regularization to the gradient, torch.optim.AdamW correctly implements decoupled weight decay. Note that PyTorch's Adam optimizer documentation explicitly warns about this distinction and recommends AdamW when you want proper weight decay behavior.
Applying Different Weight Decay to Different Parameter Groups
One practical advantage of AdamW's decoupled formulation is that you can apply different weight decay rates to different parts of the model in a principled way. In transformer models, it is common practice to apply weight decay to weight matrices but not to bias terms or layer normalization parameters. This pattern was used in the original BERT paper, GPT-2, GPT-3, and has become a near-universal convention in LLM training.
# Define parameter groups with different weight decay
def get_parameter_groups(model, weight_decay=0.01):
"""
Separate parameters into those that receive weight decay and those that don't.
Convention: apply weight decay to weight matrices, not to biases or norms.
"""
decay_params = []
no_decay_params = []
for name, param in model.named_parameters():
if not param.requires_grad:
continue
# Biases and 1D parameters (like LayerNorm scale/shift) should not be decayed
if param.ndim == 1 or name.endswith(".bias"):
no_decay_params.append(param)
else:
decay_params.append(param)
return [
{"params": decay_params, "weight_decay": weight_decay},
{"params": no_decay_params, "weight_decay": 0.0},
]
param_groups = get_parameter_groups(model, weight_decay=0.01)
optimizer_grouped = optim.AdamW(param_groups, lr=1e-3)Parameters with weight decay: 2,144 (95.7%) Parameters without weight decay: 97 (4.3%) Total parameters: 2,241
This pattern of separating weight matrices from biases and normalization parameters is standard in large language model training. The intuition is that biases and normalization scale factors serve different functional roles than weight matrices. A bias term shifts the activation function input; shrinking it toward zero would force the model to use zero-centered activations regardless of what the data requires. Layer normalization parameters control the scale and offset of normalized activations; decaying them toward zero could interfere with the normalization's intended effect.
Weight matrices, by contrast, encode learned transformations between representation spaces. Keeping these matrices from growing too large is a sensible regularization goal because large weight matrix norms are the direct mechanism by which a model can become overly sensitive to specific input patterns. The parameter group pattern lets you apply weight decay precisely where it helps and skip it where it might hurt.
Key Hyperparameters
The key AdamW hyperparameters and their practical interpretation are:
- lr (learning rate): Controls the overall step size. Typical values range from 1e-4 to 3e-4 for large language models during pretraining. A warmup schedule is always used, linearly increasing the learning rate from 0 over the first several thousand steps before following a cosine or linear decay to the end of training.
- betas=(beta1, beta2): Decay rates for the first and second moment estimates. The default (0.9, 0.999) works well for most tasks. Some LLM training recipes use beta2=0.95 or beta2=0.98 to allow faster adaptation to gradient magnitude changes over the course of training.
- eps: Numerical stability constant. The default 1e-8 is appropriate for float32 training. When training in float16 or bfloat16, increasing to 1e-6 or 1e-7 can improve numerical stability by reducing the impact of floating-point rounding on the denominator .
- weight_decay: The decoupled weight decay coefficient . Typical values are 0.01 to 0.1 for language model pretraining, with 0.01 common for fine-tuning where strong regularization can impair task-specific adaptation.