Client Drift: The Heterogeneity Problem

FedAvg works well in one condition and degrades outside it. Naming the condition precisely is the whole of this lesson.

Write the global objective as the average of client objectives, each being the loss on that client's own data. Under IID data, every client's local objective is an unbiased estimate of the global one. So a step that improves a client's local loss improves the global loss in expectation, every client is pulling in roughly the same direction, and averaging their endpoints is sound.

Under non-IID data this fails at the root. Each client's local objective has a different minimiser from the global objective. Local training does not approximate global training with noise; it optimises a genuinely different function.

The consequence has a name: client drift. As a client takes more local steps, it moves further toward its own optimum and further from where the global model should go. The very mechanism that made FedAvg communication-efficient, many local steps, is what causes it.

Heterogeneity is not one phenomenon, and the flavours have different consequences.

Label skew is the most studied. Each client holds a different mix of classes. A phone belonging to a cat owner has cat photographs; a hospital specialising in oncology sees cancers at a rate unlike a general practice.

Feature skew means the same class looks different across clients. Handwriting differs by writer. Imaging equipment differs by manufacturer, so the same pathology has a different appearance at each site.

Concept drift means the same input maps to different labels. What counts as spam depends on the recipient. A borderline diagnosis depends on local clinical convention.

Quantity skew means dataset sizes span orders of magnitude, which interacts with weighting.

A clean way to see the extreme case: give each client exactly one class. Local training then drives each client's model to predict its own class always, since that is optimal for its local loss. Averaging models that each predict a different constant produces something that predicts nothing well.

Real heterogeneity is milder, but the mechanism is identical.

The failure is geometric. Each client walks toward its own minimum, so the further they walk the more their endpoints spread. Averaging spread-out endpoints lands somewhere that is not the global optimum, and the error grows with the number of local steps.

Client drift creates a genuine dilemma rather than a bug to be patched.

Many local steps per round
  + far fewer communication rounds
  + cheap: local compute is nearly free
  - clients drift far from each other
  - averaged model can be poor or unstable

Few local steps per round
  + updates stay close, averaging is sound
  + behaves like centralised SGD
  - communication cost per unit of progress explodes

So the local-epoch count is not a free efficiency knob. It trades communication against solution quality, and in strongly heterogeneous settings the useful range can be uncomfortably narrow.

The symptoms in practice are recognisable. Training loss oscillates between rounds rather than descending smoothly. Accuracy plateaus below what centralised training on the pooled data achieves. Increasing local epochs, which should help throughput, makes results worse. Individual clients perform well on their own data and poorly on everyone else's.

The research question that follows is whether you can keep many local steps and correct the drift. That turns out to be possible.

The influential fix is SCAFFOLD, from Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh, published at ICML 2020 as "SCAFFOLD: Stochastic Controlled Averaging for Federated Learning".

The insight is that drift is a bias with predictable direction, not random noise. A client's local gradient consistently differs from the global gradient in a way that reflects that client's data. If the difference is consistent, it can be estimated and subtracted.

SCAFFOLD applies control variates, a classical variance-reduction technique. Each client maintains a control variate estimating its own gradient bias, and the server maintains one for the global direction. During local training, the client corrects each step:

local step:  w -= lr * ( grad_local(w) - c_client + c_server )

Read the correction. Subtract what this client's gradient habitually over-contributes; add back the global direction it habitually under-represents. The corrected step estimates the global gradient rather than the local one, so the client can take many steps without wandering toward its own optimum.

The paper's result is that SCAFFOLD requires significantly fewer communication rounds and is not affected by data heterogeneity or client sampling.

The correction is not free, and its costs determine where it can be deployed.

Communication doubles. Clients now send both a model update and a control variate update, so each round moves roughly twice the data. Since communication was the binding constraint, this is a real price.

Clients become stateful. Each client must remember its control variate between rounds. That is natural in cross-silo federation, where the same ten hospitals participate every round. It is a serious problem in cross-device federation, where a phone may be selected once and never again, leaving its control variate stale or undefined.

This is why the taxonomy from the previous lesson matters operationally rather than academically.

Other approaches trade differently. FedProx adds a proximal term penalising distance from the global model, which is stateless and cheap but conservative, slowing progress rather than correcting direction. Momentum-based methods offer another route, with work showing momentum benefits non-IID federated learning both simply and provably, and attractively for cross-device because the state can live on the server.

No method dominates. The choice depends on which resource, communication or client state, you can afford to spend.

There is a deeper response worth considering: if clients have genuinely different data distributions, perhaps forcing them to share a single model is the mistake.

Ask what the objective should be. If clients differ because of sampling noise around a shared truth, one global model is right and heterogeneity is an obstacle to overcome. If they differ because their situations genuinely differ, a single model is a compromise nobody wanted.

Personalised federated learning takes the second view: use collaboration to learn what is shared, then adapt locally to what is not. Common structures include training a global model then fine-tuning per client, sharing a representation while keeping per-client heads, or learning which clients should collaborate with which.

The last idea has a natural formulation as a bilevel optimisation problem, where the outer level decides collaboration structure and the inner level trains, an approach explored in work from Jaggi's lab among others.

The practical question to ask before choosing: would a single model, trained on all the pooled data, actually serve every client well? If yes, fight the drift. If no, personalisation is not a compromise, it is the correct objective.

A summary of the options against the constraints that decide between them.

Two gotchas that cause real confusion. First, more local epochs can make things worse, which inverts the usual intuition that more local work is free progress; if accuracy drops when you raise local epochs, you are seeing drift, not a bug. Second, evaluation must be per-client as well as global. A global average can look acceptable while hiding that the model fails badly on a subset of clients, and in a hospital consortium that subset is a real population.

Everything so far assumes a central server that aggregates. The next lesson removes it.