F EDERATED S ELF - SUPERVISED L EARNING FOR H ETEROGENEOUS C LIENTS
Anonymous authors
Paper under double-blind review
A BSTRACT
Federated Learning has become an important learning paradigm due to its privacy and computational beneﬁts.
As the ﬁeld advances, two key challenges that still remain to be addressed are: (1) system heterogeneity - variability in the compute and/or data resources present on each client, and (2) lack of labeled data in certain federated settings.
Several recent developments have tried to overcome these challenges independently.
In this work, we propose a uniﬁed and systematic framework,
Heterogeneous Self-supervised Federated Learning (Hetero-SSFL) for enabling self-supervised learning with federation on heterogeneous clients.
The proposed framework allows collaborative representation learning across all the clients without imposing architectural constraints or requiring presence of labeled data.
The key idea in Hetero-SSFL is to let each client train its unique self-supervised model and enable the joint learning across clients by aligning the lower dimensional representations on a common dataset.
The entire training procedure could be viewed as self and peer-supervised as both the local training and the alignment procedures do not require presence of any labeled data.
As in conventional self-supervised learning, the obtained client models are task independent and can be used for varied end-tasks.
We provide a convergence guarantee of the proposed framework for non-convex objectives in heterogeneous settings and also empirically demonstrate that our proposed approach outperforms the state of the art methods by a signiﬁcant margin.
I NTRODUCTION
Federated learning has become an important learning paradigm for training algorithms in a privacy preserving way and has gained a lot of interest in the recent past.
While traditional federated learning is capable of learning high performing models (Li et al., 2021b), two practical challenges that remain under studied are: system heterogeneity, and lack of labeled data.
In several real world scenarios, the clients involved in the training process are highly heterogeneous in terms of their data and compute resources.
Requiring each client to train identical models, like in traditional FL, may thus be severely limiting.
Similarly, assuming each client’s local data resources to be fully labeled may not be pragmatic as annotating data is time consuming and may require expertise.
Our focus in this work is to jointly address these two challenges in a systematic way so as to substantially increase the scope of FL approaches.
Prior works have studied these two issues separately, and the approaches taken do not offer a natural way to be combined so as to provide an effective heterogeneous selfsupervised FL framework.
For instance, to alleviate the scarcity of labeled data on local clients, both semi-supervised learning (Zhang et al., 2021b; Jeong et al., 2021;
Lu et al., 2022; Lubana et al., 2022) and self-supervised learning (Zhuang et al., 2021; 2022;
He et al., 2022; van Berlo et al., 2020) methods have been proposed, but they all assume identical model architectures on each client and in fact do not extend to heterogeneous settings.
Some aspects of system heterogeneity have been independently addressed for supervised learning by building personalised models on clients (Tan et al., 2022; Jiang et al., 2020; Fallah et al., 2020), but still assuming identical architectures.
Existing recent works on federated learning with independent architectures across clients consider standard supervised learning scenarios (Makhija et al., 2022) rather than self-supervised learning.
The need for self-supervised FL arises in multiple applications.
For example, consider cross-silo analytics in healthcare systems where different hospitals may possess varying amounts of private medical images.
Here, the data can neither be centralised nor can undergo extensive annotations.
Moreover, to expect each client (e.g., a hospital) to train local models of identical capacities can be highly inhibiting.
In such cases, the smaller capacity clients which are incapable of training large models of common architecture will not be accommodated in federated learning.
On the other hand, if the common model architecture size is reduced, some of the clients will not be fully engaged.
Differences in compute resources and the learning capacity demand system heterogeneity.
To the best of our knowledge, ours is the ﬁrst work that proposes a general framework for federated self-supervised learning in heterogeneous settings where non-identical clients can implement distinct, independent model architectures and obtain personalised solutions.
The proposed framework is novel and ﬂexible, and allows clients to train self-supervised models with unique (locally tuned) structures while still using the learnings from other clients in the network, in an architecture agnostic way.
To achieve this, we add a proximal term to each client’s loss function that helps in aligning the learnt lower dimensional representations (aka embeddings) across clients.
This way of transferring and acquiring global knowledge allows variability in client model architectures while keeping the entire learning process independent of labeled training data.
Furthermore, we use a kernel based distance metric for proximity calculation which provides much more ﬂexibility to the clients in deﬁning their own lower dimensional space without any constraints.
Our Contributions are summarized as follows :
Our main contribution is the new framework, Hetero-SSFL , for training heterogeneous models in a federated setting in an unsupervised way.
Hetero-SSFL allows each client to train its own customized model, using local data and computing resources while also utilising unlabeled supervision from peers.
We perform thorough experimental evaluation of the proposed approach in both image-based and text-based self-supervised learning settings.
We observe that the proposed ﬂexible approach substantially improves the accuracy of unsupervised and self-supervised federated learning in heterogeneous settings.
We also provide theoretical analysis and convergence guarantee of the algorithm for nonconvex loss functions.
Organization.
The rest of the paper is organised as follows.
Section 2 provides a brief background on Federated Learning, Self-supervised Learning and related developments.
In Section 3, we go over the preliminaries and then propose our framework.
We study the convergence guarantee of the algorithm in Section 4 and include the related proofs in the Appendix.
A thorough experimental evaluation of our method on different types of datasets is presented in Section 5 and we conclude the paper in Section 6.
R ELATED W ORK
This section provides an overview of the most relevant prior work in the ﬁelds of federated learning, self-supervised learning and federated self-supervised learning.
2. F EDERATED L EARNING (FL)
The problem of training machine learning models in distributed environments with restrictions on data/model sharing was studied by several researchers in the data mining community in the early 2000s under titles such as (privacy preserving) distributed data mining (Kargupta & Park, 2000; Gan et al., 2017; Aggarwal & Yu, 2008).
Much of this work was for speciﬁc procedures such as distributed clustering (Merugu & Ghosh, Nov, 2003) including in heterogenous settings with different feature spaces (Merugu & Ghosh, 2005), and distributed PCA (Kargupta et al., 2001), or for speciﬁc models such as SVMs (Yu et al., 2006).
Subsequently, with the re-surfacing in popularity of neural networks and proliferation of powerful deep learning approaches, the term “Federated Learning” got coined and popularized largely by an inﬂuential paper (McMahan et al., 2017) which introduced FedAvg.
Indeed FedAvg is now considered the standard distributed training method for Federated Learning, which has become a very active area of research since then.
One of the key challenges in this setting is the presence of non-iid datasets across clients.
Several modiﬁcations of the original FedAvg algorithm have been proposed to address this challenge.
Some of these approaches focus on ﬁnding better solutions to the optimization problem to prevent the divergence of the global solution (Li et al., 2020; Karimireddy et al., 2020; Zhang et al., 2021a; Pathak & Wainwright, 2020;
Acar et al., 2021; Karimireddy et al., 2021) whereas some suggest the modiﬁcation of the training procedure to incorporate appropriate aggregation of the local models (Chen & Chao,
Yurochkin et al., 2019; Wang et al., 2020b; Singh & Jaggi, 2020).
Creating local personalised models for each client to improve the overall method is also well researched (Collins et al., 2021; Ghosh et al., 2020; Smith et al., 2017a;
Li et al., 2021b; Smith et al., 2017b; Sattler et al., 2021; Yu et al., 2020).
Also, augmenting the data distributions by creating additional data instances for learning (Hao et al., 2021; Goetz & Tewari, 2020; Luo et al., 2021) has seen to help with the collaborative effort in federated learning.
2. S ELF - SUPERVISED LEARNING (SSL)
SSL aims at training high performing complex networks without the requirement of large labeled datasets.
It does not require explicitly labeled data and target pairs but instead uses inherent structure in the data for representation learning.
The two types of SSL methods include contrastive based and non-contrastive based SSL.
The contrastive methods work on triplets of the form ( x, x 0 , y ) where x and x 0 are different views of the same instance and y is a non-compatible data instance.
The key idea is to train models by adjusting the parameters in such a way that the representations(embeddings) obtained for x and x 0 are nearly identical and that of x and y are unalike.
Since obtaining these triplets from the dataset does not necessarily require presence of the labels, for example for images, x and x 0 could be different augmentations of the same image and y could be an image of a different object, these methods could learn from large volume of data.
Recent works in this ﬁeld suggest many different ways of generating the data triplets for training (Chen et al., 2020a;
He et al., 2020; Wu et al., 2018; Yan et al., 2020).
Non-contrastive methods, on the other hand, work without generating explicit triplets.
One form of the non-contrastive methods use clustering based methods to form a group of objects and use the cluster assignments as pseudo-labels for training the model (Caron et al., 2018; 2020).
Another form of non-contrastive methods use joint embedding architectures like
Siamese network but keep one of the architectures to provide target embeddings for training (Chen &
He, 2020; Grill et al., 2020; Chen et al., 2020b).
2. F EDERATED S ELF - SUPERVISED L EARNING
The interest in self-supervised federated learning is relatively new.
Although MOON (Li et al., 2021a) proposed using a contrastive loss in supervised federated learning, it was FedU (Zhuang et al., 2021) that proposed a way to leverage unlabeled data in a decentralised setting.
Subsequently, two parallel works (Zhuang et al., 2022), (He et al., 2022) proposed frameworks for training selfsupervised models in a federated setting and compared different self-supervised models.
In all of these approaches though, the clients train local models of identical architectures.
This common architecture is communicated by the server at the beginning of the training.
In each round, the clients perform update steps on the local models and send the local models to the central server at the end of the round.
The server appropriately aggregates the client models to generate the global model for that round and communicates the global model parameters back to each client.
The prominent idea in these latest developments is the divergence aware updates of the local model to the global model wherein, if the local models are very different from the global model, the local models are not updated with the global model, which can be seen as an initial step towards creating more personalised models for each client.
On the other hand, (Shi et al., 2021) does explore heterogeneous architectures across clients but uses them to train a common global model which again does not suit heterogeneity.
Additionally, some very recent works have explored self-supervised FL for speciﬁc applications like acoustics, human-activity recognition etc.
(Feng et al., 2022; Ek et al., 2022; Rehman et al., 2022).
M ETHODOLOGY
In order to enable federated learning for heterogeneous clients in an unsupervised way, we propose a new framework using self-supervised federated learning and name it Hetero-SSFL.
This framework is unique as it allows heterogeneous clients to participate in the training process, exhibits high empirical performance and is guaranteed to converge.
In this section we ﬁrst formally describe the problem setting and then elaborate on our proposed solution.
3. P ROBLEM D EFINITION
A federated learning setting consists of a set of clients k 2 | N | , with each client k
having local data instances X k of size n k drawn from the local data distribution D k .
The local objective at k th client is to solve for min W k L k
= ` ( W k ; X k ) where ` ( . ) is a suitable loss function.
After the local optimization, the federation part of the learning allows the server to access all the clients’ learnings in the form of the learned set of weights W 1 , W 2 , . .
. , W N and utilise it for achieving global objective.
In traditional FL methods, it is assumed that all W k ’s are identical in shape and in each round the server combines the local models to learn a global model ˆ W of the form
ˆ W = g ( W 1 , W 2 , . .
. , W N ) , where g is an appropriate aggregation function, and sends ˆ W back to the clients for further updates.
In contrast, our solution does not require the local models (and thus the local weights) to be of identical size, as we deploy a novel and different way of utilizing the global knowledge, W 1 , W 2 , . . .
, W N , to achieve collaboration for each client’s beneﬁt.
3. H ETERO -SSFL F RAMEWORK
We introduce the Hetero-SSFL framework in this section and provide details of all the components in the framework here.
The end-to-end pipeline is depicted in Figure 1 and the algorithm is detailed in
Algorithm 1.
General Formulation As mentioned above, the federated self-supervised learning problem involves each client k solving a local optimization problem, min W k L k
= ` ( W k ) , for some self-supervised loss, ` (
W k ) .
In Hetero-SSFL, we modify the loss function for each client to the following:
[Equation]
where d ( ., . ) is a suitable distance metric,  k ( . ; . ) is used to denote the representation learning component of the k th client,  k
( X ; W k ) denotes the representation matrix of k th client under weights W k on dataset X and ¯  denotes the aggregated representations of all clients and is computed as
[Equation]
Here, w j is the weight given to the representations obtained from the j th client.
The client weight w j could be kept higher for clients with larger resources (clients which are capable of training larger capacity models and obtain better representations), or could be set to learn in the training pipeline.
The new loss function enables collaboration across clients by aligning each client’s representations and simultaneously guiding the training on all clients.
Speciﬁcally, we modify the local loss function at each client to contain a proximal term that measures the distance between the local representations and the representations obtained on all other clients.
This framework could be used for any self-supervised learning application by suitably choosing the self-supervised loss function ` ( W k ) , the distance function d ( ., . ) , and the network architecture  k ( . ; . ) .
In this section we present the details of self-supervised learning for image recognition tasks and provide an extension for text classiﬁcation task in Appendix B, thereby providing two concrete instances of this general formulation.
Local Training We assume each client contains unlabeled datasets and employ self-supervised learning on each client.
While there are many ways to achieve self-supervision, in this work we focus on non-contrastive methods like BYOL (Grill et al., 2020) and
SimSiam (Chen & He, 2020) which involve two embedding architectures like Siamese networks for achieving self-supervision.
The two parallel networks are called the online network and target network respectively.
The target network’s task is to provide targets for regression to the online network whereas the online network is trained to learn representations.
The target network consists of an encoder and the online network is composed of an encoder and a prediction layer on top of it, whose role is to learn better representations under self-supervised loss.
The weights of the target network are just a moving average of the weights of the online network and the target network can thus be thought of as a stable version of the online network.
In federated self-supervised learning with these networks, the local set of parameters for each client k include W k = { W ok , W pk , W tk } where W ok parameterises the representation learning component of the online network with W pk being the prediction layer and W tk is the target network.
The local loss at each client ` ( . ) is similar to self-supervised loss used in central self-supervised learning and is given by -
[Equation]
i.e., it is the mean squared error between the outputs of the online network and the target network on different views(augmentations), v 0 and v 00 , of the same image v .
F ( , ; . ) and F 0 ( , ; . ) denote the online and target network functions respectively.
The representations  ( X ; . ) are obtained from the output of the client’s online network under parameters W o =
[ W o , W p ] at a layer suitable for deﬁning the representation space, and this could potentially be different from F which denotes the entire online network.
In our framework, we use the output of the predictor layer for an instance x as  ( x ; . ) .
Communication with the Server Different from traditional FL where each client sends the local models to server and the server creates a global model by performing element-wise aggregation on the local model weights, we gather representations at the server.
We assume that the server has access to an unlabeled set of data points (which could be obtained in practice from open sources like images from web) and call it Representation Alignment Dataset (RAD).
The RAD is a publicly available dataset used for the purpose of aligning representations and doesn’t possess any special properties.
i ( X ; W i ) is thus a matrix of dimensions
L x d i where L is the number of data points in
X and d i is the size of embedding for each point.
In each global training round t , server sends the
RAD to all the clients.
All the clients use the RAD, X , in the loss function as described above while simultaneously training local models.
After the clients ﬁnish local epochs, the ﬁnal representations obtained for RAD ¯ X on each client,  i ( X ; W i ) , are sent to the server.
The server then aggregates the local representations from all clients to form the global representation matrix ¯  ( X ; ( t )) , as in
Equation 2, which is then again sent to the clients for training in the next round.
This methodology is similar to what (Makhija et al., 2022) does, but for supervised learning.
Here we tackle the problem of lack of labelled data by rather doing federated self-supervised learning.
3. T HE P ROXIMAL T ERM
In this sub-section we describe the second part of the loss function and provide details on the proximal term and the distance function.
The role of the proximal term is to compare the representations obtained from various neural networks.
This requires the distance function to be able to meaningfully capture the similarity in the high dimensional representational space.
Moreover, the distance function should be such that it is sensitive to the functional behavior(e.g., outputs from initial layer versus penultimate layer) and does not change as much with the inconsequential modiﬁcations like change in random initializations.
The types of distance functions used in the literature for comparing neural network representations include Canonical Correlation Analysis(CCA), Centered Kernel ALignment(CKA) and Procustres distance based measures (Ding et al., 2021; Kornblith et al., 2019).
While the Procustres distance based metrics cannot compare the representations obtained in spaces of different dimensions, the
CCA based metrics are not suitable for training through backpropagation.
In contrast, properties of the CKA metric like ability to learn similarities between layers of different widths, invariance to random initializations, invertible linear transformations and orthogonal transformations lend it useful for use in the proximal term.
CKA measures the distance between the objects in different representation spaces by creating similarity(kernel) matrices in the individual space and then comparing these similarity matrices.
The idea is to use the representational similarity matrices to characterize the representation space.
CKA takes the activation matrices obtained from the network as inputs and gives a similarity score between 0 and 1 by comparing the similarities between all objects.
Speciﬁcally, if A i 2 R L ⇥ d i and
A j 2 R L ⇥
d j are the activation(representation) matrices for clients i and j for the RAD X of size L , the distance between A i and A j using Linear-CKA is computed as -
[Equation]
where K i and K j are kernel matrices for any choice of kernel K such that we have K i ( p, q )
=
K ( A i ( p, :) , A i ( q, :)) .
CKA also allows choice of other kernels like RBF-kernel, polynomial kernel etc.
but we observed that the results with Linear-CKA are as good as those with RBF kernel as also reported in (Kornblith et al., 2019).
Thus we focus only on Linear-CKA here.
Finally, the local loss at each client becomesmin
[Equation]
with ¯ K ( t  1) =
P
Nj =1 w j K j
( t  1) and K j ( t  1) =  j
( X ; W j ( t  1)) .  j ( X ; W j ( t  1)) T .
When µ = 0 , each client is training its own local models without any peer-supervision and when µ ! the local models try to converge on the RAD representations without caring about the SSL loss.
Algorithm 1 Hetero-SSFL Algorithm for Heterogeneous clients
Input: number of clients N , number of global communication rounds T , number of local epochs
E , parameter µ , weight vector for clients [ w 1 , w 2 , . . .
w N ]
Output: Final set of personalised models W 1 ( T ) , W 2 ( T ) . . .
W N ( T )
At ServerInitialize W 1 (0) , W 2 (0) . . .
W N (0) for t = 1 to T do
RAD =
X for each client j do
A j ( X ) =
j
( X ; W j ( t  1))
K j = K ( A j ( X ) , A j ( X )) end for ¯ K ( t  1) =
P
Nj =1 w j K j
Select a subset of clients N t for each selected client i 2 N t do
W i ( t ) =
LocalTraining ( W i ( t  1) , ¯ K ( t  1) , X , µ ) end for end for
Return W 1 ( T ) , W 2 ( T ) . . .
W N ( T )
LocalTraining ( W i ( t  1) , ¯ K ( t  1) , X , µ )
Initialize W i ( t ) with ( W i ( t  1)) for each local epoch do
Generate local representation matrix A i ( X ) =  i
( X ; W i ( t ))
Generate local kernel matrix K i =
K ( A i ( X ) ,
A i ( X ))
Update W i ( t ) after solving for loss in equation (4) end for
Return W i ( t )
C ONVERGENCE A NALYSIS
In this section we provide insights on convergence of the proposed framework.
The convergence guarantee is shown to hold under mild assumptions which are commonly made in the literature (Wang et al., 2020b; Tan et al., 2022).
Assumption 4.1 (Lipschitz Smoothness) .
The local loss function on every client k , L k , is assumed to be L 1 -Lipschitz smooth, which implies:
[Equation]
Assumption 4.2 (Stochastic Gradient) .
The local stochastic gradient for each client k , at any time t , g k,t is an unbiased estimator of the gradient, has bounded variance and the expectation of its norm is bounded above.
That is, we have
[Equation]
Assumption 4.3 (Representation Norm) .
The norm of the representations obtained for any instance on any one client is bounded above by R .
That is:
[Equation]
Assumption 4.4 (Lipschitz Continuity) .
The embedding function  ( . ; . )
for all clients is L 2 -Lipschitz continuous, which for each client k implies:
[Equation]
We start with the following lemma on the reduction of the local loss after E local epochs.
Lemma 4.5.
(Local Epochs)
In each communication round, the local loss L for each client reduces after E local epochs and is bounded as below under the Assumption 4.2.
[Equation]
where L 0 and L E denote the loss before and after E local epochs respectively and ⌘ is learning rate.
Proof of Lemma 4.5 is in Appendix A.1.
Lemma 4.5, similar to (Tan et al., 2022) , shows the bound on the client’s local loss after the completion of local epochs under one global communication round.
After the local training, the representations from all clients are sent to the server.
Our next lemma 4. shows a bound on the expected loss after every global representation update at the server.
Lemma 4.6.
(Representation Update)
In each global communication round, for each client, after the
E local updates, the global representation matrix is updated at the server and the loss function for any client k gets modiﬁed to L E 0 from L E and could be bounded as -
[Equation]
Proof of Lemma 4.6 is in Appendix A.2.
Equipped with the results from Lemmas 4.5 and 4.6, we are now ready to state our main result which provides a total deviation bound for an entire training round.
Theorem 4.7 shows the divergence in loss for any client after completion of one round.
We can guarantee the convergence by appropriately choosing µ and ⌘ which leads to a certain expected decrease in the loss function.
The result holds for convex as well as non-convex loss functions.
Theorem 4.7 (Convergence) .
After one global round, the loss function of any client decreases and is bounded as shown below:
[Equation]
Thus, if we choose ⌘ and µ such that
[Equation]
then the convergence of Algorithm 1 is guaranteed.
The proof of Theorem 4.7 follows directly from the results of Lemmas
E XPERIMENTS
In this section we describe the experimental setup and present results for our method alongside the popular baselines for the image recognition task as described above.
The experiment for text based application is provided in Appendix B. We simulate system heterogeneity by randomly choosing different architectures for local client models.
Within system heterogeneity we evaluate our model under varying statistical settings by manipulating the data distributions across clients to be IID or non-IID.
We also demonstrate the performance of our framework with increase in number of clients.
Some additional experimental results for this application are included in the Appendix C.
5. E XPERIMENTAL D ETAILS
Datasets We test our framework on image classiﬁcation task and use three datasets to compare against the baselines: CIFAR-10, CIFAR-100 and Tiny-Imagenet as suggested by several works in FL and the popular federated learning benchmark LEAF (Caldas et al., 2019).
CIFAR-10 and CIFARcontain 50,000 train and 10,000 test colored images for 10 classes and 100 classes respectively
and
Tiny-Imagenet has 100,000 images for 200 classes.
Implementation Details For the FL simulations, we explore two types of settings, non-IID setting and IID setting.
In the IID-setting, we assume each client to have access to the data of all classes and in the non-IID setting we assume each client has access to data of only a few disjoint classes.
For example, for CIFAR-10, client 1 might have access to objects of classes { 1 , 2 } versus client 2 with access to { 3 , 4 } and likewise.
The IID setting is created by dividing all the instances between the clients.
For the non-IID setting, for each client a fraction of classes is sampled and then instances are divided amongst the clients containing speciﬁc classes.
In the cases when number of clients,
N , is less than total number of classes K , each client is assumed to have access to all instances corresponding to KN classes.
The choice of encoders at each client is data and application dependent.
Since we are dealing with images, to create heterogeneous clients, for each client we randomly select between Resnet-18 and Resnet-34 models as encoders for our method.
Since all the baselines work in homogeneous client settings and hence have to use the same model, we use Resnet-34 at each client in the baselines for comparison.
The total number of global communication rounds is kept tofor all methods with each client performing E = 5 local epochs in each round.
We use SGD with momentum = 0 .
9 , batch size = 200 and learning rate = .
032 for training all models.
The weight vector used in aggregating the representations w is set to be uniform to 1 n .
All these models are trained on a machine with 4 GeForce RTX 3090 GPUs with each GPU having about 24gb of memory.
Baselines We compare our method against two very recently proposed federated self-supervised baselines: FedU and FedEMA, which reﬂect the current state-of-the-art.
Since the source code for these baselines is not publicly available, we perform these comparisons using our implementation and with the best parameters reported in these papers.
We also create additional baselines by combining central SSL algorithms like BYOL and SimSiam with the federated learning procedure of FedAvg to make FedBYOL and FedSimSiam.
The performance of supervised FedAvg and centralised BYOL is also reported for comparisons as theoretical baselines.
All of these methods however work in homogeneous settings, hence we compare the homogeneous setting of baselines with Resnetarchitecture with our heterogeneous setting having a mix of Resnet-34 and Resnet-18 architectures.
By doing this, we are comparing our framework with more powerful baselines and thus expect the improvement our baselines to be as good or better when doing an identical comparison.
One additional baseline is the SSFL algorithm, but due to the unavailability of the source code and lack of precise implementation details in the paper, we could not compare against it.
Evaluation We follow the linear evaluation protocol as mentioned in (Zhuang et al., 2021) to test the models.
Under this protocol, the local encoders are trained in self-supervised way without any labeled data and then a linear classiﬁer is trained on top of the encoder after freezing the encoder parameters.
We use a linear fully connected layer with 1000 neurons as the classiﬁer and train it for 100 epochs with Adam optimizer, batch size = 512 and learning rate = 0.003. The reported results are the accuracy of the linear classiﬁer on a separate test dataset.
Other parameters
The other parameters important for our approach are µ which controls the weight of the proximal term in the loss function, and the size of the RAD dataset.
We tune the hyperparameter µ using a validation set and found the best values of µ to be 0 . 1 for the CIFAR datasets and 1 for the Tiny-Imagenet dataset.
For the RAD size we observe that increasing the size of RAD increases performance but has an effect on the speed of simulation so we ﬁx that size to be
[Table]
5000 in our experiments.
On studying the effect of varying number of local epochs we observed that unlike the FedAvg based algorithms our method is stable with increasing local epochs.
[Table]
[Table]
R ESULTS
We show the performance of our method in comparison with baselines in non-IID settings in Tableand that for IID settings in Table 2.
We observe that our method outperforms the baselines by a signiﬁcant margin and is slightly worse than the central BYOL (requires the data to be present on the central server and is not federated) which implies that we are able to achieve successful collaboration amongst clients.
We also report the performance of our method and high-performing baselines with increase in scale in Table 3.
As the number of clients increase, in every round only a fraction of clients are sampled to perform local training, total 20 clients with 5 clients selected per round is denoted as 20clients(5) in Table 3.
[Table]
C ONCLUSION
In this work we study the problem of self-supervised federated learning with heterogeneous clients and propose a novel framework to enable collaboration in such settings.
The framework allows clients with varying data/compute resources to partake in the collaborative training procedure.
This development enables several practical settings to still gainfully employ self-supervised FL even though the different clients, for example different edge-devices or distinct organisations, might not have similar compute and data resources.
The high performance of the framework across multiple datasets indicates its promise in achieving collaboration while the theoretical results guarantee the convergence of the algorithm for fairly general loss functions.
In future work we plan to investigate the use of attention based mechanisms for adaptive weighting of different clients to generate aggregate representations based on each client’s relative importance.
Also, rather than assuming that the server has access to an unlabeled dataset and transmitting RAD from the server to clients, we can create a common generative model in a federated way like federated GANs to generate data instances required