Epoch vs. Batch Size vs. Iteration in Machine Learning: Differences, Math, and Python Breakdown

When training modern machine learning models — from gradient-boosted trees and linear regressors to deep convolutional neural networks and Large Language Models — practitioners are confronted with terminal logs displaying metrics like Epoch 3/50 - Step 120/500 - Loss: 0.042. To newcomers, the terms epoch, batch size, and iteration are frequently used interchangeably or confused with one another.

However, these three terms define the exact geometric heartbeat of numerical optimization. Selecting their values dictates whether your model converges in minutes or stalls for days, fits within GPU VRAM or crashes with an Out-Of-Memory (OOM) error, and converges to a robust, generalizable minimum or memorizes training noise.

1. Key Concepts & Mathematical Notation Glossary

Review the mathematical notation governing optimization schedules and dataset batching:

2. Mathematical Mechanics: Equations & Gradient Variance

The relationship between dataset volume, hardware batching, and training steps is governed by rigorous mathematical formulas.

2.1 The Core Training Equations

Because real-world dataset sizes $N$ are rarely integer multiples of batch size $B$ (e.g., $N = 1,000$ and $B = 64$), the number of iterations per epoch requires the ceiling division operator $\lceil \cdot \rceil$:

When $N \pmod B \ne 0$, the final batch of each epoch — termed the tail batch or remainder batch — contains fewer observations than $B$:

2.2 Gradient Variance Scaling: Why Mini-Batching Works

Let $\nabla \mathcal{L}(\theta)$ be the true, full-dataset gradient. The mini-batch gradient $g_B(\theta)$ is an unbiased estimator of the true gradient:

Where $\sigma^2$ is the per-sample gradient variance across the training distribution. This equation reveals the fundamental trade-off:

• When $B$ is small: Gradient variance is large. The optimization trajectory exhibits stochastic fluctuations. This noise prevents the optimizer from getting trapped in shallow local minima and saddle points.
• When $B$ is large: Variance drops by factor $1/B$. The gradient estimate becomes smooth and deterministic, but computing each update requires substantial memory and computational time.

2.3 The Linear Learning Rate Scaling Rule

When scaling batch size from a baseline $B_{\text{base}}$ to a larger batch size $B_{\text{new}}$ (e.g., across multi-GPU distributed clusters), empirical research (Goyal et al., 2017) proves that the learning rate must scale proportionally to preserve optimization dynamics:

If you quadruple your batch size from 32 to 128, each iteration processes 4 times more data with $\frac{1}{4}$ the gradient variance; scaling $\eta$ from $0.01$ to $0.04$ allows the model to take commensurately bolder steps per iteration.

3. The Three Gradient Descent Regimes

The choice of batch size defines three distinct optimization paradigms:

4. Hand-Worked Trace: $N = 1,000$, $B = 64$, $E = 10$

To solidify the arithmetic, walk through a concrete dataset of $N = 1,000$ samples trained with batch size $B = 64$ for $E = 10$ epochs:

1. Compute Full Batches: $\lfloor 1000 / 64 \rfloor = 15$ full batches. These 15 batches process $15 \times 64 = 960$ samples.
2. Compute Tail Batch: $1000 - 960 = 40$ remaining samples. The 16th batch contains exactly 40 observations.
3. Iterations per Epoch: $I = 15 + 1 = 16$ iterations (or $\lceil 1000 / 64 \rceil = 16$).
4. Total Parameter Updates: $T = 10 \text{ epochs} \times 16 \text{ iterations/epoch} = \mathbf{160 \text{ total iterations}}$.

During each epoch, the model undergoes 16 discrete weight adjustments; across the full 10-epoch training schedule, the model parameters are updated exactly 160 times.

5. The Training Loop Architecture

This nested control flow illustrates how epochs, batches, and iterations interact inside a deep learning training script:

6. Complete, Self-Contained Python Implementation

Below is a complete, runnable Python simulation comparing Batch Gradient Descent ($B = 600$), Mini-Batch ($B = 32$), and Pure Stochastic ($B = 1$) on the exact same synthetic linear regression dataset ($N = 600$ observations, $E = 5$ epochs):

import numpy as np

# =============================================================================
# STAGE 1: SYNTHETIC DATASET GENERATION (N = 600 samples, 2 features)
# =============================================================================
np.random.seed(42)
N = 600
X = np.random.normal(0, 1, (N, 2))
true_weights = np.array([3.5, -2.0])
true_bias = 1.2
y = X @ true_weights + true_bias + np.random.normal(0, 0.1, N)

print(f"Training Dataset Volume: N = {N} observations")

# =============================================================================
# STAGE 2: CUSTOM TRAINING LOOP (EPOCHS, BATCHES & ITERATIONS)
# =============================================================================
def train_model(X, y, batch_size, epochs, learning_rate=0.02):
    N_samples = len(X)
    weights = np.zeros(2)
    bias = 0.0
    total_steps = 0
    loss_history = []
    
    # Number of iterations required to traverse dataset once
    steps_per_epoch = int(np.ceil(N_samples / batch_size))
    
    for epoch in range(1, epochs + 1):
        # Shuffling dataset at start of every epoch
        indices = np.random.permutation(N_samples)
        X_shuffled = X[indices]
        y_shuffled = y[indices]
        
        for step in range(steps_per_epoch):
            start_idx = step * batch_size
            end_idx = min(start_idx + batch_size, N_samples)
            
            X_batch = X_shuffled[start_idx:end_idx]
            y_batch = y_shuffled[start_idx:end_idx]
            b_size = len(X_batch)
            
            # Forward pass: predictions and residual error
            y_pred = X_batch @ weights + bias
            error = y_pred - y_batch
            
            # Compute empirical mini-batch gradients
            grad_w = (2 / b_size) * (X_batch.T @ error)
            grad_b = (2 / b_size) * np.sum(error)
            
            # Parameter update step (1 ITERATION)
            weights -= learning_rate * grad_w
            bias -= learning_rate * grad_b
            total_steps += 1
            
        # Track full-dataset Mean Squared Error after each epoch
        epoch_mse = np.mean((X @ weights + bias - y) ** 2)
        loss_history.append(epoch_mse)
        
    return weights, bias, steps_per_epoch, total_steps, loss_history

# =============================================================================
# STAGE 3: EMPIRICAL BENCHMARK ACROSS 3 BATCH SIZES (E = 5 epochs)
# =============================================================================
regimes = [
    ("Batch GD (B = 600)", 600),
    ("Mini-Batch (B = 32)", 32),
    ("Stochastic (B = 1)", 1)
]

print("\n" + "="*72)
print(f"{'Regime':22s} | {'Steps/Epoch':11s} | {'Total Steps':11s} | {'Final MSE':10s} | Learned Parameters")
print("="*72)

for name, b_size in regimes:
    w_fit, b_fit, per_ep, tot_steps, hist = train_model(X, y, batch_size=b_size, epochs=5, learning_rate=0.02)
    print(f"{name:22s} | {per_ep:11d} | {tot_steps:11d} | {hist[-1]:10.4f} | w: [{w_fit[0]:.2f}, {w_fit[1]:.2f}] b: {b_fit:.2f}")
print("="*72)

7. Empirical Analysis: Interpreting the Benchmark Results

Examine the exact numerical metrics produced by executing the Python script across 5 training epochs:

Notice the dramatic operational differences across identical 5-epoch runs:

• Batch GD stalled: Because $B = 600$, the model only performed 5 total updates. It barely moved from its zero initialization, producing an unacceptable MSE of $11.81$.
• Mini-Batch ($B = 32$) balanced speed and accuracy: By taking 19 updates per epoch (95 total iterations), it rapidly converged to within $2\%$ of the true data-generating weights ($[3.42, -1.97]$ vs. $[3.50, -2.00]$).
• Stochastic ($B = 1$) over-computed: While achieving low error, it performed 3,000 individual updates. In deep learning architectures, running 3,000 sequential forward/backward passes on single images creates severe memory bus latency.

8. Production Gotchas & Deep Learning Framework Mechanics

In production PyTorch and TensorFlow systems, batch scheduling introduces three critical architectural pitfalls:

• The drop_last Dilemma in DataLoaders: If $N = 1,000$ and $B = 64$, the final batch has only 40 samples. In models utilizing Batch Normalization, a tiny tail batch creates wild variance in empirical batch mean $\mu_B$ and variance $\sigma_B^2$, corrupting running statistics. Setting DataLoader(..., drop_last=True) discards the partial batch to stabilize training.
• The Generalization Gap (Flat vs. Sharp Minima): Massive batch sizes ($B \ge 2,048$) converge to sharp local minima in the loss landscape. Smaller mini-batches ($B \in [32, 256]$) inject stochastic noise that knocks the parameters out of narrow ravines into flat minima, which generalize significantly better to unseen test distributions.
• Dataset Shuffling Every Epoch: Always ensure shuffle=True. If batches are presented in identical sequential order across epochs, the model learns spurious cyclical dependencies between adjacent mini-batches.

9. Common Misconceptions

• Believing '1 Epoch' Means '1 Step': An epoch is only equal to 1 iteration if your batch size equals your entire dataset ($B = N$). In virtually all modern neural network training, an epoch contains hundreds or thousands of iterations.
• Assuming More Epochs Always Yields Better Models: Training for excessive epochs inevitably leads to overfitting — the model begins memorizing idiosyncratic sample noise rather than structural patterns. Employ Early Stopping to halt training when validation loss stops improving.
• Confusing Batch Size with Token Sequence Length: In Natural Language Processing (NLP), models take input tensors of shape [batch_size, sequence_length]. Batch size is the number of distinct documents processed simultaneously; sequence length is the number of tokens per document.

10. Summary & Practice Exercises

1. Calculate Training Schedules: Given a dataset of $N = 50,000$ CIFAR-10 images trained for $E = 40$ epochs with batch size $B = 128$: calculate the iterations per epoch $I$, the total parameter updates $T$, and the size of the tail batch.
2. PyTorch DataLoader Experiment: Create a PyTorch DataLoader with $N = 100$ synthetic tensors and $B = 32$. Iterate through it and print the shape of each batch. Observe the final batch of size 4. Re-run with drop_last=True and observe the omission.
3. Batch Size Convergence Test: Modify the Python script from Section 6 with batch size $B = 128$. Compare its final MSE and iteration count against $B = 32$ and $B = 600$.