Federated learning

What Problem Does This Solve?

Classic centralized ML workflow:

Collect user data from millions of devices → central server
Train model on aggregated dataset
Deploy model back to devices

Problems:

Privacy: Sensitive data (health records, messages, photos) leaves user control
Bandwidth: Terabytes of raw data transmission
Regulatory: GDPR, HIPAA prohibit moving certain data
Latency: Round-trip to cloud for every training batch

Federated learning keeps data decentralized while still learning a global model.

How Federated Learning Works (Step-by-Step Derivation)

The Setup

We have:

$N$ clients (devices, hospitals, organizations), each with local dataset $\mathcal{D}_k$ of size $n_k$
Total data points: $n = \sum_{k=1}^N n_k$
Goal: Learn a shared model $\theta$ that minimizes global loss:

F(\theta) = \sum_{k=1}^N \frac{n_k}{n} F_k(\theta)

where $F_k(\theta) = \frac{1}{n_k}\sum_{i \in \mathcal{D}_k} \ell(x_i, y_i; \theta)$ is client $k$ 's local loss.

Why this weighting? Each client contributes proportionally to its data size. A client with 1000 samples should influence the model 10× more than one with 100 samples.

Algorithm: Federated Averaging (FedAvg)

The breakthrough McMahan et al. (2017) algorithm:

Server (orchestrator) steps:

Initialize global model $\theta^{(0)}$
For each round $t = 0, 1, 2, \ldots$ $t = 0, 1, 2, \dots$ :
- Sample subset $\mathcal{S}_t$ of $C \cdot N$ clients (typically $C=0.1$ , i.e., 10%)
- Broadcast current $\theta^{(t)}$ to selected clients
- Wait for client updates
- Aggregate: $\theta^{(t+1)} = \sum_{k \in \mathcal{S}_t} \frac{n_k}{n_{\mathcal{S}_t}} \theta_k^{(t+1)}$ where $n_{\mathcal{S}_t} = \sum_{k \in \mathcal{S}_t} n_k$

Client $k$ steps (when selected):

Receive $\theta^{(t)}$ from server
Local training: Run $E$ epochs of SGD on local data $\mathcal{D}_k$ : $\theta_k^{(t+1)} = \theta^{(t)} - \eta \sum_{e=1}^E \nabla F_k(\theta)$
Send back $\theta_k^{(t+1)}$ (only the updated weights, not data!)

Why this works—the math:

If each client does one SGD step ( $E=1$ , batch size $n_k$ ): $\theta_k^{(t+1)} = \theta^{(t)} - \eta \nabla F_k(\theta^{(t)})$

Server aggregation: $\theta^{(t+1)} = \sum_k \frac{n_k}{n} \left(\theta^{(t)} - \eta \nabla F_k(\theta^{(t)})\right) = \theta^{(t)} - \eta \sum_k \frac{n_k}{n} \nabla F_k(\theta^{(t)})$

This is exactly one step of gradient descent on $F(\theta)$ ! Multiple local epochs ( $E>1$ ) approximate this more efficiently.

Key Concepts Deep Dive

1. Communication Efficiency

The bottleneck: Sending model updates over networks is expensive (mobile data, slow hospital networks).

How many communication rounds vs. accuracy?

Derived from convergence analysis (simplified):

For strongly convex $F$ , after $T$ rounds: $F(\theta^{(T)}) - F(\theta^*) \leq \mathcal{O}\left(\frac{1}{T}\right) + \mathcal{O}(E^2\eta^2\sigma^2)$

First term: Faster convergence with more rounds $T$
Second term: Error from local drift (clients diverge when doing $E$ local epochs)

Trade-off:

Large $E$ → fewer rounds (less communication) but higher drift error
Small $E$ → more rounds (more communication) but lower drift

Practical choice: $E=5$ to $E=20$ epochs per round works well.

2. Non-ID Data Distribution

Real-world data is heterogeneous across clients:

Label skew: User A only types in English, User B in Spanish
Feature skew: Hospital A has older patients, Hospital B has younger
Quantity skew: Some users generate 10× more data

Mathematical impact:

Define divergence between local and global optima: $\delta_k = \|\theta_k^* - \theta^*\|^2$

where $\theta_k^* = \arg\min F_k(\theta)$ (client $k$ 's optimal) and $\theta^* = \arg\min F(\theta)$ (global optimal).

High $\delta_k$ → clients pull in conflicting directions → slower convergence.

Mitigation strategies:

More local epochs helps each client find better local minimum before aggregating
Client sampling: Sample diverse clients each round
Personalization: Allow clients to keep small personalized layers

Common mistake "Just Average the Models, It'll Work Fine"

Why this feels right: Averaging is simple and seems fair—everyone contributes equally.

Why it fails: Client with 10 samples shouldn't have equal weight to one with 10,000 samples. Unweighted averaging gives noisy clients too much influence, causing divergence.

The fix: Always weight by dataset size $n_k$ . This is the $\frac{n_k}{n}$ term in $F(\theta)$ —it's fundamental, not optional.

Proof of failure: Suppose 99 clients have near-identical data, 1 outlier client has random noise. Unweighted average: $\frac{1}{100}\sum \theta_k$ puts1% weight on noise. Weighted average: outlier gets weight $\frac{n_{\text{outlier}}}{n_{\text{total}}} \ll 1\%$ if it has little data.

3. Privacy Guarantees

Claim: "Federated learning is private because raw data never leaves devices."

Reality: More nuanced.

Attack vectors:

Gradient inversion: Reconstruct training samples from gradients
- Given $\nabla_\theta \ell(x, y; \theta)$ , solve for $x$
- Possible for small batches, high-res images
Model memorization: Neural nets memorize training examples
- Query model with prompts → it regurgitates training data

Defense: Differential Privacy

Add calibrated noise to updates before sending:

$\tilde{\theta}_k^{(t+1)} = \theta_k^{(t+1)} + \mathcal{N}(0, \sigma^2 S^2 I)$

where $S$ is clipping threshold (bound gradient norm).

Formula

$(\epsilon, \delta)$ -Differential Privacy Guarantee A mechanism $\mathcal{M}$ is $(\epsilon, \delta)$ -DP if for all neighboring datasets $\mathcal{D}, \mathcal{D}'$ (differing in one record) and all outputs $O$ : $\Pr[\mathcal{M}(\mathcal{D}) \in O] \leq e^\epsilon \Pr[\mathcal{M}(\mathcal{D}') \in O] + \delta$ Intuition: Adding/removing one person's data changes output probability by at most $e^\epsilon$ (multiplicative) plus $\delta$ (additive slack). Smaller $\epsilon$ → more privacy.

In federated learning: Typical values $\epsilon \in [1, 10]$ , $\delta = 10^{-5}$ . This provably bounds information leakage about any individual's data.

Trade-off: More noise ( $\sigma \uparrow$ ) → stronger privacy → lower accuracy.

Worked example Differential Privacy in Practice (Google Gboard)

Task: Predict next word on Android keyboards

Setup:

1 billion devices
Sensitive data: private messages
Target: $\epsilon = 2.7$ , $\delta = 10^{-9}$ (strong privacy)

Implementation:

Clip each client's gradient: $\|\nabla F_k\| \leq S = 1.0$
Add noise: $\sigma = 0.05\cdot S$ Gaussian to each weight
Aggregate: Server gets $(n_k, \text{noisy } \theta_k)$ from each client Why the clipping step? Unbounded gradients would require infinite noise forDP. Clipping at $S=1$ means we can use finite $\sigma$ while guaranteeing privacy.

Result: 94% accuracy (vs. 97% without DP). 3% accuracy cost for provable privacy.

Advanced Variants

Secure Aggregation

Problem: Even model updates can leak info if server is malicious.

Solution: Cryptographic protocol where server learns $\sum_k \theta_k$ but not individual $\theta_k$ .

How (simplified):

Each client $k$ generates random mask $r_k$ shared with other clients (via pairwise keys)
Client $k$ sends $\theta_k + r_k$ to server
Server sums: $\sum_k (\theta_k + r_k) = \sum_k \theta_k + \underbrace{\sum_k r_k}_{=0 \text{ if designed right}}$
Masks cancel out, server gets true aggregate without seeing individuals

Cost: 2-3× more communication, cryptographic overhead.

Cross-Silo vs. Cross-Device

	Cross-Device	Cross-Silo
Example	Smartphones, IoT	Hospitals, banks
# Clients	Millions	10-100
Availability	Intermittent (only when charging)	Stable
Communication	Slow (mobile networks)	Fast (datacenter links)
Data per client	Small (KBs-MBs)	Large (GBs-TBs)
Strategy	Many rounds, small $E$	Few rounds, large $E$

Why this distinction matters: Algorithm hyperparameters ( $C$ , $E$ , learning rate) must adapt to the setting.

Worked Example: Binary Classification

Setup: 3 hospitals want to predict disease from patient vitals without sharing records.

Data:

Hospital 1: 1000 patients, 60% positive
Hospital 2: 500 patients, 40% positive
Hospital 3: 2000 patients, 70% positive
Feature: Blood pressure (1D for simplicity)
Model: Logistic regression $p(y=1|x) = \sigma(\theta^T x)$

Loss: $F_k(\theta) = \frac{1}{n_k}\sum_{i=1}^{n_k} -[y_i \log \sigma(\theta^T x_i) + (1-y_i)\log(1-\sigma(\theta^T x_i))]$

Algorithm execution:

Round 0:

Server: Initialize $\theta^{(0)} = 0$ , sample all3 hospitals ( $C=1$ )
Hospital 1: Receives $\theta^{(0)}=0$
- Runs 5 epochs SGD on 1000 samples, $\eta = 0.01$
- Computes gradients: $\nabla F_1(\theta) = \frac{1}{1000}\sum_{i} (p_i - y_i)x_i$ where $p_i = \sigma(\theta^T x_i)$
- Updates: $\theta_1^{(1)} = \theta^{(0)} - \eta \nabla F_1 = 0.42$ (suppose)
- Why this step? Standard gradient descent—move weights to reduce loss on local data.
- Sends back $\theta_1^{(1)} = 0.42$
Hospital 2: Similarly gets $\theta_2^{(1)} = 0.38$
Hospital 3: Similarly gets $\theta_3^{(1)} = 0.51$
Server aggregates: $\theta^{(1)} = \frac{1000}{3500} \cdot 0.42 + \frac{500}{3500} \cdot 0.38 + \frac{2000}{3500} \cdot 0.51$ $= 0.286 \cdot 0.42 + 0.143 \cdot 0.38 + 0.571 \cdot 0.51= 0.465$

Why these weights? Total $n = 1000+500+2000=3500$ . Hospital 3 has most data (57%), so it pulls $\theta^{(1)}$ closest to its $\theta_3^{(1)}=0.51$ .

Round 1: Repeat from $\theta^{(1)} = 0.465$ ...

After 100 rounds: $\theta^{(100)} \approx 0.52$ (converged global model)

Privacy preserved: Hospitals never shared patient records, only learned model weights.

Common Pitfalls & How to Fix Them

Common mistake "More Local Epochs Always Better"

Intuition that seems right: More local training → better local models → better global model. Why it fails: Large $E$ causes client drift. Each client overfits to its local data, pulling the global model in conflicting directions.

Math: Client $k$ after $E$ epochs: $\theta_k \approx \arg\min F_k(\theta)$ . If data is non-IID, $\arg\min F_k \neq \arg\min F_j$ . Averaging contradictory models degrades performance.

The fix:

Moderate $E$ : Use $E \in [5, 20]$ , not100.
Adaptive optimization: Use FedAdam, FedYogi (momentum-based aggregation) to dampen oscillations.
Client sampling: Diverse clients each round averages out drift.

Connections to Other Concepts

Distributed Training: Federated learning is distributed training with privacy constraints and non-IID data
Differential Privacy: Formal privacy guarantee for federated updates
Optimization Theory: FedAvg is SGD variant; convergence depends on data heterogeneity
Edge Computing: Federated learning pushes computation to edge devices
Multi-Task Learning: Personalized federated learning treats each client as a task
Byzantine Robust Learning: Defense against malicious clients sending bad updates
Model Compression: Quantization + sparsity reduce communication cost per round
Privacy-Preserving ML: Federated learning is one pilar; others include homorphic encryption, secure multi-party computation

Recall Explain Federated Learning to a 12-Year-Old

Imagine you and99 friends all want to build the world's best autocorrect for texting. But here's the problem: everyone's texts are super private—you can't share them!

Old way (doesn't work): Everyone sends all their texts to one person. That person reads EVERYONE's messages (crepy!) and figures out common patterns. Then they tell everyone the autocorrect rules.

Federated learning way:

One person (the "server") writes down some basic autocorrect rules on paper.
They make100 copies and mail one to each friend.
Each friend practices on their OWN texts (privately!) and improves their copy of the rules.
Friends mail back ONLY their improved rules (not their texts!).
Server combines everyone's improvements: "80 people found that'u' usually means 'you', so let's add that."
Server sends out the new better rules. Repeat!

After 50 rounds, everyone has an amazing autocorrect that learned from 100 people's texting styles—but no one ever saw anyone else's private messages. That's federated learning: learn together, stay private.

Active Recall Questions

#flashcards/ai-ml

What is the core difference between federated learning and centralized ML? :: Federated learning brings the MODEL to the DATA (training happens on local devices, only weight updates are sent). Centralized ML brings DATA to the MODEL (raw data uploaded to central server for training).

Write the federated averaging update rule :: $\theta^{(t+1)} = \sum_{k \in \mathcal{S}_t} \frac{n_k}{\sum_{j \in \mathcal{S}_t} n_j} \theta_k^{(t+1)}$ where $\theta_k^{(t+1)}$ is client $k$ 's local model after $E$ epochs of SGD, weighted by dataset size $n_k$ .

Why do we weight clients by their dataset size

n_k

in aggregation?

To make each data point contribute equally to the global model. A client with 10× more data should influence the model 10× more. This matches the mathematical objective

F(\theta) = \sum \frac{n_k}{n} F_k(\theta)

What is the communication cost trade-off controlled by local epochs

E

Large

E

→ fewer rounds (less communication) but more client drift (local models diverge). Small

E

→ more rounds (more communication) but less drift. Optimal

E \in [5, 20]

for most tasks.

What is

(\epsilon, \delta)

-differential privacy in one sentence?

Adding or removing any one person's data changes the algorithm's output probability by at most

e^\epsilon

(multiplicatively) plus

\delta

, provably bounding information leakage.

What attack does differential privacy defend against in federated learning?

Gradient inversion attacks, where an adversary reconstructs training samples from the gradients/weight updates. DP adds calibrated noise to updates, making reconstruction impossible.

Why is non-IID data a problem for federated learning?

Clients optimize different objectives (

\theta_k^* \neq \theta_j^*

when data distributions differ). Averaging contradictory local models causes oscillation and slower convergence to the global optimum

\theta^*

What is secure aggregation?

Cryptographic protocol where the server learns the SUM of client updates

\sum_k \theta_k

but cannot see individual

\theta_k

. Protects against malicious server trying to infer info from individual updates.

Cross-device vs. cross-silo: which has more clients and less data per client?

Cross-device (millions of smartphones, KBs-MBs each). Cross-silo has 10-100 organizations with GBs-TBs each.

What is client drift and when does it happen?

When clients overfit to local data during many local epochs, their models diverge from the global optimum. Happens with large

E

and highly non-IID data. Mitigated by smaller

E

or adaptive optimizers like FedAdam.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Federated learning ka matlab simple hai: socho tumhare pas ek billion smartphones hain aur sabko ek AI model train karna hai jo next word predict kare (jaise WhatsApp mein type karte waqt suggestions ate hain). Traditional methodein kya hota? Sab logon ke messages ek central server pe upload ho jayenge, wahan model train ho, phir wapas phone pe aye. Par yahan do badi problems hain: pehla, privacy ka issue—tumhare private chats kaun padhe server pe? Dosra, bandwidth—1 billion phones se data lana toh nightmare hai.

Federated learning ka jugaad: model ko data ke pas le jao, data ko model ke paas nahi. Matlab serverek initial model bhejta hai har phone pe, har phone apne local data pe privately train karta hai (tumhare messages kabhi bahar nahi jate), phir sirf learned weights (jo model ne seekha) wapas bhejta hai server ko. Server sabke weights ko combine karke ek global model banata hai. Is tareke se, data kabhi move nahi hota, sirf knowledge move hoti hai—privacy bani rahi, bandwidth bhi bach gayi.

Core formula: