Federated learning

The objective#

Federated learning minimises a weighted sum of client objectives:

$$ \min_{w} \; F(w) = \sum_{k=1}^{K} \frac{n_k}{n} F_k(w), \qquad F_k(w) = \frac{1}{n_k}\sum_{i \in \mathcal{P}_k} \ell(w; x_i, y_i) $$

• $K$ — number of clients; $\mathcal{P}_k$ the index set held by client $k$; $n_k = |\mathcal{P}_k|$; $n = \sum_k n_k$.
• $F_k$ — the local empirical risk; $\ell$ the per-example loss.

When client data is IID, $\mathbb{E}[F_k(w)] = F(w)$ and local steps are unbiased estimates of global descent. Under non-IID data this fails, and the bias is the source of every difficulty below.

FedAvg and its convergence#

McMahan et al. (2017): each round, sample a fraction $C$ of clients, run $E$ local epochs of SGD with batch size $B$, and average the resulting weights with weights $n_k/n$.

Li et al. (2020), On the Convergence of FedAvg on Non-IID Data, give the first rigorous analysis for strongly convex, smooth $F_k$ with partial participation. The rate is $O(1/T)$, with an additive term proportional to the heterogeneity

$$ \Gamma = F^\star - \sum_k \tfrac{n_k}{n} F_k^\star $$

where $F^\star$ is the global optimum value and $F_k^\star$ the local optimum values. $\Gamma = 0$ under IID data. Crucially, the bound degrades as $E$ grows: more local work increases the drift term, which is the formal statement of the trade-off in the developer block.

Client drift and its fixes#

FedProx (Li et al., 2020) adds a proximal term $\frac{\mu}{2}|w - w^t|^2$ to each local objective, restraining clients from moving too far from the broadcast weights. Cheap, and effective mainly for mild heterogeneity.

SCAFFOLD (Karimireddy et al., 2020) maintains control variates $c_k$ per client and $c$ globally, correcting the local gradient to $g_k - c_k + c$. This provably removes the drift term and achieves convergence independent of heterogeneity, at the cost of doubling communication and requiring stateful clients — which rules it out for cross-device settings where a client is seen once.

FedNova (Wang et al., 2020) normalises updates by the number of local steps, correcting the objective inconsistency introduced when clients perform different amounts of local work.

The distinction that organises this literature is cross-silo (tens of persistent, reliable organisations, stateful methods viable) versus cross-device (millions of intermittent, stateless clients, only stateless methods viable). Kairouz et al. (2021), Advances and Open Problems in Federated Learning, is the standard survey.

Permutation invariance#

Naive weight averaging is questionable for deep networks because the loss surface has permutation symmetry: hidden units can be reordered without changing the function. Two clients converging to permuted versions of the same solution average to something worse than either.

FedAvg avoids the worst of this by re-broadcasting a common starting point each round, keeping clients in one basin. Probabilistic Federated Neural Matching (Yurochkin et al., 2019) matches neurons before averaging. Entezari et al. (2022) present evidence for the linear mode connectivity conjecture modulo permutation, which explains both why weight averaging works as well as it does and why it breaks when clients drift into different basins.

Privacy: what is and is not guaranteed#

Federated learning provides no formal privacy guarantee. The relevant attacks:

• Gradient inversion. Zhu et al. (2019) recover training images by optimising a synthetic input to match observed gradients. Geiping et al. (2020) extend this to trained networks and realistic batch sizes.
• Membership and property inference from updates. Melis et al. (2019) infer properties of a client's data from its contributions.
• Model poisoning. Bagdasaryan et al. (2020) show a small number of malicious clients can install a targeted backdoor while leaving main-task accuracy intact.

The composed defence stack:

1. Secure aggregation (Bonawitz et al., 2017) — cryptographic protocol so the server learns only $\sum_k \Delta_k$, tolerating dropouts. Removes the per-client view; does not bound what the sum reveals.
2. DP-FedAvg (McMahan et al., 2018) — clip each client update to norm $S$, add Gaussian noise to the aggregate. Gives user-level DP, a stronger and more meaningful unit than record-level for this setting.
3. Robust aggregation — Krum, coordinate-wise median, trimmed mean, against poisoning. Note that robust aggregation and secure aggregation conflict, since the former needs to inspect individual updates.

Systems constraints that shape the algorithms#

Bonawitz et al. (2019) document the production system: device eligibility requires idle, charging and unmetered network; round completion targets tolerate large dropout rates; and the same devices must not be repeatedly selected. These constraints, not the optimisation theory, dictate most design choices in cross-device deployments.

Papers#

• McMahan et al., Communication-Efficient Learning of Deep Networks from Decentralized Data, AISTATS 2017 — arxiv.org/abs/1602.05629
• Kairouz et al., Advances and Open Problems in Federated Learning, 2021 — arxiv.org/abs/1912.04977
• Li et al., On the Convergence of FedAvg on Non-IID Data, ICLR 2020 — arxiv.org/abs/1907.02189
• Karimireddy et al., SCAFFOLD, ICML 2020 — arxiv.org/abs/1910.06378
• Bonawitz et al., Practical Secure Aggregation, CCS 2017
• Geiping et al., Inverting Gradients, NeurIPS 2020 — arxiv.org/abs/2003.14053
• McMahan et al., Learning Differentially Private Recurrent Language Models, ICLR 2018 — arxiv.org/abs/1710.06963

What to learn next#

• Differential privacy — the guarantee federated learning lacks on its own.
• Adversarial attacks — what a malicious participant can do to a model.
• Edge AI — the on-device world these models run in.