Non-IID Data

"In the datacenter they shuffled my data until every shard looked the same. Out here, my whole world is one person's camera roll, and I am very confident about cats."

A Client Model Overfit to Its Owner

In Section 14.3 we built FedAvg: each round, the server broadcasts the current model, a set of clients each run several local SGD epochs on their own data, and the server averages the returned models weighted by how many examples each client holds. We mentioned that FedProx and SCAFFOLD exist as variants without yet saying what disease they cure. This section is that diagnosis. The single assumption that FedAvg silently inherited from datacenter data parallelism, that each participant's data is an independent and identically distributed (IID) sample from one global distribution, is precisely the assumption that real federated populations violate. Once we drop it, the comforting "an average decomposes, so distribution is exact" argument loses its footing, and the gap it opens is the subject of the rest of the chapter.

1. Why Shuffling Was Doing All the Work Beginner

Recall what shuffling buys you in ordinary distributed training. You take a dataset, permute it uniformly at random, and hand contiguous blocks to each worker. Because the permutation is uniform, every worker's block is a representative sample: its label histogram, its feature statistics, and its difficulty profile all match the global ones up to sampling noise. The local gradient each worker computes is therefore an unbiased estimate of the global gradient, and averaging those estimates reduces variance without introducing bias. That is the entire statistical foundation of synchronous data parallelism, and it is why a hundred workers can act as one.

Federated learning removes the permutation step and cannot put it back. The data was generated where it lives and must stay there, for reasons of privacy, regulation, bandwidth, or simple physics: you are not going to upload every phone's camera roll to a central shuffler. So the partition into clients is not a random cut through one distribution; it is a partition by source, and sources differ. Formally, write client $k$'s local distribution as $\mathcal{D}_k$ and the global mixture as $\mathcal{D} = \sum_k \frac{n_k}{n}\,\mathcal{D}_k$, where $n_k$ is client $k$'s sample count. In the IID datacenter world every $\mathcal{D}_k \approx \mathcal{D}$. In the federated world the $\mathcal{D}_k$ are different distributions, and the size of that difference is what we now have to reason about.

2. The Four Faces of Heterogeneity Beginner

"Non-IID" is an umbrella over several distinct ways client distributions can differ, and they call for different remedies, so it pays to keep them separate. Table 14.4.1 names the four that dominate practice. Factor the joint distribution as $\mathcal{D}_k(x,y) = \mathcal{D}_k(y)\,\mathcal{D}_k(x\mid y)$ on each client, and most heterogeneity types are a statement about which factor varies.

Label-distribution skew is the form most studies stress-test, because it is both common and severe: the standard benchmark sorts a dataset by label and hands each client only one or two classes, a synthetic worst case that nonetheless mirrors reality where a user generates a narrow slice of the label space. Feature skew is subtler and often invisible to the label histogram: every hospital may see the full range of diagnoses, yet each one's images carry a scanner-specific signature that a model can latch onto. Quantity skew interacts with the weighting in FedAvg, since a few enormous clients can dominate the average. Concept drift is the hardest, because it means there is no single $w^\star$ that is simultaneously right for everyone, which is the motivation for the personalized methods of Section 14.7. This last form connects directly to the distribution-shift and concept-drift machinery we built for streaming systems in Chapter 9; federated non-IID data is distribution shift across space rather than across time.

3. Client Drift, and Why Many Local Epochs Amplify It Intermediate

The mechanism that turns heterogeneity into a convergence problem is client drift: during a round, each client moves its copy of the model away from the shared starting point and toward its own local optimum, so by aggregation time the client models have spread apart. A useful scalar measure of how badly a round drifted is the average distance of the client models from their aggregate,

where $w_k$ is client $k$'s model after its local steps and $\bar{w}$ is the FedAvg aggregate. Two knobs inflate this quantity. The first is heterogeneity itself: the more the local gradients $\nabla F_k$ disagree at the shared starting point, the faster the clients diverge. The second, and the one practitioners control directly, is the number of local epochs $E$ introduced in Section 14.3. Each extra local epoch lets a client travel further down its own loss surface before the server gets a chance to average, so larger $E$ means each $w_k$ sits closer to its private $w_k^\star$ and further from its neighbors. FedAvg uses large $E$ precisely to save communication, which is the tension Section 14.5 takes up: the same local computation that cuts your network bill is what lets clients drift.

The runnable demo below makes both effects measurable on a tiny softmax classifier. It builds one synthetic dataset, partitions it two ways (an IID shuffle and a label-sorted non-IID split where each client sees essentially one class), and runs FedAvg at $E=1$ and $E=10$ local epochs. For each setting it reports the average per-round drift defined above and the final test accuracy, so you can watch heterogeneity and local epochs inflate drift while accuracy slips.

import numpy as np

rng = np.random.default_rng(0)
C, D, N = 6, 30, 12000           # classes, features, total examples
NC = 30                          # number of clients
ROUNDS = 30

means = rng.standard_normal((C, D)) * 0.45
y = rng.integers(0, C, size=N)
X = means[y] + rng.standard_normal((N, D))

def softmax(Z):
    Z = Z - Z.max(axis=1, keepdims=True)
    E = np.exp(Z); return E / E.sum(axis=1, keepdims=True)

def grad(W, Xb, yb):
    P = softmax(Xb @ W)
    Y = np.zeros_like(P); Y[np.arange(len(yb)), yb] = 1.0
    return Xb.T @ (P - Y) / len(yb)

def accuracy(W):
    return float((softmax(X @ W).argmax(1) == y).mean())

def iid_partition():
    return np.array_split(rng.permutation(N), NC)

def noniid_partition():
    # Sort by label, then chop into shards: each client sees one class
    # (extreme label-distribution skew, the standard worst case).
    return np.array_split(np.argsort(y, kind="stable"), NC)

def fedavg(parts, E, lr=1.5, batch=32):
    W = np.zeros((D, C))
    sizes = np.array([len(p) for p in parts], dtype=float)
    drift_acc = 0.0
    for _ in range(ROUNDS):
        locals_ = []
        for p in parts:
            Wl = W.copy()
            for _ in range(E):                       # E local epochs
                pp = rng.permutation(p)
                for b in range(0, len(pp), batch):
                    bi = pp[b:b + batch]
                    Wl -= lr * grad(Wl, X[bi], y[bi])
            locals_.append(Wl)
        stack = np.stack(locals_)
        Wbar = np.average(stack, axis=0, weights=sizes)   # FedAvg aggregation
        # Client drift: mean distance of each local model from the average.
        drift_acc += float(np.mean([np.linalg.norm(Wl - Wbar) for Wl in stack]))
        W = Wbar
    return drift_acc / ROUNDS, accuracy(W)

iid, non = iid_partition(), noniid_partition()
print(f"{'setting':<30}{'mean client drift':>18}{'test accuracy':>16}")
print("-" * 64)
for name, parts, E in [
    ("IID,     E=1",  iid, 1),
    ("IID,     E=10", iid, 10),
    ("non-IID, E=1",  non, 1),
    ("non-IID, E=10", non, 10),
]:
    drift, acc = fedavg(parts, E)
    print(f"{name:<30}{drift:>18.3f}{acc:>16.3f}")

setting                        mean client drift   test accuracy
----------------------------------------------------------------
IID,     E=1                               2.107           0.825
IID,     E=10                              4.932           0.824
non-IID, E=1                               2.632           0.818
non-IID, E=10                              4.971           0.805

The numbers tell the whole story of this section in one table. Non-IID data raises drift on its own, more local epochs raise it further, and the accuracy cost lands hardest exactly where both compound: many local epochs on skewed clients. The convex softmax model here is forgiving, so the accuracy gap is modest; in the non-convex deep networks of Chapter 15 the same drift can stall or destabilize training outright. The drift scalar is the early-warning signal that the final-accuracy number hides until it is too late.

4. Mitigations: From Free to Expensive Intermediate

There is a ladder of responses to client drift, ordered roughly by cost. The cheapest, demonstrated in Output 14.4.1, is to reduce the number of local steps: fewer local epochs means less room to drift, at the price of more communication rounds. The next rung adds a correction term to the local objective. FedProx, introduced in Section 14.3, augments each client's loss with a proximal penalty $\frac{\mu}{2}\lVert w - w^{\text{global}}\rVert^2$ that pulls local updates back toward the round's starting point, directly shrinking the drift scalar at the cost of a hyperparameter $\mu$ to tune. SCAFFOLD goes further with control variates: the server and each client maintain correction vectors that estimate the difference between the client's local gradient direction and the global one, and subtract it during local steps, cancelling the drift bias rather than merely penalizing it. A third rung lives on the server: applying momentum or an adaptive optimizer (FedAdam and relatives) to the sequence of aggregates smooths over the per-round noise that heterogeneity injects.

The most expensive rung relaxes the no-data-movement constraint, where policy allows. Sharing a small globally representative public dataset, or a server-held proxy set, lets clients or the server correct the skew directly; knowledge-distillation variants transmit soft predictions on shared unlabeled data instead of weights, which both sidesteps drift in weight space and shrinks the messages. These methods trade some of the privacy and bandwidth advantages that motivated federation in the first place, so they are a last resort reserved for the cross-silo settings of Section 14.2 where a small shared corpus is legally and practically available. Table 14.4.2 lays the ladder out.

import flwr as fl

# Plain FedAvg drifts under non-IID clients ...
strategy = fl.server.strategy.FedAvg(fraction_fit=0.2)

# ... swap in FedProx (proximal mu) or FedAdam (server momentum) instead.
strategy = fl.server.strategy.FedProx(fraction_fit=0.2, proximal_mu=0.5)
strategy = fl.server.strategy.FedAdam(fraction_fit=0.2)   # adaptive server optimizer

fl.server.start_server(config=fl.server.ServerConfig(num_rounds=30),
                       strategy=strategy)

5. Where Non-IID Data Reappears Beginner

The heterogeneity introduced here is not confined to this chapter; it is the constraint that shapes every later federated setting in the book. When federated learning moves to the network edge in Chapter 34, device populations are even more skewed and intermittently available, so the drift problem is compounded by clients that appear and vanish mid-round. When we reach the federated medical case study of Chapter 37, feature skew across hospital scanners and label skew across regional disease prevalence are the central engineering reality, not a textbook caricature. And the same averaging-of-divergent-updates tension recurs whenever local computation is traded against communication, including the local-SGD methods we will compare in Chapter 10; non-IID federated optimization is local SGD with the additional cruelty that you cannot fix the data.

With heterogeneity diagnosed and its mitigations laid out, the natural next question is the one we kept deferring: how few bytes can a federated round actually cost, and how does that budget interact with the local-epoch knob that drives drift? That is the communication-constraint story, and it begins in Section 14.5.