Personalized Federated Learning using Hypernetworks

In plain words

Usually, a federated system trains one shared model for many devices, and the devices may hold very different data. This paper instead trains one model that writes a separate model for each device. The writing model lives on the server. Each device has a short learned "address" that the writing model reads to produce that device's own network.

Devices train their personal network on their own data, then send back only how it changed. The server uses that change to improve the writing model. The writing model itself never leaves the server, so the large part of the system is never sent over the network.

The authors report that this beats earlier personalized methods on their benchmarks, and that it can also produce a model for a new device without retraining the whole system.

TL;DR

Personalized federated learning with one server-side hypernetwork that outputs each client's weights from a trainable client embedding. Clients send back a weight difference rather than gradients, and the server backpropagates it. The hypernetwork is never transmitted, so communication scales with the target network, not the hypernetwork. Author-reported results are strong on CIFAR10 and CIFAR100 splits but not on Omniglot against Per-FedAvg. Read the Pitfalls section before implementing: the printed update sign and the code disagree.

Problem

Clients in federated learning often have different data distributions, so a single global model fits each of them poorly §1. Personalized federated learning gives each client its own model, but previous approaches either exchange large amounts of parameters or share little across clients §1, §2. The paper states the goal as minimizing the average client loss over personal models eq. 2, with communication cost tied to the personal model size [§1, §3.2].

Method

Federated hypernetwork (pFedHN). A hypernetwork h(·; φ) on the server maps a trainable embedding v_i of client i to that client's target network weights θ_i = h(v_i; φ) §3.2. All clients share φ and the embeddings are learned per client, so the hypernetwork holds a family of personal models §3.2. The objective optimizes φ and all v_i jointly under the average client loss eq. 3.

Communication. The server sends θ_i to the client. The client runs several local optimization steps to get θ̃_i and returns the update direction Δθ_i = θ̃_i − θ_i §3.2 fig. 1. The server maps this direction back to φ and v_i with the chain rule, using the Jacobian of θ_i with respect to φ and v_i §3.2. The paper states that the hypernetwork size does not affect what is sent per round, because only target-network-sized vectors are exchanged §3.2.

Update rule. The authors motivate the update as a gradient of the surrogate ½‖θ*_i − h(v_i; φ)‖², where the unknown optimum θ*_i is replaced by the locally trained θ̃_i §3.2. Algorithm 1 gives the loop: sample a client, generate its weights, run K local steps, then update φ and v_i alg. 1. The local steps use mini-batch gradients alg. 1. The authors report that several local steps beat a single gradient step for convergence speed and accuracy §3.2 app. C.

# Personalized federated hypernetwork round  [alg. 1]
# Input:  hypernet (params phi and client embeddings), client, eta, K
# Output: updated hypernet parameters

def pfedhn_round(hypernet, client, eta, local_steps):
    # 1. Server builds the personal weights theta_i = h(v_i; phi)  [§3.2]
    theta = hypernet(client.embedding_id)
    theta_tilde = copy_of(theta)
    # 2. Client runs K local mini-batch SGD steps  [alg. 1] [§5]
    for step in range(local_steps):
        batch = client.sample_minibatch()
        theta_tilde = theta_tilde - eta * grad_loss(theta_tilde, batch)
    # 3. Difference used for the server update; sign per the code, not
    #    Alg. 1 as printed  [code: experiments/pfedhn/trainer.py#L166-L169]
    delta = theta - theta_tilde
    # 4. Chain rule through h: gradient for phi and for v_i  [§3.2]
    grads = vector_jacobian(theta, hypernet.params, delta)
    # 5. Step all hypernetwork parameters, embeddings included  [alg. 1]
    hypernet.params = hypernet.params - alpha * grads
    return hypernet

Personal classifier variant (pFedHN-PC). The hypernetwork produces only the feature extractor θ_i. Each client keeps its own output layer ω_i, which is updated locally and never sent §3.3 eq. 4. The authors use this when clients solve different tasks, for example different numbers of classes §3.3.

Analysis (not used by the experiments). For a linear hypernetwork θ_i = W v_i and linear regression with X_i^T X_i = I, the optimal W and V are given by PCA of the client-wise ERM solutions [prop. 1] app. A. The authors interpret this as denoising. A generalization bound is given in terms of the number of samples per client, the embedding dimension k and the hypernetwork size N, with a 1/n factor on the hypernetwork term [thm. 1] §4.2.

Novel clients. The hypernetwork is frozen and only a new embedding v_new is optimized on the new client's data §5.3. The paper reports better generalization than the baselines it compares to fig. 2 §5.3.

Hyperparameters

Results

Author-reported. Setup: heterogeneous CIFAR10 and CIFAR100 with 10, 50 and 100 clients, plus Omniglot with 50 clients; metric is Federated Accuracy averaged over three seeds §5 §5.1 tab. 1. The "±" values are not defined in the text.

Computational budget: 75 clients split into three equal groups with small, medium and large target networks (first conv with 8, 16 or 32 filters). One pFedHN produces all three sizes. Each baseline trains one model per group §5.2 tab. 2.

Other reported results (author-reported, no tables in the main text):

• Novel clients: accuracy gap versus total variation distance is plotted for the training and novel clients; the paper states pFedHN has the best generalization at every distance level tested fig. 2 §5.3.
• Personal classifiers: a mix-and-match test swaps feature extractors between clients in pFedHN-PC, and accuracy varies with the pairing fig. 3 §5.4.
• Client embeddings for CIFAR100 split by coarse label form clusters in a t-SNE plot fig. 4 §5.5.
• MNIST and the spectral-normalization ablation are in the appendix app. C.

Paper vs code

Official code: https://github.com/AvivSham/pFedHN, linked from the paper's introduction §1 and marked as official in its README. Cloned at commit 1bc054975232bdd6985f931030753ebbab48c56d. The repository has no license file, and the GitHub license API returns 404, so no code may be quoted. This card cites and paraphrases only. The repo contains the pFedHN and pFedHN-PC experiments for CIFAR10 and CIFAR100 only.

• Update sign. Algorithm 1 writes φ ← φ − α ∇φθᵀ Δθ with Δθ = θ̃ − θ alg. 1. The code computes delta = θ − θ̃ and takes a descent step on the gradient of that vector trainer.py trainer.py. The two give opposite signs. The code matches the surrogate-loss argument in §3.2. Inferred from the code and the text; the authors do not mention a discrepancy.
• Inner optimizer. The paper's inner update is a plain gradient step alg. 1. The code uses SGD with momentum 0.9 and weight decay 5e-5, created fresh each round, so momentum does not carry across rounds trainer.py. It also clips the global gradient norm at 50 trainer.py.
• Outer optimizer. The paper's hypernetwork update is not described beyond the equation. The code uses SGD with momentum 0.9 and weight decay 1e-3 on all hypernetwork parameters, including the embedding table, with clipping at 50 trainer.py trainer.py. An Adam option also exists trainer.py. Weight decay on a dense embedding table decays every row each step, even rows that received no gradient. This follows PyTorch SGD semantics; the paper says nothing about it.
• Update order. Algorithm 1 lists the φ update and then the v_i update alg. 1. The code computes one gradient for all hypernetwork parameters and steps them together trainer.py.
• Federated accuracy. The paper defines it as the average over clients of each client's accuracy §5. The code logs test_avg_acc as correct predictions divided by total samples, which is a sample-weighted average trainer.py. The client average is only stored as all_acc and used for the min, max and standard deviation trainer.py. Which value Table 1 reports is not stated.
• Spectral normalization. The appendix says it was applied to all hypernetwork layers app. C. The code defines --spec-norm [code: experiments/pfedhn/trainer.py#L279], but train() does not take it and does not pass it to CNNHyper trainer.py trainer.py. The flag has no effect.
• Target network. The paper says the target has two convolution and two fully connected layers §5.1. The code's CNNTarget has two convolutions and three linear layers, the last being the output models.py. The PC target stops at an 84-dimensional feature layer models.py, and a per-client linear layer produces logits models.py.
• Learning rates for PC. The PC script defaults the outer learning rate to 5e-2 trainer.py, while pFedHN defaults to 1e-2 trainer.py. The paper does not state which value was used for either.
• Test data during training. Each step computes loss and accuracy of the sent weights on the client's test loader, for logging only trainer.py. Model selection uses the validation loader trainer.py. Keep the logging out of any reimplementation's selection step.
• Client sampling and local data. One client is drawn uniformly at random with replacement each round trainer.py. Each inner step takes the first batch of a freshly shuffled loader trainer.py dataset.py.
• Partitioning. The code draws class shares from U(0.4, 0.6), normalizes them, and gives each client a fixed number of classes dataset.py dataset.py. The validation set is 10,000 images from training dataset.py, matching app. B.
• Not in the repository. Omniglot, MNIST, all baselines, the novel-client experiment and the mix-and-match and embedding analyses. The trainer raises an error for any dataset other than CIFAR10 or CIFAR100 trainer.py.
• Seeds. The paper reports three seeds §5. The code default is seed 42 and runs one seed per invocation trainer.py.

No other discrepancies found for the embedding dimension, the 3×100 hypernetwork and K = 50 trainer.py trainer.py.

Pitfalls

• Update sign. Implement the sign the code uses, which matches the surrogate loss in §3.2. Following the printed Algorithm 1 literally moves φ the wrong way (see Paper vs code).
• Unstated optimizer details. Momentum, weight decay, gradient clipping and batch size are not in the paper. Without the code, a reimplementation has no basis for these choices.
• Metric definition. The paper's formula averages per client, the code's headline number is sample-weighted §5 trainer.py. Report both and state which one you use.
• Claimed margins. The text says "2%–10%" gains over all competitors §5.1. Computed from Table 1 against the best baseline in each column, pFedHN is +1.7 to +2.7 points on CIFAR10, +3.3 to +6.8 on CIFAR100, and −4.4 on Omniglot (72.03 vs 76.46 for Per-FedAvg). The text says "4%–8%" for Table 2 §5.2. The computed gaps are +2.8 to +6.6 (smallest on CIFAR10 medium).
• Standard deviations. The ± values are given without definition; the text says accuracies are averaged over three seeds §5.
• Theorem statement. The statement omits a log(1/ε) factor. The appendix gives the corrected bound app. A [thm. 1]. The bound depends on N/n, so large hypernetworks need more samples per client by this bound.
• Cross-reference. The main text says the local-steps figure is in Appendix B §3.2, but it appears in Appendix C app. C.
• Novel clients. Steps, learning rate and embedding initialization are not stated §5.3. The source contains commented-out text about initializing from the mean of the training embeddings. That text is not in the published paper and must not be used as a stated detail.
• Learning-rate selection. The search grids are given, but the chosen values are not app. B. The code's defaults differ between pFedHN and PC.
• Local steps and communication. The authors state that K = 50 local steps beat a single step app. C. Each round still sends a full target-sized vector down and back, so communication scales with the target size and the number of rounds.
• Linear analysis. Proposition 1 needs X_i^T X_i = I for every client [prop. 1]. It does not describe the nonlinear model used in experiments.
• Ablation numbers. The hypernetwork depth ablation reports 88.38 and 88.42 for 3 and 4 hidden layers app. C, which is consistent with the default of 3.
• Consistent with code. Client embedding dimension, the 3-hidden-layer hypernetwork of width 100 and K = 50 §5 app. B all match the code defaults.

Lineage

• Builds on: hypernetworks, where a network outputs the weights of another network §2.3. The embedding-based update is said to be inspired by the lookahead method of Zhang et al. (2019) §3.2. Local-steps reasoning follows McMahan et al. (2017) and Huo et al. (2020) §3.2.
• Compared against (author-named, no cards in this wiki): FedAvg [McMahan et al. 2017], Per-FedAvg [Fallah et al. 2020], pFedMe [Dinh et al. 2020], LG-FedAvg [Liang et al. 2020], FedPer [Arivazhagan et al. 2019] §5 tab. 1.
• Superseded by: not stated.