Federated Learning: Privacy-Preserving AI

Suppose four Indian eye-care providers want to jointly train a diabetic retinopathy screening model without ever moving a patient's retinal scan off their own servers — a real constraint under India's DPDP Act, and the same constraint that pushed real multi-hospital imaging research toward federated learning in the first place (Sheller et al., 2020, Scientific Reports, showed federated brain-tumor segmentation across US and European hospitals reaching accuracy close to a model trained on pooled MRI data, without any institution sharing a scan). Picture a large teaching hospital, a mid-size hospital, a small hospital, and a rural primary health centre (PHC) network, each holding a different slice of the disease spectrum and a very different network connection. The teaching hospital sees severe, advanced-stage cases referred from elsewhere; the rural PHC sees mostly undiagnosed, often more advanced disease because patients have never been screened before; the two mid-size sites sit somewhere in between. None of their local datasets looks like the true population. And the rural PHC's link is slow enough that it can barely finish one pass over its local data before the next synchronization round starts, while the teaching hospital's GPU cluster can finish five.

This is federated learning's actual production problem, and it is sharper than "can an attacker infer something from a gradient." Two structural facts collide: client data is non-IID (not independent-and-identically-distributed — each hospital's local distribution differs from the global population), and clients differ wildly in how much local computation they can complete before a round closes. Both facts interact with the aggregation rule itself. This chapter derives that interaction from first principles, with a fully worked numeric example, and shows exactly why the standard aggregation formula — correctly weighted by data volume — still produces a biased global model when clients complete unequal amounts of local work.

FedSGD, FedAvg, and the communication budget

Start from the simplest federated scheme. A central server holds global weights w_t. It broadcasts w_t to a sampled set of clients. Each client k, holding n_k local examples, computes a single gradient step on its local loss L_k and returns the updated weight w_k. The server aggregates by a weighted average, weighting each client by its share of the total data:

w_(t+1) = Σ_k (n_k / n) · w_k,   where n = Σ_k n_k

This scheme — one local step per round — is called FedSGD. It is not a heuristic approximation of centralized training; regardless of whether client data is IID or non-IID, it is algebraically identical to one step of centralized mini-batch gradient descent on the pooled dataset. If the per-example loss is additive, the gradient of the pooled loss is exactly the n_k-weighted average of the per-client gradients, so a single weighted-average step reproduces the centralized update bit for bit. This equivalence is the entire justification for calling federated averaging "just as good" as centralized training — and it is also exactly where the equivalence stops.

McMahan, Moore, Ramage, Hampson, and Agüera y Arcas introduced FedAvg at AISTATS 2017 ("Communication-Efficient Learning of Deep Networks from Decentralized Data") precisely to break that one-step-per-round constraint, because for cross-device and cross-silo systems the communication round, not the computation, is the bottleneck: a rural PHC's uplink is scarce and expensive, so you want to extract as much useful gradient signal per round as possible before paying for another sync. FedAvg lets each client run E local epochs (over local mini-batches of size B) before reporting back, and the server still aggregates with the same n_k-weighted average. Raising E from 1 to, say, 5 can cut the number of communication rounds needed by roughly a factor of E for convex problems on IID data — a genuine, measured win in the original paper's experiments. The price is that the aggregation formula that was exact for E = 1 is no longer exact for E > 1, because each client's E-step local trajectory does not stay on the path a centralized optimizer would have taken. It walks toward that client's own local optimum instead. This is called client drift, and how much it costs you depends on two things simultaneously: how different the clients' local optima are (non-IID severity), and how much local work each client actually completes (which, in production, varies by device and network — the straggler problem).

Worked example: four hospitals, four different local optima, four different epoch counts

To see the interaction precisely, take a stripped-down but exact model. Give each client a scalar local loss that is a quadratic bowl centred on a client-specific target value, standing in for "the local data pulls the weight toward this client's population":

L_k(w) = 0.5 · (w − target_k)^2,   gradient = (w − target_k)

A local gradient-descent step with learning rate η is w ← w − η(w − target_k) = target_k + (1−η)(w − target_k). Unrolling this recursion for E steps starting from the broadcast weight w_t gives a closed form (provable by induction, since each step is an affine map with the same fixed point target_k):

w_k^(E) = target_k + (1 − η)^E · (w_t − target_k)

Set η = 0.5, so the shrinkage factor is a clean (0.5)^E, and start the round at the current global weight w_t = 0. Four clients, with data volume n_k, local optimum target_k, and a round budget of E = 5 local epochs — but two of them cannot finish the budget before the round deadline:

H1  City teaching hospital   n=800   target=+2.0   budget 5, completes 5/5 (on time)
H2  Mid-size hospital        n=500   target=+1.0   budget 5, completes 5/5 (on time)
H3  Small hospital           n=300   target=−1.0   budget 5, completes 2/5 (straggler)
H4  Rural PHC                n=400   target=−3.0   budget 5, completes 1/5 (straggler)

Run the closed form and the aggregation in code — every value below was computed, not asserted:

clients = {
    "H1": {"n": 800, "target": 2.0,  "E": 5},  # on time
    "H2": {"n": 500, "target": 1.0,  "E": 5},  # on time
    "H3": {"n": 300, "target": -1.0, "E": 2},  # straggler, 2/5 epochs
    "H4": {"n": 400, "target": -3.0, "E": 1},  # straggler, 1/5 epochs
}
eta = 0.5
w_t = 0.0
n_total = sum(c["n"] for c in clients.values())          # 2000

local_w = {}
for name, c in clients.items():
    shrink = (1 - eta) ** c["E"]                          # (0.5)^E
    local_w[name] = c["target"] + shrink * (w_t - c["target"])

w_weighted = sum(clients[k]["n"] * local_w[k] for k in clients) / n_total
w_naive    = sum(local_w.values()) / len(local_w)

print(local_w)
print("weighted FedAvg:", round(w_weighted, 4))
print("naive average  :", round(w_naive, 4))

Tracing it by hand confirms the code: H1's shrink factor is 0.5⁵ = 0.03125, so w_H1 = 2.0 + 0.03125·(0 − 2.0) = 1.9375. H2 works out to 0.96875 the same way. H3 only gets 2 epochs, so its shrink factor is 0.5² = 0.25, giving w_H3 = −1.0 + 0.25·(0 − (−1.0)) = −0.75. H4 gets a single epoch, shrink 0.5, giving w_H4 = −3.0 + 0.5·3.0 = −1.5. The printed dictionary is exactly {'H1': 1.9375, 'H2': 0.96875, 'H3': -0.75, 'H4': -1.5}.

The correctly n_k-weighted aggregate is (800·1.9375 + 500·0.96875 + 300·(−0.75) + 400·(−1.5)) / 2000 = 1209.375 / 2000 = 0.6047. A common production bug — averaging the four returned models with equal weight 1/4 each, ignoring how much data or how many patients backed each update — gives (1.9375 + 0.96875 − 0.75 − 1.5)/4 = 0.1641 instead. That single unweighted-mean bug moves the aggregate by 0.4406, a 73% relative swing, because it lets H3 (300 patients) and H4 (400 patients) outvote H1 (800 patients) exactly as much as H1 outvotes them — the textbook reason FedAvg's weighting by n_k exists at all, and the textbook reason to check, in any real federated pipeline, that the aggregation step is actually reading batch sizes off the wire rather than just counting participants.

Why even the correctly weighted formula still overshoots

Here is the sharper, less obvious result. If every client had finished its full E = 5 budget — no stragglers at all — the weighted FedAvg output would be 0.2906 (run the same code with H3 and H4 both at E = 5: their local weights become −0.96875 and −2.90625, and the n_k-weighted average comes out to 93/320 = 0.290625). Compare that to the true population-weighted optimum — the value a single model trained on all 2000 pooled records would converge toward, Σ n_k·target_k / n = (1600 + 500 − 300 − 1200)/2000 = 0.30. With equal epoch counts, weighted FedAvg lands within about 3.1% of the pooled optimum after just one round. That is the equivalence the n_k-weighting formula is designed to guarantee, and on equal footing it delivers.

But the moment H3 and H4 straggle — completing 2 and 1 epochs instead of 5 — the same n_k-weighted formula jumps to 0.6047, roughly double the true target and now on the opposite side of it from where equal-epoch FedAvg landed. The weighting by n_k never changed; what changed is that a straggling client's returned weight has barely moved from the broadcast starting point w_t, because it only got through a fraction of its assigned local descent. Decompose the aggregate algebraically: since w_t = 0, every client's contribution to w_t+1 is (n_k/n)·(1 − 0.5^E_k)·target_k. The factor multiplying each client's data share is not just n_k/n, it is n_k/n times (1 − 0.5^E_k) — call this the client's effective weight. For H1 and H2 (E = 5), effective weight equals fair data share almost exactly (0.3875 against a fair 0.40; 0.2422 against a fair 0.25) because 1 − 0.5⁵ = 0.96875 is nearly 1. For H4 (E = 1), effective weight collapses to 0.10 against a fair share of 0.20 — the straggler's real influence on the round is cut in half, not by the aggregation rule discriminating against it, but because its own unfinished local descent never got the chance to move toward its data. H3 loses even more of its influence in relative terms (0.1125 effective against 0.15 fair). The well-resourced, well-connected clients end up dominating the round not because the formula favours them by design, but because only they finished the work the formula assumes everyone finished.

This is precisely the failure mode Wang, Liu, Liang, Joshi, and Poor formalize as objective inconsistency in "Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization" (NeurIPS 2020) — the paper that introduces FedNova, whose step-count-normalized averaging directly targets this exact arithmetic gap: clients that complete different numbers of local steps get their updates rescaled before aggregation, rather than mixed in raw. A related, widely cited fix is Li, Sahu, Zaheer, Sanjabi, Talwalkar, and Smith's FedProx ("Federated Optimization in Heterogeneous Networks," MLSys 2020). Its fix adds a proximal penalty μ/2·‖w − w_t‖² to every client's local loss, which caps how far any client — straggler or not — can drift from the broadcast point regardless of how many local steps it manages to complete, making E_k variability far less consequential to the final aggregate. Together, the two papers give a direct, citable production answer to the exact arithmetic gap this worked example exposes.

Client availability is not a random sample

Bonawitz, Eichner, and colleagues describe the actual production system Google runs for on-device federated learning (Bonawitz et al., "Towards Federated Learning at Scale: System Design," MLSys 2019). Devices are only eligible to participate in a round when they are charging, idle, and on an unmetered network — a sensible policy for not draining a user's battery or data plan, but one with a systematic side effect: the devices that are available to train are not a representative sample of the fleet. Applied to the hospital scenario, the rural PHC is in double jeopardy. Its local case-mix is already the furthest from the population average (non-IID severity), and its poor connectivity means it is also the client most likely to straggle or be dropped from a round entirely (availability bias) — the two effects compound rather than cancel, and Kairouz, McMahan, and dozens of co-authors survey exactly this compounding as one of the open problems in "Advances and Open Problems in Federated Learning" (2021). The practical mitigations in production systems follow directly from the arithmetic above: bound client drift with a proximal term (FedProx), oversample more clients per round than you need so dropouts do not silently shrink your effective sample, and — critically — never let "finished its round" silently stand in for "well represented," because the worked example shows a fully honest, correctly n_k-weighted aggregation can still be off by 2× the true target purely from unequal completed local work.

Common misconception

Students frequently reason: "federated learning keeps raw data on-device, so besides the extra network rounds, the resulting model is mathematically the same as training centrally — privacy is free." This is true only in the narrow FedSGD case (E = 1), regardless of whether client data is IID or non-IID, where the weighted-average step is provably identical to a centralized gradient step. The instant you introduce multiple local epochs — which is the entire point of FedAvg, since it is what makes federated learning communication-efficient enough to deploy at all — and the instant clients are non-IID or complete unequal amounts of local work, the aggregate model is a biased estimate of the centralized solution, not an equivalent one. The worked example above is not an edge case; unequal client compute and non-IID data are the default condition of every real cross-silo or cross-device deployment. Privacy-preservation and statistical fidelity to a hypothetical pooled model are separate properties, and FedAvg buys the first without a free guarantee of the second.

The aggregation mechanism

Read the two bars per client: the blue bar is what each hospital's data volume alone entitles it to (n_k/n); the orange bar is what it actually got, once incomplete local epochs are folded in. H1 and H2 barely lose anything — their orange bars sit almost under their blue ones. H4's orange bar is roughly half its blue bar's height: a rural PHC with a fifth of the total patient volume ends up shaping the global model about as much as a client with a tenth. That gap, not any attacker, is the everyday production failure mode of federated averaging.

Active recall

Attempt each question before reading its answer.

1. Why does FedSGD (E = 1), even on non-IID client data, produce an aggregate that is provably identical to one step of centralized mini-batch gradient descent, while FedAvg with E > 1 does not?

2. Using the four-hospital setup, suppose no client straggles at all — every hospital completes its full budget of E = 5 local epochs. Compute the new weighted FedAvg global weight, and compare it to the true population-weighted optimum of 0.30.

3. (Ripple-effect.) H3's rural link actually holds this round, and it completes all 5 assigned epochs instead of only 2 (H4 still only manages 1/5). Recompute: (a) H3's local w_k, (b) the new weighted FedAvg w_t+1, and (c) the effective-weight table for all four clients. Is H4 still the most under-represented client?

4. Why does even the correctly n_k-weighted FedAvg formula from McMahan et al. (2017) still overshoot the true target when clients complete unequal numbers of local epochs? What does FedProx (Li et al., 2020) change to address this specifically?

5. A junior engineer implements aggregation as a plain arithmetic mean of the four returned client weights, without reading batch sizes off the wire. Using the original straggler scenario, what direction and magnitude of bias does this introduce relative to the correct weighted result?

6. Google's production cross-device system (Bonawitz et al., 2019) only admits a device to a round when it is charging, idle, and on unmetered WiFi. Explain one concrete failure this creates for the rural-PHC scenario, and propose one mitigation.

Answers

1. With an additive loss L(w) = Σ_k (n_k/n)·L_k(w), the gradient of the pooled loss is exactly the n_k-weighted average of the per-client gradients: ∇L(w) = Σ_k (n_k/n)·∇L_k(w). A single local gradient step at each client, followed by the n_k-weighted average of the resulting weights, is algebraically the same operation as taking the weighted-average gradient first and then stepping once — because a single step is a linear (affine) function of the gradient, and weighted averaging commutes with an affine map. Once E > 1, each client's local trajectory is a nonlinear function of its own gradient sequence (each step depends on where the previous step landed), and weighted-averaging nonlinear trajectories is no longer the same as weighted-averaging one global trajectory — this is exactly where client drift enters.

2. With every client at E = 5, the shrink factor 0.5⁵ = 0.03125 applies to all four: local weights become 1.9375, 0.96875, −0.96875, and −2.90625 for H1–H4 (H3 and H4 now shrink almost all the way toward their own targets, same as H1 and H2 do). The n_k-weighted aggregate is (800·1.9375 + 500·0.96875 + 300·(−0.96875) + 400·(−2.90625))/2000 = 93/320 = 0.2906, within about 3.1% of the true 0.30 population optimum. This confirms the earlier claim directly: when local computation is equal across clients, the n_k-weighting formula alone is enough to track the centralized solution closely; the large error in the original scenario (0.6047) comes entirely from the unequal epoch counts, not from non-IID data by itself.

3. (a) H3 now uses shrink factor 0.5⁵ = 0.03125: w_H3 = −1.0 + 0.03125·(0 − (−1.0)) = −0.96875 (H4 is unchanged at −1.5, since it still only completes 1 epoch). (b) The new weighted aggregate is (800·1.9375 + 500·0.96875 + 300·(−0.96875) + 400·(−1.5))/2000 = 183/320 = 0.5719 — closer to the E=5-for-everyone answer of 0.2906 than the original 0.6047 was, but still far from it, because H4 alone is still stragglers. (c) Effective weights recompute to H1 = 0.3875 (unchanged), H2 = 0.2422 (unchanged), H3 = 0.15·(1 − 0.03125) = 0.1453 (up from 0.1125, since H3 is no longer straggling), H4 = 0.10 (unchanged, still the worst-off). Yes — H4 remains the single most under-represented client, at exactly half its 0.20 fair share, because it is the only client whose epoch count did not change between the two scenarios; fixing one straggler does not fix the others.

4. The n_k-weighting formula assumes every client's local update reflects a comparable amount of "descent progress" toward its own target — it weights by data volume, not by how far each client's local trajectory actually moved. A straggler that completes only 1 of 5 assigned epochs returns a weight that has barely left the broadcast point w_t, so its true influence on the aggregate collapses (effective weight = fair share × (1 − 0.5^E_k)) even though its n_k weighting stayed the same on paper. FedProx (Li et al., 2020, MLSys) adds a proximal term μ/2·‖w − w_t‖² to every client's local objective, which explicitly bounds how far any client — fast or slow — is allowed to drift from w_t in a single round. This shrinks the gap between a 5-epoch client and a 1-epoch client's local movement, making the n_k-weighted aggregate far less sensitive to exactly how many local steps each client happened to finish.

5. The unweighted mean gives (1.9375 + 0.96875 − 0.75 − 1.5)/4 = 0.1641, versus the correct weighted result of 0.6047 — a difference of 0.4406, i.e. the naive result is about 73% lower than the correct one. The direction of the bias comes from H1 (800 patients, largest positive pull toward +2.0) being outvoted 3-to-1 by H2, H3, and H4 combined under equal weighting, when in the true data H1 alone represents 40% of all patients — unweighted averaging always erases exactly this kind of volume advantage, regardless of which direction it happens to push the number in any one round.

6. Concrete failure: a rural PHC's devices or edge servers may rarely or never satisfy "unmetered WiFi and charging and idle" simultaneously — mobile-first or intermittently-powered sites can be excluded from rounds not because their data is uninformative, but because their infrastructure never meets the eligibility gate. Over many rounds this is a systematic, not random, exclusion — exactly the availability bias that compounds with non-IID severity, since the excluded sites are also the ones furthest from the population average. One mitigation: relax or widen the eligibility window for chronically under-represented client segments (e.g., allow metered-but-off-peak connections for PHC-tier clients, or schedule dedicated rounds that specifically wait for rural clients rather than always taking whichever clients respond fastest), combined with oversampling more clients per round than needed so that dropouts from any one segment do not zero out its influence entirely.

Think About It

Think about this: How would you explain federated learning: privacy-preserving ai to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where federated learning: privacy-preserving ai is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting federated learning: privacy-preserving ai to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind federated learning: privacy-preserving ai, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.