Momentum › Momentum in PyTorch
Language AI HandbookAll chapters Momentum › Momentum in PyTorchOne section · Sections in this chapter9 / 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's torch.optim.SGD supports both classical and Nesterov momentum through the momentum and nesterov parameters.
import torch.nn as nn
import torch.optim as optim
# A small neural network for demonstration
model = nn.Sequential(
nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 1)
)
# Classical (Polyak) momentum
optimizer_sgd = optim.SGD(
model.parameters(),
lr=0.01, # Base learning rate
momentum=0.9, # Momentum coefficient beta
dampening=0, # Dampening for the gradient (0 = pure momentum)
nesterov=False, # Classical momentum
weight_decay=1e-4, # Optional L2 regularization
)
# Nesterov momentum
optimizer_nag = optim.SGD(
model.parameters(),
lr=0.01,
momentum=0.9,
dampening=0, # Must be 0 when nesterov=True
nesterov=True, # Enable Nesterov momentum
)The dampening parameter is a lesser-known option that scales down the gradient contribution to the velocity update. With dampening=d, the update becomes:
Setting dampening=0 gives the standard PyTorch momentum formulation. Nesterov momentum requires dampening=0. The dampening parameter is rarely tuned in practice and should be left at its default of 0 for standard momentum training.
The Training Loop
Using momentum in a training loop is identical to using plain SGD; the optimizer handles the velocity state internally. PyTorch stores the velocity buffer as a stateful parameter in the optimizer, so it persists correctly across batches and epochs without any manual management.
import numpy as np
import torch
torch.manual_seed(42)
np.random.seed(42)
# Synthetic regression dataset
n_samples = 500
n_features = 10
X = torch.randn(n_samples, n_features)
true_weights = torch.randn(n_features, 1)
y = X @ true_weights + 0.5 * torch.randn(n_samples, 1)
# Split data
X_train, y_train = X[:400], y[:400]
X_test, y_test = X[400:], y[400:]
# Re-initialize model for a clean run
model_momentum = nn.Sequential(
nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 1)
)
optimizer = optim.SGD(model_momentum.parameters(), lr=0.01, momentum=0.9)
loss_fn = nn.MSELoss()
losses = []
for epoch in range(100):
optimizer.zero_grad() # Clear accumulated gradients
preds = model_momentum(X_train)
loss = loss_fn(preds, y_train)
loss.backward() # Compute gradients via backprop
optimizer.step() # Apply momentum update
losses.append(loss.item())Initial loss: 3.7185 Final loss: 0.1944 Loss reduction: 94.8%
The training loop is straightforward: zero_grad() clears old gradients, backward() computes new gradients, and step() applies the momentum-based update. The optimizer maintains the velocity buffer internally, so you do not need to track it manually. Notice that the interface is identical to plain SGD; momentum is configured once at optimizer creation and then operates transparently during training.
The zero_grad() call clears the accumulated parameter gradients but does not touch the velocity buffer. The velocity persists across iterations because that is the whole point: it accumulates information from previous steps. If you want to reset the momentum, you must either re-create the optimizer or manually clear the state_dict entries corresponding to the momentum buffers.
Comparing SGD, Momentum, and Nesterov
To make the difference concrete, let us train the same network on the same data with three optimizers and compare their training loss curves.
def train_with_optimizer(optimizer_fn, n_epochs=100):
"""Train a fresh model with the given optimizer constructor."""
torch.manual_seed(42)
model = nn.Sequential(
nn.Linear(10, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, 1),
)
opt = optimizer_fn(model.parameters())
loss_fn_inner = nn.MSELoss()
epoch_losses = []
for _ in range(n_epochs):
opt.zero_grad()
preds = model(X_train)
loss = loss_fn_inner(preds, y_train)
loss.backward()
opt.step()
epoch_losses.append(loss.item())
return epoch_losses
losses_sgd = train_with_optimizer(lambda p: optim.SGD(p, lr=0.01))
losses_momentum = train_with_optimizer(
lambda p: optim.SGD(p, lr=0.01, momentum=0.9)
)
losses_nesterov = train_with_optimizer(
lambda p: optim.SGD(p, lr=0.01, momentum=0.9, nesterov=True)
)SGD final loss: 0.6077 Momentum final loss: 0.1952 Nesterov final loss: 0.1934 Speedup (momentum vs SGD): 3.11x lower loss Speedup (Nesterov vs SGD): 3.14x lower loss
Both momentum variants achieve largely lower loss than plain SGD within the same number of epochs. This shows the convergence acceleration. The speedup is meaningful: momentum changes how effectively the optimizer uses each gradient computation rather than providing only a marginal improvement.