Federated Self-supervised Learning for Heterogeneous Clients
Authors: Disha Makhija
Nhat Ho
Joydeep Ghosh
Venue: Submitted to ICLR 2023 (2023)
Keywords: Federated Learning
Self-supervised Learning

Abstract: Federated Learning has become an important learning paradigm due to its privacy and computational benefits. As the field 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 unified and systematic framework, \emph{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 significant margin. 

Under review as a conference paper at ICLR 2023
FEDERATED SELF-SUPERVISED LEARNING FOR HET- EROGENEOUS CLIENTS
Anonymous authors Paper under double-bling review
ABSTRACT
Federated Learning has become an important learning paradigm due to its privacy and computational benefits. As the field 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 chal- lenges independently. In this work, we propose unified and systematic framework, Heterogeneous Self seeffsupervieed Federated Learning (Hetero-SSHe cef-ssupervieed learning with federation on heterogeneous clients. The proposed framework allows collaborative representation learning across all the clients with- out imposing architectural constraints or requiring presence of labeled data. The key idea in Hetero-SSFL. is let each client train its unique self-supervised model and enable the joint learning across clients by aligning the lower dimensional rep- 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 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 significant margin.
for enabling
resentations
INTRODUCTION
Federated learning has become an important learning paradigm for training algorithms in privacy preserving way and has gained a lot of interest in the recent past. While traditional federated learning capable of learning high performing models (Li al., 202Tb), 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 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 taken do not offer a natural way to be combined so as to provide an effective heterogeneous self- supervised FL framework. For instance, alleviate the scarcity of labeled data on local clients, both semi-supervised learning (Zhang et al. 2021 Jeong et al. 2021 Lu et al. 2022: Lubana et al. 2022) and supervisee learning (Zhuang al. 2021 2022 He et al. 2022 van Berlo 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 clients TTan jt al., 2022, Jiang et al., 2020 Fallah et al.) (200),, but still assuming identical architectures. Existing recent works federated learning with independent architectures across clients consider standard supervised learning scenarios (Makhija et al. 2022) rather than self-supervised learning.
issues separately, and the approaches
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., hospital) to train local models of identical capacities can

---

Under review as a conference paper at ICLR 2023
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 first work that proposes 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 flexible, and allows clients to train self- seff-supervieed 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 '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 labeled training data. Furthermore, we use a kernel based distance metric for proximity calculation which provides much more flexibility to the clients in defining their own lower dimensional space without any constraints.
Our Contributions are summarized follows:
1. 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.
2. 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 flexible approach substantially improves the accuracy of unsupervised and self-supervised federated learning heterogeneous settings.
3. We also provide theoretical analysis and convergence guarantee of the algorithm for non- convex loss functions.
Organization. The rest of the paper is organised as follows. Section 2 provides brief background on Federated Learning, Self-supervised Learning and related developments. In Section 3 we go over the preliminaries and then propose framework. We study the convergence guarantee of the algorithm in Section 4 and include the related proof' in the Appendix. A thorough experimental evaluation of our method on different types of datasets is presented Section5aant we conclude the paper in Section 6
2 RELATED WORK
This section provides an overview of the most relevant prior work in the fields of federated learning, self-supervised learning and federated self-supervised learning.
2.1
FEDERATED LEARNING(FL)
The problem of training machine lcarning 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 al. 2017 Aggarwal & Yu 2008). Much of this work was for specific 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., (000), or for specific models such as SVMs YY 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 influential paper (McMahan et al,, (017)) which introduced FedAvg. Indeed FedAvg is now considered the standard distributed training method for Federated Learning, which has become very active area of research since then. One of the key challenges in this setting the presence of non-iid datasets across clients. Several modifications of the original FedAvg algorithm have been proposed to address this challenge. Some of these approaches focus on finding better solutions to the optimization problem to prevent the divergence of the global solution (Li et al. 2020 Karimireddy et al. 2020 Zhang al. 2021a| Pathak & Wainwright 2020 Acar et al.
2021: Karimireddy et al.,, 2021) whereas some suggest the modification of the traiiin procedure to
incorporate appropriate aggregation of the local models (Chen & Chao 2021 Wang et al. 2020a
2

---

Under review as a conference paper at ICLR 2023
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. [022 Ghosh et al. 2020; Smith et al. 2017a Li ct al. 2021b: Smith et al. 2017b /attler et al. [021] Yu et al.
2020). Also, augmenting the data distributions by creating additional data instances for learning (Hao
al. 202T Goetz& & Tewari 2020 Luo al. 202T) has seen to help with the collaborative effort in
federated learning.
2.2 SELF-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', y) where x and x' are different views of the same instance and is non-compatible data instance. The key idea is to train models by adjusting the parameters in such a way that the appresenationssemmedddigge obtained for and z' are nearly identical and that of 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' could be different augmentations of the same image and y could be an image of different object, these methods could learn from large volume of data. Recent works in this field suggest many different ways of generating the data triplets for training (Chen et al./ 2020al He al. 2020) Wu it al. 2018 Yan et al. (020). 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 group of objects and use the cluster assignments as pseudo-labels for training the model (Caron ct 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.3
FEDERATED
LEARNING
The interest in self-suue federated learning is relatively new. Although MOON (Li et al. 2021a proposed using contrastive loss in supervised federated learning, it was FedU (Zhuang et al. 2021 that proposed way to leverage unlabeled data in a decentralised setting. Subsequently,
Two parallel works (Zhuang et al., 2022). (He al. 2022) proposed frameworks for training self- supervised 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 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 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 common global model which again does not suit heterogeneity. Additionally, some very recent works have explored self-supervised FL for specific applications like acoustics, human- anttaitactivity recognition etc. (Feng al. 2022 Ek et al. 2022
Rehman et al. 2022
3 METHODOL.OGY
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 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 first formally describe the problem setting and then elaborate on our proposed solution.
3.1 PROBLEM DEFINITION
A federated learning setting consists of set of clients k E [N), with each client k having local data instances X's of size nę drawn from the local data distribution D&. The local objective at t client is to solve for minw. Lt &(Wr; X&) where f(.) is 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 leamed set of weights Wi, Wz...., .., Wy and utilise it for achieving global objective.
3

---

Under review as a conference paper at ICLR 2023
In traditional FL methods, it is assumed that all W3 's are identical in shape and in each round the server combines the local models learn global model W of the form W g(Wi, W2,..., Wx). 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, W1, W2,.... ....................... WN, to achieve collaboration for each client's benefit.
3.2 HETERO-SSFL FRAMEWORK
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 [and the algorithm is detailed in Algorithm
General Formulation As mentioned above, the federated self-supervised learning problem involves each client & solving local optimization problem, minw, La ((W)), for some self-supervised loss, E(W).. In Hetero- SSFL, we modify the loss function for each client to the following:
min Lu = k(W)) + pdl t (x: W.). (X: (t 1 W4
(1)
where .(...) is suitable distance metric, (=:.:) is used denote the representation learning component of the 1tth client, (X; Wi) denotes the representation matrix of 14th client under weights We on dataset X and & denotes the aggregated representations of all clients and is computed
as
N
{ (t)) = Σ wy (x;(((
5=1
Here, wj is the weight given to the representations obtained from the jth client. The client weight wy could be kept higher for clients with larger resources (clients which are capable of training larger capacity models and obtain better representations), 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. Specifically, 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)). the distance function d(.,.), and the network architecture (.: .). In this section we present the details of self-supervised learning for image recognition tasks and provide an extension for text classification 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 many ways achieve self-supervision, in this work we focus on non-contrastive methods like BOL (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 targel network respectively. The target network's task is provide targets for regression the online network whereas the online network 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 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 include W₁ = (W2, WP, W where W2 parameterises the representation learning component of the online network with W being the prediction layer and W & is the target network. The local loss at each client E(.) is similar to self--spe loss used in central self-super learning and is given by-
((W+) = [[F('': [Wo. WE)) - -F''
ie., it is the mean squared error between the outputs of the online network and the target network on different views(augmentations)), s' and v", of the same image v. F(...) :) and (.: 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 (Wo, WP) at layer suitable for
4

---

Under review as a conference paper at ICLR 2023
defining the representation space, and this could potentially be different from FF which denotes the entire online network. In our framework, we use the output of the predictor layer for an instance 2: as ((z;)).
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-wiss 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 publicly available dataset used for the purpose of aligning representations and doesn't possess any special properties. (X; W;) is thus matrix of dimensions I. x d; where I. is the number of data points in X and de 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 described above while simultaneously training local models. After the clients finish local epochs, the final representations obtained for RAD X on each client, 42 (X; W:), the local representations from all clients to form the global representation matrix ((X: (t)), as in Equation 2 which is then again sent the clients for training in the next round. This methodology similar to what (Makhija el al.) (022) does, but for supervised learning. Here we lackle the problem of lack of labelled data by rather doing federated supervueee learning.
sent the server. The server then aggregates
3.3 THE PROXIMAL TERM
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 modifications like change in random initializations. The types of distance functions used in the literature for comparing neural network representations include Canonical Correlation Analysis(CCA) (CCA). Centered Kernel (CKg) 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 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 by comparing the similarities between all objects. Specifically, if A; TK. lxxd, and A5 E RLdd, are the acdivation(reeresennttiio matrices for clients and j for the RAD of size L,, the distance between A, and A, using Linear-CKA is computed as
114 1
Linear CKA(A,A) AA43) Linear CKA(K,,K)) ||A A 1147 Aglle
(3)
where K; and Kj are kernel matrices for any choice of kernel K. such that we have K2(p,)) q) (p,),, A (g.:)) :)). CKA also allows choice of other kernels like RBF-kemnel, 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 becomes-
min L. = E(WW)) + J dLimeur-CKA Kg, K(t -t)- We
(4)
with (1-1)= W. K (t-1) 1) and Kイ 1) (, (1-1))) (X; 1))a. When f = O, each client is training its own local models without any peer-supervision and when μ + co the local models try to converge on the RAD representations without caring about the SSL loss.
5

---

Under review as a conference paper at ICLR 2023
Ondine Parsent
Pred cer
Chert
6
Loco
Taggl Breode 1erez10
Representation Aggreganor
RAD
Orai
Chort N
8
Local
Peare150
Clients
Server
Loss
Loss
SSL
Dissimlarity Loss
Representation
58L
ngo Upoad
Aggregaaed
Look BP LEIT comboned bee
Figure Overview of the proposed framework: Hetero-SSFL. In each training round, the server sends the RAD and the aggregated representation matrix to all clients, each client trains local self- supervised models by optimising the combined loss, the representations obtained using local model on RAD are sent back to the server which aggregates the representations and sends it back to all clients for next training round. Note: the local models are named Resnet- 18 and Resnet-50 as examples to depict heterogeneous client models.
Algorithm 1 Hetero-SSFL. Algorithm for Heterogeneouss clients
Input: number of clients N, number of global communication rounds T, number of local epochs E. parameter J4, weight vector for clients [w), W2, WN]
Output: Final set of personalised models (T),, W2(T) (T)
At Server- Initialize (1()), W2(0) ... Wx(()
for f. = 1 to 7 do
RAD=X
for each client j do
(X) = 1 (X: W3 (t (t- 1)) (X)
K, = -
end for ( 1(1-1)=E, === wy Kj
Select subset of clients M₁ for each selected client E Nr do
Wt( LocalTraining W. 1), R(t M.)
end for end for
Return (T),) Wy(T).. WN(T)
LocalTraining W. - 1), K(t 1) Initialize (t) with (1,,(t 1))
for each local epoch do Generate local representation matrix A;(X) (X: Generate local kernel matrix K, AX,
Update (t(t) after solving for loss in equation end for Return W,(t)
4 CONVERGENCE ANALYSIS
In this section we provide insights on convergence the proposed framework. The convergence guarantee is shown hold under mild assumptions which commonly made in the literature (Wang et al. 20206 lan et al. 2022).

---

Under review as a conference paper at ICLR 2023
Assumption 4.1 (Lipschitz Smoothness). The local loss function every client k, Lx. is assumed to be Li-Lisschitz smooth, which implies:
L1 (x(a) (b) (b) (Wa(a) (a( + - - Va,b> (5)
Assumption 4.2 (Stochastic Gradient). The local stochastic gradient for each client k at any time t, I is unbiased estimator of the gradient, has bounded variance and the expectation of its norm is bounded above. That is, we have
[[gx.)] = VL and Var(su,)) S o²,
Assumption 4.3 (Representation Norm). The norm of the representations obtained for any instance on any one client bounded above by R. That is:
[[ Wa) R.
Assumption 4.4 (Lipschitz Continuity). The embedding function ((:-) for all clients is L--Liischitz continuous, which for each client k implies:
(.; (-; Wa) ( Wb)ll2 S T2llWa - Wallz-
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 C for each client reduces after E local epochs and is bounded as below under the Assumption 4.2
L1n2 >
11 4, 0².
E[C s Co -n
(6)
2
where Co and LE denote the loss before and after É local epochs respectively and 1j is learning rate.
Proof of Lemma 4.5 is in Appendix A. Lemma 4.5 similar to (Tan el al. 2022). shows the bound on the client's local loss after the completion of local epochs under global communication round.
After the local training, the representations from all clients are sent to the server. Our next lemma 4.6 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, each client, after the E local updates, the global representation matrix is updated at the server and the loss function for any client lo gets modified to Lg from Lg and could be bounded as
E < E[ + 2pmL2PR L
Proof of Lemma 4.6 is in Appendix
Equipped with the results from Lemmas 4.5 and 4.6
ready to state main result which
provides a total deviation bound for an entire training round. Theorem (.7mmm4./hhww 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:
Lun²
E-1
En
Bi S Lo
2
0²
P
(8)
2
==0
Thus, if we choose 17 and μ such that
2(EE-1 VL 2nL2PR8L22
V
1=
L LE V 2 + Eo²)
2L
PR
then the convergence of Algorithm guaranteed.
The proof of Theorem 4.7 follows directly from the results of Lemmas
4.5 and 4.6
7

---

Under review as a conference paper at ICLR 2023
5 EXPERIMENTS
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.1 EXPERIMENTAL DETAILS
Datasets We test our framework on image classification 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 CIFAR-100 contain 50,000 train and 10,000 test colored images for 10 classes and 100 classes respectively and Tiny-Imagend has 100,000 images for 200 classes.
Implementation Details For the FL simulations, we explore two types of settings, non-IlI 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, 10, client 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 divided amongst the clients containing specific 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 TV 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 to 200 for 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. 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 reflect the current state-of-the-att 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 B BOL 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 Resnet-34 architecture with our heterogeneous setting having a mix of Resnet-34 and Resnet-18 architectures. By doing this, we 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. 202 to test the models. Under this protocol, the local encoders are trained in supe-supe way without any labeled data and then a linear classifier is trained on top of the encoder after freezing the encoder parameters. We use linear fully connected layer with 1000 neurons as the classifier 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 classifier 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 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 we fix that size to be
8

---

Under review as a conference paper at ICLR 2023
Table 1: Test Accuracy comparison of method with baselines under the linear evaluation protocol in non-lID settings.
Method
CIFAR-10
CIFAR-100
Tiny-Imagenet
FedBYoL
78.1 + 1.1
58.1 0.3
31.8 = 2,0
FedSimSiam
76.7 + 1.5
29.1 1.2
FedU
79.6 + 0.5
58.9 0.12
58.23 + 2.3
FedEMA
81.2 + 1.6
61.8 0.31
58.2 = 3.1
Hetero-SSFL
90.29 0.3
66.1 + 0.21
61.5 1.8
Supervised FedAvg (Supervised upper bound) Central-BYOL (IID upper bound)
68.5 + 1.7
44.29 + 1.34
27.2 3.2
94.3 +: 0.2
74.2 + 0.67
78.7 + 1.3
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 2: Test Accuracy comparison of our method with baselines under the linear evaluation protocol in IID settings(with equal number of classes all clients).
Dataset
FedB Yol.
FedSimSiam
FedU
FedEMA
Hetero-SSFL.
CIFAR-10
82.24 + 0.8
80.4 + 1.2
81.6 + 1.8
85.9 + 0.3
88.5 + 1.3
CIFAR-100
23.8 + 0.14
22.19 + 1.8
20.4 + 0.91
25.7 + 0.4
34.8 + 0.7
5.2 RESULTS
We show the performance of our method in comparison with baselines in non-llD settings in Table ! and that for IID settings Table [ We observe that our method outperforms the baselines by sigmificant 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 achieve successful collaboration amongst clients. We also report the performance of our method and high-performing baselines with increase scale 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 3: Test Accuracy of various federated self-supervised methods with change in scale CIFAR- 100 dataset for the non-IID setting.
Dataset
Sclients(5)
20clients(S)
100clients(10)
FedU
58.9 + 0.12
48.29 1.8
41.1 2.3
FedEMA
61.8 0.31
52.5 0.8
42.0 + 2.7
Hetero- -SFL
66.1 66.1 0.21
57.5 0.9
43.3 1.9
FedAvg (Supervised Upper Bound)
65.6 + 0.5
60.68 + 0.23
55.5 + 1.74
6 CONCLUSION
In this work we study the problem of self-supervised federated learning with heterogeneous clients and propose 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 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 s relative importance. Also, rather than assuming that the server has access to unlabeled dataset and transmitting RAD from the server to clients, common generative model in federated way like federated GANs to generate data instances required for RAD.
create

---

Under review as a conference paper at ICLR 2023
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. International Conference on Learning Representations, 2021.
Charu C Aggarwal and Philip S Yu. A general survey of privacy-preserving data mining models and algorithms. In Privacy-preserving data mining, pp. 11-52. Springer. 2008.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, Brendan McMahan, Virginia Smith, and Ameet Talwalkar, Leaf: A benchmark for federated settings, 2019.
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsuper- vised learning of visual features. Vittorio Ferrari, Martial Hebert, Cristian Sminchisescu, and Yair Weiss (eds.), Computer Vision ECCV 2018, pp. 139-156 156. Springer International Publishing, 2018.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Aisuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Process- ing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Hong-You You Chen and Wni-Lun Chao. Fed[be]: Making bayesian model ensemble applicable to federated learning. In International Conference on Learning Representations, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020a.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. arXiv preprint UrXiv:2011.10566 2020.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning, 2020b. cite arxiv:2003.08 arxiv:2003.04297Comm w4w..70mm9mmmt Tech report, pages + references.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared represen- tations for personalized federated learning. arXiv preprint arXiv:2102.07078, 2021.
Frances Ding, Sean-Stanislas Denain, and Jacob Steinhardt. Grounding representation similarity through statistical testing. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Sannara Ek, Romain Rombourg, François Portet, and Philippe Lalanda. Federated sepf-supe learning in heterogeneous settings: Limits of baseline approach on har. In 2022 IEEE Interna- tional Conference on Pervasive Computing and Communications Workshops and other Affiliated 2022. 2022.
Events (PerCom Workshops), pp. 557-562,
Alireza Fallah, Aryan Mokhtari, and Asuman F. Ozdaglar. Personalized federated learning: A meta-lcarning approach. CoRR, hass/2002.07944. 2020.
Meng Feng, Chich-Chi Kao, Qingming Tang. Ming Sun, Viktor Rozgic, Spyros Matsoukas, and Chao Wang. Federated self-supervised learning for acoustic event classification. In ICASSP 2022 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 481-485, 2022.
Wensheng Gan, Chun-Wad Lin, Han-Chich Chao, and Justin Zhan. Data mining in distributed environment: a survey. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery, 7:e1216, 07 2017. doi: 10. w0t.0002ww..1222 1216.
10

---

Under review as a conference paper at ICLR 2023
Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ranchandran. An efficient framework for clustered federated learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tice Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Jack Goetz and Ambuj Tewari. Federated learning via synthetic data. CoRR, abs/2008.04489 2020.
lean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukcuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent new approach to self-supervised learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 21271-21228 Curran Associates, Inc., 2020.
Weituo Hao, Mostafa El- El-Khamy, Jungwon Lee, Jianyi Zhang, Kevin J. Liang, Changyou Chen, and Lawrence Carin. Towards fair federated learning with zero-shot data augmentation. In IEEE Conference on Computer Vision and Pattern Recognition Workshops, CVPR Workshops 2021, Computer Vision Foundation IEEE, 2021.
virtual, June 19.25, 2021, pp. 1310-33199
Chaoyang Hc, Zhengyu Yang. Erum Mushtag, Sunwoo Lee, Mahdi Soltanolkotabi, and Salman Avestimehr. SSFL: Tackling label deficiency in federated learning via personalized self-supervision, 2022.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9726-9735, 2020. doi: 10. 10.109/CVR222000.200000000
Wonyong Jeong, Jaehong Yoon, Eunho Yang. and Sung Ju Hwang. Federated semi-supervised learning with inter-client consistency disjoint learning. In International Conference on Learning Representations, 2021.
Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning, 2020.
H. Kargupta and B.-H. Park. Collective data mining: A new perspective toward distributed data mining. In H. Kargupta and P. Chan (eds.), Advances in Distributed and Parallel Knowledge Discovery. AAAI/MIT Press, 2000.
H. Kargupta, W. Huang, Krishnamoorthy, and E. Johnson. Distributed clustering using collective principal component analysis. Knowledge and Information Systems Journal Special Issue on Distributed and Parallel Knowledge Discovery, 3:422 448, 2001.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 1 119 of Proceedings of Machine Learning Research, pp. 5132-5143. PMLR, 13-18 Jul 2020.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning, 2021. URL
Simon Korbblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3519-35211 PMLR, 2019.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: lite bert for self-supervised learning of language representations. In ICLR. OpenReview.net, 2020.
Qinbin Li, Bingsheng He, and Dawn Song. conel-contrastive federated learning. In Proceedings of the IEFE/CVF Conference Computer Vision and Pattern Recognition, 2021a.

---

Under review as a conference paper at ICLR 2023
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Inderjit S. Dhillon, Dimitris S. Papailiopou- los, and Vivienne Sze (eds.). Proceedings of Machine Learning and Systems 2020, MLSys 2020, Austin, TX, USA, March 2-4, 2020. mlsys.org. 2020.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 6357-6368. PMLR, 2021b.
Nan Lu, Zhao Wang, Xiaoxiao Li, Gang Niu, Qi Dou, and Masashi Sugiyama. Federated learning from only unlabeled data with paas-cooditionaa--sariig clients. In International Conference on Learning Representations, 2022.
Ekdeep Singh Lubana, Chi Ian Tang, Fahim Kawsar, Robert P. Dick, and Akhil Mathur. Orchestra: Unsupervised federated learning via globally consistent clustering. In ICML, pp. 14461-14484, 2022.
Mi Luo. Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, and Jiashi Feng. No fear of heterogeneity: Classifier calibration for federated learning with non-iid data, 2021.
Disha Makhija, Xing Han, Nhat Ho, and Joydeep Ghosh. Architecture agnostic federated learning for neural networks. In ICML. 2022.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--fficeenn Learning of Deep Networks from Decentralized Data. In Aarti Singh and Jerry Zhu (eds.), Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54 of Proceedings of Machine Learning Research, pp. 1273-1282, PMLR, 20-22 Apr 2017.
S. Merugu and J. Ghosh. distributed learning framework for heterogeneous data sources. In Proc. KDD, pp. 208- 2005.
S. Merugu and J. Ghosh. Privacy perserving distributed clustering using generative models. In Proc. ICDM, pp. 211-218, Nov, 2003.
Reese Pathak and Martin J Wainwright. Fedsplit: an algorithmic framework for fast federated optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 7057-7066. Curran Associates, Inc., 2020.
Yasar Abbas Ur Rehman, Yan Gao, Jiajun Shen, Pedro Porto Buarque de Gusmao, and Nicholas Lane. Federated self-supervised learning for video understanding, 2022.
Felix Sattler, Klaus-Robert Müller, and Wojeiech Samek. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. IEEE Trans. Neural Networks Learn. Syst. 32(8):3710-3222 2021.
Haizhou Shi, Youcai Zhang, Zijin Shen, Siliang Tang, Yaqian Li, Yandong Guo, and Yueting Zhuang. Towards communication--fficii and privacy-preserving federated representation learning, 2021.
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33, 2020.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017a.
Virginia Smith, Chao-Kai Kai Chiang, Maziar Sanjabi, and Ameet S. Talwalkar. Federated multi-task learning. In Isabelle Guyon, Ulrkee von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 4424 4434, 2017b.
12

---

Under review as a conference paper at ICLR 2023
Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across heterogeneous clients. In AAAI Conference on Artificial Intelligence, 2022.
Bram van Berlo, Aaqib Saeed, and Tanir Ozcelebi. Towards federated unsupervised representation learning. In Proceedings of the Third ACM International Workshop on Edge Systems, Analytics and Networking, pp. 31-36, New York, NY. USA, 2020. Association for Computing Machinery. ISBN 9781450371322.
3394333443
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Fed- erated learning with matched averaging. International Conference on Leurning Representations, 2020a.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, Curran Associates, Inc., 2020b.
volume 33, pp. 7611-7623.
Zhirong Wu, Yuanjun Xiong, Stella X. Yu, and Dahua Lin. Unsupervised feature learning via non-parametric instance discrimination. 2018 IEEE/CVF Conference on Computer Vision and Pultern Recognition, pp. 3733-3742, 2018.
Xueting Yan, Ishan Mísra, Abhinav Gupta, Deepti Ghadiyaram, and Dhruv Mahajan. Clusterfit: Improving generalization visual representations. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pp. 6508-6517. Computer Vision Foundation IEEE, 2020.
Hwanjo Yu, Xiaogian Jiang, and Jaideep Vaidya. Privacy-preserving svm using nonlinear kernels on horizontally partitioned data. In Proceedings of the 2006 ACM Symposium on Applied Computing, SAC '06, pp. 603-610, New York, NY, USA. 2006. Association for Computing Machinery. ISBN 1595931082, doi: 10. 1414//41277...1144555
Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. CoRR, abs/2002.04758, 2020.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald. Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 7252-7261. PMLR, 2019.
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. Fedpd: federated learning framework with adaptivity to non-iid data. IEEE Transactions Signal Processing, 69:6055-6070, 2021a. doi: 00.110997SP222211111992.
Zhengming Zhang, Yaoqing Yang, Zhewei Yao, Yujun Yan, Joseph E Gonzalez, Kannan Ramchan- dran, and Michael W Mahoney. Improving semi-supervised federated learning by reducing the gradient diversity of models. IEEE International Conference on Big Data (Big Data), 2021b.
Weiming Zhuang, Xin Gan, Yonggang Wen, Shuai Zhang, and Shuai Yi. Collaborative unsupervised visual representation learning from decentralized data. CoRR, abs/2108.0640 2021.
Weiming Zhuang, Yonggang Wen, and Shuai Zhang. Divergence- federated self-supervised learning. In International Conference Learning Representations, 2022.
13

---

Under review as a conference paper at ICLR 2023
Supplement for "Federated Self-supervised Learning for Heterogeneous Clients"
In this supplementary material, we provide proofs for the key results in the paper in Appendix A along with the intuitive explanations of the results. Then, we present a setting to illustrate the framework's utility in the text classification setting in Appendix B and provide additional experiments in Appendix C
A PROOFS
We first provide proof of Lemma 4.5 in Appendix A.I Then, the proof of Lemma i.6 given in Appendix A.2
A.1 PROOF OF LEMMA 4.5
Let Lg denote the loss function after E' local epochs. Then, from the Lipschitz smoothness assumption we obtain that
in Assumption 4.1
Ln
L1 S Lo gol +
11
Therefore, we have
E[L]] & B[Co] - 7 goو + !,12
E[11go|1
2
E[Co] - 7/ /. 11 + L1
2
L
B|[o] V 112 +
(V
(gr( + E|ll )11
2
2 2
-
B[Co] - (n
|| + 2
By summing the above bounds over E number of local epochs, we find that
E
E-1
Lun E-1
E-1
In
Enj
Σ
EEC - Σ
[|C,] - (1
| 2 +
2
1=0
L
L, En
E[Lg] < Co - (1
1|
+
o².
2
2
1=0
2(> im 22
Therefore, if 1)
we obtain LE S Lo and the local loss reduces after
L <> 0
L + Eoz)
E local epochs. As consequence, obtain the conclusion of the lemma.
A.2 PROOF OF LEMMA 4.6
For any arbitrary client after the global representation update step, if the loss function gets modified to Lg from Lg, have
L = LE + LE LE:
LE + ( (d(::t) - d(;;t d(:;-11) -1))
As for client i, d(.:t) (Linaa-CKK((K(()) have,
LE LE (Lnear-CKA(KK())) K(t)) Linear-CKA(KK((), K(t R(t-1))
14

---

Under review as a conference paper at ICLR 2023
Since this holds for all clients, we drop the client index going forward and obtain that
E LF + p (tace(((t)((tt) (race(((t)R((t
11)
LE + 11
(())))(( R(t - 1)1)
N
= LE + 1
trace K (1)
wk kx(t) - Σ
wk K - 1))
R=1
k=1
LE + p
K(t)
Σ (t) Σ Wa - 1)bm))
= LE + 1
Σ
WiR
((0(s(Ka(0)) - Kull - 1)ss)))
(x((-1)(s))
k=1
L
LE + p
Wik
K(D),((butz ( WE) - t (1:s 1) w/ (11
=1
Given the above equations, we find that
LE = =L= + 14
Wk
k(t)as ()(c:WW)((0((s)) - *2(23; W 1))
+ (O - W/ 1))
Taking norm on both the sides, we get
||
LE + 1
K(t)ss
(Bl(:: ( - D,(1;; W/
k=
1=1
+ (zzz W/¹ ( (x) WP) (e(
< |Coel| + u Σ Wk K(t)es
1))
+ W (
( (2x; Wi) To (x (x:
LE + 11 Wk K( (I) WX) o,(ay; w
+ + W (x) W 2a(
< LE + 1 Σ Wk K(t)asil (I0A((T: [1) W IDIl
+ 11 / ( -
N
L +
[(t),,)
(10 W + (2 W (
1=1
15

---

Under review as a conference paper at ICLR 2023
CIFAR-10, 10, n=10, 2 classes per client
4.0
3.5
Local Client Loss
3.0
2.5
2.0
1.5
10
0.5
50
100
150
200
250
300
350
400
Iteration
Figure Local loss at the end of each training round for randomly selected client.
Taking expectation on both sides of the above bounds, we have
N
E[Gg S E(LE] 11
WF
< E(CE] 11
||<())=))
[RLnnP PL21P)
WFkk
1=
E[CE] + 2pm PR
|K(()s
L L
A <B[LE] + 2 PR
Σ
R?
==
[[CE] + 2
Thus, we have,
B[CF] < B[C=] 2umL2
which concludes the proof.
A.3 CONVERGENCE DEMONSTRATION IN EXPERIMENTS
We monitor the local loss at the end of each training round and we see that the decreasing loss theorem holds in practice well. The Figure 2 shows plot of the local loss for one randomly selected client with the number of global round for the proposed method on CIFAR-I0 dataset.
B APPLICATIONS IN OTHER DOMAINS
While the experiments with image recognition tasks demonstrated the potential of the approach in popular federated self-supervised settings, note that the Hetero-SSFL framework is indeed extensible to other domains and applications as well since the key idea the framework (i.e., enhance local self-supervised learning heterogeneous clients with limited data by obtaining peer supervision from other clients) is broadly applicable. We now briefly illustrate how the proposed method can be applied for the text classification task.
16

---

Under review as a conference paper at ICLR 2023
Consider each client to be learning transformer for language modeling using the self- self-supervised Masked Language Modeling(MLMM objective. This objective tries to reconstruct the sentence by predicting the masked words in the sentence. Thus, the local self-supervised loss function for client k, E(W),, corresponds to the MLM loss. The peer supervision is achieved as usual, by server obtaining the embeddings on the RAD from each client, aggregating those embeddings and sending it to the clients for the next round of training. The local loss function at each client is similar to Equation (4) but with the new ((Ws). After pre-training using self-supervision, we evaluate the performance of training by following linear evaluation protocol using sentiment classification task on Sentiment| 40 dataset from LEAF (Caldas al. 2019).
Experimental Setup For local self-supervised learning we use ALBERT transformer (Lan et al. 2020). We partition the Sentiment149 data non-iid partitions for each client and fine-tune the transformer for each client on locally available data for 200 rounds, using Adam W optimizer and a range μa values between .01 and 100 and report the best performance. After pre-training, the parameters of the transformer are frozen and linear classifier is trained on the sentence embeddings obtained from the transformer. The test accuracy the linear classifier is reported as the result.
For baselines, consider the client's 's local model training local data, FedAvg version of ALBERT transformer and a divergence aware update method as shown in the baselines FEDEMA and FedU. Since the other self-supervised federated methods like FedEMA and FedU consider only the image classification task, we are unable to compare against them.
The Table below shows the test accuracy of Hetero-SSFF in comparison with the above mentioned baselines.
Table Test accuracy for different methods for text classification on Sentiment 40 dataset.
Method
FedAvg
Local
Divergence-Awer Aware
Hetero-SSFl
Test Accuracy 61.53 + 7.2 61.7 = 2.2
58.19 + 9.2 62.79 1.7
C ADDITIONAL EXPERIMENTS
Table Top-1 test accuracy for different methods under semi-supervised learning for CIFAR100.
Method
1% labeled data 10% labeled data
FedU
20.5 0.5
46.3 0.57
FedEMA
23.11 = 0.65
48.6 = 0.47
Hetero-SSFL
37.4 + 0.95
57.4 0.43
Table 6: Top-1 test accuracy for different methods under semi-supervised learning for CIFARIO.
Method
1% labeled data 10% labeled data
FedU
72.9 + 2.25
86.8 + 0.5
FedEMA
75.8 + 1.4
87.4 + 0.8
Hetero-SSFL.
79.3 + 1.85
88.7 + 0.25
Table 7: Average effect of collaboration on different clients. Column depicts the test accuracy of clients with local models and column 3 the accuracy under Hetero-SSFL.
Client Model Architecture
Local SSL
Hetero-SSFL.
Resnet-18
71.9
85.4
Resnet-34
77.8
93.7
We also test our framework on the semi-supervised learning task, as described in FedU (Zhuang et al. 202T) and FedEMA (Zhuang et al. 2022) papers. Under this protocol, we let the federated SSFT.
17

---

Under review as a conference paper at ICLR 2023
methods train the local models for representation learning. For the evaluation, we assume that only a small (raction11%% or 1%) of labeled data is available. Then, after the federated representation learning of local models clients, we fine-tune the local models along with linear classification layer on top for 100 epochs. We evaluate our method and the baselines, FedU and FedEMA, in this setting CIFAR1O and CIFAR 100 datasets and observe that our method outperforms the baselines. For fine-tuning the models, we use SGD with Nesterov momentum as the optimizer with momentum 0.9, learning rate 0.05 and batch size 512 and 256 for 10% labeled data and 1% labeled data respectively. The results reported in Table 5 and Table 6 are for federated setting with non-III clients with 5 users, 20 classes per user for CIFARIO0 and 2 classes per user for CIFARIO.
For additional insights, we show the average change in test accuracy of clients with different architectures when performing standalone training versus collaborative learning in the heterogeneous setting in Table 7
D DETAILS ABOUT RAD
The Representation Alignment Dataset (RAD) is important for training in Hetero-SSFL, Functionally, the RAD is used to guide the simultaneous training on the peer clients by serving as a common dataset for representation alignment. Also, since the models are not transmitted and only the RAD representations are transmitted from clients to the server, there is no privacy loss due to involvement of RAD in the training process. The only assumption about this dataset is to be of the same domain as the training data, i.e., for image classification tasks it could be images sourced from the web and for the clxsilaasiiiiatoon tasks could be publicly available text like from Wikipedia, Reddit etc. Because the only requirement for RAD is to come from the same domain, it does not pose practical constraints. Several applications can make use of publicly available datasets to use for training in Hetero-SSFL, for example, learning self-supervised models for image datasets could use MNIST or CIFAR datasets for training, similarly acoustic applications can use audio datasets provided by Huggingface, and likewise.
we subset part of the dataset (5000 data points) and set it aside to be used as RAD before beginning the training process. Since the RAD is considered generic with no bearing on clients, it is created before the clients are intialised and data is partitioned into clients. We further analysed the effect of using simple MNIST dataset as an RAD and saw similar results. The change in performance with respect the size of RAD is given below in the Figure 3 When the RAD size is very low, the local models performance is closer to the performance in standalone training. As the size of the RAD grows, the performance gets better but because the gain is small after 3000 points, we keep 7. = 3000 for all the experiment simulations. Also, in cases when μ = O, the performance of the Hetero-SSFL is same as standalone training ensuring that the performance gain is not due to additional data in the form of RAD,
In the experiments described in Table 1 and Table 2
E COMMUNICATION COST ANALYSIS
In Hetero-SSt SSFI, each global round involves the server sending the common dataset, RAD, to each client and the clients sending back the representations obtained on the RAD. The size of the RAD is given by L, thus the communication cost in each round is O(L). In all our experiments, we set the RAD size between 3000 and 5000. Thus the communication cost in each round of Hetero-SSFL corresponds to transferring these many (3000-5000) parameters twice. This is orders of magnitude less than the conventional FL algorithms that require sharing of the model parameters which run in millions. For federated training of the Resnet18 models, having 1 11M number of parameters, Hetero-SSFL. with I. = 5000 will require 2000 times less communication cost than conventional methods in each round, making it very practical.
F ABLATION STUDIES
18

---

Under review as a conference paper at ICLR 2023
50
Arcuray
55
45-
Hetero-SSFL
40
1000
1500
2000
2300
3200
RAD
Figure 3:
Change performance with change in RAD size for CIFAR-100 dataset in non-IID setting.
Table 8: Test Accuracy comparison of our method with increasing heterogeneity for CIFAR-100 datasets and non-IlI setting.
Setting
Small Models
Resnet Models
20 % small
39.4 + 1.55
64.2 = 2.9 63.7 +: 1.89
40 % small
37.5 + 1.96
Standalone Training
28.8 + 0.9
55.3 + 1.3
F.1 INCREASED HETEROGENEITY IN CLIENTS
To demonstrate the effectiveness of the framework with diverse types of heterogeneous clients, experiment by adding clients with models of much smaller capacities like CNN and MLP along with Resnet18 and Resnet34 models. We use Hetero-SSFL. in this setup on CIFAR-100 dataset achieve federated training of self-supervised models on CIFAR- 100 dataset and use linear evaluation protocol, as described in Section [5. to measure the test accuracy. We show the average performance of the small capacity clients, and average performance of the large capacity clients in settings with 20% and 40% small capacity models in Table as compared the standalone training.
F.2 EFFECT OF μ
The parameter regulates the weight of the proximal term (used for peer alignment) against the local self-supervised loss, shown Equation Thus, the value of μ plays an important role in the entire training process. 0 corresponds to standalone training where each client is minimizing the local supervisdd loss on its own without any federation. As the μ increases 0.001, 0.01, some federation starts to happen and the performance gets better than the standalone training. μ 0.1 achieves a good balance the two loss functions for CIFAR datasets, and μ = 1 for Tiny-Imagenet. But the starts increasing beyond 1, the weight of the proximal term becomes higher than the self-supervised loss leading to a drop in performance. This effect is more pronounced when the heterogeneous models are diverse, setting j 10 in those settings inhibits learning the more powerful clients while still improving the performance on the less powerful clients. With this knowledge, we can potentially vary y a client with less number of data points and lower training capacity could benefit more from proximal term by setting higher μ.
clients depending their local resources, for example,
19


Reviews:
[{"review_id": "uj1q05GYwn", "reviewer": "Reviewer_N6gv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents Hetero-SSFL, a framework for SSL in FL that enables diverse clients (heterogeneous in terms of their compute capabilities and/or data distribution) to collaboratively learn a good SSL model. At the end of N rounds, the SSL model gets personalised (via linear evaluation protocol) to each client's data. Different from the majority of previous FL works, is the fact that clients in Hetero-SSFL do not communicate their locally-trained model with the server. Instead, what clients send are the embeddings outputted by the local model when using a common (i.e. same in all clients) dataset. These representations are then aggregated in the server and used in the following round to align (via a distance loss term) the representations each client learns during SSL. Hetero-SSFL is primarily evaluated on vision classification tasks but also text classification.", "summary_of_the_review": "I am happy with this work but there are some open questions that need to be addressed ( see both the Strengths and Weaknesses lists) specially in regards to the RAD and results in Table1&2. I think the Authors can address these during the rebuttal period easily.", "strength_and_weaknesses": "### Strengths\n\n*    The paper is fairly well written, although it would benefit if Figure 1 was presented earlier in the text (maybe at the top of page 4?) Making it more self-contained would be good too.\n\n*    The results in Tables 1,2,3 look super good compared to the other baselines but. Could the Authors explain why the figures in Table 2 are lower than in Table 1? It seems like all the baselines do better in the IID case (as you'd normally expect) for CIFAR-10 but Hetero-SSFL does worse... But then for CIFAR-100 in Table 2 (IID) all methods do worse than in Table 1 (non-IID). Could the Authors elaborate? Maybe the numbers got mixed up adding them to the table?\n\n### Weaknesses\nSome weaknesses and clarifications needed:\n\n*    The Authors present this work as enabling collaboratively SSL in FL when clients are of different compute capabilities. However, only two types of networks ResNet-18/34 are considered and both train with batch size 200. That leaves the setting to just two types of clients. If the Authors considered a few more models to cover a wider spectrum of system capabilities (specially towards more lightweight models -- e.g mobilenet, efficientnet or similar?) and different batch sizes (which directly impacts on memory peak at training time), then I would be convinced that Hetero-SSFL is indeed addressing the client device heterogeneity problem.\n\n*    The RAD is one of the key ingredients in Hetero-SSFL but the Authors do not provide details about it. Is it also a vision dataset? which one? How does the quality of the SSL change with different RADs?\n\n*    The Authors state that the RAD \"doesn't posses any special properties\", could this be elaborated? I'm inclined to size and diversity in the RAD would be important, but to what extent? is MNIST a good RAD?\n\n*    How the SSL training on the clients happens is clear. However, how is the linear evaluation protocol implemented in this work? Is there labelled \"validation\" set to train the classifier and then the evaluation is done on a, yet another, sub-set of data that the client has (i.e. a labelled test set)? Do these validation and test sets follow the same label distribution as the (unlabelled) training set each client has? \n\nIn Table 3, the Authors present some results on how Hetere-SSFL extrapolates from the cross-silo (5 total clients and 5 clients per round -- as considered for Table1 and, I believe, Table 2 also) to a cross-device setting. I understand the cross-device setting is more challenging, specially when it comes to ensuring the alignment of representations helps the training process. \n*    Could the Authors comment upon why SSL suffers such a sharp drop in performance in cross-device settings?\n*    It would be good to include in Table 3 the standard (non-SSL) FL baseline, which should be something in the high 50%s, given that FedAdam(https://arxiv.org/abs/2003.00295) achieves 52.5% in a setup with 500 clients (sampling 10 per round I believe). Adding that extra row is good for putting the results in context: SSL is really important and clearly the way to go for FL, but we aren't there yet.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "qGKNTtq9qG", "reviewer": "Reviewer_GXVZ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a new method for federated self-supervised learning with clients who may train on local tasks using heterogeneous model architectures. The authors adopt a Siamese structure for local self-supervised learning, implemented to BYOL or SimSiam, and validate it using both image and text classification tasks. To allow semantic aggregation of local clients' representation, they introduce a public dataset named Representation Alignment Dataset (RAD) that all local models utilize for implicitly transferring representations across clients by minimizing the CKA distance of RAD representation. Therefore, the proposed model doesn't need direct weight aggregation and allows communication with heterogeneous clients.\n", "summary_of_the_review": "The paper proposes a practical framework allowing model heterogeneity during federated self-supervised learning. The method is reasonable and provides convergence analysis, but there is a lack of analyses to understand the behavior of the proposed methods and the issue of RAD construction.", "strength_and_weaknesses": "**strength**\n- propose a new method to transfer the knowledge across clients without direct weight aggregation, so that allows participation in heterogeneous self-supervised models \n- validate the model over text- and image-based classification problems\n- provide the convergence analysis for the proposed model\n\n**weakness**\n- Missing communication cost analysis.\n- Missing analyses of the RAD size and the construction of the RAD. To avoid the data privacy issue, the distribution of the RAD should be different from the distributions of local datasets (i.e., Non-IID). Therefore, the construction of the RAD can highly impact the model performance.\n- Only two ResNet-variants are used for architecture heterogeneity. More diverse architectures are required for backbones, such as Transformers, simple CNNs, and MLPs.\n- Performance gap with baselines decreases for larger participating clients at each round.\n\n----\n+ local clients from baselines do not train on the samples constructed to RAD?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LuMHpl_nTZ", "reviewer": "Reviewer_9iyo", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Summary: This work proposes an FL paradigm for SSL under client heterogeneity. The proposed method, Hetero-SSFL allows each client to train personalized and different self-supervised models. It then enables joint learning across clients by aligning the lower dimensional representation with those of a common dataset assumed to be available on the server. The work is an extension of Makhija et al., 2022 by changing the supervised learning task to SSL. Also, Herero-SSFL has the impractical assumption of having a common dataset on the server, which may limit its application in some real settings. The author compared Hetero-SSFL with other FL+SSL methods on benchmark datasets.", "summary_of_the_review": "The paper aims to address the interesting question, but the novelty and the written quality downgrade it from the other submissions. ", "strength_and_weaknesses": "Pros: Overall, the idea of allowing different network architecture across clients in FL for SSL is very interesting. The selected technique foundation (Makhija et al., 2022) is reasonable. \n\nCons: \n1.Limited methodology novelty. Despite of the different learning paradigm, Hetero-SSL is very similar to Makhija et al., 2022 (for supervised learning tasks). The core techniques discussed in Sec 3, including general formulation, RAD, CKA, etc., seems directly borrow from Makhija et al., 2022. The authors may want to highlight the differences if there is any or give sufficient credits to the previous work. \n\n2.The assumption of requiring a common public dataset available in the server will raise many practical concerns. For example, what kind of data can be considered as the common dataset, how many data points need to be used, is it easy to find the privacy-free public data to serve this purpose, etc.\n\n3.What do a and b mean in Assumption 4.1. I assume they are two iterations. Then don’t you want a,b > 0?\n\n4.Lemma 4.5 is similar to Tan et al.  The authors should give proper reference here.\n\n5.Are the results in Table 3 from IID or non-IID setting?\n\n6.How to select the common (global) data is missing in Section 5.\n\n7.Grammar and format issues should be corrected.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Gf2oYcO97Ve", "reviewer": "Reviewer_X5at", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposed a new approach called Hetero-SSFL, to conduct federated learning under self-supervised tasks with heterogeneous clients. Specifically, besides local training, all clients learn to align the lower dimensional representations on a common dataset. The paper provides a convergence guarantee theoretically and demonstrate the improved performance empirically.", "summary_of_the_review": "The paper has some important details missing and the experimental part was not comprehensive enough. The current version is not ready for publication.", "strength_and_weaknesses": "Strengths:\n- The modification by adding an objective to make the representation on a common dataset aligned is intuitive and clear. \n- Authors gave a convergence guarantee for the proposed method, making the approach more thorough and convincing from the theoretical side.\n\nWeaknesses:\n- Experiments were not very comprehensive. Some important ablation studies were missing. For example, how will the coefficient $\\mu$ will affect the performance? And what is the influence of the size of the common dataset RAD? It is not clear whether the performance gain was due to the additional data.\n- It was not clear how the server utilized the representations from different local clients, given that the dimensions were not the same. In Eq. (1), the second term was the distance of two representations, while in Eq. (4), it was changed to two kernels. It seems that the paper did not mention how to aggregate representations.\n- Comparison of communication efficiency in federated learning is necessary in terms of metrics such as the size of transmitted data in one round.\n- Some implementation details were not complete. For example, how to obtain the RAD dataset? The authors also stated that weight vector $[w_1, w_2, \\dots, w_N]$ can be learned, but how to achieve it? Will learning the weight vector require the communication between the server and clients just as FedAvg to average weights from different clients?\n- There is a lack of discussion about privacy in federated learning and the authors should talk about potential privacy problems and whether the method can preserve differential privacy.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject