Interactive Explainer

Federated Learning, Round by Round

Six clients with private data train a shared classifier without ever sending the data to a server. Watch FedAvg succeed on IID clients and collapse on non-IID ones; flip to FedProx and watch it recover. Then the privacy + security stack (secure aggregation, differential privacy).

~20 min Distributed ML · Privacy · FedAvg

Step 1

The federated setting

Data lives on $K$ clients (phones, hospitals, edge devices, sensor networks). Each client has a private dataset $\mathcal{D}_k$ it will not share. A central server coordinates training of a shared model $\theta$. Each round:

Server sends $\theta$ to a random sample of clients.
Each client runs $E$ local epochs of SGD on its data.
Clients send back updated $\theta_k$.
Server averages: $\theta \leftarrow \sum_k \frac{n_k}{n}\,\theta_k$ (FedAvg, McMahan et al., 2017).

FedAvg is one of the most cited DL papers of the decade and one of the brittlest. It works beautifully when client datasets are IID (random partition of a global dataset); it collapses when they aren't (which is most of the time in production).

Step 2

FedAvg vs FedProx — see them break and heal

Six clients, three-class 2-D classification. The non-IID slider controls how skewed each client's label distribution is. At 0, every client sees roughly uniform classes; at 1, each client sees only one or two classes. FedProx (Li et al., 2020) adds a proximal term $\frac{\mu}{2}\|\theta_k - \theta\|^2$ to each client's local objective, gently pulling local updates back toward the global model.

algorithm non-IID α 0.50 local epochs E 3 round 0 global acc —

Global classifier on test data

Per-client class distribution

Accuracy vs rounds

What to watch. With non-IID α near 0, FedAvg converges within ~10 rounds. Push α to 1 and FedAvg oscillates or diverges — each client's local SGD overfits to its own classes and the average ends up worse than any client's local model. Switch to FedProx; the proximal term clips drift and accuracy climbs again.

Step 3

The privacy + security stack

FedAvg by itself isn't private. Sharing gradients leaks information; reconstruction attacks (Zhu et al., 2019) can recover original images from a single client update. Three layers fix this:

Secure aggregation (Bonawitz et al., 2017). Cryptographic protocol so the server only sees the sum $\sum_k \theta_k$, not individual updates. Information-theoretically secure under a t-honest-but-curious threat model.
Differential privacy (DP-FedAvg, Abadi et al., 2016). Each client clips its update to a bounded norm and adds Gaussian noise; the global model satisfies $(\varepsilon, \delta)$-DP. Trade accuracy against privacy budget.
Homomorphic encryption / TEE. The server computes the average on encrypted values; or runs in a Trusted Execution Environment (Intel SGX, AMD SEV). Heavier, used in regulated industries.

Practical setups in 2026: FedML, Flower, NVFlare, OpenFL all ship FedAvg/FedProx/SCAFFOLD, secure aggregation, and DP-SGD as drop-in components.

Step 4

Algorithm zoo — what to use when

Algorithm	Key idea	Best when
FedAvg	Plain weight average	IID-ish clients, baseline
FedProx	+ proximal term on client objective	Non-IID, system heterogeneity
SCAFFOLD	Control variates to cancel client drift	Strongly non-IID, large E
FedAvgM / FedAdam	Server-side momentum / Adam	Slow / unstable convergence
FedDF	Distillation on a server-side public set	Heterogeneous client architectures
FedBN	Keep BatchNorm stats local, average everything else	Image clients with different statistics
Personalised FL (pFedMe, Ditto)	Per-client head, global trunk	Per-client distribution drift; medical, mobile keyboards
Federated MoE	Mixture-of-experts where each expert is local	Strongly heterogeneous clients; 2024+

Step 5

Practical pitfalls

Stragglers. One slow client blocks every round. Standard fix: drop the slowest 20% per round; use asynchronous variants like FedBuff.
Heterogeneous compute. A phone client can't run 10 epochs in the time the server waits. Use FedProx (insensitive to E) or per-client local-step budgets.
Catastrophic forgetting. If a client joins late with a new class, the global model can lose old classes. FedDF / replay buffers / continual-FL papers address this.
Reconstruction attacks. Gradients reveal data. Always combine FedAvg with secure aggregation and DP for any privacy-sensitive deployment.
Free-rider / Byzantine clients. A malicious client can poison the model. Robust aggregation (Krum, median, trimmed mean) defends against a bounded fraction; cryptographic verifiable aggregation goes further.
Evaluation. Global accuracy hides per-client disparity. Always report worst-client and 10th-percentile accuracy too.

The slogan. Federated learning works when clients differ a little; collapses when they differ a lot; needs FedProx / SCAFFOLD / personalised variants in the middle; and always needs privacy primitives if the data is private. The algorithmic part is small; the systems and threat-model parts dominate.