Training at Scale

Building the transformer architecture is only half the battle. To train it effectively, we need techniques that make training stable, prevent overfitting, and work within the constraints of hobby-scale hardware. This section covers two critical techniques: gradient accumulation and validation splits.

The Challenge: Small Batches, Noisy Training

What is a batch? During training, we process multiple examples together in a “batch.” The model makes predictions for all examples, we compute the average loss, then we calculate gradients and update weights. Larger batches give us more stable gradient estimates because we’re averaging over more examples.

The problem with small batches: On hobby hardware (like an M1 Mac or consumer GPU), we’re limited to small batches—typically just 8 sequences at a time. Small batches lead to noisy gradients: each batch gives a slightly different signal about which direction to update the weights, causing erratic training.

Gradient Accumulation: Large Batches Without the Memory Cost

The key insight: We don’t need to process all examples simultaneously! Gradient accumulation lets us simulate large batch sizes by accumulating gradients over multiple small batches before updating weights.

How it works:

Process batch 1: Forward pass → Loss → Backward pass → Store gradients (don’t update yet!)
Process batch 2: Forward pass → Loss → Backward pass → Add gradients to stored ones
Repeat for N batches (e.g., 16 times)
Update weights: Use the accumulated (averaged) gradients

Gradient accumulation comparison

Why this works mathematically: Gradients are linear, so averaging gradients from N separate batches gives the same result as computing the gradient on one large batch containing all N×batch_size examples. The key formula:

∇(L₁ + L₂ + … + Lₙ) = ∇L₁ + ∇L₂ + … + ∇Lₙ

By accumulating gradients over 16 batches of 8 sequences each, we get gradients equivalent to a batch of 128 sequences—16× more stable!—while only ever holding 8 sequences in memory at once.

Implementation

# Without accumulation (noisy)
for batch in dataloader:
    loss = compute_loss(batch)
    loss.backward()           # Compute gradients
    optimizer.step()          # Update every batch (noisy!)
    optimizer.zero_grad()

# With accumulation (stable)
accumulation_steps = 16
for i, batch in enumerate(dataloader):
    loss = compute_loss(batch)
    loss = loss / accumulation_steps  # Scale for correct averaging
    loss.backward()                   # Accumulate gradients

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()              # Update every 16 batches (stable!)
        optimizer.zero_grad()

Validation: Detecting Overfitting

The Problem: Memorization vs. Learning

Imagine a student preparing for an exam. They could:

Memorize answers to practice problems → Fails on new problems (overfitting)
Learn concepts from practice problems → Succeeds on new problems (good generalization)

The same happens with neural networks. As training progresses, the model might start memorizing the training data instead of learning general patterns. This is called overfitting.

The Solution: Validation Split

We set aside 10% of our data that the model never sees during training. After each epoch, we evaluate the model on this “validation” data. If the model is truly learning patterns (not memorizing), it should perform well on both training and validation data.

Training vs validation loss curves

How to Interpret the Curves

✓ Good Training

Train: 5.0 → 4.0 → 3.0
Val:   5.2 → 4.2 → 3.2

Both losses decreasing together. Model is learning general patterns that work on new data!

⚠ Underfitting

Train: 5.0 → 4.8 → 4.7
Val:   5.2 → 5.0 → 4.9

Both losses barely improving. Model is too simple or needs more training epochs.

⚠ Overfitting

Train: 5.0 → 3.0 → 1.5
Val:   5.2 → 3.5 → 4.0

Training loss decreasing but validation increasing. Model is memorizing training data!

Implementation

# Training with validation
for epoch in range(num_epochs):
    # Training phase
    model.train()
    for batch in train_dataloader:
        # ... forward, backward, update ...

    # Validation phase (no weight updates!)
    model.eval()
    with torch.no_grad():
        for batch in val_dataloader:
            val_loss = compute_loss(batch)
            # Just measure, don't update

    print(f"Train loss: {train_loss:.2f}, Val loss: {val_loss:.2f}")

    # Check for overfitting
    if val_loss > train_loss * 1.3:
        print("Warning: Possible overfitting!")

Expected Improvements

With gradient accumulation and validation:

20-30% lower final loss due to stable training
Smoother training curves that are easier to debug
Confidence in generalization by monitoring validation
Early stopping when validation stops improving
Works on hobby hardware without expensive GPUs

Full Code

See the full implementation: