AI.Cruz - Learn AI Through Interactive Visualizations

FEDERATED LEARNING

FedAvg: The Federated Averaging Algorithm

Each client runs several local SGD steps; the server computes a weighted average of all updates proportional to each client's dataset size.

1Why Averaging Weights Works

Before FedAvg, the naive FL approach was FedSGD: each client computes one gradient step and sends it; the server accumulates gradients. This requires one round of communication per gradient — enormously expensive at scale. FedAvg's insight: let each client run multiple local SGD steps before communicating. The resulting weight deltas are then averaged. More local compute per round = fewer rounds = drastically less communication.

Why does weight averaging work? Neural network loss surfaces are approximately convex in the region of convergence. If two clients both start from the same global weights and move towards local minima via SGD, the midpoint of their weights is typically also a lower-loss point than the global start — so weighted averaging produces a globally better model. This breaks down with very non-IID data or too many local steps.

2The FedAvg Algorithm

FedAvg has two levels: the server loop and the client update. The server samples a fraction C of clients each round, sends them the global model, waits for their updated weights, then computes a weighted average. The weight given to client k is proportional to its local dataset size n_k divided by the total data n.

Server loop (T rounds)

For each round t = 1 to T:
S_t ← random sample of C·K clients
For each client k in S_t (in parallel):
w_k ← ClientUpdate(k, w_t)
w_{t+1} ← Σ_k (n_k / n) · w_k

ClientUpdate(k, w) — E local epochs

Split local data into batches of size B
For epoch e = 1 to E:
For each batch b:
w ← w − η · ∇L(w; b)
Return w

K — total clients available. Only C·K participate per round (partial participation).

E — local epochs per round. Increasing E reduces communication but increases divergence risk with non-IID data.

B — local batch size. B=∞ means full-batch gradient on local data per epoch.

3Convergence and Hyperparameters

FedAvg converges to a good solution under IID data — clients effectively pool their data through the average. Under non-IID data, convergence degrades because clients diverge towards their local optima during the E local epochs before re-synchronising. This is called client drift. The remedy is either reducing E, reducing C (more clients per round), or using FedProx which adds a proximal term penalising deviation from the global model.

✓ Good hyperparameter choices

✓ E=1 for highly non-IID; E=5–20 for near-IID — reduces client drift
✓ C=0.1 to 0.3 — enough clients for stable aggregation
✓ η (learning rate) lower than centralised baseline — averaged updates amplify instability

✗ Common pitfalls

✗ E too large with non-IID data — catastrophic client drift, divergence
✗ Uniform weighting (instead of n_k/n) — large-data clients under-represented
✗ No learning rate decay — oscillation around global optimum

4Beyond FedAvg

FedAvg spawned an entire research field. Key extensions: FedProx adds a proximal term to each client's loss to limit divergence from the global model. SCAFFOLD uses variance reduction to correct for client drift. FedNova normalises local updates by the number of steps taken to avoid biasing towards clients with more compute. FedMA uses neuron-level matching to average dissimilar network architectures.

FedProx
Adds ‖w − w_global‖² penalty to client loss. Limits how far clients drift. Hyperparameter μ controls penalty strength.

SCAFFOLD
Uses control variates to estimate and subtract client drift. Theoretically optimal convergence even with non-IID data.

Personalised FL
Rather than one global model, maintain a personalised model per client. Trades global generalisation for local fit — key for heterogeneous tasks.

FedAvg from scratch — pure Python simulation

import random

# ── Toy linear regression: y = 2x + noise ──────────────────────────
def make_data(n, slope=2.0, noise=0.3, seed=None):
    rng = random.Random(seed)
    xs = [rng.uniform(-1, 1) for _ in range(n)]
    ys = [slope * x + rng.gauss(0, noise) for x in xs]
    return list(zip(xs, ys))

def local_sgd(w, data, lr=0.01, epochs=5):
    """One client: run SGD for 'epochs' passes over local data."""
    for _ in range(epochs):
        for x, y in data:
            pred = w * x
            grad = -2 * x * (y - pred)
            w -= lr * grad
    return w

def fed_avg(clients_w, clients_n):
    """Weighted average of client weights."""
    total_n = sum(clients_n)
    return sum(w * n / total_n for w, n in zip(clients_w, clients_n))

# ── Setup: 4 clients with different dataset sizes ─────────────────
client_data = [
    make_data(n=200, slope=2.0, seed=1),
    make_data(n=100, slope=2.0, seed=2),
    make_data(n=300, slope=2.0, seed=3),
    make_data(n=150, slope=2.0, seed=4),
]
client_n = [len(d) for d in client_data]

w_global = 0.0  # start at 0 (true value: 2.0)

print(f"{'Round':<8} {'w_global':<12} {'Error'}")
print("-" * 35)
for rnd in range(1, 11):
    # Each client trains locally starting from global weights
    updated_ws = [local_sgd(w_global, data) for data in client_data]
    # Server aggregates
    w_global = fed_avg(updated_ws, client_n)
    err = abs(w_global - 2.0)
    print(f"{rnd:<8} {w_global:<12.4f} {err:.4f}")

print(f"\nTrue slope: 2.0000 | Final estimate: {w_global:.4f}")