Forward: compute outputs and an objective for a mini-batch.
Backward: compute gradients of that objective with respect to trainable parameters.
Update: let the optimizer combine gradients with its state to change parameters.
PyTorch accumulates gradients by default, so most loops clear them before the next step:
for batch in dataloader:
optimizer.zero_grad()
pred = model(batch.x)
loss = loss_fn(pred, batch.y)
loss.backward()
optimizer.step()
This is the core skeleton. Real training may add mixed precision, gradient accumulation, distributed synchronization, clipping, schedules, checkpointing, and evaluation. Those systems can change numerical behavior and are not merely cosmetic wrappers.
Choose an Optimizer with the Full Setup
SGD with momentum is a strong baseline and remains competitive, especially in well-tuned vision training.
Adam uses moving estimates of first and second moments to adapt parameter-wise step sizes.
AdamW applies decoupled weight decay. It is common for Transformers, but it is not universally the best default.
Learning rate, batch size, warmup, schedule, precision, and regularization interact with optimizer choice. “Adam works with little tuning” is not a reliable production rule.
Forward computes values, backward computes sensitivities, and the optimizer changes state. Separating those roles tells you where to inspect a stalled or unstable run.
Code
The complete training loop·python
import torch
import torch.nn as nn
import torch.optim as optim
# Tiny dataset
X = torch.randn(100, 4)
y = (X.sum(dim=1) > 0).long() # binary classification
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 2))
optimizer = optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
# The five-line training loop
for epoch in range(50):
optimizer.zero_grad()
pred = model(X)
loss = loss_fn(pred, y)
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print(f"epoch {epoch}: loss = {loss.item():.4f}")
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
# Tiny dataset — same shape as the PyTorch version
mx.random.seed(0)
X = mx.random.normal(shape=(100, 4))
y = (X.sum(axis=1) > 0).astype(mx.int32)
class Net(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(4, 16)
self.fc2 = nn.Linear(16, 2)
def __call__(self, x):
return self.fc2(nn.relu(self.fc1(x)))
net = Net()
optimizer = optim.Adam(learning_rate=0.01)
# Functional loss — takes the model as an argument
def loss_fn(model, X, y):
return nn.losses.cross_entropy(model(X), y, reduction="mean")
# nn.value_and_grad wraps loss_fn into a (loss, grads) returning function
loss_and_grad = nn.value_and_grad(net, loss_fn)
for epoch in range(50):
loss, grads = loss_and_grad(net, X, y)
optimizer.update(net, grads)
mx.eval(net.parameters(), optimizer.state) # materialize lazy ops
if epoch % 10 == 0:
print(f"epoch {epoch}: loss = {loss.item():.4f}")
Run the training loop above. Then change lr=0.01 to lr=10.0 (way too high) and lr=0.0001 (way too low). Observe: divergence vs. crawling. The learning rate is the most important hyperparameter.
Hint
High lr: loss explodes or oscillates. Low lr: loss barely moves. Adam's default 0.001-0.01 range is usually a good starting point.
Progress
Progress is local-only — sign in to sync across devices.