Federated Learning with Matched Averaging

In plain words

In federated learning, many devices each train a copy of a model on their own data, and a server combines the copies into one shared model without ever seeing the data. The standard way to combine them, FedAvg, simply averages the weights position by position. That assumes neuron number 5 on one device does the same job as neuron number 5 on every other device, but nothing guarantees it: the order of neurons in a layer is arbitrary, so two networks can learn the same features in a shuffled order. Averaging a "horse leg detector" with an unrelated filter that happens to sit in the same slot ruins both.

FedMA first lines the neurons up, then averages. For each layer the server looks for pairs of neurons (or convolution channels, or LSTM hidden units) across devices whose weights look alike, treats them as the same feature, and averages only those. A neuron that has no good partner anywhere is kept as a new neuron instead of being forced into a bad match, so the shared layer can grow a little when devices have seen different kinds of data.

Matching a whole deep network at once works poorly, so FedMA goes one layer at a time: match the first layer, send it back, let every device retrain the remaining layers with that first layer fixed, then match the second layer, and so on. The last layer is averaged with weights based on which classes each device has. On the authors' experiments with a 9-layer CNN on CIFAR-10 and an LSTM on Shakespeare text, this beat FedAvg and FedProx while sending less data in total, and it kept improving when devices trained longer locally, where FedAvg got worse.

TL;DR

• Weight-space averaging is broken by permutation invariance of hidden units; FedMA solves a per-layer matching problem so that only matched units are averaged §2 eq. 2.
• The matching allows a growing global layer (unmatched local units become new global units) and is solved with BBP-MAP, the Beta-Bernoulli-process MAP solver from PFNM, via iterated Hungarian assignments §2.1 eq. 3.
• Extends matching to CNN channels (flatten C^{in}\times w\times h) and LSTM hidden states (match on stacked input-to-hidden weights, permute hidden-to-hidden on both sides) §2.2 eq. 5 eq. 6.
• Matching all layers in one shot fails on VGG-9; FedMA matches one layer per communication round and has clients retrain the layers above with the matched layer frozen §2.3 alg. 1.
• Author-reported: VGG-9 / CIFAR-10, 16 heterogeneous clients: 87.53% vs FedAvg 86.29%, FedProx 85.32%; LSTM / Shakespeare, 66 clients: 49.07% vs 46.63% / 45.83% tab. 1 tab. 2.

Problem

FedAvg averages client parameters coordinate by coordinate. For a hidden layer with L units there are L! functionally equivalent orderings, so two clients trained from different starts generally end up in different orderings, and their coordinate-wise average is not a permuted copy of either solution §2 eq. 1. The authors argue this hurts accuracy and forces many communication rounds §1. PFNM (Yurochkin et al., 2019) fixed this by matching neurons, but only for fully connected networks; the authors report that a direct extension to CNNs works on LeNet/MNIST but breaks down on VGG-9/CIFAR-10 §1 §2.2 fig. 1. Assumptions: all clients train the same architecture, and in experiments the server samples every client in every round §3. FedAvg and FedProx baselines are run without a shared random initialization app. A.

Method

Matched averaging of one layer

Let w_{jl} be the l-th unit (incoming weight vector) of client j, \theta_i the i-th global unit and c(\cdot,\cdot) a similarity/cost. The matching problem is eq. 2:

\min_{\{\pi^j_{li}\}} \sum_{i=1}^{L} \sum_{j,l} \min_{\theta_i} \pi^j_{li}\, c(w_{jl},\theta_i) \quad \text{s.t.}\ \sum_i \pi^j_{li}=1\ \forall j,l;\ \ \sum_l \pi^j_{li}=1\ \forall i,j.

With the resulting permutations \Pi_j, the federated weights of a one-hidden-layer net are W_1=\frac1J\sum_j W_{j,1}\Pi_j^T and W_2=\frac1J\sum_j \Pi_j W_{j,2} §2.1. With squared Euclidean cost this is k-means with permutation constraints; with equal sizes everywhere it is a Wasserstein barycenter §2.1.

It is solved by coordinate descent over clients: fix all other clients' assignments, compute the global units from them, and run the Hungarian algorithm for client j' §2.1. To handle heterogeneous data the global size L is free, with \max_j L_j \le L \le \sum_j L_j; a local unit opens a new global unit when its best match costs more than \epsilon plus a size penalty f §2.1 eq. 3:

\min_{\{\pi^{j'}_{li}\}_{l,i}} \sum_{i=1}^{L+L_{j'}}\sum_{j=1}^{L_{j'}} \pi^{j'}_{li} C^{j'}_{li} \quad\text{s.t.}\ \sum_i \pi^{j'}_{li}=1\ \forall l;\ \ \sum_l \pi^{j}_{li}\in\{0,1\}\ \forall i, \qquad C^{j'}_{li}=\begin{cases} c(w_{j'l},\theta_i), & i\le L\\ \epsilon+f(i), & L<i\le L+L_{j'}.\end{cases}

The new size is the largest used column index; local layers smaller than the global one are padded with zero "dummy" units that must not count in the average §2.1. The paper does not choose c, \epsilon, f by hand: it uses PFNM's Beta-Bernoulli-process model, where c is a Gaussian posterior term and \epsilon, f follow from the Indian Buffet Process prior; this solver is called BBP-MAP §2.1. The paper gives no formula for these terms; the code computes them as below pfnm.py.

Architectures

• Deep FC: x_n=\sigma(x_{n-1}\Pi_{n-1}^T W_n \Pi_n), \Pi_0 and \Pi_N identity eq. 4. Joint matching of all layers is NP-hard, so layers are matched recursively: apply the previous layer's permutation to the rows of W_{j,n}, then match §2.2.
• CNN: channels play the role of neurons: x_n=\sigma(\mathrm{Conv}(x_{n-1},\Pi_{n-1}^T W_n\Pi_n)) eq. 5. Each output channel becomes a vector of size D=C^{in}_n\cdot w\cdot h (after permuting input channels by the previous layer's \Pi). Pooling is compatible because it acts per channel §2.2.
• RNN/LSTM: h_t=\sigma(h_{t-1}\Pi^T H\Pi + x_t W\Pi) eq. 6. Matching H directly is a quadratic assignment problem (NP-hard), so permutations are found from input-to-hidden weights only, and H=\frac1J\sum_j\Pi_j H_j\Pi_j^T §2.2. For LSTMs the input-to-hidden weights of all S gates are stacked into an SD\times L matrix for matching; the embedding layer is handled like an FC layer; deep LSTMs are processed layer by layer §2.2.

FedMA (layer-wise matching with retraining)

# FedMA: one pass over an N-layer network  [alg. 1] [§2.3]
# Input:  client_weights[j][n] for J clients, layers n = 1..N
# Output: global_weights[1..N]

def fedma_pass(client_weights, class_share):
    for n in range(1, N + 1):
        if n < N:
            # match layer n across clients with BBP-MAP  [alg. 1] [eq. 2]
            perms = bbp_map([client_weights[j][n] for j in clients])
            # average matched units  [alg. 1]
            global_weights[n] = mean_over_j(
                client_weights[j][n] @ perms[j].T)
        else:
            # last layer: per-class weighted average  [alg. 1] [§2.3]
            # class_share[j][k] = share of class-k data held by client j
            # (resolution from code: main.py#L1538-L1550)
            global_weights[n] = sum_over_k_j(
                class_share[j][k] * client_weights[j][n][:, k])
        for j in clients:
            # permute next layer's input dimension to the matched order
            # [alg. 1]
            client_weights[j][n + 1] = perms[j] @ client_weights[j][n + 1]
            # client sets layers 1..n to the global ones, frozen, and
            # trains layers n+1..N on local data  [alg. 1] [§2.3]
            client_weights[j] = train_local(
                client_weights[j], frozen_upto=n)
    return global_weights

One pass costs as many communication rounds as the network has layers §2.3.

FedMA with communication §2.3: after a pass, each client rebuilds a local model of its original size from the global model using its own matching from the previous round, retrains it, and a new pass begins. The code rebuilds by selecting the global units its own units were matched to main.py, retrains the whole network with nothing frozen main.py, then runs a full layer-wise pass with retraining main.py.

BBP-MAP as implemented

The code works in natural parameters: each global unit stores a precision- weighted sum of its members plus the prior, and the final weight is that sum divided by the accumulated precision pfnm.py pfnm.py.

# BBP-MAP for one layer  [code: matching/pfnm.py#L213-L277]
# units[j]: (L_j, D+1) incoming weights with bias appended, per client
# s = 1/sigma2 (local precision), s0 = 1/sigma0_2 (prior precision)
# prior mean: 0 for weights, 0.1 for bias  [code: matching/pfnm.py#L838-L839]

def bbp_map(units, gamma, sigma2, sigma0_2, iterations=5):
    natural = [u / sigma2 for u in units]
    prior_nat = prior_mean / sigma0_2
    # start from the client with most units  [code: pfnm.py#L216-L228]
    order = sort_clients_by_unit_count_desc()
    sums = prior_nat + natural[order[0]]
    precisions = s0 + s                        # per row and dimension
    counts = ones(len(sums))
    for j in order[1:]:
        assign[j] = match_client(j, sums, precisions, counts)
    # refine: remove a client's units, re-match them  [pfnm.py#L241-L273]
    for it in range(iterations):
        for j in random_permutation(clients):
            remove_client(j, sums, precisions, counts)  # drop empty units
            assign[j] = match_client(j, sums, precisions, counts)
    # posterior mean of each global unit  [code: matching/pfnm.py#L1009]
    return sums / precisions, assign

def match_client(j, sums, precisions, counts):
    # gain of adding unit l to existing global unit i
    # [code: matching/pfnm.py#L14-L56]
    gain_old[l, i] = (
        sum_d((natural[j][l] + sums[i]) ** 2 / (s + precisions[i]))
        - sum_d(sums[i] ** 2 / precisions[i])
        + log(min(counts[i], 10) / (J - min(counts[i], 10))))
    # gain of opening new unit number k (k = 1..K)
    # [code: matching/pfnm.py#L61-L67]
    gain_new[l, k] = (
        sum_d((natural[j][l] + prior_nat) ** 2 / (s0 + s))
        - sum_d(prior_nat ** 2 / s0)
        - 2 * log(k) + 2 * log(gamma / J))
    # K = min(L_j, max(700 - L, 1))            # [code: pfnm.py#L62]
    # maximise total gain with a dense LAP solver  [code: pfnm.py#L91]
    rows, cols = lapsolver_solve_dense(-hstack(gain_old, gain_new))
    # matched: add natural params and precision; new: append unit
    # [code: matching/pfnm.py#L100-L111]
    return cols_mapped_to_global_indices

Complexity per layer: worst case (nothing matched) O(D\,(JL)^2) to build the cost matrix and O((JL)^3) for the Hungarian step; best case (everything matched) O(D L^2 + L^3) app. E. The authors use lapsolver (shortest-augmenting-path Hungarian), reporting orders-of-magnitude speedup over a vanilla implementation app. E.

Hyperparameters

VGG-9 architecture (all activations ReLU; conv shapes are (C_{in}, C_{out}, c, c), stride 1, padding 1) tab. 4 model.py:

LSTM: encoder (embedding) 80\times8, one LSTM layer with hidden size 256 (weight_ih 1024\times8, weight_hh 1024\times256, both biases 1024), decoder 80\times256 plus bias 80; 80 classes; 293k parameters tab. 5 tab. 3.

Results

All numbers are author-reported. Setup: simulated federation on AWS p3.2xlarge nodes, PyTorch, all clients sampled every round §3. Each FedMA "round" in the comparisons below counts one pass; for a fixed budget the paper allows 11 FedMA rounds versus 99 (VGG-9) and 33 (LSTM) rounds for FedAvg/FedProx §3.

Final global model sizes: VGG-9 grows from 3,491,530 to 3,900,235 parameters (conv channel widths 47, 79, 143, 143, 271, 271; FC widths 527, 527); the LSTM grows from 293,584 to 310,160 (embedding 8 to 21, hidden 256 to 257) tab. 6 tab. 7. Growth is largest in the first layers (conv1 1.47x, conv2 weights 1.81x) tab. 7.

Other findings, given only as plots:

• Communication: FedMA reaches higher accuracy than FedAvg and FedProx for the same number of rounds and, more clearly, for the same total message size in GB fig. 2.
• Local epochs: with E\in\{10,\dots,150\}, FedMA (6 rounds) improves with longer local training, FedAvg (54 rounds) degrades, FedProx degrades less §3 fig. 3.
• One-shot matching (PFNM extended to CNNs, all layers matched at once) does well on LeNet/MNIST but poorly on VGG-9/CIFAR-10; layer-wise FedMA does well on both and on LSTM §2.2 §2.3 fig. 1.
• Data bias: CIFAR-10 with 5 classes 95% grayscale and 5 classes 5% grayscale, split into a color-dominated and a grayscale-dominated client; FedMA with communication beats centralized training on the skewed data, FedAvg, FedProx and oversampling, while a color-balanced dataset does well §3 fig. 4. Centralized baselines with SGD vs AMSGrad: entire data 78.71% vs 75.91%, no-bias 91.23% vs 88.5%, color-balanced 87.79% vs 87.81% app. F. The FedMA number itself is only in the figure.
• Data efficiency: CIFAR-10 cut into 25 heterogeneous shards of about 2k points, added 5 clients at a time; one FedMA pass improves as clients are added while FedAvg with 9 rounds gets worse §3 fig. 6.
• Interpretability: matched first-layer filters (client 1 filter 0, client 2 filter 23) and their global filter respond to the same image feature, while the FedAvg filter 0 (average of both clients' filter 0) loses it §3 fig. 5.

Paper vs code

Checked at 4b586a5, mainly the CIFAR-10 path (main.py, matching/pfnm.py); the LSTM code was only spot-checked.

• Averaging rule. Paper averages matched units uniformly, \frac1J\sum_j W_{j,n}\Pi_j^T alg. 1 §2.1. Code returns the BBP posterior mean: (prior mean / \sigma_0^2 + sum of members / \sigma^2) / (1/\sigma_0^2 + count/\sigma^2) pfnm.py pfnm.py. With \sigma^2=\sigma_0^2=1, a unit matched by m clients is shrunk by m/(m+1) toward the prior mean (0, or 0.1 for biases); an unmatched unit is halved. The average is also unweighted by data size.
• \gamma in FedMA with communication. fedma_comm sets gamma = 5.0 main.py, but it calls BBP_MAP, which resets gamma = 7.0 main.py; the effective value matches the paper's 7 app. D.2.
• Local training weight decay. Paper lists \ell_2 10^{-4} for all CNN training tab. 3; the initial local training in code uses SGD with momentum and no weight decay on CIFAR-10 (decay only for CINIC-10) and a constant lr main.py. FedAvg/FedProx retraining in code does use 10^{-4} main.py main.py.
• Last-layer lr. Paper: 10^{-4} app. A and 0.001 tab. 3. Code: retrain_lr / 10 = 0.001 with the shipped --retrain_lr=0.01 main.py run.sh.
• Epochs and rounds. Paper's best FedMA E is 150 tab. 1; run.sh uses --epochs=20 --retrain_epochs=20 --comm_round=10 and the README example uses 50 rounds run.sh.
• LSTM lr. Paper: constant 0.8 tab. 3; code: 0.25 with gradient clip 0.25 lstm_fedma_with_comm.py.
• Shakespeare filter. Paper keeps clients with at least 10k data points §3; the README preprocessing sets a minimum of 9K samples per user.
• Heterogeneous partition. Paper: sample \mathrm{Dir}_J(0.5) per class §3. Code also zeroes the proportion for any client that already holds at least N/J points and resamples the whole split until every client has at least 10 points utils.py.
• Matching order and caps (absent from the paper): clients initialized largest first, then 5 random-order remove-and-rematch sweeps pfnm.py; popularity counts capped at 10 in the cost and at most about 700 global units per layer pfnm.py pfnm.py.

Pitfalls

• "Rounds" is ambiguous. One FedMA pass needs one round per layer §2.3, and the budget "11 rounds for FedMA vs 99/33 for FedAvg/FedProx" §3 equals 11 passes times 9 (VGG-9) or 3 (LSTM) layers. So read FedMA "rounds" in tab. 1 tab. 2 fig. 2 as passes. This is our inference from the numbers, not stated.
• Algorithm 1 has typos. The retraining line trains up to W_{j,L} where N is meant; the last-layer formula indexes W_{jl,n} and calls the weight p_k while summing p_{jk} alg. 1. The RNN average writes H_h for H_j §2.2. In code the last-layer weight for client j and class k is client j's count of class k divided by the total count of class k, applied to class k's output column including its bias main.py main.py.
• Global layers grow, so every later layer changes shape. Clients must pad or permute the input dimension of layer n+1 to the matched width. For the first FC layer after a conv stack this is a block permutation of spatial positions per channel, not a simple row permutation pfnm.py. Zero "dummy" units must be left out of the averages §2.1.
• Clients train a model larger than their own after each matched layer. With communication, models are cut back to the original size from the client's own assignment before the next round §2.3; without that, the global model would keep growing §2.3.
• Posterior-mean shrinkage (see Paper vs code): implementing plain averages per alg. 1 gives different weights from the released code (factor m/(m+1) for a unit shared by m clients when both variances are 1); decide which one to reproduce.
• Hidden cost. Each FedMA pass runs N-1 local retraining jobs per client alg. 1, plus a full retrain per round in the with-communication variant main.py. Matching cost is cubic in JL in the worst case app. E; a fast LAP solver is needed in practice app. E.
• No batch norm. All BN layers were removed from VGG-9 and left to future work §3; it is not stated how running statistics would be matched.
• LSTM matching ignores hidden-to-hidden weights in the cost; they are only permuted and averaged afterwards §2.2.
• BBP-MAP terms are not written out in the paper; the cost in the Method pseudocode comes from the code only pfnm.py and follows PFNM.
• Hyperparameters were tuned per method: E is picked per method by best final test accuracy on the global test set §3; FedMA's best E (150) is much larger than the baselines' (20) tab. 1.
• Last-layer retraining is unstable on VGG-9 at the initial lr; the authors lowered it by 10x and trained 3x longer app. A.

Lineage

• Builds on: PFNM (Yurochkin et al., 2019; Bayesian nonparametric neuron matching for fully connected nets), whose BBP-MAP solver is reused unchanged and extended to CNN channels and LSTM hidden states, then wrapped in layer-wise matching with retraining §1 §2.1.
• Compared against: FedAvg (McMahan et al., 2017, arXiv:1602.05629), FedProx (Sahu et al., 2018, arXiv:1812.06127), uniform ensemble of client models, and centralized training §3.
• Related: concurrent model fusion via Wasserstein barycenters (Singh and Jaggi, 2019) §2.1.