Week 10: Large-Batch and Distributed Training

Minibatch as Parallelism#

A minibatch of size $B$ provides a gradient estimator:

$\hat{g}_B = \frac{1}{B}\sum_{i \in \mathcal{B}} \nabla f_i(x)$

where $\mathcal{B}$ is a random subset of $\{1, \ldots, n\}$ of size $B$. The variance of this estimator is:

$\operatorname{Var}(\hat{g}_B) = \frac{1}{B} \cdot \frac{n - B}{n - 1} \cdot \sigma^2 \approx \frac{\sigma^2}{B}$

for $B \ll n$, where $\sigma^2$ is the per-example gradient variance. The variance falls as $1/B$, but the computational cost rises as $B$ — a classic trade-off.

In a data-parallel setup with $K$ workers, each worker processes a local minibatch of size $B_{\text{local}}$ and computes a local gradient. The global batch size is $B = K \cdot B_{\text{local}}$, and the gradient is the average of the $K$ local gradients. Workers operate independently on disjoint data shards — no model-level synchronisation during the forward/backward passes.

Data Parallel Loop#

Each iteration of data-parallel SGD:

1. Broadcast: The current model parameters $x$ are sent from the parameter server (or replicated via all-gather) to all $K$ workers.
2. Local forward/backward: Each worker $k$ samples a minibatch of size $B_{\text{local}}$, computes the local gradient $\hat{g}_k = \frac{1}{B_{\text{local}}}\sum_{i \in \mathcal{B}_k} \nabla f_i(x)$.
3. All-reduce: Workers collectively sum (or average) their gradients. The all-reduce primitive computes $\sum_{k=1}^K \hat{g}_k$ and makes it available to every worker.
4. Update: Each worker computes the global gradient $\hat{g} = \frac{1}{K}\sum_k \hat{g}_k$ and applies the optimiser step $x \leftarrow x - \eta \hat{g}$ (or communicates the averaged gradient to a parameter server).

The all-reduce operation (e.g., ring all-reduce) is the communication bottleneck. For a gradient vector of size $d$ (number of parameters), the communication time per iteration is roughly $2(K-1)d/K$ worth of data transfer per worker in a bandwidth-optimal ring reduce. With $K = 8$ workers and $d = 10^7$ (40 MB per gradient), this is about 80 MB per worker per step — a few milliseconds on high-bandwidth interconnects (NVLink, InfiniBand), tens of milliseconds on Ethernet.

Learning Rate Scaling Folklore#

Linear scaling rule: If you multiply the batch size by $r$ (by adding workers), multiply the learning rate by $r$ to maintain the same expected parameter update per step:

$\eta_{\text{new}} = r \cdot \eta_{\text{old}}$

Why it sometimes works. The expected gradient is the same regardless of $B$: $\mathbb{E}[\hat{g}_B] = \nabla f(x)$. The total displacement after one step with batch $B$ and learning rate $\eta$ is $-\eta \hat{g}_B$. If you use $B' = rB$ and $\eta' = r\eta$, the expected displacement per step scales as $r$, but the variance scales as $(r\eta)^2 \cdot (\sigma^2/rB) = r \eta^2 \sigma^2 / B$ — variance grows linearly with $r$.

When it breaks:

• Large $r$ causes instability: The variance increase can cause the iterates to diverge early in training before the gradient shrinks.
• The noise is beneficial: SGDStochastic Gradient Descent noise helps escape sharp minima (Week 12). Reducing noise by increasing batch size can lead to worse generalisation.
• Warmup is essential: Starting with the full scaled learning rate can cause immediate divergence. A linear warmup (gradually increasing $\eta$ from a small value) is standard practice.

The linear scaling rule is a heuristic, not a theorem. It is a starting point; hyperparameter tuning (warmup duration, momentum correction, gradient clipping) is nearly always needed.

Generalization Gap Stories#

A consistent empirical finding: large-batch training often produces models with worse test accuracy than small-batch training, even when the training loss is matched. This "generalisation gap" has several proposed explanations:

1. Sharp minima hypothesis: Large batches produce lower-variance gradients, which converge to sharper minima — points where the loss increases rapidly in some directions. Small-batch SGD noise helps escape these sharp regions and land in flatter minima that generalise better (Week 12).

2. Effective learning rate: For a fixed total number of epochs, larger batches mean fewer parameter updates (fewer iterations). The model sees fewer noisy gradient evaluations, effectively exploring less of the loss landscape.

3. Regularisation effect of noise: SGDStochastic Gradient Descent noise acts as implicit regularisation, similar to dropout or weight decay. Reducing noise by scaling the batch removes this regularisation.

Counterpoint. Subsequent work has shown that with careful training recipes (longer training, adjusted regularisation, SAMSharpness-Aware Minimization), large-batch models can match small-batch generalisation. The "generalisation gap" is not a law of nature but a failure of naive scaling — it signals that the optimisation procedure (not just the batch size) must be adapted.

What This Course Does Not Cover#

Distributed training beyond data parallelism is a rich field that this course only gestures at:

• Model parallelism: Different layers of the network reside on different workers (pipeline parallelism). Used when a single model is too large for one GPU.
• Tensor parallelism: Individual matrix multiplications are split across workers.
• ZeRO / FSDP: Sharding optimiser states, gradients, and parameters to reduce per-GPU memory.
• Gradient compression: Quantisation and sparsification to reduce all-reduce communication volume.

These are implementation concerns, not optimisation theory — they belong to systems for ML, not an optimisation curriculum. The conceptual takeaway: data-parallel scaling is mathematically clean (averaging independent gradient estimates) until communication cost or generalisation intervenes.

Browser lab#

Simulate $K$ workers each producing a noisy gradient estimate. Demonstrate that the variance of the averaged gradient falls as $1/K$, and that the expected gradient is unbiased regardless of $K$.

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(123)
d = 100  # dimensional
true_grad = np.random.randn(d)

# Generate K workers, each with noisy gradient: true + noise
def simulate_workers(K, noise_std=1.0, n_trials=2000):
    estimates = np.zeros((n_trials, d))
    for t in range(n_trials):
        grads = np.array([true_grad + noise_std * np.random.randn(d) for _ in range(K)])
        estimates[t] = grads.mean(axis=0)
    # Variance of averaged gradient (averaged over dimensions and trials)
    var = np.mean(np.var(estimates, axis=0))
    bias = np.linalg.norm(estimates.mean(axis=0) - true_grad)
    return var, bias

K_values = [1, 2, 4, 8, 16, 32, 64]
variances = []
biases = []

for K in K_values:
    var, bias = simulate_workers(K, noise_std=1.0)
    variances.append(var)
    biases.append(bias)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

ax1.loglog(K_values, variances, "bo-", markersize=6, label="Empirical variance")
ax1.loglog(K_values, [variances[0]/k for k in K_values], "r--", linewidth=1, label=r"$\propto 1/K$ (theory)")
ax1.set_xlabel("Number of workers $K$")
ax1.set_ylabel("Variance of averaged gradient")
ax1.set_title("Gradient variance vs $K$")
ax1.legend()
ax1.grid(True, alpha=0.3)

ax2.plot(K_values, biases, "go-", markersize=6)
ax2.axhline(0, color="gray", linestyle="--", alpha=0.5)
ax2.set_xlabel("Number of workers $K$")
ax2.set_ylabel("$\|\|\mathbb{E}[\hat{g}] - \\nabla f\|\|$")
ax2.set_title("Bias of averaged gradient vs $K$ (unbiased)")
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print("K    | Variance       | Bias")
print("-" * 35)
for K, var, bias in zip(K_values, variances, biases):
    print(f"{K:3d}  | {var:.6f}       | {bias:.6f}")