Federated Learning from Pre-Trained Models: A Contrastive Learning Approach
Abstract
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 ﬁxed pre-trained models rather than training a large-scale model from scratch.
To capture more client-speciﬁc and class-relevant information from the pre-trained models and jointly improve each client’s ability to exploit those off-theshelf models, we design
a Fed erated P rototypewise C ontrastive L earning (FedPCL) approach which shares knowledge across clients through their class prototypes and builds client-speciﬁc representations in a prototype-wise contrastive manner.
We perform a thorough evaluation of the proposed FedPCL in the lightweight framework, measuring its ability to fuse various pre-trained models on popular FL datasets.
1. Introduction
Federated Learning (FL) is a new promising ﬁeld of machine learning that allows multiple clients to train together without sharing their private data (McMahan et al., 2017).
Vanilla
FL aims to train a single global model over all participating clients by periodically synchronizing their model parameters.
However, the learned model usually does not perform well on all clients due to the statistical heterogeneity among local datasets (Kairouz et al., 2019; Long et al., 2020).
Personalized Federated Learning (PFL) is proposed to solve this problem by training a personalized model for each client.
Recent studies on PFL leverage various techniques to enable
First Workshop of Pre-training: Perspectives, Pitfalls, and Paths Forward ICML 2022 , Baltimore, Maryland, USA, PMLR 162, more common underlying information shared across different clients (Zhang et al., 2021a; Fallah et al., 2020; T Dinh et al., 2020; Oh et al., 2022; Long et al., 2022).
So far, FL and PFL have been widely used in real-world tasks, including computer vision (Park et al., 2021), natural language processing (Lin et al., 2021), and graph data mining (Xie et al., 2021; Chen et al., 2022; Liu et al., 2022).
However, the models in real-world applications are usually large-scale neural networks which incur high computation costs and require high communication bandwidth when trained from scratch.
This can make it infeasible to train such models in some practical FL scenarios, e.g., lowresource device-based federated learning.
To alleviate the above issues, we propose a lightweight FL framework that uses multiple ﬁxed pre-trained backbones as the encoder, followed by learnable layers to fuse the representations generated by the backbones for each client.
The proposed framework is capable of fusing the representations generated by pre-trained models with various architectures or obtained from various source data, expanding the scope of federated learning by integrating off-the-shelf foundation models.
Also, it makes it possible to utilize large-scale pre-trained models (Liu et al., 2021b), e.g., VisionTransformer (Dosovitskiy et al., 2020) and Swin Transformer (Liu et al., 2021c), in a computation resource-constrained case to enhance the overall performance.
Using the pre-trained foundation models as the ﬁxed encoder can efﬁciently reduce costs because neither complicated backward propagation computation nor large-scale neural network transmission between the server and clients is needed during the training stage.
To enable a better personalized representation ability for each client under this lightweight FL framework, we need to select an appropriate information carrier to share common underlying knowledge across clients.
Motivated by (Snell et al., 2017; Luo et al., 2021; Tan et al., 2022), class-wise prototypes, deﬁned as “a representative embedding” for a speciﬁc class, can be an effective information carrier for communication between the server and clients.
Sharing prototypes allows for better knowledge sharing across various learning domains, which has been proved in transfer learning (Quattoni et al., 2008) and multi-task learning (Kang et al., 2011) scenarios.
To efﬁciently extract the useful shared information learned from the pre-trained models via prototypes, we design an algorithm called Fed erated
P rototype-wise C ontrastive L earning (FedPCL) where both local prototypes and global prototypes are used for knowledge sharing in a supervised contrastive manner (Khosla et al., 2020; Liu et al., 2021a).
By maximizing the agreement between the fused representation and its corresponding prototypes with contrastive learning, class-relevant information and semantically meaningful knowledge are captured by each client.
Concretely, global prototypes force the fused representation to be closer to the global class center, while local prototypes force clients to share more higher-level feature information in a pairwise way.
Using prototypes to realize inter-client communication allows for knowledge sharing in a latent space and brings additional beneﬁts.
Compared with directly transmitting the representation of a sample, prototypes can eliminate bias from a single sample and protect clients’ privacy.
Compared with model parameters, prototypes are much smaller in size, which signiﬁcantly reduces communication costs.
We quantitatively evaluate the performance of FedPCL and state-of-the-art FL algorithms under the lightweight framework based on benchmark foundation models and datasets.
We also perform extensive experiments to validate the effectiveness of FedPCL in fusing representations output by different backbones and its capability to integrate knowledge from backbones with different model architectures.
Our main contributions are summarized as follows:
2. Problem Formulation
We formulate our proposed lightweight FL framework which integrates off-the-shelf pre-trained models as ﬁxed backbones and learns to fuse them adaptively.
The Proposed Lightweight FL Framework.
Similar to most FL frameworks, there are m clients and a central server involved in the multiple pre-trained backbone-based framework.
Each client i ∈ [1 , m ] owns K shared and ﬁxed backbones and a private dataset D i that cannot be shared with each other.
Each learning model can be seen as a combination of at least two parts: (i) Feature Encoder r ( · ; Φ ∗ ) :
R d → R K × d e , comprising K ﬁxed pretrained backbones, each of which maps the raw sample x of size d to a representation vector of size d e .
K representation vectors are concatenated together as the output r ( x ; Φ ∗ ) , denoted as r x for short.
(ii) Projection Network h ( · ; θ i ) :
R K × d e → R d h , which fuses the K representation vectors and maps r x from one latent space to another for further representation learning.
The formal deﬁnition is provided as follows.
Deﬁnition 2.1.
Let φ ∗ k , where k ∈ [1 , K ] , denote the optimal parameter of the k -th backbone pre-trained on a speciﬁc dataset, r k be the embedding function of the k -th backbone, and x denote an instance sampled from a local dataset.
We deﬁne the concatenated representation output by the feature encoder as
[Equation]
For client i , the projection network h , parameterized by θ i , aims to fuse the representations output by multiple backbones into another abstract space.
The output of the projection network is computed as
[Equation]
Based on the above deﬁnition, we formulate the global objective of the lightweight FL framework as
[Equation]
The target of the framework is to learn the personalized projection network { θ i } mi =1 for each client.
Given an input sample x , the representations output by pre-trained backbones are concatenated together as r ( x ; Φ ∗ ) .
Then, the projection network h
( θ i ) optimized for each client converts the concatenated representation to z ( x ) .
3. FedPCL
We elaborate our proposed algorithm Fed erated P rototypewise C ontrastive L earning (FedPCL), which is illustrated in Figure 1.
In each client, z ( x ) is generated by the projection network which fuses the representations from multiple backbones.
Then, to share the common underlying knowledge across clients, we employ a prototype-based communication scheme to transmit and aggregate prototype sets between the server and clients.
With the prototypes returned from the server, we perform local optimization via well-crafted prototype-wise contrastive loss function, which extracts class-relevant information while sharing more interclient knowledge in the latent space.
We provide detailed illustration for each procedure in the rest of this section.
Prototype as the Information Carrier.
To capture more class-relevant information and semantically meaningful knowledge, we propose to transmit prototypes between the server and clients.
Compared with transmitting the learnable model parameters, there are several advantages brought by transmitting class-wise prototypes.
Firstly, prototype is more compact in form, which signiﬁcantly decreases communication costs required during the training process.
Secondly, non-parametric communication allows each client to learn a more customized local model without synchronizing parameters with others.
Thirdly, prototypes are high-level statistic information rather than raw features, which raise no additional privacy concerns to the system and are robust to gradient-based attacks.
Backbones 𝑟 ! ...
[Figure]
Server ...
|ℂ| { avg ... ...
𝐿 !
Eq.
(6) 𝑪 " 𝑪 # avg ... ...
𝑪 " 𝑪 # ...
ℂ ...
𝐿 $
Client
[Figure]
𝐿 ! ... ...
𝑪 " 𝑪 # ...
ℂ ...
𝐿 $ Backbones 𝑟 !
Projection Network ℎ(&;𝜃 " )
Projection Network ℎ(&;𝜃 # ) 𝑧(𝐱) 𝑧(𝐱)
[Figure]
To extract useful class-relevant information from the local side, we construct a local prototype as the information carrier for knowledge uploading.
Speciﬁcally, it is deﬁned in the latent space of projection network’s output by the mean of the fused representations within the same class j ,
C ( j ) i
[Equation]
where D i,j refers to the subset of D i composed of all instances belonging to class j , and C i denotes the local prototype set of the i -th client.
After the above computation, the local prototype set of each client is sent to the central server for knowledge aggregation, which shares the local class-relevant information extracted on each speciﬁc client based on its local dataset.
Server Aggregation.
After receiving local prototype sets { C i } mi =1 from all participating clients, the server ﬁrst computes the global prototype as
[Equation]
where N j denotes the set of clients that own instances of class j , and N j denotes the number of instances belonging to class j over all clients.
The global prototype set is denoted as C = { ¯ C (1) , ¯ C (2) , . . .
With such an aggregation mechanism, the global prototype set summarize coarse-grained class-relevant knowledge shared by all clients, which provides a high-level perspective for representation learning.
Local Training.
After receiving the prototype sets from the server, the main target of local training is to efﬁciently extract useful knowledge from the local and the global prototypes, respectively, so as to maximally beneﬁt local representation learning.
To achieve that, we propose a prototypewise supervised contrastive loss that consists of two terms, i.e., global term and local term.
To force the fused representation z ( x ) generated by the local projection network to be closer to its corresponding global class center so as to extract more class-relevant but clientirrelevant information, we deﬁne the global prototype-based loss term as
[Equation]
where z x represents z ( x ) for short, A ( y ) :
= { y a ∈ [1 , | C | ] : y a (cid:54) =
y } is the set of labels distinct from y , τ is the temperature that can adjust the tolerance for feature difference (Zhang et al., 2021b;
Chen et al., 2020; Khosla et al., 2020).
For a speciﬁc instance x sampled from D i , we use an inner dot product to measure the similarity between the fused representation z x and prototypes.
Apart from the global term, to align z ( x ) with each client’s local prototypes by alternate client-wise contrastive learning in the latent space and enable more inter-client knowledge sharing, we deﬁne the local prototype-based loss term as
[Equation]
For the i -th client, the local objective function in Eq.
(3) is deﬁned as a combination of L g and L p in the following form,
[Equation]
4. Experiments 4.1. Experimental Setup
Datasets and Non-IID Settings.
We evaluate our proposed framework on the following three benchmark datasets:
Digit-5 (Zhou et al., 2020), Ofﬁce-10 (Gong et al., 2012),
[Table]
and DomainNet dataset (Peng & Saenko, 2018).
To mimic non-IID scenarios in a more general way, we investigate three different non-IID settings.
Details about the non-IID settings can be found in Appendix A.1.
Baselines and Implementation.
We compare our proposed method with popular FL algorithms including FedAvg (McMahan et al., 2017), pFedMe (Dinh et al., 2020),
PerFedAvg (Fallah et al., 2020), FedRep (Collins et al., 2021), FedProto (Tan et al., 2022), and Solo, i.e., training independently within each client.
Details about the model and each dataset can be found in Appendix A.1.
4.2. Performance Comparison
[Table]
Table 1 and Table 2 report the results of our method and baselines in mean (std) format over clients with three independent runs.
The results suggest that: (1) compared with single backbone cases, multiple pre-trained backbones lead to higher test accuracy in most cases and about a 1% − 4% test accuracy improvement for our method; (2) apart from locally training each client (Solo), our method achieves relatively smaller deviation across different runs compared with most baselines, which demonstrates that FedPCL is able to fuse the representations in a more stable way; (3) the number of communicated parameters in prototype-based method is much lower than that of the model parameterbased methods.
[Table]
4.3. Integrating Backbones with Various Architectures
In Table 3, we also show that our proposed lightweight framework is capable of (i) fusing representations generated by backbones with different architectures, i.e., a two-layer
CNN, AlexNet (Krizhevsky et al., 2012), and VGGNet (Simonyan & Zisserman, 2014).
(ii) integrating large-scale pre-trained models, i.e., ViT (Dosovitskiy et al., 2020), to enhance the performance without huge computation resources to train or ﬁne-tune them.
5. Conclusion
To address the problems on excessive computation and communication demands in current FL frameworks, we propose a lightweight framework that leverages multiple neural networks as ﬁxed pre-trained backbones to replace the learnable feature extractor.
To customize the general representations generated by these backbones for each client, class-wise prototypes are shared across the clients and the server.
To efﬁciently extract the shared knowledge from the prototypes, we develop FedPCL algorithm that uses contrastive learning at the client-side during the local update.
Extensive experiments are conducted to show the superiority of FedPCL under the proposed lightweight framework.
A. Experimental Details and Additional Results A.1. Experimental Details
A.1.1.
N ON -IID S ETTINGS .
[Figure]
To mimic non-IID scenarios in a more general way [Figure] we investigate three different non-IID settings as follows and visualize them in Figure 2.
• Feature shift non-IID:
The datasets owned by clients have the same label distribution but different feature distributions.
The number of classes and the number of samples per class are the same across clients.
• Label shift non-IID: The datasets owned by clients have the same feature distribution but different label distributions.
Similar to existing works (Lin et al., 2020; Wang et al., 2020), we use Dirichlet distribution with parameter α to allocate examples for this kind of non-IID setting.
• Feature & Label shift non-IID:
The datasets owned by clients are different in both label distribution and feature distribution, which is more common but challenging in real-world scenarios. [Figure] USPS Synth
[Figure]
USPS Synth [Figure] A.1.2. V ISUALIZATIN OF R AW S AMPLES .
Some examples of raw instances can be found in Figure 3.
A.1.3. M ODEL A RCHITECTURE .
For all baseline methods, the backbone module is followed by two fully connected layers, corresponding to projection network and classiﬁer, respectively, whereas for our method, the backbone module is followed by only one fully connected layer as the projection network.
The output dimension of each backbone and the projection network are 512 and 256, respectively.
Note that for a fair comparison, all baselines use the same network architecture on top of the frozen backbones as FedPCL except their essential classiﬁer.
We use a batch size of 32 , and an Adam (Kingma & Ba, 2014) optimizer with weight decay 1 e − 4 and learning rate 0 . 001 .
The default setting for local update epochs is E =
and the temperature τ is 0 .
We implement all the methods using
PyTorch and conduct all experiments on one NVIDIA Tesla V100 GPU.
For the single backbone cases, we use ResNet18 pre-trained on Quickdraw as the backbone.
For the multiple backbone cases, we use three pre-trained ResNet18 as the backbones.
They are pre-trained on Quick Draw (Jongejan et al., 2016),
Aircraft (Maji et al., 2013), and CU-Birds (Liu et al., 2020) public dataset, respectively.
The model architecture following the backbone module is shown in Table 4.
[Table]
A.1.4. T RAINING D ETAILS .
We provide the detailed settings for the experiments conducted in Section 4.2.
Table 5, 6, 7 show the data partitioning details in feature shift non-IID, label shift non-IID, and feature & label shift non-IID, respectively.
We run each algorithm till the convergence of its loss.
[Table]
USPS
[Table] [Table] [Table]
[Table]
A.2. Additional Experiments A.2.1. P ERFORMANCE UNDER T HREE N ON -IID S ETTINGS .
Table 8 shows the performance of FedPCL and baseline methods under label shift non-IID scenarios.
[Table]
Following the metrics in (Li et al., 2019)
[Table]
we verify the advantage of FedPCL in fairness and report the results over 40 andclients in Table 9 and Table 10, respectively.
We list the average, the worst 10% , the worst 20% , the worst 40% , the best 10% of test accuracy over all clients across three runs in Digit-5 dataset.
We also report the variance of the accuracy distribution across clients (the smaller, the fairer) (Li et al., 2019).
The results demonstrate that in such a multi-backbone scenario, it is easier to obtain a fair solution using prototype-based communication rather than model parameter/gradient-based communication, because the information conveyed by prototypes is more local data distribution-independent while the information conveyed by model parameters tends to be affected by local data distribution.
Detailed experimental results on fairness can be found in Table 9 and
[Table] [Table]
[Table]
A.2.3. A BLATION S TUDY
In this section [Table] we present the results for various ablation experiments to test the effect of global/local prototypes and the contrastive loss.
More results can be found in Appendix A.2.3.
Effect of the contrastive loss.
We test the performance of the local loss function in different forms: (1) Cross Entropy loss; (2) Cross Entropy loss + ProtoDist term, which takes the distance between prototype and fused representation as an additional term to regularize the cross entropy loss; (3) Supervised Contrastive loss that only uses local embedding for contrastive learning; (4) Our prototype-wise contrastive loss that uses both global and local prototypes for local contrastive learning.
As shown in Table 11, our prototype-wise supervised contrastive loss outperforms others.
Effect of prototypes.
To verify the effectiveness of different kinds of prototypes used for local training in our proposed
FedPCL, we compare the following three cases: (i) Only global prototypes computed at the server are used for local training;
[Table]
(ii) Only local prototypes aggregated at the server are used for local training; (iii) Both global and local prototypes are used for local training.
All these three cases are implemented under three non-IID settings.
As shown in Table 12, without global or local prototypes being used for local supervised contrastive learning, the performance drops 0 .
3% - 2% , indicating that the knowledge conveyed by global prototypes and local prototypes can beneﬁt the local learning framework from different perspectives
[Table]
[Table]
Effect of the Number of Backbones.
The number of pre-trained backbones can be adjusted for a speciﬁc task correspondingly.
To study its effect, we compare the performance and parameter size when different numbers of ﬁxed backbones are used.
The results in Table 13 indicate that more pre-trained backbones can lead to better performance but consume more computing resources and memory space [Figure] [Table] [Figure]
[Table]
A.2.4. V ISUALIZING THE F USING R ESULTS OF F ED PCL
[Figure]
To better understand how FedPCL fuses the representation output by backbones, we visualize the normalized similarity scores under the feature shift non-IID setting.
The heatmap in Figure 4 summarizes the values of the normalized similarity scores, computed by the inner product between local prototypes in the single backbone case and the multi-backbone case.
For example, in Figure 4(a), the element 0 .
5 on row 1 and column 2 represents the similarity score between the local prototype computed when only one backbone pre-trained in Quickdraw is available and the local prototype computed when three backbones pre-trained in Quickdraw, Aircraft, and Birds are available.
We ﬁnd that the backbone pre-trained on Quickdraw contributes more to the fused prototype, while the backbone pre-trained on Birds contributes the least.
Clients who own a speciﬁc dataset show different levels of utilization when fusing the representations from various backbones.
A.2.5. P RIVACY P ROTECTION .
We incorporate FedPCL with privacy-preserving techniques to observe its variation in performance.
Concretely, we add random noise of various distributions into the communicated prototypes and the original images, respectively.
Tableshows that the performance of FedPCL remains high after injecting noise to the prototypes.
Figure 5 visualizes an original image of a bike from Ofﬁce-10 dataset and what it looks like after the noise injection operation.
It is hard to tell the objects in the right column of Figure 5(c), but the accuracy only drops about 2% , compared to vanilla FedPCL.
In conclusion,
FedPCL has the potential to combine with privacy-preserving techniques without an obvious decrease in performance
[Table]
[Table]
[Figure]