Adam Optimizer › PyTorch Adam Implementation
Language AI HandbookAll chapters Adam Optimizer › PyTorch Adam ImplementationOne section · Sections in this chapter12 / 18 ↓
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
PyTorch provides torch.optim.Adam as a built-in optimizer. Understanding how to use it correctly, including how to set per-layer learning rates and integrate it with learning rate scheduling, is essential for practical deep learning.
Basic Setup
First, install the required libraries.
Define a simple model and configure the Adam optimizer with default hyperparameters.
torch.manual_seed(42)
# Simple two-layer network for binary classification
model = nn.Sequential(
nn.Linear(20, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, 1),
nn.Sigmoid(),
)
# Standard Adam configuration
optimizer = torch.optim.Adam(
model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0
)Model parameters: 3,457 Optimizer: Adam lr: 0.001 betas: (0.9, 0.999) eps: 1e-08
Training Loop
A complete training loop using Adam over a synthetic binary classification dataset.
# Generate synthetic dataset
np.random.seed(42)
X = torch.randn(1000, 20)
w_true = torch.randn(20)
y = ((X @ w_true + torch.randn(1000) * 0.5) > 0).float().unsqueeze(1)
# Split into train/test
X_train, X_test = X[:800], X[800:]
y_train, y_test = y[:800], y[800:]
criterion = nn.BCELoss()
train_losses = []
test_losses = []
for epoch in range(200):
model.train()
optimizer.zero_grad()
y_pred = model(X_train)
loss = criterion(y_pred, y_train)
loss.backward()
optimizer.step()
train_losses.append(loss.item())
model.eval()
with torch.no_grad():
test_pred = model(X_test)
test_loss = criterion(test_pred, y_test).item()
test_losses.append(test_loss)Final train loss: 0.0159 Final test loss: 0.2141 Train accuracy: 1.000 Test accuracy: 0.910
The model converges within 200 epochs using Adam's default learning rate of 0.001. This shows that the default configuration works well without any tuning on this type of task. Notice the standard pattern: optimizer.zero_grad() clears gradients from the previous step, loss.backward() computes new gradients, and optimizer.step() applies the Adam update. Forgetting zero_grad() is one of the most common bugs in PyTorch training: gradients accumulate across steps, effectively multiplying the gradient by the number of steps since the last reset, which causes wildly incorrect updates.
Accessing Optimizer State
You can inspect Adam's internal state to understand what the optimizer has learned about each parameter.
# Inspect the Adam state for the first layer's weight matrix
first_group = optimizer.state_dict()["state"]
first_param_key = list(first_group.keys())[0]
param_state = first_group[first_param_key]
step_count = param_state["step"].item()
m_mean = param_state["exp_avg"].abs().mean().item()
v_mean = param_state["exp_avg_sq"].abs().mean().item()Steps taken: 200 Mean |m_t| (first moment): 0.000243 Mean |v_t| (second moment): 0.000000 Mean sqrt(v_t) (scale): 0.000666 Effective LR (approx): 1.502169
The exp_avg and exp_avg_sq tensors correspond to and in Adam's equations. After training, tells you how large the gradients have been for each parameter. Parameters with large exp_avg_sq values have been actively updated throughout training. This state information can be very useful for debugging: if exp_avg_sq is near zero for many parameters, those parameters have barely been updated, which often signals a problem in the gradient flow (dead neurons, disconnected components, or missing gradients through a non-differentiable operation).
When you save and restore a model checkpoint, you should save the optimizer state alongside the model weights. Resuming training from a checkpoint without the optimizer state means Adam loses all its accumulated moment estimates and must restart the calibration process from zero, causing a temporary dip in training efficiency and sometimes a brief spike in loss.
Per-Layer Learning Rates
Adam supports different learning rates for different parameter groups, which is useful when fine-tuning pretrained models where you want the base model to update slowly and the new head to update quickly.
# Different learning rates for each layer
optimizer_multilr = torch.optim.Adam(
[
{"params": model[0].parameters(), "lr": 1e-4},
{"params": model[2].parameters(), "lr": 1e-3},
{"params": model[4].parameters(), "lr": 1e-2},
],
lr=1e-3,
)Group 0: lr=0.0001, params=1,344 Group 1: lr=0.001, params=2,080 Group 2: lr=0.01, params=33
This is a common pattern for transfer learning, where the pretrained backbone should update at a small fraction of the head's learning rate to preserve learned representations. The intuition is that the backbone already encodes useful features; updating it too aggressively will destroy that knowledge before the head has a chance to use it. A common heuristic is to set backbone learning rates to 1/10th or 1/100th of the head learning rate, though optimal values vary by task and model architecture.
Learning Rate Warmup with Adam
Learning rate warmup is nearly universal in transformer training. The most common schedule is linear warmup followed by cosine or inverse-square-root decay.
def get_linear_warmup_cosine_decay(
optimizer, warmup_steps, total_steps, min_lr_ratio=0.1
):
"""Linear warmup then cosine decay schedule."""
def lr_lambda(step):
if step < warmup_steps:
return step / warmup_steps
progress = (step - warmup_steps) / (total_steps - warmup_steps)
cosine_decay = 0.5 * (1 + np.cos(np.pi * progress))
return min_lr_ratio + (1 - min_lr_ratio) * cosine_decay
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
torch.manual_seed(42)
model_sched = nn.Sequential(nn.Linear(20, 1))
optimizer_sched = torch.optim.Adam(model_sched.parameters(), lr=1e-3)
total_steps = 1000
warmup_steps = 100
scheduler = get_linear_warmup_cosine_decay(
optimizer_sched, warmup_steps, total_steps
)
lr_schedule = []
for step in range(total_steps):
lr_schedule.append(optimizer_sched.param_groups[0]["lr"])
scheduler.step()LR at step 0: 0.00e+00 LR at step 50: 5.00e-04 LR at step 100: 1.00e-03 (end of warmup) LR at step 500: 6.28e-04 LR at step 999: 1.00e-04
The warmup ramps the learning rate from zero to the target value . During this phase, Adam's moment estimates are also near zero (from their zero initialization), so the combined effect of small and small moment magnitudes means the first warmup steps are very conservative. As both the schedule and the moment estimates warm up together, training becomes stable.
Key Parameters
The key parameters for PyTorch's Adam optimizer are:
- lr: Global learning rate . Default 0.001. This is the most impactful hyperparameter.
- betas: Tuple for first and second moment decay. Default (0.9, 0.999). Rarely need changing.
- eps: Numerical stability constant . Default 1e-8. Increase to 1e-6 for float16 training.
- weight_decay: L2 regularization coefficient. Default 0. For true decoupled weight decay, use AdamW instead (covered in the next chapter).
- amsgrad: Whether to use AMSGrad variant. Default False. Rarely needed in practice.