PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Learning from Pre-Trained Models: A Contrastive Learning Approach

Yue Tan, Guodong Long, Jie Ma, Lu Liu, Tianyi Zhou, Jing Jiang

arXiv:2209.10083v1cs.CRcs.AIcs.LG

TL;DR

Large-scale models make federated learning costly in computation and communication, motivating a lightweight framework based on fixed pre-trained backbones. FedPCL fuses their representations through client-specific projection networks and prototype-wise contrastive learning, with experiments showing improved personalization and knowledge integration over baselines.

  • Problem

    Large neural networks trained from scratch impose excessive computation and communication demands, while heterogeneous client data can limit a single global model’s local performance.

  • Method

    The framework fixes multiple pre-trained backbones, learns client-specific projection networks to fuse their representations, and uses class prototypes with prototype-wise contrastive learning for federated knowledge sharing.

  • Results

    FedPCL outperforms baselines with better personalization and knowledge integration across experiments evaluating representation fusion from varied pre-trained models and datasets.

  • Takeaways & Limitations

    Sharing prototypes instead of learnable model parameters keeps shared knowledge compact while enabling clients to personalize fused representations under the lightweight framework.

  • Takeaways & Limitations

    The study mainly focuses on vision tasks, leaving methods for language tasks for future work.

Abstract

from arXiv · show

Federated Learning (FL) is a machine learning paradigm that allows decentralized clients to learn collaboratively without sharing their private data. However, excessive computation and communication demands pose challenges to current FL frameworks, especially when training large-scale models. To prevent these issues from hindering the deployment of FL systems, we propose a lightweight framework where clients jointly learn to fuse the representations generated by multiple fixed pre-trained models rather than training a large-scale model from scratch. This leads us to a more practical FL problem by considering how to capture more client-specific and class-relevant information from the pre-trained models and jointly improve each client's ability to exploit those off-the-shelf models. In this work, we design a Federated Prototype-wise Contrastive Learning (FedPCL) approach which shares knowledge across clients through their class prototypes and builds client-specific representations in a prototype-wise contrastive manner. Sharing prototypes rather than learnable model parameters allows each client to fuse the representations in a personalized way while keeping the shared knowledge in a compact form for efficient communication. We perform a thorough evaluation of the proposed FedPCL in the lightweight framework, measuring and visualizing its ability to fuse various pre-trained models on popular FL datasets.

1 Introduction

The paper proposes a lightweight federated learning framework that fuses representations from fixed pre-trained backbones, reducing training costs while enabling personalized, class-relevant knowledge sharing through prototypes and contrastive learning.

  • FL enables collaborative training without sharing private client data, but statistical heterogeneity can make a single global model perform poorly across clients.
  • Large neural networks trained from scratch impose high computation and communication costs, limiting practical deployment on resource-constrained devices.
  • The lightweight framework fixes multiple pre-trained backbones and learns client-specific layers to fuse their representations, avoiding large-scale backpropagation and parameter transmission.
  • Class-wise prototypes serve as compact communication carriers for sharing common, class-relevant knowledge while allowing personalized representation learning across clients.
  • Compared with training ResNet18 from scratch, the fixed-backbone setting requires less training time per round and fewer learnable parameters under the stated Digit-5 evaluation.
  • FedPCL applies prototype-wise contrastive learning during local updates using local and global prototypes, improving personalization and knowledge integration over baselines.

2 Related Work

Related work covers personalized federated learning, contrastive learning, prototype learning, and pre-trained foundation models as foundations for the proposed approach.

  • Personalized Federated Learning: Personalized federated learning trains client-specific models to address local performance differences caused by heterogeneous private datasets.
  • Personalized Federated Learning: Existing personalized FL methods use personalized objectives, partial model sharing, hypernetworks, adaptive aggregation, and meta-learning.
  • Contrastive Learning: Contrastive learning brings embeddings of similar samples closer and pushes embeddings of different samples apart, supporting representation learning.
  • Prototype Learning: Prototypes are commonly used across transfer, multi-task, and few-shot learning, often representing the mean feature vector of samples in a class.
  • Pre-Trained Foundation Model: Pre-trained foundation models leverage large parameter counts and broad training data to capture downstream knowledge and support green AI.

3 Problem Formulation

The framework uses fixed pre-trained backbones to produce concatenated client representations, then learns client-specific projection networks that fuse them into task representations within federated optimization.

  • The global objective follows federated learning by aggregating client losses over private datasets, with model parameters potentially synchronized fully or partially.
  • Each client owns K shared, fixed pre-trained backbones and a private dataset, while the framework separates feature encoding from representation projection.
  • The optimization target is a personalized projection network for every client rather than a newly trained large feature extractor.
  • The feature encoder concatenates the representations generated by the pre-trained backbones for each input sample.
  • Each client’s projection network h(θ_i) maps the concatenated backbone representation into another abstract representation space.
  • The multi-backbone framework can accommodate FedAvg and personalized FL methods by sharing or adapting the learnable projection parameters.

4 Federated Prototype-wise Contrastive Learning (FedPCL)

FedPCL uses compact prototype-based communication and prototype-wise contrastive learning to personalize representation fusion from multiple fixed pre-trained backbones. Clients learn from global and local class information while avoiding parameter synchronization.

  • Framework: FedPCL fuses representations from multiple fixed pre-trained backbones through client-specific projection networks.The projection network generates fused representations while the backbones remain fixed.
  • Prototype-based Communication: Class-wise prototypes replace learnable model parameters as the communication carrier, reducing communication costs and supporting personalized models.Prototypes are compact, non-parametric, high-level statistics that allow customized local models without parameter synchronization.
  • Prototype-based Communication: Each client computes local class prototypes from the mean fused representation of samples sharing class j and uploads them for server aggregation.The server aggregates local prototype sets into global class prototypes that summarize shared class-relevant knowledge.
  • Prototype-based Communication: The server returns global and complete local prototype sets to clients, padding missing classes when non-IID clients contain only a few classes.This procedure ensures local prototype sets cover all classes before local training.
  • Local Training: The local objective combines global and local prototype-wise supervised contrastive losses.Both terms pull fused representations toward positive class prototypes and push them away from negative prototypes, while providing different information granularities.
  • Prototype-based Inference: At inference, each client predicts the class whose local prototype has the greatest similarity to the projected test representation.The local prototypes contain compact client-specific information while absorbing general knowledge from participating clients.

5 Experiments

Experiments evaluate FedPCL across benchmark datasets, non-IID settings, backbone configurations, client counts, and ablations. Results support robust, stable fusion of pre-trained representations with compact prototype communication.

  • Experimental Setup: Experiments use Digit-5, Office-10, and DomainNet under feature shift, label shift, and combined feature-and-label shift settings.Digit-5, Office-10, and DomainNet comprise multiple source-domain datasets, while label shift is simulated with a Dirichlet parameter α.
  • Experimental Setup: The evaluation compares FedPCL with FedAvg, pFedMe, PerFedAvg, FedRep, FedProto, and independently trained Solo models.Client counts are generally five for feature and label shift, with 5, 4, and 6 clients for Digit-5, Office-10, and DomainNet under combined shifts.
  • Performance Comparison: Multiple pre-trained backbones improve test accuracy in most cases, yielding about 1%−4% improvement for FedPCL over single-backbone cases.The results are reported as mean (std) over clients across three independent runs; prototype-based communication also sends far fewer parameters than model-parameter methods.
  • Performance Comparison: FedPCL achieves the best performance across all tested heterogeneity levels and has the smallest deviation across runs in the label shift setting.Heterogeneity varies with α ∈ {0.5, 1, 2, 5, 10} over 10 clients, where smaller α indicates more heterogeneous label distributions.
  • Performance Comparison: FedPCL’s average test accuracy is about 15% higher than FedAvg as the number of clients increases from 5 to 80.Its deviation across three runs is also much lower, especially with more participating clients.
  • Integrating Backbones with Various Architectures: FedPCL integrates backbones with different architectures and large-scale models, including two-layer CNNs, AlexNet, VGGNet, and ViT.The framework uses fixed pre-trained models and can select models based on a small data set before large-scale training.
  • Ablation Study: Prototype-wise supervised contrastive loss outperforms cross-entropy, ProtoDist-regularized cross-entropy, and local-embedding supervised contrastive loss.The prototype-wise loss uses both global and local prototypes for local contrastive learning.
  • Ablation Study: Removing either global or local prototypes reduces performance by 0.3%-2% across three non-IID settings.The ablation compares global-only, local-only, and combined prototype use during local supervised contrastive learning.

6 Conclusion

The paper presents a lightweight framework using fixed pre-trained backbones, shared class-wise prototypes, and client-side contrastive learning to personalize representations while reducing FL demands. It evaluates three non-IID settings and notes that the study focuses on vision tasks, leaving language methods for future work.

  • Fixed pre-trained backbones replace learnable feature extractors, reducing computation and communication demands in federated learning.
  • FedPCL shares class-wise prototypes and applies client-side contrastive learning to extract shared knowledge while customizing representations for each client.
  • The study focuses on vision tasks, while methods for language tasks remain an area for future research.
  • The evaluation considers feature shift, label shift, and combined feature-and-label shift non-IID settings.

A.1.2 Visualizatin of Raw Samples.

Figure 5 presents examples of raw instances.

  • Figure 5 contains examples of raw instances.
  • The section directs readers to Figure 5 for raw-instance examples.
  • Raw instances are illustrated through examples in Figure 5.

A.1.3 Model Architecture.

The experiments use fixed pre-trained ResNet18 backbones, with three different pre-training datasets in multiple-backbone cases and a client-side architecture containing learnable layers.

  • Single-backbone experiments use ResNet18 pre-trained on Quickdraw.
  • Multiple-backbone experiments use three ResNet18 models pre-trained on Quick Draw, Aircraft, and CU-Birds.
  • The client architecture includes fully connected and BatchNormalization layers, with K denoting the number of available pre-trained backbones and K=3 in the experiments.

A.1.4 Training Details.

Training uses held-out validation data, convergence-based stopping, grid-search hyperparameter tuning, and benchmark partitions for three non-IID settings.

  • Approximately 10% of each dataset is used for training, with 20% of that training portion temporarily held out for validation before retraining.
  • Each algorithm is run until its loss converges, using benchmark partitions for feature shift, label shift, and combined feature-and-label shift settings.
  • Grid search tunes learning rate, weight decay, and the output dimension of each backbone.
  • The learning-rate search space is {0.1, 0.01, 0.001, 0.0001}, weight decay is {1e-3, 5e-4, 1e-4, 5e-5, 1e-5}, and backbone output dimension is {128, 256, 512, 1024}.
  • Table 11 reports FedPCL and baseline performance under label shift non-IID scenarios.

A.2.2 Fairness across Clients.

FedPCL is evaluated for fairness across clients under label-shift and feature-and-label-shift non-IID settings. The evaluation reports accuracy distributions and examines how backbone count affects performance and resource use.

  • Fairness evaluation: Fairness evaluation reports average, worst 10%, worst 20%, worst 40%, and best 10% test accuracy, plus cross-client accuracy variance.Lower variance indicates a fairer solution.
  • Fairness evaluation: The fairness experiments cover 40- and 80-client settings on the Digit-5 dataset across three runs.
  • Fairness evaluation: Prototype-based communication produces fairer solutions than model parameter or gradient communication in the multi-backbone scenario.The paper attributes this to prototypes being more independent of local data distributions than model parameters.
  • Backbone count: The number of fixed pre-trained backbones can be adjusted for a task, with more backbones improving performance while increasing computing and memory requirements.

A.2.4 Privacy Protection.

FedPCL is evaluated with privacy-preserving noise applied to communicated prototypes and images. The reported results indicate that prototype noise preserves high performance, while the framework is limited to settings with available pre-trained models.

  • Privacy-preserving noise: FedPCL remains high-performing after noise is injected into communicated prototypes.The privacy experiment uses multiple backbones on Office-10 under label-shift non-IID conditions.
  • Privacy-preserving noise: The authors conclude that FedPCL can combine with privacy-preserving techniques without an obvious decrease in performance.
  • Privacy-preserving noise: Gaussian and Laplace noise are visualized on an Office-10 bike image using perturbation coefficients p = 0.1 and p = 0.2.The perturbed training image is defined as e_x = (1 − p)x + e.
  • Scope: The framework is limited to cases where pre-trained models are available.For cases without them, the authors suggest multiple fixed encoders as an alternative direction for future exploration.

C Proof of Generalization Bound

The proof establishes a generalization bound using concentration and complexity arguments. It applies these tools to prototype-related quantities and combines the resulting bounds across clients and classes.

  • Generalization bound: The bound depends on the hypothesis class H, its VC-dimension d, and the total sample count N.The proof text also identifies σ_i as independent uniform random variables taking values in {−1, +1}.
  • Proof strategy: The proof begins with McDiarmid’s inequality to control deviations between a function of samples and its expectation.
  • Prototype bound: For a prototype, the proof substitutes the relevant function into the concentration argument to obtain a high-probability bound for a specific prototype.
  • Union over prototypes: Because there are (m + 1)|C| prototypes, the proof applies Boole’s inequality to obtain a simultaneous high-probability statement.
Loading 2209.10083v1…