Distributed and Federated Optimization

Data spread across devices or data centers forces training to compress, average, and tolerate stragglers, trading some exactness for parallelism.

Local steps on each worker followed by averaging trade communication for exactness

Why Does This Exist?

One GPU cannot hold modern training, and private data cannot move to one server. So optimization runs where the data sits: data centers split batches across workers, phones train on local keystrokes. The math problem becomes coordination under constraints: limited bandwidth, heterogeneous data, and machines that answer late or never.

Naive distribution (ship every gradient everywhere) drowns in communication. The field answers with local steps, compressed messages, and drift corrections. See single-machine rates in non-convex SGD convergence and the stochastic foundations in stochastic approximation.

The one idea you need first

Data parallelism splits each batch across $K$ workers that compute gradients on their shard and average them. Ideal speedup is $K$ times, paid for with $K$ gradient transfers per step. Synchronous averages are exact; asynchronous ones are fast but stale.

Think of It Like This

Cooks in separate kitchens sharing one recipe

Ten cooks in ten kitchens each taste their own local ingredients and adjust the shared recipe card by mail. Mail is slow, so each cook tries several tweaks before posting an update, and shorthand keeps letters short.

The recipe improves fast but drifts toward local tastes: the coastal kitchen oversalts, the mountain kitchen overspices. Regular full-text syncs and drift corrections pull everyone back to a shared dish.

Where the analogy stops: real gradients are millions of numbers with provable averaging properties, not subjective taste notes.

How It Actually Works

FedAvg protocol: server broadcasts weights $w_t$ . Each of $K$ clients runs $E$ local SGD steps on its own data, producing $w_t^k$ . Server averages $w_{t+1} = \sum_k (n_k/n) w_t^k$ . Communication drops by factor $E$ versus per-step sync.

Cost model with model size $M$ bytes and $R$ rounds: traffic $\approx 2RM$ per client (down plus up). With $M = 100$ MB, $R = 1000$ , $E = 10$ versus $E = 1$ : FedAvg moves about $200$ GB per client instead of $2$ TB. Client drift under non-IID data adds a penalty scaling with gradient dissimilarity $G^2$ and $E^2\eta^2$ ; SCAFFOLD-style control variates subtract it.

Worked example

Two clients, true minimum at $w = 0$ . Client A loss $(w-2)^2$ , client B loss $(w+2)^2$ . Start $w_0 = 0$ , $\eta = 0.1$ , $E = 5$ local steps each round:

Client A gradient at $0$ : $-4$ . Each step keeps $0.8$ of the distance to the target, so five steps land at $(1 - 0.8^5) \times 2 \approx 1.34$ .
Client B mirrors to $-1.34$ . Average returns to $0$ . Center holds.
Same setup from $w_0 = 1$ : A pulls toward $2$ , B toward $-2$ , and the average after one round lands near $0.33$ instead of the exact-sync $0.8$ . The $0.47$ gap is client drift, growing with $E$ .

Code

1import numpy as np2
3
4def fedavg_round(w: float, eta: float = 0.1, local: int = 5) -> float:5    def local_run(start: float, target: float) -> float:6        x = start7        for _ in range(local):  # E local SGD steps on (x - target)^28            x = x - eta * 2.0 * (x - target)9        return x10    return (local_run(w, 2.0) + local_run(w, -2.0)) / 2.0  # server average11
12
13print(round(fedavg_round(1.0), 4))  # -> 0.3277, drift of about 0.47 vs exact-sync 0.8

Watch Out For

Running many local epochs on non-IID clients

Large $E$ saves bandwidth but each client overfits its own distribution. Averaging then lands far from any client's good region and accuracy collapses.

Cap $E$ (single digits is typical), use proximal terms (FedProx) or control variates, and watch per-client validation, not just the averaged loss.

Compressing gradients without error feedback

Top-k or 1-bit quantization without memory biases the average permanently. Training converges to the wrong point with full confidence.

Keep error-feedback buffers that add back the compression residual each step. The extra memory is small and the bias fix is large.

The Quick Version

Distribution splits data across workers; federation keeps raw data on devices.
FedAvg runs local steps then averages, cutting communication by the local-epoch factor.
Non-IID data causes client drift that grows with local steps and dissimilarity.
Compression needs error feedback or the average converges biased.
Budget traffic as rounds times model size, and monitor per-client quality.