Federated Learning from Pre-Trained Models: A Contrastive Learning Approach
Anonymous Author(s)
AfﬁliationAddress
Abstract
Federated Learning (FL) is a machine learning paradigm that allows decentralizedclients 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 theseissues from hindering the deployment of FL systems, we propose a lightweightframework where clients jointly learn to fuse the representations generated bymultiple ﬁxed pre-trained models rather than training a large-scale model fromscratch.
This leads us to a more practical FL problem by considering how tocapture more client-speciﬁc and class-relevant information from the pre-trainedmodels and jointly improve each client’s ability to exploit those off-the-shelfmodels.
In this work, we design a Fed erated P rototype-wise C ontrastive L earning(FedPCL) approach which shares knowledge across clients through their classprototypes and builds client-speciﬁc representations in a prototype-wise contrastivemanner.
Sharing prototypes rather than learnable model parameters allows eachclient to fuse the representations in a personalized way while keeping the sharedknowledge in a compact form for efﬁcient communication.
We perform a thoroughevaluation of the proposed FedPCL in the lightweight framework, measuring andvisualizing its ability to fuse various pre-trained models on popular FL datasets.
Introduction
Federated Learning (FL) is a new promising ﬁeld of machine learning that allows multiple clients totrain together without sharing their private data [1].
Vanilla FL aims to train a single global modelover all participating clients by periodically synchronizing their model parameters.
However, thelearned model usually does not perform well on all clients due to the statistical heterogeneity amonglocal datasets [2].
Personalized Federated Learning (PFL) is proposed to solve this problem bytraining a personalized model for each client.
Recent studies on PFL leverage various techniques toenable more common underlying information shared across different clients [3, 4, 5, 6].
So far, FLand PFL have been widely used in real-world tasks, including computer vision [7], natural languageprocessing [8], and graph data mining [9].
However, the models in real-world applications are usually large-scale neural networks which incurhigh computation costs and require high communication bandwidth when trained from scratch.
Thiscan make it infeasible to train such models in some practical FL scenarios, e.g., low-resource devicebased federated learning.
To alleviate the above issues, we propose a lightweight FL framework thatuses 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 fusingthe 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 mod- els.
Also, it makes it possible to utilize large-scale pre-trained models, e.g., VisionTransformer
[10]and Swin Transformer [11], in a computation resource-constrained case to enhance the overallperformance.
Using the pre-trained foundation models as the ﬁxed encoder can efﬁciently reducecosts because neither complicated backward propagation computation nor large-scale neural networktransmission between the server and clients is needed during the training stage.
As shown in Table 1,compared to training a ResNet18 [12] from scratch using FedAvg [1], a much smaller number of parameters are communicated per round when learning to fuse the representations generated by ﬁxed pretrained models.
Besides, higher performance can be achieved when using pre-trained models after thesame communication rounds.
Compared with using a single pre-trained model, using multiple modelspre-trained on diverse datasets can provide a more comprehensive view for an input sample [13].
[Table]
To enable a better personalized representation ability for each client under thislightweight FL framework, we need toselect an appropriate information carrierto share common underlying knowledgeacross clients.
Motivated by [17, 18, 19],class-wise prototypes, deﬁned as “a representative embedding” for a speciﬁc class,can be an effective information carrierfor communication between the serverand clients.
Sharing prototypes allowsfor better knowledge sharing across various learning domains, which has been proved in transfer learning [20] and multi-task learning [21]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 asupervised contrastive manner.
By maximizing the agreement between the fused representation andits corresponding prototypes with contrastive learning, class-relevant information and semanticallymeaningful knowledge are captured by each client.
Concretely, global prototypes force the fusedrepresentation to be closer to the global class center, while local prototypes force clients to sharemore 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.
Comparedwith directly transmitting the representation of a sample, prototypes can eliminate bias from a singlesample and protect clients’ privacy.
Compared with model parameters, prototypes are much smallerin size, which signiﬁcantly reduces communication costs.
We quantitatively evaluate the performance of FedPCL and state-of-the-art FL algorithms underthe lightweight framework based on benchmark foundation models and datasets.
We also performextensive experiments to validate the effectiveness of FedPCL in fusing representations output bydifferent backbones and its capability to integrate knowledge from backbones with different modelarchitectures.
Our main contributions are summarized as follows:
• We take the ﬁrst step towards integrating pre-trained models into federated learning to forma lightweight FL framework, which substantially reduces computation and communicationcosts of current FL frameworks.• We further propose FedPCL, a novel algorithm that uses class prototypes as the informationcarrier and conducts contrastive learning during local updates, which allows clients to sharemore class-relevant knowledge from both local and global prototypes.• Experiments are conducted on a variety of benchmark foundation models and datasets tomeasure FedPCL’s ability to fuse various pre-trained models for each client.
The results
indicate that FedPCL outperforms baselines with a better personalization and knowledgeintegration ability.
Related Work
Personalized Federated Learning.
Personalized Federated Learning (PFL) aims to train a personalized model for each client in the FL framework so that the model can achieve better performanceon the local dataset.
Existing approaches for PFL are based on various techniques.
[22, 23, 24] addan additional term to the original loss function of each client to produce better personalized models according to the private data.
[16, 25] share part of the model and keep personalized layer privateto achieve personalization.
[26] proposes to use a central hypernetwork to generate personalizedmodels for clients.
[27] enables a more ﬂexible personalization by adaptively weighted aggregation.[4, 28] study PFL from a meta-learning perspective where a meta-model is learned to generate theinitialized local model for each client.
Contrastive Learning.
Contrastive learning has been widely applied to self-supervised learningscenarios in recent years, achieving state-of-the-art performance in the unsupervised training ofdeep image models [29, 30].
A number of works are focused on learning an encoder where theembeddings of the same sample are pulled closer and those of different samples are pushed apart [31,32, 33].
[34] extends contrastive learning from self-supervised settings to fully supervised settings,enabling us to better exploit label information with contrastive learning.
There are also some worksincorporate contrastive learning into federated learning to assist local training to achieve higher modelperformance [35, 36, 37].
Prototype Learning.
Prototypes have been widely used in a variety of tasks in transfer learning [20,38], multi-task learning [21, 39], and meta-learning [17, 40].
It is usually deﬁned as the meanfeature vectors of samples within the same class [41, 42].
The authors in [43] represent task-agnosticinformation by prototypes for distributed machine learning systems and propose a multi-task modelfusion method that integrates prototypes for a new task.
Since the prototype has the ability togeneralize semantic knowledge from similar samples, it is used to assist the local training in federatedlearning by several studies [44, 19, 36, 35].
Pre-Trained Foundation Model.
Due to the huge number of parameters and the broad data availablefor training, pre-trained foundation models can better capture knowledge when being applied todownstream tasks [45].
Recently, pre-trained models (e.g., ViT [10], DETR [46], BERT [47]) havebeen widely investigated in both vision and natural language processing (NLP) tasks [48].
Extractingtask-speciﬁc knowledge from the pre-trained models has the potential to achieve the state-of-the-artperformance due to their generality and adaptability to different tasks [49, 50, 51].
However, how tointegrate the pre-trained models into a federated learning paradigm and keep the whole frameworklightweight still remains an open problem.
Problem Formulation
In this section, we formulate our proposed lightweight FL framework which integrates off-the-shelfpre-trained models as ﬁxed backbones and learns to fuse them adaptively.
We start from the generalframework of FL, and then explicitly deﬁne the problem and explain the proposed global objective.
General FL Framework.
Formally, the global objective of general FL across m clients is
[Equation]
where L i and w i are the local loss function and model parameters for client i , respectively.
D i is theprivate dataset of the i -th client.
N is the total number of samples among all clients.
The objective of general FL, such as vanilla FedAvg and its variants, is to learn an optimal globalmodel w across m clients, where w = w 1 = w 2 = · · · = w m [1].
This can be achieved byperiodically synchronizing the model parameters of all clients at the server.
However, parametersynchronization sometimes deteriorates the performance on local datasets in the presence of dataheterogeneity [25, 52].
Some recent studies [5, 3] also explore personalized FL by applying variousconstraints and regularization terms, which allows clients to keep different models w i , i ∈ [1 , m ] toachieve higher performance on their local datasets.
The Proposed Lightweight FL Framework.
Similar to most FL frameworks, there are m clientsand a central server involved in the multiple pre-trained backbone-based framework.
Each clienti ∈ [1 , m ] owns K shared and ﬁxed backbones and a private dataset D i that cannot be shared witheach 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 pre-trained backbones, each of which mapsthe raw sample x of size d to a representation vector of size d e .
K representation vectors areconcatenated together as the output r ( x ; Φ ∗ ) , denoted as r x for short.
(ii) Projection Networkh ( · ; θ i ) :
R K × d e → R d h , which fuses the K representation vectors and maps r x from one latentspace to another for further representation learning.
The formal deﬁnition is provided as follows.
[Figure]
Deﬁnition 3.1.
Let φ ∗ k , where k ∈ [1 , K ] , denote the optimal parameter of the k -th backbone pretrained on a speciﬁc dataset, r k be the embedding function of the k -th backbone, and x denote aninstance sampled from a local dataset.
We deﬁne the concatenated representation output by thefeature encoder as
[Equation]
For client i , the projection network h , parameterized by θ i , aims to fuse the representations output bymultiple backbones into another abstract space.
The output of the projection network is computed asz
[Equation]
x
[Equation]
h
[Equation]
Based on the above deﬁnition, we formulate the global objective of the lightweight FL framework as
[Equation]
where Φ ∗
[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 concatenatedtogether as r ( x ; Φ ∗ ) .
Then, the projection network h
( θ i ) optimized for each client converts theconcatenated representation to z ( x ) .
The proposed multi-backbone setting is friendly to most FL methods.
The learnable parameter θin the projection network can be fully or partly shared across clients for different purposes.
Forexample, FedAvg [1] and its variants [53, 54, 55, 56] can synchronize θ every communication round,and personalized FL methods can also adapt to this framework without much modiﬁcation [22, 25].
Federated Prototype-wise Contrastive Learning (FedPCL)
In this section, we elaborate our proposed algorithm Fed erated P rototype-wise C ontrastive L earning(FedPCL), which is illustrated in Figure 1.
In each client, z ( x ) is generated by the projection networkwhich fuses the representations from multiple backbones.
Then, to share the common underlyingknowledge across clients, we employ a prototype-based communication scheme to transmit andaggregate prototype sets between the server and clients.
With the prototypes returned from theserver, we perform local optimization via well-crafted prototype-wise contrastive loss function, whichextracts class-relevant information while sharing more inter-client knowledge in the latent space.
Weprovide detailed illustration for each procedure in the rest of this section.
Prototype as the Information Carrier.
To capture more class-relevant information and semanticallymeaningful knowledge, we propose to transmit prototypes between the server and clients.
Comparedwith transmitting the learnable model parameters, there are several advantages brought by transmittingclass-wise prototypes.
Firstly, prototype is more compact in form, which signiﬁcantly decreasescommunication costs required during the training process.
Secondly, non-parametric communication allows each client to learn a more customized local model without synchronizing parameters withothers.
Thirdly, prototypes are high-level statistic information rather than raw features, which raiseno additional privacy concerns to the system and are robust to gradient-based attacks.
To extract useful class-relevant information from the local side, we construct a local prototype asthe information carrier for knowledge uploading.
Speciﬁcally, it is deﬁned in the latent space ofprojection network’s output by the mean of the fused representations within the same class j , [Equation] ( j ) [Equation]
i
[Equation]
(cid:88)
[Equation]
x
[Equation] i [Equation] ) where D i,j refers to the subset of D i composed of all instances belonging to class j , and C i denotesthe local prototype set of the i -th client.
After the above computation, the local prototype set of eachclient is sent to the central server for knowledge aggregation, which shares the local class-relevantinformation 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, theserver ﬁ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 numberof 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 coarsegrained class-relevant knowledge shared by all clients, which provides a high-level perspective forrepresentation learning.
After aggregation, the server sends both the global prototype set and the full local prototype setscollected from all clients back to every clients.
For the situation where only a few classes in eachclient in some non-IID cases, we introduce a prototype padding procedure in the server to ensure thateach local prototype set contains prototypes corresponding to all classes:
[Equation]
Prototype-based communication and aggregation allow each client to own a unique projection networkthat is able to fuse the general representations in a customized way.
The returned local prototype setscan encourage a mutual learning from a client-relevant perspective while the global prototype set,where each element indicates a class center in the overall data, provides a chance to learn from ahighly summarized client-irrelevant perspective.
Local Training.
After receiving the prototype sets from the server, the main target of local trainingis to efﬁciently extract useful knowledge from the local and the global prototypes, respectively, soas to maximally beneﬁt local representation learning.
To achieve that, we propose a prototype-wisesupervised 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 itscorresponding global class center so as to extract more class-relevant but client-irrelevant 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 fromy , τ is the temperature that can adjust the tolerance for feature difference [57, 33, 34].
For a speciﬁcinstance x sampled from D i , we use an inner dot product to measure the similarity between the fusedrepresentation z x and prototypes.
Apart from the global term, to align z ( x ) with each client’s local prototypes by alternate client-wisecontrastive learning in the latent space and enable more inter-client knowledge sharing, we deﬁne thelocal prototype-based loss term as
[Equation]
Given the ground-truth label y of x , a prototype set can be divided as one positive prototype C ( y )and a set of negative prototypes C ( y a ) .
Both L g and L p force the representation of a sample to becloser to those positive prototypes and apart from those negative prototypes. ¯ C ( y ) in L g and C ( y ) in
L p summarize abstract class-relevant information to different granularity, which provides guidancefor optimizing the local projection network from different perspectives.
For the i -th client, the localobjective function in Eq.
(4) is deﬁned as a combination of L g and L p in the following form,
L (cid:16)
[Equation]
i
[Equation]
z ( [Equation] ) , y, C , {
[Equation]
p [Equation] (cid:17)
[Equation] [Equation] [Equation] Input: D i , θ i , i = 1 , · · · , m , and K pre-trained backbones with parameters φ ∗ 1 , φ ∗ 2 , · · · , φ ∗ K , respectively.
Server executes: 1: Initialize prototype sets { C p } mp =1 .
2: for each round T = 1 , 2 , ... do 3: for each client i in parallel do 4:
C
i ← LocalUpdate ( i, C , { C p } mp =1 ) 5: end for 6:
Update global prototype by Eq.
7: end for
LocalUpdate ( i, C , { C p } mp =1 ) : 1: for each local epoch do 2: for each batch in D i do 3:
Compute L g by Eq.
(8) with global prototypes.
Compute L p by Eq.
(9) with local prototypes.
Update θ i by Eq.
6: end for 7: end for 8: Compute local prototypes by Eq.
9: return C i
At the end of the local training in eachround, clients upload their local prototypeset C i to the server.
We present the FedPCL algorithm in detail in Algorithm 1.
Prototype-based Inference.
After adequate optimization, in each client, the localprototypes generated by projection networknot only contain compact local class-wiseinformation but also absorb general knowledge from all participating clients.
Considering their powerful representation capability, we leverage the local prototypes to predict unknown labels in the inference stage.
Concretely, for a test sample, we ﬁrst calculate the similarity scores between the output of projection network and every localprototypes.
Then, the predicted result is theclass with the maximum similarity score.
Generalization Bound.
We provide insights into the performance analysis of FedPCL.
A detailed description and derivationscan be found in Appendix C.
Theorem 4.1.
(Generalization Bound of FedPCL.)
Consider an FL system with m clients.
Let
D 1 , D 2 , · · · , D m be the true data distribution and (cid:98) D 1 , (cid:98) D 2 , · · · , (cid:98) D m be the empirical data distribution.
Denote the projection network h as the hypothesis from H and d be the VC-dimension of H .
The totalnumber of samples over all clients is N .
Then, with probability at least 1 − δ :
[Equation]
≤
[Equation]
Theorem 4.1 indicates that compared with the model trained on an ideal data distribution, theperformance of the empirically trained model is associated with the VC-dimension of H and the sizeof the prototype set C . An expected performance gap can be achieved by using appropriate projectionnetwork and number of classes.
The generalization bound is also important to ensure a satisfyingperformance of FedPCL, especially when applying it to some safety-critical scenarios [3, 58, 59].
Experiments
5.1 Experimental Setup
Datasets and Non-IID Settings.
We evaluate our proposed framework on the following threebenchmark datasets: Digit-5 [15], Ofﬁce-10 [60], and DomainNet dataset [61].
Digit-5 [15] is abenchmark dataset for digit recognition, including ﬁve benchmark datasets, namely SVHN, USPS,
SynthDigits, MNIST-M, and MNIST.
Ofﬁce-10 [60] is a standard benchmark dataset consistingof four datasets, namely Amazon, Caltech, DSLR and WebCam.
DomainNet [62] is a large-scaledataset consisting of six datasets, namely Clipart, Info, Painting, Quickdraw, Real, and Sketch. [Table]
To mimic non-IID scenarios in a more general way, we investigate three different non-IID settings:(i)
Feature shift non-IID:
The datasets owned by clients have the same label distribution but differentfeature distributions.
(ii) Label shift non-IID:
The datasets owned by clients have the same featuredistribution but different label distributions which is simulated by Dirichlet distribution with parameterα [58, 55].
(iii) Feature & Label shift non-IID:
The datasets owned by clients are different in bothlabel distribution and feature distribution, which is more common but challenging in real-worldscenarios.
Details about the non-IID settings can be found in Appendix A.1.
Baselines and Implementation.
We compare our proposed method with popular FL algorithmsincluding FedAvg [1], pFedMe [22], PerFedAvg [4], FedRep [25], FedProto [19], and Solo, i.e.,training independently within each client.
In feature shift and label shift non-IID setting, the numberof clients is 5.
In feature & label shift setting, the number of clients is 5, 4, 6 for Digit-5, Ofﬁce-10,and DomainNet, respectively.
Similar to [13, 14], we use the ResNet18 architecture [12] pre-trained on Quick Draw, Aircraft, and
CU-Birds [14] as the backbones.
For all baseline methods, the backbone module is followed by twofully connected layers, corresponding to projection network and classiﬁer, respectively, whereas forour method, the backbone module is followed by only one fully connected layer as the projectionnetwork.
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 topof the frozen backbones as FedPCL except their essential classiﬁer.
We use a batch size of 32 , and an Adam [63] optimizer with weight decay 1e-4 and learning rate0 .
The default setting for local update epochs is E =
and the temperature τ is 0 .
Weimplement all the methods using PyTorch and conduct all experiments on one NVIDIA Tesla V
GPU.
Details about the model and each dataset can be found in Appendix A.1.
5.2 Performance Comparison
[Table]
Table 2 and Table 3 report the results of our methodand baselines in mean (std) format over clients withthree independent runs.
Table 2 shows the results oftwo cases: (1) s : a single backbone is available; (2) m :multiple backbones are available.
The results suggest that: (1) compared with singlebackbone cases, multiple pre-trained backbones lead tohigher test accuracy in most cases and about a 1% − 4%test accuracy improvement for our method; (2) apartfrom locally training each client (Solo), our methodachieves relatively smaller deviation across differentruns compared with most baselines, which demonstrates that FedPCL is able to fuse the representations in a more stable way; (3) the number ofcommunicated parameters in prototype-based method is much lower than that of the model parameterbased methods. [Figure] (b)
Robustness to varying levels of heterogeneity.
In the label shift non-IID setting, we use Dirichlet distribution withparameter α > 0 to simulate the heterogeneity level over 10 clients.
As α becomes smaller, the label distribution becomes more heterogeneous.
We evaluateour proposed FedPCL with FedRep and
Solo by varying α ∈ { 0 . 5 , 1 , 2 , 5 , 10 } on10 clients.
Figure 2(a) suggests that FedPCL achieves the best performance forall values of α .
Moreover, FedPCL hasthe smallest deviation across several runs,which implies the stable convergence androbustness of FedPCL.
Effect of varying numbers of participating clients.
We also compare FedPCL with the vanilla FedAvg by varying the number of clients ( m ) from 5 to 80 as shown in Figure 2(b).
The non-IIDlevel remains the same when more clients participate in the training procedure.
The results showthat the average test accuracy of FedPCL is higher than that of FedAvg by about 15% .
Furthermore,the deviation of FedPCL across three runs is much lower than that of FedAvg, especially when mbecomes larger.
This indicates the scalability of the lightweight framework and the potential andability of FedPCL when adapting to a large-scale FL system.
5.3 Integrating Backbones with Various Architectures
[Table]
[Table]
The above experiments use backbones with thesame architecture but pre-trained on differentsource data, showing the capability of FedPCLon integrating knowledge from different sourcedomains.
In Table 4, we also show that ourproposed lightweight framework is capable of(i) using pre-trained models with different architectures, i.e., a two-layer CNN, AlexNet [64],and VGGNet [65]; (ii) integrating large-scalepre-trained models, i.e., ViT [10, 66], to enhance the performance without the need of hugecomputation resources to train or ﬁne-tune them.
Pre-trained models with different representationabilities should be adaptively leveraged for different tasks to achieve expected performance.
They can be selected based on a small set ofdata before large-scale training.
5.4 Ablation Study
In this section, we present the results for various ablation experiments to test the effect of global/localprototypes and the contrastive loss.
More results can be found in Appendix A.2.3.
[Table]
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 termto regularize the cross entropy loss; (3) Supervised
Contrastive loss that only uses local embedding forcontrastive learning; (4) Our prototype-wise contrastive loss that uses both global and local prototypes for local contrastive learning.
As shown in
Table 5, our prototype-wise supervised contrastiveloss outperforms others.
Effect of prototypes.
To verify the effectiveness of different kinds of prototypes used for localtraining in our proposed FedPCL, we compare the following three cases: (i) Only global prototypescomputed at the server are used for local training; (ii) Only local prototypes aggregated at the serverare used for local training; (iii) Both global and local prototypes are used for local training.
All thesethree cases are implemented under three non-IID settings.
As shown in Table 6, without global orlocal 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 thelocal learning framework from different perspectives.
[Table]
5.5 Visualizing the Fusing Results of FedPCL
[Figure]
To better understand how FedPCL fusesthe representation output by backbones,we visualize the normalized similarityscores under the feature shift non-IIDsetting.
The heatmap in Figure 3 summarizes the values of the normalizedsimilarity scores, computed by the inner product between local prototypes inthe single backbone case and the multibackbone case.
For example, in Figure 3(a), the element 0 .
5 on row 1 andcolumn 2 represents the similarity scorebetween the local prototype computed when only one backbone pre-trained in Quickdraw is availableand 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, whilethe backbone pre-trained on Birds contributes the least.
Clients who own a speciﬁc dataset showdifferent levels of utilization when fusing the representations from various backbones.
5.6 Incorporating with Privacy-Preserving Techniques
To further eliminate concerns about privacy leakage, we also incorporate FedPCL with privacypreserving techniques to observe the variation in performance.
Concretely, we add various randomnoise to the prototypes to be communicated and the original images, respectively.
The results show thatthe performance of FedPCL remains high after noise injection.
Details are given in Appendix A.2.4.
Conclusion
To address the problems on excessive computation and communication demands in current federatedlearning 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 theclients 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.
This research has great potential to bring existing FL frameworks to a new paradigm and providemore options on local neural network architectures.
However, we mainly focus on the vision taskin this work.
Methods for language tasks can be further explored in future studies.
Considering theevident motivation and effectiveness of the lightweight framework and FedPCL, we believe this workis intriguing for future studies.
Checklist
For all authors
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes] See Section 1.
(b) Did you describe the limitations of your work?
[Yes] See Section 5.6 and Section 6.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix C.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See Section 5.
and source code in supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Section 5.1 and Appendix A.1.
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Section 5.2.
(d) Did you include the total amount of compute and the type of resources used (e
g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Section 5.1.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Section 5.1.
(b) Did you mention the license of the assets?
[Yes] See Section 5.1.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
We do not use any new datasets in this paper.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[No] The datasets used in this paper are publicly available.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content? [N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c)
Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
A Experimental Details and Additional Results
A.1 Experimental Details
[Figure]
A.1.1 Non-IID Settings
[Figure]
To mimic non-IID scenarios in a more general way, we investigate three different non-IID settings asfollows and visualize them in Figure 4.
• Feature shift non-IID:
The datasets owned by clients have the same label distribution but differentfeature distributions.
The number of classes and the number of samples per class are the sameacross clients.
• Label shift non-IID:
The datasets owned by clients have the same feature distribution butdifferent label distributions.
Similar to existing works [58, 55], we use Dirichlet distributionwith 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]
Synth
[Figure]
SVHN
USPS
Synth
[Figure]
Amazon
[Figure]
Painting
[Figure]
Synth
[Figure]
A.1.2 Visualizatin of Raw Samples.
Some examples of raw instances can be found in Figure 5.
A.1.3 Model Architecture.
For the single backbone cases, we use ResNet18 pre-trained on Quickdraw as the backbone.
For themultiple backbone cases, we use three pre-trained ResNet18 as the backbones.
They are pre-trainedon Quick Draw [67], Aircraft [68], and CU-Birds [13] public dataset, respectively.
The modelarchitecture following the backbone module is shown in Table 7.
[Table]
A.1.4 Training Details.
[Table]
We provide the detailed settings for the experiments conducted in Section 5.2.
For data splitting, we use a portion of data as training samples ( 10% ) and the rest as test samples.
We ﬁrst take out a 20% subset of the training set for validation and return the validation set back tothe training set and retrain the model after selecting the optimal hyperparameters.
Table 8, 9, 10 show the data partitioning details in feature shift non-IID, label shift non-IID
[Table]
andfeature & label shift non-IID, respectively.
We run each algorithm till the convergence of its loss.
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
For hyperparameter tuning, we use grid search to ﬁnd the optimal hyperparameters including learningrate, weight decay, and the output dimension of each backbone.
Concretely, the grid search is carriedout in the following search space:
• learning rate: { 0 . 1 , 0 . 01 , 0 . 001 , 0 . 0001 }• weight decay: { 1e-3, 5e-4, 1e-4, 5e-5, 1e-5 }• output dimension of each backbone: { 128 , 256 , 512 , 1024 }
A.2 Additional Experiments
A.2.1 Performance under Three Non-IID Settings.
Table 11 shows the performance of FedPCL and baseline methods under label shift non-IID scenarios.
A.2.2 Fairness across Clients.
Following the metrics in [69], we verify the advantage of FedPCL in fairness and report the resultsover 40 and 80 clients in Table 12 and Table 13, respectively.
We list the average, the worst 10% ,
[Table]
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) [69].
The results demonstrate that in such a multi-backbone scenario, it is easier to obtaina fair solution using prototype-based communication rather than model parameter/gradient-basedcommunication, because the information conveyed by prototypes is more local data distributionindependent while the information conveyed by model parameters tends to be affected by local data
[Table]
distribution.
Detailed experimental results on fairness can be found in Table 12 and 13.
[Table]
[Table]
A.2.3 Effect of the Number of Backbones.
The number of pre-trained backbones can be adjusted for a speciﬁc task correspondingly.
To study itseffect, we compare the performance and parameter size when different numbers of ﬁxed backbones areused.
The results in Table 14 indicate that more pre-trained backbones can lead to better performancebut consume more computing resources and memory space
[Table]
A.2.4 Privacy Protection.
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 theoriginal images, respectively.
Table 15 shows that the performance of FedPCL remains high afterinjecting noise to the prototypes.
Figure 6 visualizes an original image of a bike from Ofﬁce-10dataset and what it looks like after the noise injection operation.
It is hard to tell the objects in theright column of Figure 6(c), but the accuracy only drops about 2% , compared to vanilla FedPCL.
[Table]
In conclusion
[Table]
FedPCL has the potential to combine with privacy-preserving techniques without anobvious decrease in performance.
[Table] ss i a n s =0. s =0. s =0.
L a p l a c e (b) p = 0 .
G a u ss i a n s =0.
s =0.
s =0.
L a p l a c e (c) p = 0 .
B Discussion
B.1 Applying FedPCL to a Wider Range of Scenarios
Since there are many off-the-shelf foundation models available nowadays and there is a trend to usethem in CV and NLP communities, it is a promising research direction to integrate those pre-trainedfoundation models into an FL framework.
Compared with training from scratch, using pre-trainedmodels can save more computation and communication resources.
Our proposed framework is limited to the cases where pre-trained models are available.
This ismainly due to the fact that most types of data, e.g., image, text, graph, have corresponding pre-trainedmodels nowadays.
For the cases without pre-trained models, using multiple ﬁxed encoders can be analternative solution, which is another interesting problem and can be explored in the future.
We do think incorporating pre-trained models into existing learning frameworks is a promising trendin deep learning, especially when models are becoming larger and larger in scale and hard to trainfrom scratch.
So far, the idea to utilize pre-trained foundation models has been proposed for CV and
NLP tasks and achieved certain improvement [45, 70, 71].
B.2 Comparison to Related Works
This paper has provided novel contributions in terms of (1) integrating pre-trained models intofederated learning, (2) proposing a novel algorithm FedPCL which allows clients to share knowledgevia prototype-based local contrastive learning.
Although prototypical learning and contrastive learning exist in prior work, they are still based on thelearnable parameter aggregation scheme [35, 36, 44] or just use prototypes/contrastive learning toregularize the original local training [19].
Instead of synchronizing learnable parameters, our methodallows each client to keep their own local parameters while extracting shared knowledge only bycontrastive learning.
Prototype is used as the information carrier to achieve that.
Some state-of-the-art PFL methods do not work well under our proposed lightweight framework.
Their performance is mainly due to the following two reasons.
• Most of these baselines are designed based on the training-from-scratch framework where
there are a large number of learnable parameters.
The proposed new setting where mostparameters are ﬁxed is not friendly to some PFL methods.
• When there are feature/label shift non-IID across clients, the local performance might be
further deteriorated after parameter aggregation.
The deterioration can be more signiﬁcant
when only a small number of parameters are locally trained.
Since most PFL methods are
still
based on the parameter aggregation scheme, their performances are inevitably affected.
Performance might be enhanced by decreasing the number of shared parameters, optimizing localtraining schemes, and introducing ﬁne-tuning procedures during local training, etc.
There is also a branch of FL studying ﬁne-tuning techniques for FL [72, 73].
These works mainly focuson how to adjust the local/global model to improve its representation ability on biased/generic datadistribution, while our work utilizes ﬁxed pre-trained foundation models and focuses on improvingthe fusing ability.
C Proof of Generalization Bound
Theorem C.1.
(Generalization Bound of Alg. 1).
Consider an FL system with m clients.
Let
D 1 , D 2 , · · · , D m be the true data distribution and (cid:98) D 1 , (cid:98) D 2 , · · · , (cid:98) D m be the empirical data distribution.
Denote the projector head h as the hypothesis from H and d be the VC-dimension of H .
The totalnumber of samples over all clients is N .
Then, with probability at least 1 − δ :
[Equation]
Proof.
We start from the McDiarmid’s inequality that
[Equation]
when
[Equation]
[Equation]
Eq.
13 equals to
[Equation]
which means that with probability at least 1 − exp (cid:16) −
[Equation]
[Equation]
Let δ = exp
[Equation]
−
[Equation]
i
[Equation]
, the above can be rewritten as with probability at least 1 − δ ,
[Equation]
≤
[Equation]
For prototype, by substituting g ( · ) with
[Equation]
we can obtain that with probability at least 1 − δ , the following holds for a speciﬁc prototype,
[Equation]
Considering there are ( m + 1) | C | prototypes in total
[Equation]
by using Boole’s inequality, with probability atleast 1 − δ , the following holds,
[Equation]
where N is the total number of samples over all clients.
[Equation]
where H is the hypothesis set of projector head h , d is the VC-dimension of H , N is the total numberof samples over all clients.
(a) follows from the deﬁnition of Rademacher complexity
[Equation]
[Equation]
where m is the number of samples, σ i refers to independent uniform random variable taking value in{− 1 , +1 } , and (b) follows from Jensen’s inequality.
So,
[Equation]