PaperlayerView exact revision on arXiv

Source-linked AI summary

FedProto: Federated Prototype Learning across Heterogeneous Clients

Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, Chengqi Zhang

arXiv:2105.00243v4cs.LGcs.DC

TL;DR

Heterogeneous federated learning clients may have incompatible model architectures, data distributions, and input/output spaces, limiting gradient-based aggregation. FedProto instead aggregates class prototypes and regularizes local training with global prototypes, achieving higher test accuracy with lower communication costs across heterogeneous settings.

Abstract

from arXiv · show

Heterogeneity across clients in federated learning (FL) usually hinders the optimization convergence and generalization performance when the aggregation of clients' knowledge occurs in the gradient space. For example, clients may differ in terms of data distribution, network latency, input/output space, and/or model architecture, which can easily lead to the misalignment of their local gradients. To improve the tolerance to heterogeneity, we propose a novel federated prototype learning (FedProto) framework in which the clients and server communicate the abstract class prototypes instead of the gradients. FedProto aggregates the local prototypes collected from different clients, and then sends the global prototypes back to all clients to regularize the training of local models. The training on each client aims to minimize the classification error on the local data while keeping the resulting local prototypes sufficiently close to the corresponding global ones. Moreover, we provide a theoretical analysis to the convergence rate of FedProto under non-convex objectives. In experiments, we propose a benchmark setting tailored for heterogeneous FL, with FedProto outperforming several recent FL approaches on multiple datasets.

Introduction

The paper proposes prototype aggregation-based federated learning to address statistical and model heterogeneity without aggregating model parameters or gradients. It introduces a heterogeneous-FL benchmark, emphasizes communication efficiency, and provides a convergence guarantee.

Motivation

  • Federated learning trains a global model from private data distributed across local clients, which cannot be freely transmitted for privacy or communication concerns.Source: Introduction · p. 1

Limitations of existing methods

  • Gradient-based aggregation incurs high communication costs and relies heavily on homogeneous local models, limiting many existing heterogeneous and personalized FL methods.Source: Introduction · p. 1
  • Knowledge-distillation FL supports different architectures but requires an extra public dataset to align outputs, increasing computation costs.Source: Introduction · p. 1

FedProto

  • The proposed framework transmits only prototypes, allowing clients with different model architectures and input/output spaces to exchange information without aggregating parameters or gradients.Source: Introduction · pp. 1–2
  • Each class prototype is the mean representation of same-class observed samples, and prototype aggregation enables efficient communication across heterogeneous clients.Source: Introduction · p. 2

Contributions

  • The paper contributes a heterogeneous-FL benchmark, a prototype aggregation method intended to improve communication efficiency, and a theoretical convergence guarantee.Source: Introduction · p. 2

Related Work

Prior federated learning work addresses statistical heterogeneity through local regularization, personalized models, or multiple global models, but generally relies on gradient-based aggregation. Prototypes, commonly class-level means of feature vectors, are well suited to cross-client federated learning because clients typically have limited training data.

Federated learning methods

  • FedProx uses local regularization, while other methods train personalized models or provide multiple global models to address heterogeneous federated learning.These approaches aim to combine globally shared information with client-specific components.Source: Heterogeneous Federated Learning · p. 2

Limitations of prior work

  • Most existing methods target only one heterogeneous scenario and use gradient-based aggregation, raising communication-efficiency and gradient-based-attack concerns.The cited concerns include communication efficiency and vulnerabilities to gradient-based attacks.Source: Heterogeneous Federated Learning · p. 2

Prototype learning

  • Prototypes represent classes in image classification by averaging feature vectors within each class and are also used as averaged representations in action recognition.The passage describes prototypes as means of multiple features across several learning tasks.Source: Prototype Learning · p. 2
  • Prototypes are widely used when training samples are limited, matching the cross-client federated learning assumption that each client has too few instances to train a high-performing model independently.This limited-data assumption is supported by federated learning benchmark datasets and applications such as healthcare.Source: Prototype Learning · p. 2

Problem Setting

Heterogeneous federated learning allows clients to differ in data distributions, input/output spaces, and model architectures, making conventional global-model aggregation problematic. FedProto addresses this challenge by communicating and aggregating class prototypes instead of model parameters.

Heterogeneity

Aggregation challenge

  • Different model architectures can give local parameters different formats and sizes, preventing optimization of a global parameter through averaging.This limitation motivates communicating and aggregating prototypes in federated learning.Source: Heterogeneous Federated Learning Setting · p. 3

Prototype definition

Prototype-based model inference

  • Prototype-based inference predicts an instance’s label by measuring the L2 distance between its representation and each class prototype.The instance representation is produced by the embedding function f(φ; x).Source: Prototype-Based Aggregation Setting · p. 3

Methodology

FedProto addresses heterogeneous federated learning by exchanging and aggregating class prototypes rather than gradients or model parameters. Clients jointly minimize local classification loss while regularizing local prototypes toward global prototypes, with convergence analyzed for non-convex objectives.

Prototype communication

  • FedProto exchanges class prototypes between clients and server, averaging overlapping but non-identical local prototype sets by class.Clients may recognize different class subsets, such as digits {2, 3, 4} and {4, 5}.Source: Methodology · p. 3
  • Prototype-based communication avoids exchanging gradients or model parameters, supporting heterogeneous model architectures and class spaces.Clients can share an embedding space when label spaces align, even when their optimal model parameters differ.Source: Methodology · p. 4, Global Prototype Aggregation · p. 4

Optimization objective

  • FedProto minimizes each client’s supervised classification loss plus a regularizer measuring distance between local and corresponding global prototypes.The regularization term uses L2 distance in the objective, while the distance metric can also be L1 or earth mover’s distance.Source: Optimization Objective · p. 4, Local Model Update · p. 4

Optimization procedure

  • Training alternates local supervised updates with server aggregation of local prototypes, then returns updated global prototypes to clients for subsequent rounds.Each client computes prototypes and loss over local batches, updates its model, and returns its prototype set.Source: Optimization Objective · p. 4, Local Model Update · p. 4

Convergence analysis

  • Under four smoothness, stochastic-gradient, boundedness, and embedding-continuity assumptions, FedProto’s loss decreases per communication round and converges for suitable η and λ.Theorem 2 bounds the expected gradient norm by any ϵ after selecting communication rounds and hyperparameters; smaller ϵ requires more rounds.Source: Convergence Analysis · pp. 4–5

Discussion

FedProto’s prototype-based communication offers advantages in model inference, communication efficiency, and privacy preservation. New clients can adapt pretrained representations using global class prototypes, while prototypes are substantially smaller than model parameters and harder to reverse into raw data.

Model inference

  • New clients can initialize representation layers from a pretrained model, download relevant global class prototypes, and fine-tune local models with random decision layers.The example uses a ResNet18 pretrained on ImageNet.Source: Model Inference · p. 5

Communication efficiency

  • 50 prototype dimensions per class versus 21,500 model parameters on MNIST illustrates FedProto’s communication-efficiency advantage.FedProto transmits prototypes rather than model parameters between clients and server.Source: Communication Efficiency · p. 5

Privacy preserving

  • Averaged low-dimensional prototype vectors support privacy preservation because the averaging process is irreversible and attackers cannot reconstruct raw data without local models.Prototypes are 1D-vectors generated from representations of samples in the same class.Source: Privacy Preserving · p. 5

Experiments

Experiments evaluate FedProto in heterogeneous federated settings using benchmark datasets, varied client tasks, and comparisons with established FL baselines. FedProto achieves the highest accuracy and least variance in most cases, supporting uniformity among heterogeneous clients.

Datasets and local models

  • Experiments use MNIST, FEMNIST, and CIFAR10 with CNN models for the first two datasets and ResNet18 for CIFAR10.All clients own local data and communicate with a central server.Source: Training Setups · p. 6

Local tasks

Baselines of FL

Implementation Details

  • Experiments use 20 clients for every dataset, sample all clients in each communication round, and set the average size of each class per client to 100.Implementations use PyTorch; CIFAR10 initializes ResNet18 with ImageNet pre-training.Source: Training Setups · p. 6

Results

  • FedProto achieves the highest average test accuracy and least variance in most cases, ensuring uniformity among heterogeneous clients.The average test accuracy is reported over all clients in Table 1.Source: Performance in Non-IID Federated Setting · p. 7

Conclusion

FedProto addresses heterogeneous federated learning by aggregating and exchanging class prototypes instead of gradients. Its effectiveness is analyzed theoretically and experimentally.

Conclusion

  • FedProto tackles heterogeneous input/output spaces, data distributions, and model architectures through prototype aggregation rather than gradient exchange.The method collaboratively trains intelligent models by communicating prototypes between clients and server.Source: Conclusion · p. 7
  • The proposed method’s effectiveness is comprehensively evaluated from both theoretical and experimental perspectives.Source: Conclusion · p. 7

Experimental Details and Extra Results

Experiments evaluate FedProto and FedProto-mh on MNIST, FEMNIST, and CIFAR10 under heterogeneous settings, showing higher accuracy with lower communication costs than existing FL methods. Additional results indicate applicability to model-heterogeneous scenarios and similar performance to homogeneous settings.

Experimental setup

  • Clients use SGD with learning rate 0.01, momentum 0.5, one local epoch, and batch size 8; the best λ values are 1 for MNIST, 1 for FEMNIST, and 0.1 for CIFAR10.The best λ is selected by grid search from a limited candidate set, and heterogeneity is controlled by the standard deviation of n.Source: Experimental Details · p. 10

Overall results

  • FedProto and FedProto-mh outperform existing FL methods in test accuracy while reducing communication costs across MNIST, FEMNIST, and CIFAR10.The framework also supports model-heterogeneous scenarios with performance similar to homogeneous scenarios.Source: Extra Results · p. 10, Table 2 · p. 10, Table 3 · p. 11

MNIST

  • On MNIST, FedProto achieves higher local average test accuracy than FeSEM, FedProx, FedPer, and FedAvg in all evaluated settings with low communication cost.Table 2 compares methods under a non-IID split over clients.Source: Extra Results · p. 10, Table 2 · p. 10

FEMNIST

  • On FEMNIST, FedProto has much smaller accuracy variance across clients than other FL methods and communicates around 0.025% of the total parameters.The evaluation considers client heterogeneity levels with standard deviation of n equal to 1 and 2.Source: Extra Results · p. 10, Table 3 · p. 11

CIFAR10

  • On CIFAR10, FedProto converges faster in most heterogeneous cases, while FedProto and FedProto-mh communicate far fewer parameters per round than baseline methods.Table 4 reports comparisons under a non-IID split over clients.Source: Extra Results · p. 10, Table 4 · p. 11

Convergence Analysis for FedProto

FedProto’s convergence analysis models local embedding and decision functions, then derives convergence under smoothness, stochastic-gradient, gradient-norm, and embedding-continuity assumptions. The results establish monotonic loss decrease and a non-convex convergence rate across communication rounds.

Setup

  • The analysis represents each client with an embedding function and a shared-form decision function, while local loss combines classification error with prototype regularization.The embedding maps inputs to prototypes, and the regularizer penalizes distance from the global prototype.Source: Additional Notation · p. 12

Assumptions

  • The convergence proof assumes Lipschitz-smooth local objectives, unbiased stochastic gradients with bounded variance, bounded expected gradient norms, and Lipschitz-continuous embeddings.The embedding-continuity assumption is acknowledged as strong but is used only over a narrow E-step SGD interval.Source: Assumptions · p. 12

Lemmas

  • Lemma 1 bounds an arbitrary client’s loss during local updates by applying smoothness, stochastic-gradient assumptions, expectation, and telescoping over E local steps.This lemma analyzes the interval from the beginning of communication round t + 1 through the final local update.Source: Key Lemmas · p. 13
  • Lemma 2 bounds post-aggregation loss by relating changes in prototype regularization to global and local prototype definitions and embedding continuity.The bound is stated after server-side prototype aggregation for an arbitrary client.Source: Key Lemmas · pp. 13–14

Convergence Results

  • Theorem 1 gives an arbitrary client’s one-round deviation under Assumptions 1 to 4, while Corollary 1 states that FedProto’s non-convex loss monotonically decreases when its condition holds.The corollary concludes that the loss function converges.Source: Theorems · pp. 14–15
  • Theorem 2 states a non-convex convergence rate for FedProto under Assumptions 1 to 4, parameterized by the initial-to-optimal loss gap Δ = L0 − L∗ and any ϵ > 0.Its proof telescopes expected losses across communication rounds and local timesteps.Source: Theorems · p. 15, Completing the Proof of Theorem 2 · p. 15

Related papers

Loading 2105.00243v4…