Empirical Risk Minimization at Scale

"They told me my job was to minimize the average loss. Nobody mentioned the average would be scattered across two thousand machines that all wanted to go first."

An Optimizer With Too Many Workers

Training a model, stripped of its framework and its jargon, is the act of choosing parameters that make a chosen loss small on the data you have. You do not get to measure loss on the true distribution of all possible inputs, because you never see that distribution; you see a finite sample. So you settle for the next best thing, the average loss over the examples in front of you, and you search for the parameters that minimize it. That average has a name, empirical risk, and minimizing it, empirical risk minimization or ERM, is the optimization problem underneath linear regression, logistic regression, support vector machines, gradient-boosted trees, and every deep network in this book. Chapter 1 used the gradient of this objective to show that data parallelism is exact; this chapter takes that same objective and builds the full distributed-optimization toolkit around it.

1. The Objective Behind All Training Beginner

Fix a model with parameters $w$ and a per-example loss $\ell(w; x_i, y_i)$ that measures how badly the model with parameters $w$ predicts the label $y_i$ for input $x_i$. Given $N$ training pairs, the empirical risk is the average of that loss over the sample, and ERM is the problem of choosing $w$ to make it as small as possible:

The optional term $\lambda R(w)$ is regularization, a penalty on the size or complexity of $w$ (most often $R(w) = \tfrac{1}{2}\lVert w \rVert^2$, ridge or weight decay) that discourages overfitting the finite sample. Different models simply plug in different losses: squared error $\ell = \tfrac{1}{2}(w^\top x_i - y_i)^2$ gives linear regression, the logistic loss gives logistic regression, the hinge loss gives a linear support vector machine, and the cross-entropy of a deep network's output gives a neural classifier. The shape of $\ell$ changes everything about how hard the search is, but the form of the objective, an average over examples plus a penalty, does not change at all. That invariance is what lets one chapter on optimization serve models as different as a linear regressor and a transformer.

To minimize $L$ we follow its gradient downhill, and here the averaging structure carries straight through differentiation. Because the derivative of a sum is the sum of the derivatives, the gradient of the empirical risk is itself an average of per-example gradients:

This is the identity Section 1.1 used to prove data parallelism is exact (Section 1.1), and it is the hinge on which this entire chapter swings. Read it as a recipe for distribution: to compute $\nabla L$, you sum a contribution from every example, and a sum can be partitioned across machines and recombined by addition. The regularizer $\lambda \nabla R(w)$ is a function of $w$ alone, not of the data, so it costs nothing to distribute; every worker can add its share locally.

2. Why the Sum Is Separable, Exactly Intermediate

The property that makes ERM distributable is that its objective is a separable sum: a total formed by adding one independent term per example, where no term depends on which other examples happen to share the dataset. Separable sums decompose along any partition of the index set. Split the $N$ examples into $K$ disjoint shards $S_1, \dots, S_K$, hand shard $S_k$ to worker $k$, and let that worker compute the unnormalized partial sum over only its own rows,

Then $\sum_{i=1}^{N} \nabla \ell = g_1 + g_2 + \cdots + g_K$, because addition is associative and commutative and does not care how the terms were grouped. Dividing the summed partials by $N$ (and adding the cheap regularizer gradient) recovers $\nabla L(w)$ exactly. The combine step, summing one vector held on each worker and sharing the result, is the collective operation called all-reduce, introduced in Chapter 4 and turned into a training loop in Chapter 15. Figure 10.1.1 traces exactly this path from the single objective at the top to the recombined gradient at the bottom.

3. Convex and Non-Convex Objectives Intermediate

Separability tells us the gradient distributes, but it says nothing about whether following that gradient actually finds a good minimum. That question turns on the shape of $L$, and the decisive distinction is convexity. An objective is convex if its graph is bowl-shaped: the line segment between any two points on the surface lies on or above the surface, so the function curves the same way everywhere. Linear regression with squared loss, logistic regression, and the linear SVM are all convex in their parameters. Convexity is a gift, because a convex objective has no spurious local minima; every local minimum is a global minimum, and gradient descent from any starting point converges to the same optimal loss. For distributed optimization this means the convergence guarantees are clean and provable: synchronous distributed SGD on a convex objective converges to the global optimum at rates we can write down, and those rates are what Chapter 3's cost models combine with communication time to predict end-to-end training speed.

Deep networks break convexity. Composing linear maps with nonlinear activations produces an objective with many local minima, saddle points, flat regions, and sharp valleys; the bowl becomes an egg carton. Non-convex ERM has no guarantee that gradient descent reaches the global optimum, and different runs from different initializations land in different basins. Yet the same averaged-gradient machinery still applies, because separability is a property of the sum, not of its curvature. We still compute $\nabla L$ as an average of per-example gradients, still split that average across shards, still all-reduce it. What changes is the theory we can offer: for non-convex objectives we prove convergence only to a stationary point (a place where the gradient is near zero), not to a global minimum, and the empirical fact that this works well enough to train every modern model is one of the happy surprises of the field. The good news for this chapter is that the distribution mechanics are identical in both regimes; convexity affects what you can prove about the destination, not how the workers cooperate to take each step.

4. Computing the Gradient as a Sum of Shard Gradients Intermediate

The claim that the full-data gradient equals the sum of per-shard gradients is easy to state and worth verifying in code, because the rest of the chapter rests on it. The program below sets up a regularized logistic-regression ERM objective on synthetic binary-classification data, then computes its gradient two ways: once the single-machine way over all $N$ rows, and once as $K$ independent shard computations whose partial sums are added and divided by $N$, exactly the all-reduce-then-normalize path of Figure 10.1.1. It reports how far apart the two answers are.

import numpy as np

rng = np.random.default_rng(7)
N, d, K = 120_000, 40, 6          # examples, features, workers (shards)
lam = 1e-2                        # L2 regularization strength

# Synthetic binary-classification data for regularized logistic regression.
X = rng.standard_normal((N, d))
w_star = rng.standard_normal(d)
p = 1.0 / (1.0 + np.exp(-(X @ w_star)))
y = (rng.random(N) < p).astype(float)        # labels in {0, 1}
w = rng.standard_normal(d) * 0.1             # the point where we evaluate the gradient


def reg_logistic_grad(Xb, yb, w):
    """Gradient contribution of a block of rows to the FULL-data objective.
    The data term is summed (not averaged) so blocks compose by addition; the
    regularizer is the per-row share lam*w scaled by the block's row count."""
    sigma = 1.0 / (1.0 + np.exp(-(Xb @ w)))
    data_term = Xb.T @ (sigma - yb)          # sum over this block's rows
    reg_term = lam * w * Xb.shape[0]         # this block's share of N * lam * w
    return data_term + reg_term


# Single-machine gradient: one process sees all N rows.
full = reg_logistic_grad(X, y, w) / N

# Data-parallel gradient: K workers, each blind to the others' rows.
shards = np.array_split(np.arange(N), K)
partials = [reg_logistic_grad(X[s], y[s], w) for s in shards]   # K workers
allreduced = np.sum(partials, axis=0) / N                       # all-reduce, then /N

print("workers K              :", K)
print("shard sizes            :", [len(s) for s in shards])
print("max abs difference     :", f"{np.max(np.abs(allreduced - full)):.2e}")
print("relative error         :", f"{np.linalg.norm(allreduced - full) / np.linalg.norm(full):.2e}")
print("gradients identical?   :", bool(np.allclose(allreduced, full)))

workers K              : 6
shard sizes            : [20000, 20000, 20000, 20000, 20000, 20000]
max abs difference     : 1.65e-15
relative error         : 6.33e-15
gradients identical?   : True

The difference is not merely small; it is zero up to floating-point rounding, and np.allclose confirms the two gradients are identical for every practical purpose. This is the same exactness Section 1.1 demonstrated for linear regression, now shown to survive the addition of a nonlinear logistic loss and an $L^2$ regularizer, because none of that disturbs the separable-sum structure. The only subtlety the code makes explicit is bookkeeping: each shard sums its data-term contributions rather than averaging them, and the regularizer is apportioned by row count, so that the partial sums add back to exactly $N \nabla L$ before the single division by $N$. Get that bookkeeping wrong (for instance by averaging within each shard and then averaging across shards of unequal size) and the exactness breaks, a failure you will provoke deliberately in Exercise 10.1.2.

# Run with: torchrun --nproc_per_node=6 thisfile.py
import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist

dist.init_process_group("nccl")                     # join the group of K workers
model = DDP(nn.Linear(40, 1).cuda())                # logistic regression as a 1-layer net
loss_fn = nn.BCEWithLogitsLoss()                    # mean cross-entropy over the shard

for xb, yb in this_workers_shard_loader:            # each worker sees only its rows
    loss = loss_fn(model(xb).squeeze(), yb)
    loss.backward()      # autograd builds the per-shard gradient AND all-reduces it
    optimizer.step()     # every worker now steps with the exact full-batch direction

5. Setting Up the Chapter's Attack on This Objective Beginner

We now have the target. Every method in this chapter is a different way of minimizing the empirical risk $L(w)$ across machines, and they differ in how they handle two costs that Code 10.1.1 hid because everything ran in one process. The first cost is computing the full sum at all: with millions or billions of examples, evaluating $\nabla L$ exactly every step is wasteful, so Section 10.2 replaces it with the gradient of a small random mini-batch, an unbiased sample of the full average, giving stochastic gradient descent. The second cost is communication: once the gradient is split across workers, the partial sums must be combined every step, and that exchange takes real network time. Section 10.3 performs the combine synchronously so every worker steps in lockstep with an exact averaged gradient, while Section 10.4 lets workers run asynchronously, trading exactness for the elimination of waiting, at the cost of the stale gradients that Chapter 11 manages with bounded staleness.

Communication is the thread that runs through all of it. Because the combine step moves a vector the size of the model every step, its cost grows with both model size and worker count, and for large models it can dominate the computation it serves. The communication-efficient methods of Section 10.7 shrink that cost by compressing, quantizing, or sparsifying the messages, or by communicating less often (local SGD), and Section 10.9 asks how far that can be pushed in principle through communication-complexity lower bounds. The federated setting of Chapter 14 takes the same ERM objective to its communication-starved extreme, where data never leaves the devices that hold it. Throughout, the object being optimized never changes; what changes is how cleverly the workers cooperate to evaluate and follow its averaged gradient.

The plan for the chapter, then, is set. We have the objective (empirical risk), the structural fact that makes it distributable (a separable sum whose gradient is an average), the distinction that governs what we can prove about the result (convex versus non-convex), and a verified demonstration that the distributed gradient equals the single-machine one exactly. The next section takes the first and most consequential step away from the exact full-data gradient: sampling it.