Federated Learning Explained: FedAvg, Privacy and a Python Demo

Imagine a restaurant chain with five branches across Indore: Vijay Nagar, Palasia, Rajwada, Bhawarkua and Rau. Head office wants a better masala recipe. Each branch knows what its own customers like, but its customer feedback and order history must stay in that branch.

So head office sends the same base recipe to every branch. Each kitchen adjusts it based on a week of its own feedback and sends back only the adjusted recipe, never the customer records. Head office blends the five versions, giving more weight to busier branches, and sends the blend out again.

After a few rounds, the chain has one recipe shaped by all its customers, and no customer list ever left a branch. That is federated learning.

What is federated learning?

Map the restaurant chain to machine learning:

• Head office is the server that coordinates training.
• Each branch is a client: a phone, a hospital, a bank or a factory.
• The base recipe is the global model, and each branch's adjustment is a local update.
• Blending the recipes is aggregation, and one send-train-return cycle is a round.

Federated learning is a decentralised training approach in which many clients train a shared model on their own local data and send only model updates to a central server, which combines them into an improved global model, so the raw data never leaves its owner.

It follows transfer learning in this series. Transfer learning moves knowledge from one task to another; federated learning gathers knowledge from many data owners into one model.

Why not just collect all the data?

Traditional machine learning copies every row to one server and trains there. That often fails in practice:

• Privacy and law: health, message and financial data are sensitive, and laws such as India's Digital Personal Data Protection Act, 2023 restrict how personal data is collected and moved.
• Competition: banks or hospitals may want a shared model but will not hand each other raw data.
• Volume and cost: millions of phones produce more data than is practical to upload.

In centralised training the data moves to the model. In federated learning the model moves to the data, and only updates come back.

How does federated learning work?

The most widely used algorithm is Federated Averaging (FedAvg), introduced by researchers at Google in 2017.

One FedAvg round: broadcast the global model, train locally, return the weights, and average them in proportion to each client's data size.

Each round follows the same steps:

1. The server picks a set of available clients.
2. It sends them the current global model weights.
3. Each client trains the model for a few epochs on its own data.
4. Each client sends back its new weights (or the change in weights).
5. The server averages them, weighting each client by its number of training examples.
6. The average becomes the new global model, and the next round begins.

The simulation below has five hospitals in one city with 20 to 90 diabetes patients each, too few for a good model alone. It compares each hospital training alone, FedAvg across all five, and the central model you would get if pooling data were allowed. The model is plain logistic regression trained with gradient descent.

import numpy as np

rng = np.random.default_rng(0)
true_w = rng.normal(0, 1, 12)                    # how 12 test results relate to risk

def hospital(n):
    X = rng.normal(0, 1, (n, 12))
    y = (rng.random(n) < 1 / (1 + np.exp(-(X @ true_w)))).astype(float)
    return X, y
sizes = [20, 30, 45, 60, 90]                     # patients per hospital
clients = [hospital(n) for n in sizes]
X_test, y_test = hospital(5000)

def train(w, X, y, epochs, lr=0.5):              # logistic regression by gradient descent
    for _ in range(epochs):
        w = w - lr * X.T @ (1 / (1 + np.exp(-(X @ w))) - y) / len(y)
    return w
acc = lambda w: np.mean(((X_test @ w) > 0) == y_test)

w_global = np.zeros(12)
for rnd in range(30):                            # 30 communication rounds
    local = [train(w_global.copy(), X, y, epochs=10) for X, y in clients]
    w_global = np.average(local, axis=0, weights=sizes)   # FedAvg: weight by data size

alone = [acc(train(np.zeros(12), X, y, epochs=300)) for X, y in clients]
X_all = np.vstack([c[0] for c in clients]); y_all = np.concatenate([c[1] for c in clients])
print("each hospital alone:", np.round(alone, 3))
print(f"federated (FedAvg): {acc(w_global):.3f}")
print(f"central (pooled data, not allowed here): {acc(train(np.zeros(12), X_all, y_all, 300)):.3f}")

Alone, the hospitals reach 0.70 to 0.79 accuracy. FedAvg reaches 0.81, matching the central model, and no patient row left its hospital. The smallest hospitals gain the most.

Types of federated learning

Cross-device training is closely tied to edge AI, since the training happens on the device.

Challenges in federated learning

Non-IID data and personalisation

In a central dataset you can shuffle rows so every batch looks like the whole. In federated learning each client's data reflects its own users. A rural clinic sees different patients from a city hospital, and one person's phone reflects one person's typing. This is called non-IID data (not independent and identically distributed).

One global model then fits the average client and may serve unusual clients badly. A common fix is personalisation: after federated training, each client fine-tunes the global model briefly on its own data.

import numpy as np

rng = np.random.default_rng(3)
w_city = rng.normal(0, 1, 6)
w_rural = w_city.copy(); w_rural[0] *= -1    # one test behaves differently in the rural clinic

def clinic(n, w):
    X = rng.normal(0, 1, (n, 6))
    y = (rng.random(n) < 1 / (1 + np.exp(-3 * (X @ w)))).astype(float)
    return X, y

def train(w, X, y, epochs, lr=0.5):
    for _ in range(epochs):
        w = w - lr * X.T @ (1 / (1 + np.exp(-(X @ w))) - y) / len(y)
    return w

rules = [w_city] * 4 + [w_rural]
names = ["city A", "city B", "city C", "city D", "rural"]
train_sets = [clinic(200, w) for w in rules]
test_sets = [clinic(2000, w) for w in rules]
acc = lambda w, X, y: np.mean(((X @ w) > 0) == y)

w = np.zeros(6)
for rnd in range(20):                         # FedAvg with equal-sized clients
    w = np.mean([train(w.copy(), X, y, 5) for X, y in train_sets], axis=0)

for name, (X, y), (Xt, yt) in zip(names, train_sets, test_sets):
    personal = train(w.copy(), X, y, 30)      # fine-tune the global model locally
    print(f"{name:>6}: global model {acc(w, Xt, yt):.2f} | after local fine-tuning {acc(personal, Xt, yt):.2f}")

The global model scores 0.92 in the city clinics but only 0.64 in the rural one. Thirty local steps lift the rural clinic to 0.91 without hurting anyone else. Research methods such as FedProx and SCAFFOLD also reduce this "client drift" during training.

Communication, stragglers and dropouts

Sending millions of weights to thousands of phones every round is expensive. Systems train several local epochs per round, compress updates and sample a fraction of clients each time. Phones also go offline mid-round, so the server must tolerate missing replies.

Privacy: updates can still leak

Updates are not raw data, but they are not harmless: researchers have shown that individual gradients can sometimes be inverted to reconstruct training examples. Two defences are standard.

Secure aggregation lets the server learn only the sum of updates, not any single one. In the toy version below, every pair of clients agrees on a random mask; one adds it and the other subtracts it. Each masked update looks like noise, but the masks cancel in the sum.

import numpy as np

rng = np.random.default_rng(7)
updates = {"A": np.array([0.20, -0.10, 0.05]),       # each hospital's real model update
           "B": np.array([0.15, -0.05, 0.10]),
           "C": np.array([0.25, -0.20, 0.00])}
names = list(updates)

# every pair of hospitals agrees on a random mask; one adds it, the other subtracts it
masked = {n: updates[n].copy() for n in names}
for i in range(len(names)):
    for j in range(i + 1, len(names)):
        mask = rng.normal(0, 10, 3)
        masked[names[i]] += mask
        masked[names[j]] -= mask

for n in names:
    print(f"server receives from {n}: {np.round(masked[n], 2)}")   # looks like noise
print("sum of masked updates:", np.round(sum(masked.values()), 2))
print("true sum of updates:  ", np.round(sum(updates.values()), 2))

Differential privacy goes further: each update is clipped to a maximum size and random noise is added, so no single person's data can noticeably change the result. Real protocols also handle clients that drop out after masking, which this toy skips.

Where is federated learning used?

• Phone keyboards: Google has described training next-word and emoji suggestions in its Gboard keyboard with federated learning, so typing data stays on phones.
• Healthcare: hospitals train tumour-detection or risk models together without sharing patient scans.
• Banking and fraud: banks learn fraud patterns from one another's transactions without exposing customer data, which links to anomaly detection.

Open-source frameworks include Flower, TensorFlow Federated and NVIDIA FLARE. They handle client selection, communication and aggregation, so you write mostly ordinary training code.

Common federated learning mistakes

• Assuming federated means private. Without secure aggregation or differential privacy, updates can leak information.
• Testing only on IID simulations. Real clients are uneven. Simulate skewed splits before trusting a design.
• Averaging without weights. A clinic with 20 patients should not count as much as one with 2,000.
• Too many local epochs on skewed data. Clients drift towards their own optimum; keep local training short or use a drift-correcting method.
• Ignoring evaluation per client. A good global average can hide a client where the model fails, as the rural clinic showed.

Federated setups combine ML, systems and privacy thinking, the kind of cross-topic reasoning practised in Debugshala's ML course for graduates and career-switchers. For revision, check your federated learning basics with MCQs on Debugshala Test.

Key takeaways

• Federated learning trains one model across many data owners while raw data stays where it is.
• FedAvg loops: broadcast, train locally, return weights, average by data size.
• Non-IID data is the central difficulty; per-client evaluation and local fine-tuning help.
• Updates can leak information, so combine federated learning with secure aggregation and differential privacy.
• Cross-device (phones) and cross-silo (organisations) setups have very different engineering needs.