Federated learning (FedAvg family) and non-IID sensors

"You told me to learn from a thousand people without meeting any of them. So I asked each to teach me privately, then quietly averaged what they said. The trouble was, no two of them lived in the same world."

A Diplomatic AI Agent

FedAvg: ship the model, average the updates

The global objective is an average of per-client losses. With \(K\) clients, client \(k\) holding \(n_k\) local samples and \(n = \sum_k n_k\), federated training minimizes

The what of FedAvg (McMahan et al., AISTATS 2017) is a communication-round loop. In each round the server broadcasts the current weights \(w^t\) to a sampled subset of clients. Each selected client runs \(E\) local epochs of SGD on its own data, producing \(w_k^{t+1}\). The clients send back only those updated weights (or the delta), and the server forms the next global model as a size-weighted average,

The why is communication. Naive distributed SGD would exchange a gradient after every minibatch, which over a cellular link is prohibitive (recall the radio-energy argument of Chapter 59). By letting each client take many local steps (\(E\) epochs) before synchronizing, FedAvg trades computation, which is cheap and local, for communication, which is expensive and shared. The when: use it whenever the data is partitioned across many devices that cannot pool it, and a single global model is still the goal. That is exactly the shape of a wearable fleet, a network of factory sensors, or a set of phones.

Why sensor fleets are aggressively non-IID

Averaging assumes the things being averaged point roughly the same way. That holds when every client draws from the same distribution (independent and identically distributed, IID). Sensor fleets violate it along several axes at once, and understanding which axis you face determines the fix.

Feature skew is the most physical. The same activity produces different signals depending on sensor placement, body, and hardware: a phone in a trouser pocket, a smartwatch, and a chest strap record walking as three distinct accelerometer signatures (Chapter 26). Label skew is quantity imbalance: an office worker's wearable logs almost no swimming, an athlete's logs little sitting, so each client's label histogram \(P_k(y)\) is wildly different. Concept drift across clients means \(P_k(y \mid x)\) itself differs, for instance the same heart-rate pattern signaling exertion in one user and arrhythmia risk in another. On top of statistical heterogeneity sits systems heterogeneity: devices differ in speed, battery, and availability, so clients complete different numbers of local steps and some drop out mid-round.

The consequence has a name: client drift. When \(P_k\) differ, each client's local optimum \(w_k^\star\) sits somewhere different, and the FedAvg average of far-run local models is biased away from the true global optimum \(w^\star\). Empirically this shows up as an accuracy curve that climbs, then oscillates or plateaus well below what centralized training on the pooled data would reach. Non-IID data does not merely slow federated learning; past a point it caps the achievable accuracy.

The FedAvg family: correcting client drift

The family keeps FedAvg's ship-model-average-updates skeleton and adds a term that fights drift. FedProx (Li et al., MLSys 2020) is the gentlest: each client minimizes its local loss plus a proximal penalty \(\tfrac{\mu}{2}\lVert w - w^t \rVert^2\) that anchors the local model to the broadcast global weights, so no client wanders too far in \(E\) epochs. It also tolerates stragglers by accepting partial local work. SCAFFOLD (Karimireddy et al., ICML 2020) is more surgical: it maintains control variates, per-client and global correction vectors that estimate the drift direction and subtract it from each local update, in effect variance-reducing the client updates. It converges in far fewer rounds on heterogeneous data at the cost of communicating and storing an extra vector the size of the model. FedAvgM (Hsu et al., 2019) simply applies momentum on the server when aggregating updates, which smooths the noisy round-to-round direction that skew induces. The how you choose: start with FedProx for a cheap, robust default; reach for SCAFFOLD when communication rounds are the binding cost and you can afford the extra state; add server momentum when the global curve is jittery.

The listing below simulates FedAvg versus FedProx on a deliberately skewed two-client problem so you can see drift and its correction without any networking. It is referenced in the discussion that follows.

import numpy as np

rng = np.random.default_rng(0)
# Two clients, opposite feature skew: same task, shifted inputs (non-IID).
def make_client(shift, n=400):
    X = rng.normal(shift, 1.0, size=(n, 2))
    y = (X[:, 0] + X[:, 1] > 2 * shift).astype(float)   # local decision boundary
    return X, y

clients = [make_client(-1.5), make_client(+1.5)]

def local_train(w, X, y, w_glob, mu, epochs=5, lr=0.1):
    for _ in range(epochs):
        z = X @ w
        p = 1 / (1 + np.exp(-z))
        grad = X.T @ (p - y) / len(y) + mu * (w - w_glob)  # mu>0 => FedProx
        w = w - lr * grad
    return w

def federated(mu, rounds=30):
    w = np.zeros(2)
    for _ in range(rounds):
        updates = [local_train(w.copy(), X, y, w, mu) for X, y in clients]
        w = np.mean(updates, axis=0)          # size-weighted avg (equal here)
    return w

def acc(w):
    correct = sum(((X @ w > 0).astype(float) == y).mean() for X, y in clients)
    return correct / len(clients)

print(f"FedAvg  (mu=0.0): global acc = {acc(federated(0.0)):.3f}")
print(f"FedProx (mu=0.5): global acc = {acc(federated(0.5)):.3f}")

As Listing 64.2.1 makes concrete, the only structural change from FedAvg to FedProx is one extra term in the gradient. That is the family's design signature: the communication loop is sacred; the correction lives inside the local objective or the aggregation step. Note also that this simulation, like all federated evaluation, must be read against a per-client, subject-disjoint test set, or the reported accuracy is an illusion, a theme Section 64.7 makes rigorous.

import flwr as fl

class SensorClient(fl.client.NumPyClient):
    def get_parameters(self, config):        return get_weights(model)
    def fit(self, parameters, config):
        set_weights(model, parameters)
        train_one_round(model, local_loader)   # your normal training loop
        return get_weights(model), len(local_loader.dataset), {}
    def evaluate(self, parameters, config):
        set_weights(model, parameters)
        return local_loss(model), len(val_loader.dataset), {"acc": local_acc(model)}

fl.simulation.start_simulation(
    client_fn=lambda cid: SensorClient(),
    num_clients=1000,
    strategy=fl.server.strategy.FedProx(proximal_mu=0.1),  # swap to FedAvg/FedAvgM freely
)

What's Next

In Section 64.3, we take seriously the admission this section ended on: when clients are heterogeneous enough, no single averaged model is the right answer. We turn from forcing everyone onto one global model to personalization and local adaptation, giving each device a model that starts from the shared knowledge but bends to fit its own wearer, placement, and hardware.