Most machine learning still assumes a convenient fiction: all the data sits in one place, ready to be shuffled, split, and trained on. Real data rarely cooperates. It lives on phones, in hospitals, in banks, and in factories, where moving it is expensive, slow, or simply not allowed. Federated learning is the attempt to train anyway, without centralising the data.
The core idea is disarmingly simple. Instead of sending data to the model, send the model to the data. Each client trains locally, sends back only an update, and a server averages those updates into a new global model. That averaging step, called Federated Averaging or FedAvg, looks almost too naive to work. Understanding why it usually does work, and exactly where it breaks, is the key to using it well.
The problem FedAvg solves
Suppose there are K clients. Client k holds n_k examples, and n is the total number of examples across all clients. The global objective is a weighted average of local objectives:
Here w is the model, and is the loss on client k's data alone. The server never sees the data behind . It only sees updates derived from it.
In a centralised setting, stochastic gradient descent would repeatedly compute a gradient on a random batch drawn from the whole population. In federated learning, each client can only draw batches from its own slice. If every client held a representative sample of the whole population, this would barely matter. In practice, clients are heterogeneous: one hospital sees different patients from another, one phone user types different words from the next, and one sensor sees a different operating regime. That heterogeneity, usually called non-IID data, is the central difficulty of the field.
FedAvg in one round
A round of FedAvg proceeds as follows:
- The server sends the current global model to a subset of clients.
- Each selected client runs several steps of SGD on its own data, starting from , and obtains a local model .
- Each client sends its update, or its new model, back to the server.
- The server forms the next global model as a weighted average:
Two design choices matter enormously here. First, clients do multiple local steps rather than one. That reduces communication, which is often the real bottleneck when clients are phones on slow networks or devices that are only intermittently available. Second, the averaging is weighted by dataset size, so a client with more data pulls the global model more strongly. Both choices are sensible defaults, and both can mislead.
Why should averaging models work at all? Averaging parameters is not the same as averaging predictions. For a linear model, the two coincide in a useful way. For a deep network, averaging two models trained from different starting points usually produces a poor model, because the same function can be represented by permuted hidden units. FedAvg avoids the worst of this by starting every client from the same each round and only letting them drift for a few local steps. The local models stay in the same basin, so their average remains a sensible model. Stretch the local training too far, and that assumption fails.
A minimal FedAvg loop
The following sketch captures the structure, stripped of networking, serialisation, and client selection. It is useful precisely because it shows how little code the core idea needs, and how much real systems add around it.
# w: global model parameters (a flat vector for clarity)
# clients: list of local datasets, each with its own loss and gradient
def client_update(w, data, steps, lr):
local_w = w.copy()
for _ in range(steps):
batch = sample_batch(data)
local_w = local_w - lr * gradient(local_w, batch)
return local_w
def federated_round(w, clients, weights, steps, lr):
updates = []
for client, weight in zip(clients, weights):
updates.append(weight * client_update(w, client, steps, lr))
return sum(updates) # weights sum to 1, so this is the weighted average
In a real deployment, client_update runs on the device, federated_round runs on the server, and only the returned vector crosses the network. Everything else - dropout handling, secure aggregation, versioning, retries for clients that vanish mid-round - is engineering around this loop.
Client drift: the failure mode to watch
When data are heterogeneous, each client's local optimum is different. Let ParseError: KaTeX parse error: Expected group after '^' at position 4: w_k^̲ be the minimiser of . Local SGD pulls client k toward ParseError: KaTeX parse error: Expected group after '^' at position 4: w_k^̲ , not toward the global minimiser. After several local steps, the average of the local models can drift away from a good global solution, even if each local step reduced its own client's loss.
A simple picture helps. Imagine two clients, one whose data are mostly cats and one whose data are mostly dogs. Each local model becomes excellent at its own class and worse at the other. Their average may be mediocre at both. This is not a bug in the averaging formula. It is what the formula does when the local objectives disagree.
Several symptoms point to client drift in practice:
- Training loss falls on each client but global validation accuracy stalls or oscillates.
- More local steps make things worse, not better, even though they reduce communication.
- A small subset of clients dominates, because they have more data or are selected more often.
- Performance varies wildly across clients after training, even when average accuracy looks acceptable.
The standard remedies trade off against each other. Fewer local steps reduce drift but increase communication. A proximal term that penalises movement away from , as in FedProx, keeps local models closer to the shared starting point. Momentum-like corrections and control variates, as in SCAFFOLD, explicitly estimate and subtract the drift. None of these removes heterogeneity. They manage it, at the cost of extra computation, extra state, or extra tuning.
Communication, computation, and stragglers
Federated learning inverts the usual resource picture. In a datacentre, compute is the bottleneck and communication is cheap. On the edge, communication is expensive, devices are slow and heterogeneous, and some clients simply disappear. A round waits for the slowest selected client unless the system imposes a deadline, and a deadline means some updates are discarded.
This creates three practical trade-offs:
- Local steps vs. rounds. More local work per round saves bandwidth but increases drift and straggler delay.
- Clients per round vs. stability. Sampling more clients gives a better estimate of the global gradient but costs more coordination.
- Model size vs. participation. A larger model may be more accurate, but every extra parameter must be sent down and back up, which excludes slow or constrained clients.
Compression, quantisation, and partial participation all help, but they change the optimisation dynamics. A compressed update is a noisy update. The question is never whether compression loses information. It does. The question is whether the lost information mattered for the global objective, and that depends on how redundant the clients' updates already are.
Privacy: what averaging does and does not guarantee
It is tempting to say that federated learning is private because raw data never leaves the device. That claim is too strong. Model updates are functions of the data, and in some settings they leak a surprising amount about it. Gradient inversion attacks can reconstruct recognisable inputs from gradients for small batches and certain architectures. Membership inference can reveal whether a particular example was in a client's training set. A curious server, or a compromised aggregator, is a real threat model, not a theoretical one.
This is why production systems layer additional protections on top of averaging:
- Secure aggregation lets the server learn only the sum of updates, not any individual client's update, using cryptographic masking that cancels out in the aggregate.
- Differential privacy adds calibrated noise to updates, so that the influence of any single example is bounded. The cost is accuracy, and the cost grows when many rounds are run.
- Data minimisation on the client, such as training on features rather than raw inputs where possible, reduces what any leaked update can reveal.
None of these is free, and none should be treated as a checkbox. Differential privacy in particular forces a concrete choice: how much accuracy are you willing to trade for a stated privacy guarantee, measured over the full training run, not just one round? If a proposal cannot answer that, its privacy claim is marketing, not engineering.
When federated learning is the right tool
Federated learning earns its complexity when at least one of these holds:
- Data cannot legally or contractually leave its location, as with health records or financial transactions.
- Data are too large or too dispersed to centralise, as with sensor networks.
- The deployment itself is distributed, so a model trained centrally would face a distribution shift on the devices that actually use it.
It is a poor fit when data could be centralised with modest effort, when clients are nearly homogeneous, or when the team cannot afford the systems work: client management, versioning, monitoring per-client performance, and debugging a training run where the data are, by design, invisible to the developers. In those cases, centralised training with strong access controls is simpler, faster, and easier to audit.
A useful test before committing: simulate federation on data you can already centralise, by partitioning it artificially across clients in a way that mimics the real heterogeneity. If FedAvg underperforms centralised training badly in simulation, it will not magically improve when the data become genuinely inaccessible. If it performs close to centralised training, the remaining work is systems work, not algorithmic hope.
Common pitfalls
- Treating the average as the whole story. Report per-client performance, not just the global average. A model that works for the majority and fails for a minority is a deployment failure, not a success.
- Tuning on IID partitions. Randomly shuffling centralised data across clients hides exactly the heterogeneity that makes federation hard. Partition by natural client boundaries, or by label skew, to get an honest estimate.
- Ignoring client availability. Clients that are online at night differ systematically from clients online during the day. Selection bias in federated learning is not a sampling detail. It changes the objective being optimised.
- Adding privacy after the fact. Noise and secure aggregation interact with optimisation. They need to be part of the training design and evaluation, not bolted on before launch.
- Assuming more rounds fix drift. If local objectives disagree, more rounds of the same procedure can converge to a compromise that satisfies nobody. Diagnose drift first, then choose a correction, rather than simply training longer.
Takeaway
Federated averaging works because it exploits a narrow but powerful regime: clients start from a shared model, move only a little on their own data, and are then pulled back together by averaging. That regime gives communication efficiency without centralising data. It breaks when local steps go too far, when clients disagree too much, or when privacy is assumed rather than engineered.
The art of federated learning is therefore not in the averaging formula itself. It is in choosing how far clients may drift before they are reined in, how much noise privacy requires, and how honestly per-client performance is measured. Get those choices right, and averaging becomes a practical way to learn from data that can never be gathered in one place. Get them wrong, and the average is just a polite way of describing a model that works nowhere in particular.
Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial
Top comments (0)