The Training Loop: Forward, Backward, Step

The Three-Step Dance

Modern model training is the same loop, repeated millions of times:

Forward pass: feed batch through the network, compute loss.
Backward pass: backprop computes gradients for every parameter.
Step: optimizer updates parameters using the gradients.

Repeat until the loss stops dropping. That's it. That's how every neural network in the world is trained.

The Standard PyTorch Idiom

You'll write this pattern thousands of times in deep learning:

for batch in dataloader:
    optimizer.zero_grad()      # clear stale gradients
    pred = model(batch.x)      # forward pass
    loss = loss_fn(pred, batch.y)
    loss.backward()            # backward pass — backprop runs here
    optimizer.step()           # update parameters

Five lines. The whole training loop. Memorize this rhythm; everything fancy in deep learning is built on top of it.

Common Optimizers

SGD — vanilla. $θ \leftarrow θ - η \nabla L$ .
SGD + momentum — accumulates past gradients. Smoother trajectories.
Adam — per-parameter adaptive learning rates. Modern default. Works with little tuning.
AdamW — Adam with proper weight decay. Best default for large transformer models.

Forward → Backward → Step. That's the universal training loop. Everything fancy is wrapping or extending these three steps.

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}")

MLX flavor — value_and_grad + optimizer.update + mx.eval·python

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}")

Exercise

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.

The Training Loop: Forward, Backward, Step

The Three-Step Dance

The Standard PyTorch Idiom

Common Optimizers

Code

External links

Exercise

Progress

Comments 0