Personalized Federated Learning

"They averaged a thousand of us into one model and called it consensus. It fit nobody, least of all the night-shift phone in Jakarta that only ever sees keyboards in the dark."

A Global Model That Pleased No Client

The previous sections built a single global model and worked to defend it: FedAvg averaged client updates (Section 14.3), non-IID skew threatened its convergence (Section 14.4), communication budgets constrained how often it could be refreshed (Section 14.5), and secure aggregation protected the updates on the way in (Section 14.6). All of that assumed one model is the right answer. Now we question the assumption. If your thousand phones split into a group that types in English and a group that types in Hindi, no shared keyboard model serves both well, and forcing one is a modeling error, not a privacy or systems problem. Personalized federated learning accepts that each client may need its own model while insisting that those models be learned together, so a client with little data still benefits from everyone else's.

1. When One Global Model Is the Wrong Goal Beginner

Federation, as built so far, minimizes the average loss across clients with one shared parameter vector $w$. Write client $k$'s local risk as $F_k(w)$ and its share of the data as $p_k$; FedAvg targets

When every $F_k$ is minimized at roughly the same place, this single $w^\star$ is excellent for all of them, and there is nothing to personalize. The trouble is the non-IID regime, where the client minimizers $w_k^\star = \arg\min_w F_k(w)$ scatter. A point that minimizes the weighted average can then sit far from every individual optimum, fitting the centroid of the clients rather than any client. The averaged keyboard predicts a blend of English and Hindi that neither typist wants. This is not a bug in FedAvg; it is FedAvg correctly solving a problem whose answer no longer matches what each client needs.

The honest alternative at the other extreme is to give up on sharing: let client $k$ train only on its own data and keep $w_k$ entirely private. That removes all bias toward other clients, but on a phone that has typed a few hundred words, the local minimizer is estimated from so little data that it is dominated by noise. The model overfits its own scarce, noisy sample and generalizes poorly to the same client's future inputs. Personalization is the search for a middle point: a per-client model that inherits the shared structure all clients have in common, then adjusts for what makes the client different.

2. Five Ways to Personalize Intermediate

The literature offers a handful of recurring strategies; they differ in how much of the model is shared and where the personalization enters. The simplest, and a startlingly strong baseline, is to fine-tune the finished global model locally: run a few steps of gradient descent on the client's own data starting from $w^\star$. The global model supplies a good initialization estimated from all clients, and a short local pass bends it toward the individual without enough steps to overfit. This is the method we measure in Section 4, and it often recovers most of the achievable gain for a few lines of code.

Meta-learning makes that initialization explicit. Instead of training $w$ to be good as-is, Per-FedAvg (Fallah et al., 2020) trains $w$ so that one or a few gradient steps from $w$ are good on each client, the federated cousin of MAML. The server optimizes for adaptability rather than immediate accuracy, producing a launch point from which every client reaches its own model quickly. Multi-task and clustered federated learning take a structural view: if clients fall into groups (the English typists, the Hindi typists), discover the groups and train one model per cluster, so each client shares strength only with its similar peers rather than the whole population. Clustered FL (Sattler et al., 2020) recursively splits clients whose update directions disagree.

Parameter decoupling splits the model itself. A shared global backbone (the lower layers, a feature extractor) is federated across all clients, while a small personalized head (the final layer or two) is kept and trained locally on each device and never averaged. FedPer and FedRep build on this idea: the representation is common because good features are universal, the classifier is private because the decision boundary is client-specific. The diagram in Figure 14.7.1 shows this split. Finally, interpolation methods keep both a global and a local model and mix them, predicting with $\alpha \, w_{\text{local}} + (1-\alpha)\, w_{\text{global}}$ or regularizing the local model to stay near the global one. The mixing weight $\alpha$ is the bias-variance knob made literal, tuned per client from how much its data justifies trusting itself over the crowd.

3. The Cost of Personalizing: Generalizing to New Clients Intermediate

Personalization is not free, and its price is paid by clients who were not in the training population. A single global model has a clear story for a brand-new device: hand it $w^\star$ and it works at average quality on the first request, before the device has produced a single gradient. A clustered or decoupled scheme has no head for a client it has never seen, and a meta-learned initialization is good only after the new client takes its adaptation steps, which requires that the client has some data and the compute to use it. The more sharply a method tailors to seen clients, the less it has to say about unseen ones; this is the generalization-to-new-clients axis along which personalization methods are fairly compared, and it is why meta-learning, which yields a deployable initialization, is often preferred over per-cluster models when the client population churns.

There is a systems cost too. Decoupling and interpolation mean the server now manages per-client state (heads, mixing weights, cluster assignments), turning a stateless aggregation into something closer to the per-client model store of a recommendation system, an idea we meet again with sharded embeddings in Chapter 11. The reward for these costs is concrete, and the next section measures it: on non-IID clients with scarce data, a global-then-fine-tuned model beats both the shared global model and the purely local model at the same time.

4. Personalization Versus Global and Local, Measured Intermediate

The code below builds the bias-variance picture concretely with pure Python and NumPy. Six clients share a large common linear concept and add a small per-client tilt, the mild non-IID setting where a backbone is worth sharing but no single model is exactly right for anyone. Each client holds only fourteen noisy training examples, far too few to trust alone, and four hundred held-out test examples. We train a global model by FedAvg, then produce a personalized model per client by fine-tuning that global model on the client's own fourteen points, and compare both against a purely local model trained from scratch.

import numpy as np

rng = np.random.default_rng(7)
K, d = 6, 12                  # clients, features
n_per, n_test = 14, 400       # FEW noisy train points per client, many test points
lr, noise = 0.03, 1.0

# Clients share a large COMMON backbone and add a small per-client tilt -> mildly
# non-IID. The shared part dominates, so borrowing strength pays off, but the
# tilt means no single global model is exactly right for any one client.
common = 2.0 * rng.standard_normal(d)
client_w = []
for k in range(K):
    offset = np.zeros(d)
    offset[k % d] = 1.5 * (1 if k % 2 == 0 else -1)   # client-specific tilt
    client_w.append(common + offset)

def make(k, n):
    X = rng.standard_normal((n, d))
    return X, X @ client_w[k] + noise * rng.standard_normal(n)

train = [make(k, n_per) for k in range(K)]
test  = [make(k, n_test) for k in range(K)]
grad = lambda w, X, y: (2.0 / len(y)) * (X.T @ (X @ w - y))
mse  = lambda w, X, y: float((X @ w - y) @ (X @ w - y) / len(y))

# FedAvg: each round, every client takes a few local steps; the server averages.
w_global = np.zeros(d)
for _ in range(400):
    updated = []
    for k in range(K):
        wk = w_global.copy()
        for _ in range(5):
            wk -= lr * grad(wk, *train[k])
        updated.append(wk)
    w_global = np.mean(updated, axis=0)          # the single shared model

global_test = [mse(w_global, *test[k]) for k in range(K)]

# Personalized: start from the global model, fine-tune on each client's own data.
pers_test = []
for k in range(K):
    wk = w_global.copy()
    for _ in range(200):
        wk -= lr * grad(wk, *train[k])
    pers_test.append(mse(wk, *test[k]))

# Purely local: train from scratch, no shared learning at all.
local_test = []
for k in range(K):
    wk = np.zeros(d)
    for _ in range(200):
        wk -= lr * grad(wk, *train[k])
    local_test.append(mse(wk, *test[k]))

print("client |  global  |  local   | personalized")
for k in range(K):
    print(f"   {k}   | {global_test[k]:8.3f} | {local_test[k]:8.3f} | {pers_test[k]:11.3f}")
print(f"  mean | {np.mean(global_test):8.3f} | {np.mean(local_test):8.3f} | {np.mean(pers_test):11.3f}")
print(f"  worst| {np.max(global_test):8.3f} | {np.max(local_test):8.3f} | {np.max(pers_test):11.3f}")

client |  global  |  local   | personalized
   0   |    3.763 |    4.314 |       1.913
   1   |    3.154 |    3.699 |       3.505
   2   |    3.014 |    4.031 |       1.622
   3   |    2.463 |    2.026 |       2.447
   4   |    2.546 |    5.203 |       2.402
   5   |    4.186 |    7.982 |       3.112
  mean |    3.188 |    4.543 |       2.500
  worst|    4.186 |    7.982 |       3.505

Three numbers tell the whole story. The global model's mean error (3.188) is the bias floor: estimated from everyone, it is stable but tuned to a population centroid, so it never fits a client tightly. The local model's mean error (4.543), and especially its worst-client error (7.982), is the variance ceiling: with fourteen noisy points a client cannot estimate its own model reliably and overfits. Personalization (mean 2.500, worst 3.505) beats both at once, exactly the bias-variance middle the key insight promised. It does so for the price of a short local fine-tune, no new server protocol, and no extra communication beyond the FedAvg rounds that produced the global model. That is why fine-tuning is the baseline every more elaborate method must beat.

# pip install flwr ; the server runs standard FedAvg, unchanged
import flwr as fl

class PersonalizedClient(fl.client.NumPyClient):
    def fit(self, global_params, config):
        self.model.set_weights(global_params)         # start from the shared model
        self.model.train(self.local_data, epochs=1)   # contribute to the global average
        return self.model.get_weights(), self.n, {}

    def evaluate(self, global_params, config):
        self.model.set_weights(global_params)
        self.model.train(self.local_data, epochs=2)    # FINE-TUNE locally, then score
        loss, acc = self.model.test(self.local_test)   # per-client personalized metric
        return loss, self.n, {"personalized_accuracy": acc}

5. The Tension That Outlives the Server Advanced

The global-versus-local tension this section navigates is not specific to having a central server. It is the same conflict that organizes the rest of the chapter and the federated thread of the book. In decentralized learning (Section 14.8), there is no server to hold a global model at all; clients average only with gossip neighbors, and "how much should I trust my neighbors versus my own data?" is precisely the interpolation knob $\alpha$ of Section 2, now set by the mixing matrix of the communication graph. At the edge (Section 14.9), on-device compute and intermittent connectivity make local adaptation cheap and global synchronization expensive, tilting the same trade toward more personalization. The federated arc continues into Chapter 34 on federated edge learning and culminates in the federated medical case study of Chapter 37, where every hospital is a client whose patient distribution differs and a personalized model per site is not a luxury but a clinical requirement.

We now have the personalization toolkit (fine-tune, meta-learn, cluster, decouple, interpolate), the bias-variance reason it works, the new-client cost it incurs, and a measured win for its simplest form. The same global-versus-local trade is about to outlive the server entirely: the next section removes the central aggregator and lets clients reach agreement by talking only to neighbors. That decentralized world begins in Section 14.8.