FedProto: Federated Prototype Learning across Heterogeneous Clients

In plain words

Federated learning trains a model across many devices (phones, sensors, hospitals) without collecting their data in one place. The usual recipe sends each device's model weights to a server, averages them, and sends the average back. That breaks down when devices see different classes, hold very different amounts of data, or run different network architectures: weights of different shapes cannot be averaged at all.

FedProto sends something much smaller instead. Each device runs its own network up to an embedding layer and computes, for every class it knows, the average embedding of its examples: a "prototype", like each person's mental picture of a dog. The server averages the prototypes per class across devices and sends these global prototypes back. Each device then keeps training its own classifier, with an extra penalty that pulls its embeddings toward the shared prototype of the right class.

Because only one vector per class travels over the network, communication drops from hundreds of thousands of parameters per round to a few thousand, and every device can use its own architecture as long as the embedding size matches. In the authors' experiments on MNIST, FEMNIST and CIFAR10 with 20 clients, this matched or beat weight-averaging methods while sending far less. The price is that there is no shared global classifier: every client keeps its own model.

TL;DR

Clients share class prototypes (the mean embedding per class, eq. 3) instead of gradients. The server forms global prototypes by per-class aggregation (eq. 6); each client minimizes cross-entropy plus λ times the distance between its prototypes and the global ones (eq. 7). Clients may differ in label space, data size and architecture, provided the embedding dimension is shared. Author-reported results with 20 clients: best or near-best accuracy on MNIST, FEMNIST and CIFAR10 with far lower variance across clients and about 4×10³ (MNIST, FEMNIST) or 41×10³ (CIFAR10) parameters communicated per round tab. 1. The official code differs from the paper in several material ways (aggregation weights, per-sample regularizer, evaluation bookkeeping); read "Paper vs code" before reproducing.

Problem

Weight-averaging FL (FedAvg, eq. 1) assumes all clients share one architecture and that a single parameter vector suits all of them §Problem Setting eq. 1. Under statistical heterogeneity each client's distribution P_i differs, for example each client sees only a subset of classes; under model heterogeneity each client has its own F_i, so the ω_i have different shapes and cannot be averaged eq. 2 §Problem Setting. Knowledge-distillation FL handles different architectures but needs a public dataset and degrades when that dataset diverges from the non-IID client data §Introduction. FedProto targets both kinds of heterogeneity without a public dataset and with low communication §Introduction.

Method

Model split. Each client's model is F_i = g_i(ν_i) ∘ f_i(φ_i): representation layers f_i map input x to an embedding h_i = f_i(φ_i; x), and decision layers g_i produce the prediction §Problem Setting. The embedding dimension d_c must be the same for all clients; f_i itself may differ app. B.

Local prototype. For client i and class j, the prototype is the mean embedding of the client's training examples of that class eq. 3:

C_i^{(j)} = \frac{1}{|D_{i,j}|} \sum_{(x,y)\in D_{i,j}} f_i(\phi_i; x) eq. 3

Prototype inference. A sample can be classified by its nearest prototype under L2 distance eq. 4:

\hat{y} = \arg\min_j \lVert f(\phi; x) - C^{(j)} \rVert_2 eq. 4

Global objective. Minimize the size-weighted supervised loss over clients plus a regularizer between local and global prototypes, weighted by class counts eq. 5:

\arg\min_{\{\bar{C}^{(j)}\}} \sum_{i=1}^m \frac{|D_i|}{N} \mathcal{L}_S(\mathcal{F}_i(\omega_i;x), y) + \lambda \sum_{j=1}^{|\mathbb{C}|} \sum_{i=1}^m \frac{|D_{i,j}|}{N_j} \mathcal{L}_R(\bar{C}_i^{(j)}, C_i^{(j)}) eq. 5

N is the total number of instances, N_j the number of class-j instances over all clients, and L_R is the L2 distance §Methodology. It is solved by alternating minimization: (1) each client updates ω_i with the global prototypes fixed; (2) the server updates the global prototypes with all ω_i fixed §Methodology.

Global aggregation. With 𝒩_j the set of clients that have class j eq. 6:

\bar{C}^{(j)} = \frac{1}{|\mathcal{N}_j|} \sum_{i \in \mathcal{N}_j} \frac{|D_{i,j}|}{N_j} C_i^{(j)} eq. 6

Local loss. Cross-entropy plus λ times the prototype regularizer eq. 7 eq. 8:

\mathcal{L}(D_i, \omega_i) = \mathcal{L}_S(\mathcal{F}_i(\omega_i;x), y) + \lambda \cdot \mathcal{L}_R(\bar{C}_i, C_i), \qquad \mathcal{L}_R = \sum_j d(C_i^{(j)}, \bar{C}_i^{(j)}) eq. 7 eq. 8

d can be L1, L2 or earth mover's distance §Methodology; the main text uses L2 §Methodology, while the appendix analysis uses the squared L2 distance between each sample's embedding and the global prototype, λ‖f_i(φ_i; x) − C̄‖²₂ app. B. The code uses the latter form, as a mean squared error (see "Paper vs code").

Algorithm, as written in the paper alg. 1, with the code's concrete choices noted:

# FedProto  [alg. 1]
# Input: per-client data D_i and local models (f_i, g_i), i = 1..m
# Output: per-client models and global prototypes {class: vector of size d_c}

def server(clients, num_rounds):
    # Alg. 1 initializes global prototypes; the code starts with none
    # [alg. 1] [code: exps/federated_main.py#L39]
    global_protos = {}
    for round in range(num_rounds):
        local_protos = {}
        # every client takes part in every round  [§Experiments]
        for client in clients:
            local_protos[client.id] = client.local_update(global_protos)
        # per-class aggregation over clients that have the class  [eq. 6]
        global_protos = aggregate(local_protos)
    return global_protos

def aggregate(local_protos):
    # code: unweighted mean over clients holding class j, not the
    # |D_ij|/N_j weights of eq. 6  [eq. 6] [code: lib/utils.py#L221-L240]
    by_class = group_by_class(local_protos)
    return {j: mean(vectors) for j, vectors in by_class.items()}

def local_update(self, global_protos):
    # SGD, lr 0.01, momentum 0.5; 1 local epoch  [app. A]
    for epoch in range(local_epochs):
        for images, labels in batches(self.data):
            # log-probs and embeddings from the same forward pass
            # [code: lib/models/models.py#L54-L61]
            log_probs, embeddings = self.model(images)
            # supervised term L_S: negative log-likelihood  [eq. 7]
            loss_ce = nll(log_probs, labels)
            # targets: the global prototype of each sample's class; a sample
            # whose class has no global prototype targets itself (zero term)
            # [code: lib/update.py#L162-L171]
            targets = copy(embeddings).detach()
            for row, label in enumerate(labels):
                if label in global_protos:
                    targets[row] = global_protos[label]
            # regularizer L_R as mean squared error  [code: lib/update.py#L171]
            loss_proto = mse(embeddings, targets)
            # total loss with importance weight lambda  [eq. 7]
            loss = loss_ce + lam * loss_proto
            loss.backward()
            optimizer.step()
            # code keeps these training-pass embeddings for the local
            # prototypes  [code: lib/update.py#L177-L181]
            store_embeddings_by_label(embeddings, labels)
    # mean embedding per class  [eq. 3] [code: lib/utils.py#L205-L219]
    return mean_per_class(stored_embeddings)

Convergence. Under four assumptions (L1-smooth local objective, unbiased stochastic gradients with variance ≤ σ², expected gradient norm ≤ G, L2-Lipschitz embedding function) eq. 9–eq. 14, the expected one-round change of a client's loss obeys eq. 15:

\mathbb{E}[\mathcal{L}_{(t+1)E+1/2}] \le \mathcal{L}_{tE+1/2} - \left(\eta - \frac{L_1\eta^2}{2}\right) \sum_{e=1/2}^{E-1} \lVert \nabla \mathcal{L}_{tE+e} \rVert_2^2 + \frac{L_1 E \eta^2}{2}\sigma^2 + \lambda L_2 \eta E G eq. 15

The last term is what the prototype update adds, through the change in global prototypes between rounds app. B. The loss decreases every round if η and λ meet the bounds of Corollary 1 eq. 16 eq. 17, and after

T = \frac{2\Delta}{E\epsilon(2\eta - L_1\eta^2) - E\eta(L_1\eta\sigma^2 + 2\lambda L_2 G)} eq. 18

rounds the averaged squared gradient norm is below ε eq. 19, provided η < 2(ε − λL₂G) / (L₁(ε + σ²)) and λ < ε / (L₂G) §Methodology. Δ = L₀ − L* §Methodology.

New clients. A new client can start from a pretrained representation (for example an ImageNet ResNet18) with random decision layers, download the global prototypes of its classes and fine-tune on the local objective §Discussion.

Hyperparameters

Results

Author-reported. Setup: 20 clients, all sampled every round, non-IID class subsets drawn per client with average n classes and about 100 samples per class §Experiments. Metric: mean ± std of test accuracy over clients (%), plus communication rounds and parameters sent per round (×10³) tab. 1. FedProto-mh uses heterogeneous architectures; all other rows use one architecture. Main table tab. 1:

Bold as in the paper. The appendix adds a second heterogeneity level per dataset tab. 2 tab. 3 tab. 4; rows for that level (rounds and parameters are the same as above):

The FedPer CIFAR10 n=5 entry is printed as "72,43" in the source tab. 4.

Other findings:

• λ sweep (FEMNIST, n=3, k=100, λ in [0, 4]): best λ = 1; accuracy rises sharply from λ = 0 to λ = 1, then falls by about 6 points as λ grows further, while the prototype-distance loss keeps decreasing fig. 4.
• Fewer samples per class (MNIST): accuracy of both FedProto and FedAvg drops as k shrinks, FedProto more slowly fig. 3.
• t-SNE (MNIST, 20 clients, n=3): under FedProto, samples of one class form several nearby but separable clusters, one per client, around the global prototype; FedAvg merges each class into one cluster and FedPer separates by client fig. 2.
• FEMNIST communication: about 0.025% of the parameters communicated by the baselines app. A.

Paper vs code

Checked against 97e42e2.

• Aggregation weights. Eq. 6 weights each client's prototype by |D_ij|/N_j and also divides by |𝒩_j| eq. 6. The code takes a plain unweighted mean over clients that hold the class utils.py.
• Regularizer target. Eq. 7 and Alg. 1 compute the client's local prototype and compare it with the global one eq. 7 alg. 1. The code compares each sample's embedding with the global prototype of its label via MSELoss, i.e. squared error averaged over batch and embedding dimension update.py; this matches the appendix form λ‖f_i(φ_i; x) − C̄‖²₂ up to the averaging app. B, not the main-text L2 distance between prototypes §Methodology.
• When local prototypes are computed. Eq. 3 is a mean over the whole local class subset eq. 3. The code averages the embeddings from the training forward passes of the last local epoch, collected while weights change and in train mode, so the Dropout2d after conv2 also perturbs them update.py models.py. With drop_last=True some samples never contribute update.py.
• Initialization. Alg. 1 initializes the global prototypes alg. 1; the code starts empty, so round 1 trains with cross-entropy only update.py.
• Batch size. Paper: 8 for all datasets app. A. Code default: 4 options.py; the README's CIFAR10 command uses 32.
• ImageNet initialization. Paper: ResNet18 pretrained on ImageNet §Experiments. The code loads ImageNet weights but keeps random weights for conv1, bn1 and fc federated_main.py.
• Per-client accuracy. The evaluation keeps correct and total as running sums across all clients and across the "without prototypes" and "with prototypes" passes, so each client's printed accuracy is a cumulative average over itself and all earlier clients update.py. The ± values reported in the tables are standard deviations over clients tab. 1; with this code, they would be computed over running averages, which shrinks them. Which of the two printed accuracies feeds the tables is not stated.
• Heterogeneity noise. Paper: random noise with a given standard deviation added to n and k §Experiments. Code: n is a uniform integer in [max(2, n − s), min(C − 1, n + s)] and k a uniform integer in a narrow range around 100 (MNIST: [k − s + 1, k + s − 2]; FEMNIST and CIFAR10: [k − s + 1, k + s]) federated_main.py.
• Baseline numbers. Tab. 1 lists FedAvg on MNIST (stdev 2, n=3) as 95.04±6.48, while appendix tab. 2 lists the same cell as 91.40±6.48 tab. 1 tab. 2.
• The code also contains baseline and FedProx helpers that are not used by the FedProto entry point; the FedProx term there is computed between state-dict copies and does not reach the gradient update.py.

Pitfalls

• Embedding dimension must match across clients, even in the model-heterogeneous setting app. B. In the code the heterogeneous CNNs differ only in conv2 channels and resize fc1's input, keeping a 50-dim prototype models.py.
• The prototype is a feature map for ResNet. The code takes layer4's output before average pooling resnet.py. On 32×32 CIFAR10 inputs with the strides used, this is 512×1×1; with other input sizes the "prototype" becomes a spatial tensor.
• Eq. 6 is not a convex combination. The weights |D_ij|/N_j already sum to 1 over 𝒩_j, so the extra 1/|𝒩_j| shrinks every global prototype toward zero by the number of contributing clients eq. 6. Using it literally would make the regularizer pull embeddings toward a scaled-down target; the code ignores both factors utils.py.
• Classes without a global prototype contribute zero regularization but still count in the MSE denominator, so the effective λ depends on the share of covered samples in a batch update.py.
• λ scale depends on the distance form. Mean squared error per dimension (code), squared L2, and L2 give losses of very different magnitudes; the reported λ values (1, 1, 0.1) app. A are tied to the code's MSELoss.
• No global classifier. The served "global model" is a set of prototypes; each client keeps its own model, and test sets are drawn from each client's own classes §Discussion sampling.py. Accuracies are therefore personalized, per-client numbers, not accuracy on the full test set.
• Prototype-based prediction is restricted to local classes in the code: classes outside the client's set get a fixed distance of 100 update.py.
• Communication counts are per round and depend on how many classes clients hold; the exact counting rule behind the 4×10³ and 41×10³ entries is not stated tab. 1.
• Privacy claims are informal. The paper argues prototypes are hard to invert without the local model §Discussion; no attack or formal guarantee is evaluated.
• Theory vs practice. The proof's Lemma 2 uses an unsquared norm when differencing the loss across a prototype update, while the local loss it analyses uses the squared norm app. B. The bounds involve unknown constants (L₁, L₂, G, σ²), so they do not yield usable η or λ.
• Learning-rate schedule, weight decay, data augmentation for CIFAR10 and the seeds behind the ± values are not stated in the paper.

Lineage

• Builds on: FedAvg (arXiv:1602.05629) for the federated protocol and default hyperparameters; prototypical networks (arXiv:1703.05175) for class prototypes as mean embeddings.
• Compared against: Local training, FedAvg, FedProx (arXiv:1812.06127), FeSEM, FedPer (arXiv:1912.00818) and FedRep (arXiv:2102.07078) tab. 1.