Training loops from scratch

What is a training loop?

The training loop is the iterative process that transforms random weights into a trained model. At its core, it implements empirical risk minimization: find parameters $\theta$ that minimize average loss over training data.

WHY these components? Each step solves a specific problem:

Forward pass: "What does my current model predict?"
Loss: "How wrong am I?" (scalar measure of error)
Backward: "Which weights caused this error?" (gradient flow)
Update: "How should I adjust weights to reduce error?"
Iteration: "Am I generalizing, or just memorizing?"

Deriving the training loop from first principles

Starting point: The optimization objective

We want to minimize expected risk: $\mathcal{R}(\theta) = \mathbb{E}_{(x,y) \sim \mathcal{D}}[\mathcal{L}(f(x;\theta), y)]$

WHY this form? We care about performance on unseen data from distribution $\mathcal{D}$ , not just training examples.

Since we don't know $\mathcal{D}$ , we approximate with empirical risk: $\hat{\mathcal{R}}(\theta) = \frac{1}{N}\sum_{i=1}^{N} \mathcal{L}(f(x_i;\theta), y_i)$

WHY the average? Raw sum grows with dataset size; averaging makes loss magnitude dataset-independent.

From objective to iterative updates

Direct minimization is intractable for neural networks (non-convex, millions of parameters). Instead, we use gradient descent:

$\theta_{t+1} = \theta_t - \eta \nabla_\theta \hat{\mathcal{R}}(\theta_t)$

WHY this works: Gradient $\nabla_\theta \hat{\mathcal{R}}$ points toward stepest increase in loss; moving opposite (negative gradient) decreases loss locally.

The problem: Computing $\nabla_\theta \hat{\mathcal{R}}$ requires full-dataset pass—expensive!

The solution: Stochastic Gradient Descent (SGD) approximates the gradient using a mini-batch:

$\nabla_\theta \hat{\mathcal{R}} \approx \frac{1}{B}\sum_{i \in \text{batch}} \nabla_\theta \mathcal{L}(f(x_i;\theta), y_i)$

WHY batches? Single examples are too noisy; full dataset is too slow; mini-batches balance variance and computation.

Formula Complete Training Loop Update Rule

For batch $\mathcal{B}_t$ at step $t$ : $\theta_{t+1} = \theta_t - \frac{\eta}{|\mathcal{B}_t|} \sum_{(x,y) \in \mathcal{B}_t} \nabla_\theta \mathcal{L}(f(x;\theta_t), y)$ Derivation:

Start with gradient descent: $\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{R}$
Replace population gradient with mini-batch estimate
Apply chain rule: $\nabla_\theta \mathcal{L} = \frac{\partial \mathcal{L}}{\partial \hat{y}} \frac{\partial \hat{y}}{\partial \theta}$ (backpropagation!)
Average over batch to reduce variance

Implementation from scratch: Step-by-step

Worked example Minimal training loop (no framework)

Let's train a 2-layer network on synthetic data.

import numpy as np
# Generate data: y = x1*w1 + x2*w2 + noise
np.random.seed(42)
N, D = 1000, 2
X = np.random.randn(N, D)
true_w = np.array([3.0, -2.0])
y = X @ true_w + 0.1*np.random.randn(N)
 
# Initialize model: 2 → 4 → 1 network
W1 = np.random.randn(D, 4) * 0.1  # WHY small init? Prevent saturation
b1 = np.zeros(4)
W2 = np.random.randn(4, 1) * 0.1
b2 = np.zeros(1)
# Hyperparameters
lr = 0.01  # Learning rate
batch_size = 32
epochs = 50
# Training loop
for epoch in range(epochs):
    # Shuffle data (WHY? Breaks correlation between consecutive batches)
    indices = np.random.permutation(N)
    X_shuffled = X[indices]
    y_shuffled = y[indices]
    
    epoch_loss = 0.0
    for i in range(0, N, batch_size):
        # 1. Extract batch
        X_batch = X_shuffled[i:i+batch_size]
        y_batch = y_shuffled[i:i+batch_size]
        
        # 2. FORWARD PASS
        z1 = X_batch @ W1 + b1  # Linear layer 1
        a1 = np.maximum(0, z1)   # ReLU activation
        # WHY ReLU? Non-linearity without vanishing gradients
        z2 = a1 @ W2 + b2        # Output layer
        y_pred = z2.squeeze()
        
        # 3. LOSS COMPUTATION
        loss = np.mean((y_pred - y_batch)**2)  # MSE
        # WHY MSE? Differentiable, penalizes large errors quadratically
        epoch_loss += loss
        
        # 4. BACKWARD PASS (manual gradients)
        B = len(X_batch)
        dloss_dy = 2 * (y_pred - y_batch) / B  # d(MSE)/d(y_pred)
        
        # Output layer gradients
        dW2 = a1.T @ dloss_dy.reshape(-1, 1)
        # WHY this form? Chain rule: dL/dW2 = dL/dz2 * dz2/dW2 = dL/dy * a1
        db2 = np.sum(dloss_dy)
        
        # Hidden layer gradients
        da1 = dloss_dy.reshape(-1, 1) @ W2.T
        dz1 = da1 * (z1 > 0)  # ReLU derivative: 1 if z>0, else 0
        # WHY multiply by (z1>0)? Backprop through ReLU activation
        dW1 = X_batch.T @ dz1
        db1 = np.sum(dz1, axis=0)
        
        # 5. PARAMETER UPDATE
        W2 -= lr * dW2
        b2 -= lr * db2
        W1 -= lr * dW1
        b1 -= lr * db1
    
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {epoch_loss/(N//batch_size):.4f}")
# Final weights approximate true_w via learned representation
print(f"True weights: {true_w}")
print(f"Learned mapping effective weights: ~3.0, ~-2.0 (hidden in W1, W2)")

WHY each step:

Shuffle: Prevents model from learning batch order artifacts
Batch extraction: Enables parallel gradient computation
Forward: Transforms input to prediction via learned parameters
Loss: Quantifies prediction quality (drives learning signal)
Backward: Computes how each parameter contributed to error
Update: Adjusts parameters to reduce error on this batch

Worked example Real-world training loop (PyTorch style)

Modern frameworks automate gradient computation:

import torch
import torch.nn as nn
# Same synthetic data
X_torch = torch.randn(1000, 2)
y_torch = X_torch @ torch.tensor([3.0, -2.0]) + 0.1*torch.randn(1000)
# Define model
model = nn.Sequential(
    nn.Linear(2, 4),
    nn.ReLU(),
    nn.Linear(4, 1)
)
# Loss and optimizer
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
 
# Training loop
batch_size = 32
for epoch in range(50):
    perm = torch.randperm(len(X_torch))
    for i in range(0, len(X_torch), batch_size):
        indices = perm[i:i+batch_size]
        X_batch, y_batch = X_torch[indices], y_torch[indices]
        
        # Forward
        y_pred = model(X_batch).squeeze()
        loss = criterion(y_pred, y_batch)
        
        # Backward + Update
        optimizer.zero_grad()  # WHY? Clear old gradients
        loss.backward()        # Autograd computes all gradients
        optimizer.step()       # Apply update rule
    
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

Key differences from scratch:

loss.backward() replaces manual gradient computation
optimizer.step() replaces manual parameter updates
Same underlying math, cleaner code

Common mistakes and fixes

Common mistake Mistake 1: Forgetting to zero gradients

The error:

for batch in dataloader:
    loss = compute_loss(batch)
    loss.backward()
    optimizer.step()
    # Missing: optimizer.zero_grad()

Why it feels right: "backward() computes gradients, step() uses them, done!"

Why it's wrong: PyTorch accumulates gradients by default (useful for gradient accumulation). Without zero_grad(), gradients from previous batches add up, causing:

Incorrect gradient magnitudes (sum instead of current batch)
Exploding gradients over many steps
Model doesn't learn the current batch pattern

The fix: Always call optimizer.zero_grad() or model.zero_grad() before backward().

When accumulation is intentional: Large effective batch size:

accumulation_steps = 4
for i, batch in enumerate(dataloader):
    loss = compute_loss(batch) / accumulation_steps
    loss.backward()  # Accumulate gradients
    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()  # Now clear

Common mistake Mistake 2: Wrong loss scaling in gradient accumulation

The error:

for i, batch in enumerate(dataloader):
    loss = compute_loss(batch)  # No division!
    loss.backward()
    if (i+1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Why it feels right: "I'm accumulating, so gradients will average out naturally."

Why it's wrong: Gradients sum over accumulation steps. If accum_steps=4, you get 4× larger gradients than intended. This is effectively a 4× larger learning rate.

The fix: Divide loss by accumulation_steps to normalize gradient magnitude, or scale learning rate down.

Common mistake Mistake 3: Training mode vs eval mode confusion

The error:

# Training
for batch in train_loader:
    outputs = model(batch)  # model.train() never called
    loss = criterion(outputs, targets)
    ...
# Evaluation
with torch.no_grad():
    for batch in val_loader:
        outputs = model(batch)  # model.eval() never called

Why it feels right: "no_grad() disables gradients, that's enough for eval."

Why it's wrong: Layers like Dropout and BatchNorm behave differently during training vs. inference:

Dropout randomly zeros activations during training, but keeps all during eval
BatchNorm uses batch statistics during training, running during eval

Forgetting model.eval() means validation uses training-mode behavior (wrong statistics, random drops), giving pessimistic metrics.

The fix:

model.train()  # Before training loop
# ... training code ...
model.eval()   # Before evaluation
with torch.no_grad():
    # ... evaluation code ...

Common mistake Mistake 4: Not shuffling training data

The error:

for epoch in range(100):
    for batch in dataloader:  # dataloader has shuffle=False
        # training code

Why it feels right: "SGD is stochastic anyway, order doesn't matter." Why it's wrong: Without shuffling, the model sees identical batch sequences every epoch:

Early batches might all be class 0, late batches class 1 → catastrophic forgetting
Corelated samples in consecutive batches → biased gradient estimates
Model learns batch order artifacts instead of true patterns

The fix: DataLoader(..., shuffle=True) or manual permutation each epoch.

Advanced: Learning rate scheduling and early stopping

A fixed learning rate often underperforms. Learning rate scheduling adapts $\eta$ during training.

WHY schedule? Early training needs large steps to escape initialization; late training needs small steps to settle into minima.

Early stopping prevents overfitting by monitoring validation loss:

best_val_loss = float('inf')
patience = 10  # Epochs to wait for improvement
patience_counter = 0
 
for epoch in range(max_epochs):
    train_loss = train_one_epoch(model, train_loader, optimizer)
    val_loss = evaluate(model, val_loader)
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        save_checkpoint(model, 'best_model.pth')
        patience_counter = 0
    else:
        patience_counter += 1
    if patience_counter >= patience:
        print(f"Early stopping at epoch {epoch}")
        break

WHY this works: Validation loss increases when model starts memorizing training noise → stop before overfitting worsens.

Diagram: Training loop flow

Recall Explain to a 12-year-old

Imagine you're learning to shoot basketball free throws, but you're blindfolded and can only feel how far you missed.

Training loop is like this practice routine:

Shoot the ball (forward pass) - you take a shot with your current technique
Feel where it landed (compute loss) - did you miss left, right, short, or long?
Figure out what went wrong (backward pass) - was your elbow too high? Did you use too much power?
Adjust your form (update parameters) - bend your elbow a little less, ease up on the power
Shoot again (next iteration) - try with your adjusted technique

Each shot teaches you something small. After hundreds of shots (epochs), your muscle memory (weights) gets really good at hitting the basket. The "learning rate" is like how much you adjust each time - big changes early when you're way off, tiny tweaks later when you're close. The magic is that you never need to see the basket - just feeling the error and adjusting step-by-step eventually makes you a great shooter!

Connections

3.1.02-Backpropagation-algorithm - How gradients flow backward (the math behind loss.backward())
3.3.01-PyTorch-fundamentals - Tensors and autograd that power training loops
3.3.06-Custom-loss-functions - What goes into the loss computation step
3.4.03-Learning-rate-schedules - Advanced strategies for the update step
3.4.01-Batch-normalization - Layers that behave differently in train vs eval mode
3.5.02-Data-loaders-and-batching - How batches are prepared for the loop
4.2.01-Gradient-descent-variants - SGD, Adam, RMSprop - different update rules
5.1.01-Overfitting-and-underfitting - Why we monitor train vs validation loss

#flashcards/ai-ml

What are the five core components of a training loop? :: 1. Forward pass (compute predictions), 2. Loss computation (measure error), 3. Backward pass (compute gradients), 4. Parameter update (apply optimizer), 5. Iteration (repeat over batches/epochs)

Why do we use mini-batch SGD instead of full-batch gradient descent?

Full-batch requires computing gradients over entire dataset (slow and memory-intensive). Mini-batches approximate the gradient with much less computation while providing less noisy updates than single-example SGD. Balances variance and efficiency.

What is empirical risk minimization?

Minimizing the average loss over training data as a proxy for minimizing expected risk over the true data distribution. Formula:

\hat{\mathcal{R}}(\theta) = \frac{1}{N}\sum_{i=1}^{N} \mathcal{L}(f(x_i;\theta), y_i)

Why must you call optimizer.zero_grad() before backward()?

PyTorch accumulates gradients by default. Without zeroing, gradients from previous batches add to current gradients, causing incorrect gradient magnitudes and exploding gradients. Only intentionally skip when implementing gradient accumulation.

What happens if you forget model.eval() during validation?

Layers like Dropout and BatchNorm use training-mode behavior (random drops, batch statistics instead of running statistics), giving incorrect and typically pessimistic evaluation metrics. Always call model.eval() before validation and model.train() before training.

What is the update rule for SGD with mini-batches?

\theta_{t+1} = \theta_t - \frac{\eta}{|\mathcal{B}_t|} \sum_{(x,y) \in \mathcal{B}_t} \nabla_\theta \mathcal{L}(f(x;\theta_t), y)

where

\mathcal{B}_t

is the mini-batch,

\eta

is learning rate, and we average the gradient over the batch.

Why do we shuffle training data each epoch?

Prevents model from learning batch order artifacts and correlations between consecutive batches. Without shuffling, corelated samples bias gradient estimates and the model may learn the sequence pattern instead of true underlying patterns.

What is early stopping and why does it work?

Monitoring validation loss and stopping training when it stops improving (after patience period). Works because validation loss increases when model starts overfitting (memorizing training noise), so stopping before this point preserves generalization.

How do you implement gradient accumulation correctly?

Divide loss by accumulation_steps before backward(), accumulate gradients over multiple batches without zero_grad(), then call optimizer.step() and zero_grad() after accumulation_steps batches. This simulates a larger effective batch size.

What is the purpose of learning rate scheduling?

Adapts learning rate during training - large steps early to escape initialization quickly, small steps later to settle into minima without overshooting. Common schedules: step decay, exponential decay, cosine annealing, reduce on plateau.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Training loop deep learning ka dil hai, dost. Socho tumhare pas ek model hai jo bilkul nayi hai, kuch nahi janta. Training loop woh process hai jisse model seekhta hai data se.

Pehle forward pass hota hai - matlab model apna current guess deta hai (prediction). Fir loss function bata hai kitna galat tha (error). Ab magic shuru hoti hai: backpropagation se har weight ko pata chalta hai ki usne kitna contribution diya is error mein. Yeh gradients hain. Phir optimizer (jaise SGD ya Adam) un gradients ka use karke weights ko thoda adjust karta hai, taki agle baar error kam ho.

Yeh sab ek batch ke liye hota hai (32-256 examples). Pura dataset khatam hone tak yeh repeat karta rehta hai - iskoek epoch kehte hain. Kai epochs ke bad model expert ban jata hai. Lekin dhyaan rakho - agar validation loss badhne lage toh early stopping lagao, warna overfitting ho

What is a training loop?

Deriving the training loop from first principles

Starting point: The optimization objective

From objective to iterative updates

Implementation from scratch: Step-by-step

Common mistakes and fixes

Advanced: Learning rate scheduling and early stopping

Diagram: Training loop flow

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Deep Learning Frameworks

Connections