FedAvg, live.
Eight clients each keep a private slice of handwritten digits and train one shared model by exchanging only weights. Drag the skew slider to watch non-IID data hurt.
How it works
Everything runs in this tab. The data is the UCI optical digits set (scikit-learn's load_digits, 8×8 greyscale): 1,297 training images split across 8 clients, 500 held out for testing.
- Partition. Each client "owns" two digits. With probability skew an image goes to an owner of its digit, otherwise to a random client. Skew 0 is IID; skew 1 means every client sees exactly two digits.
- Round. The server sends the global weights to the sampled clients. Each runs mini-batch SGD (batch 16) on its own slice for E local epochs and sends the weights back.
- Average. The server replaces the global model with the mean of the returned weights, weighted by each client's sample count — that is Federated Averaging.
- Baselines step in lockstep from the same initial weights: a centralized model trained E epochs per round on pooled data, and eight local-only models that never talk to anyone.
The model is a 64→32→10 MLP (ReLU, softmax cross-entropy) with backprop written by hand in about 100 lines — 2,410 parameters, 9.4 KB as float32. No libraries, no build step.
Measured results
bench(): 30 rounds, 2 local epochs, 8/8 clients, lr 0.1, mean over seeds 1–3, run in headless Chrome. Test accuracy on the 500 held-out digits after round 30.
| label skew | fedavg | centralized | local-only (mean) |
|---|---|---|---|
| 0.0 · iid | 96.7% | 97.3% | 91.4% |
| 1.0 · 2 digits / client | 86.1% | 97.3% | 19.9% |
Seed ranges — fedavg: 96.0–97.6% (iid), 84.4–87.0% (skew 1); centralized: 97.2–97.4%. Weights are 55.7× the raw pixels. Raw output: docs/bench-chrome.json.
What I learned
Federated learning buys privacy, not bandwidth. Over 30 rounds the clients exchanged 4,519 KB of weights to protect 81 KB of pixels — 55.7× the raw data. With a tiny dataset and a full model sent every round, shipping the data would have been far cheaper. The trade only pays off when data is large and the model is small, or when the data legally cannot move.
Non-IID is the real story. With IID data FedAvg lands within 0.6 points of centralized training (96.7% vs 97.3%) and passes 90% by round 7. Give every client just two digits and it drops to 86.1% and, on seed 1, never reaches 90% in 30 rounds. Local-only models collapse to 19.9% — they can only ever recognise the 2 digits they have seen, so 20% is their ceiling.
Skew and partial participation compound. Sampling 2 of 8 clients per round barely matters with IID data (95.8%) but at full skew it falls to 63.4%: each round the global model is pulled toward whichever four digits happened to be sampled.
More local work did not fix it. At full skew, 5 local epochs scored 85.1% against 86.5% for 1 epoch. Longer local training lets each client drift further toward its own two digits before averaging, which is exactly the problem FedProx and SCAFFOLD were designed for.
Limits
- It is a single-tab simulation: no real network, no stragglers, no dropped clients, no clocks.
- No secure aggregation and no differential privacy. The server sees every client's raw update, and updates can leak training data. "Only weights leave the device" is a weaker promise than it sounds.
- The model and dataset are tiny. 8×8 digits are easy, so the gaps here are smaller than on real non-IID data.
- Plain FedAvg only — no FedProx, SCAFFOLD, momentum or server-side optimiser, which are the usual fixes for skew.
- Three seeds is enough to see the effect, not to put tight error bars on it.