DIVERGENCE-AWARE FEDERATED SELF-SUPERVISED LEARNING
Year: 2022

Abstract:
Self-supervised learning (SSL) is capable of learning remarkable representations from centrally available data. Recent works further implement federated learning with SSL to learn from rapidly growing decentralized unlabeled images (e.g., from cameras and phones), often resulted from privacy constraints. Extensive attention has been paid to SSL approaches based on Siamese networks. However, such an effort has not yet revealed deep insights into various fundamental building blocks for the federated self-supervised learning (FedSSL) architecture. We aim to fill in this gap via in-depth empirical study and propose a new method to tackle the nonindependently and identically distributed (non-IID) data problem of decentralized data. Firstly, we introduce a generalized FedSSL framework that embraces existing SSL methods based on Siamese networks and presents flexibility catering to future methods. In this framework, a server coordinates multiple clients to conduct SSL training and periodically updates local models of clients with the aggregated global model. Using the framework, our study uncovers unique insights of FedSSL: 1) stop-gradient operation, previously reported to be essential, is not always necessary in FedSSL; 2) retaining local knowledge of clients in FedSSL is particularly beneficial for non-IID data. Inspired by the insights, we then propose a new approach for model update, Federated Divergence-aware Exponential Moving Average update (FedEMA). FedEMA updates local models of clients adaptively using EMA of the global model, where the decay rate is dynamically measured by model divergence. Extensive experiments demonstrate that FedEMA outperforms existing methods by 3-4% on linear evaluation. We hope that this work will provide useful insights for future research.

heading: 1 INTRODUCTION
text: Self-supervised learning (SSL) has attracted extensive research interest for learning representations without relying on expensive data labels. In computer vision, the common practice is to design proxy tasks to facilitate visual representation learning from unlabeled images (Doersch et al., 2015; Noroozi & Favaro, 2016; Zhang et al., 2016; Gidaris et al., 2018). Among them, the state-of-the-art SSL methods employ contrastive learning that uses Siamese networks to minimize the similarity of two augmented views of images (Wu et al., 2018; Chen et al., 2020a; He et al., 2020; Grill et al., 2020; Chen & He, 2021). All these methods heavily rely on the assumption that images are centrally available in cloud servers, such as public data on the Internet.
However, the rapidly growing amount of decentralized images may not be centralized due to increasingly stringent privacy protection regulations (Custers et al., 2019). The increasing number of edge devices, such as street cameras and phones, are generating a large number of unlabeled images, but these images may not be centralized as they could contain sensitive personal information like human faces. Besides, learning representations from these images could be more beneficial for downstream tasks deployed in the same scenarios (Yan et al., 2020). A straightforward method is to adopt SSL methods for each edge, but it results in poor performance (Zhuang et al., 2021a) as decentralized data are mostly non-independently and identically distributed (non-IID) (Li et al., 2020a).
Federated learning (FL) has emerged as a popular privacy-preserving method to train models from decentralized data (McMahan et al., 2017), where clients send training updates to the server instead
of raw data. The majority of FL methods, however, are not applicable for unsupervised representation learning because they require fully labeled data (Caldas et al., 2018), or partially labeled data in either the server or clients (Jin et al., 2020a; Jeong et al., 2021). Recent studies implement FL with SSL methods that are based on Siamese networks, but they only focus on a single SSL method. For example, FedCA (Zhang et al., 2020a) is based on SimCLR (Chen et al., 2020a) and FedU (Zhuang et al., 2021a) is based on BYOL (Grill et al., 2020). These efforts have not yet revealed deep insights into the fundamental building blocks of Siamese networks for federated self-supervised learning.
In this paper, we investigate the effects of fundamental components of federated self-supervised learning (FedSSL) via in-depth empirical study. To facilitate fair comparison, we first introduce a generalized FedSSL framework to embrace existing SSL methods that differ in building blocks of Siamese networks. The framework comprises of a server and multiple clients: clients conduct SSL training using Siamese networks — an online network and a target network; the server aggregates the trained online networks to obtain a new global network and uses this global network to update the online networks of clients in the next round of training. FedSSL primarily focuses on the cross-silo FL where clients are stateful with high availability (Kairouz et al., 2019).
We conduct empirical studies based on the FedSSL framework and discover important insights of FedSSL. Among four popular SSL methods (SimCLR (Chen et al., 2020a), MoCo (He et al., 2020), BYOL (Grill et al., 2020), and SimSiam (Chen & He, 2021), FedBYOL achieves the best performance, whereas FedSimSiam yields the worst performance. More detailed analysis uncover the following unique insights: 1) Stop-gradient operation, essential for SimSiam and BYOL, is not always essential in FedSSL; 2) Target networks of clients are essential to gain knowledge from online networks; 3) Keeping local knowledge of clients is beneficial for performance on non-IID data.
Inspired by the insights, we propose a new approach, Federated Divergence-aware Exponential Moving Average update (FedEMA) 1 , to address the non-IID data problem. Specifically, instead of updating online networks of clients simply by the global network, FedEMA updates them via exponential moving average (EMA) of the global network, where the decay rate of EMA is measured by the divergence of global and online networks dynamically. Extensive experiments demonstrate that FedEMA outperforms existing methods in a wide range of settings. We believe that important insights from this study will shed light on future research. Our main contributions are threefold:
• We introduce a new generalized FedSSL framework that embraces existing SSL methods based on Siamese networks and presents flexibility catering to future methods.
• We conduct in-depth empirical studies of FedSSL based on the framework and discover deep insights of the fundamental building blocks of Siamese networks for FedSSL.
• Inspired by the insights, we further propose a new model update approach, FedEMA, that adaptively updates online networks of clients with EMA of the global network. Extensive experiments show that FedEMA outperforms existing methods in a wide range of settings.

heading: 2 RELATED WORK
text: Self-supervised Learning In computer vision, self-supervised learning (SSL) aims to learn visual representations without any labels. Discriminative SSL methods facilitate learning with proxy tasks (Pathak et al., 2016; Noroozi & Favaro, 2016; Zhang et al., 2016; Gidaris et al., 2018). Among them, contrastive learning (Oord et al., 2018; Bachman et al., 2019) has become a promising principle. It uses Siamese networks to minimize the similarity of two augmented views (positive pairs) and maximize the similarity of two different images (negative pairs). These methods are either contrastive or non-contrastive ones: contrastive SSL methods require negative pairs (Chen et al., 2020a; He et al., 2020) to prevent training collapse; non-contrastive SSL methods (Grill et al., 2020; Chen & He, 2021) are generally more efficient as they maintain remarkable performances using only positive pairs. However, these methods do not perform well on decentralized non-IID data (Zhuang et al., 2021a). We analyze their similarities and variances and propose a generalized FedSSL framework.
Federated Learning Federated learning (FL) is a distributed training technique for learning from decentralized parties without transmitting raw data to a central server (McMahan et al., 2017).
1Intuitively, FedSSL is analogous to a SuperClass in object-oriented programming (OOP), then FedEMA is a SubClass that inherits FedSSL and overrides the model update method.
Among many studies that address the non-IID data challenge (Zhao et al., 2018; Li et al., 2020b; Wang et al., 2020; Zhuang et al., 2021c), Personalized FL (PFL) aims to learn personalized models for clients (Tan et al., 2021). Although some PFL methods interpolate global and local models (Hanzely et al., 2020; Mansour et al., 2020; yuyang deng et al., 2021), our proposed FedEMA differ in the motivation, application scenario, and measurement of the decay rate. Besides, the majority of existing works only consider supervised learning where clients have fully labeled data. Although recent works propose federated semi-supervised learning (Jin et al., 2020b; Zhang et al., 2020b; Jeong et al., 2021) or federated domain adaptation (Peng et al., 2020; Zhuang et al., 2021b), they still need labels in either the server or clients. This paper focuses on purely unlabeled decentralized data.
Federated Unsupervised Learning Learning representations from unlabeled decentralized data while preserving data privacy is still a nascent field. Federated unsupervised representation learning is first proposed by van Berlo et al. (2020) based on autoencoder, but it neglects the non-IID data challenge. Zhang et al. (2020a) address the non-IID issue with potential privacy risk for sharing features. Although Zhuang et al. (2020) address the issue based on BYOL as our FedEMA, they do not shed light on why BYOL works best. Since SSL methods are evolving rapidly and new methods are emerging, we introduce a generalized FedSSL framework and deeply investigate the fundamental components to build up practical guidelines for the generic FedSSL framework.

heading: 3 AN EMPIRICAL STUDY OF FEDERATED SELF-SUPERVISED LEARNING
text: This section first defines the problem and introduces the generalized FedSSL framework. Using the framework, we then conduct empirical studies to reveal deep insights of FedSSL.

heading: 3.1 PROBLEM DEFINITION
text: FedSSL aims to learn a generalized representation W from multiple decentralized parties for downstream tasks in the same scenarios. Each party k contains unlabeled data Dk = {Xk} that cannot be transferred to the server or other parties due to privacy constraints. Data is normally non-IID among decentralized parties (Li et al., 2020a); each party could contain only limited data categories (e.g., two out of ten CIFAR-10 classes) (Luo et al., 2019). As a result, each party alone is unable to obtain a good representation (Zhuang et al., 2021a). The global objective function to learn from multiple parties is minw f(w) := ∑K k=1 nk n fk(w), where K is the number of clients, and n = ∑K k=1 nk is the total data amount. For client k, fk(w) := Exk∼Pk [f̃k(w;xk)] is the expected loss over data distribution Pk, where xk is the unlabeled data and f̃k(w;xk) is the loss function.

heading: 3.2 GENERALIZED FRAMEWORK
text: We introduce a generalized FedSSL framework that empowers existing SSL methods based on Siamese networks to learn from decentralized data under privacy constraints. Figure 1 depicts the
end-to-end training pipeline of the framework. It comprises of three key operations: 1) Local Training in clients; 2) Model Aggregation in the server; 3) Model Communication (upload and update) between the server and clients. We implement and analyze four popular SSL methods — SimCLR (Chen et al., 2020a), MoCo (V1 (He et al., 2020) and V2 (Chen et al., 2020b)), SimSiam (Chen & He, 2021), and BYOL (Grill et al., 2020). Variances in Siamese networks of these methods lead to differences in executions in these three operations 2.
Local Training Firstly, each client k conducts self-supervised training on unlabeled data Dk based on the same global model W og downloaded from the server. Regardless of SSL methods, clients train with Siamese networks — an online network W ok and a target network W t k for E local epochs using cooresponding loss functions L. We classify these SSL methods with two major differences (Figure 8 in Appendix A): 1) Only SimSiam and BYOL contain a predictor in the online network, so we denote their online network W ok = (Wk,W p k ), where Wk is the online encoder and W p k is the predictor; As for SimCLR and MoCo, W ok = Wk. 2) SimCLR and SimSiam share identical weights between the online encoder and the target encoder, so W tk = Wk. In contrast, MoCo and BYOL update the target encoder with EMA of the online encoder in every mini-batch: W tk = mWk + (1−m)W tk, where m is the momentum value normally set to 0.99. Model Communication After local training, client k uploads the trained online network W ok to the server and updates it with the global model W og after aggregation. Considering the differences of SSL methods, we upload and update encoders and predictors separately: 1) we upload and update the predictor when it presents in local training; 2) we follow the communication protocol Zhuang et al. (2021a) to upload and update only the online encoder Wk when encoders are different.
Model Aggregation When the server receives online networks from clients, it aggregates them to obtain a new global model W og = ∑K k=0 nk n W o k . W o g = (Wg,W p g ) if predictor presents, otherwise W og = Wg , where Wg is the global encoder. Then, the server sent W o g to clients to update their online networks. The training iterates these three operations until it meets the stopping conditions. At the end of the training, we use the parameters of W og as the generic representation W for evaluation.

heading: 3.3 EXPERIMENTAL SETUP
text: We provide basic experimental setups in this section and describe more details in Appendix B.
Datasets We conduct experiments using CIFAR-10 and CIFAR-100 datasets (Krizhevsky et al., 2009). To simulate federated settings, we equally split a dataset into K clients. We simulate nonIID data with label heterogeneity, where each client contains limited classes — l = {2, 4, 6, 8, 10} number of classes for CIFAR-10 and l = {20, 40, 60, 80, 100} for CIFAR-100. The setting is IID when each client contains 10 (100) classes for CIFAR-10 (CIFAR-100).
Implementation Details We implement FedSSL in Python using popular deep learning framework PyTorch (Paszke et al., 2017). To simulate federated learning, we train each client on one NVIDIA V100 GPU. These clients communicate with the server through NCCL backend. We use ResNet-18 (He et al., 2016) as default network for the encoders and present results of ResNet-50 in Appendix C. The predictor is a two-layer multi-layer perceptron (MLP). By default, we train for R = 100 rounds with K = 5 clients, E = 5 local epoches, batch size B = 128, learning rate η = 0.032 with cosine decay, and non-IID data l = 2 (l = 20) for CIFAR-10 (CIFAR-100).
Linear Evaluation We evaluate the quality of representations following linear evaluation (Kolesnikov et al., 2019; Grill et al., 2020) protocol. We first learn representations from the FedSSL framework. Then, we train a new linear classifier on the frozen representations.

heading: 3.4 ALGORITHM COMPARISONS
text: We benchmark and compare the SSL methods using the FedSSL framework. To denote the implementation of an SSL method, We add a prefix Fed to the name of the SSL method. For example, FedBYOL denotes using BYOL in the FedSSL framework.
2Intuitively, FedSSL is analogous to a SuperClass in OOP, then implementation of each method is a SubClass that inherits FedSSL and overrides methods of local training, model communication, and aggregation.
Table 1 compares linear evaluation results of these methods under the non-IID setting of CIFAR datasets. On the one hand, contrastive FedSSL methods obtain similar performances. As SimCLR is previously reported to need a large batch size (e.g., B = 4096) (Chen et al., 2020a), it is surprising that FedSimCLR obtains competitive results using the same batch size B = 128 as the others. On the other hand, the results of non-contrastive FedSSL methods have large variances: FedBYOL achieves the best performance, whereas FedSimSiam yields the worst performance. Since SimSiam is capable to learn as powerful representations as BYOL (Chen & He, 2021), as well as considering that noncontrastive methods are conceptually simpler and more efficient (Tian et al., 2021), we focus on non-contrastive methods and further investigate the effects of their fundamental components.

heading: 3.5 IMPACT OF FACTORS OF NON-CONTRASTIVE METHODS
text: This section analyzes the impact of fundamental components of non-contrastive FedSSL methods. From empirical studies, we obtain the following insights: 1) predictor is essential; 2) EMA and stopgradient improves performances; 3) Local encoders should retain local knowledge of the non-IID data; 4) Target encoder should gain knowledge from the online encoder. Details are as followed.
Predictor is essential. Figure 2 (left plot) presents the kNN testing accuracy as a monitoring process for FedBYOL and FedSimSiam with and without predictors. Without predictors, both methods can barely learn due to collapse in local training. It affirms the vital role of predictor (Chen & He, 2021; Tian et al., 2021) even when learning from decentralized data.
Stop-gradient operation is previously indicated as an essential component for SimSiam and BYOL (Tian et al., 2021), but it is not essential for FedBYOL. Stop-gradient prevents stochastic gradient optimization on the target network. Figure 2 shows that FedSimSiam without stop-gradient collapses, whereas FedBYOL without stop-gradient still achieves competitive performance. It is because online and target encoders are significantly different in FedBYOL as the online encoder is updated by the global encoder every communication round. In contrast, SimSiam or FedSiamSiam share weights between online and target encoders, so removing stop-gradient leads to collapse.
Exponential Moving Average (EMA) is not essential, but it helps improve performance. Table 2 (first and second rows) shows that FedBYOL outperforms FedSimSiam at different levels of nonIID data, which is represented by {2, 4, 6, 8, 10} classes per client of CIFAR-10. EMA is the main difference between SimSiam and BYOL, indicating that EMA is helpful to improve performance. Based on these results, we further analyze the underlying impact of EMA on the encoders below.
Encoders that retain local knowledge of non-IID data helps improve performance. EMA in FedBYOL allows the parameters of the online encoder to be different from the target encoder. As a result, the global encoder only updates the online encoder, not the target encoder. We hypothesize that retaining such local knowledge of data in the target encoder is beneficial especially when the data distribution is highly skewed. For comparison, we remove such local knowledge by updating both online and target encoders with the global model. Table 2 shows that FedBYOL with both encoders updated leads to lower performance than FedBYOL; It achieves results close to FedSimSiam when the data distribution is more skewed (2 or 4 classes per client). These results demonstrate the importance of keeping local knowledge in the encoders. Besides, the results of {6, 8, 10} classes per client also indicate the benefit of EMA.
Target encoder is essential to gain knowledge from the online encoder. Figure 3 shows that FedBYOL without EMA can merely learn, and FedBYOL without EMA and stop-gradient (sg) degrades in performance. In both cases, the target encoder is either never updated (w/o EMA) or is updated only through backpropagation (w/o EMA and sg) — not updated by the online encoder. On the other hand, we also identify that FedSSL methods in Table 1, which achieve competitive results, all update target encoders with knowledge of online encoders (the global encoder is the aggregation of the online encoder). We argue that target encoder is crucial to gain knowledge from the online encoder to provide contrastive targets. We further validate it by using the global encoder to update both online and target encoders when removing both EMA and stop-gradient. Figure 3 shows that such method improves performance and achieves comparable results.

heading: 4 DIVERGENCE-AWARE DYNAMIC MOVING AVERAGE UPDATE
text: Built on the FedSSL framework, we propose Federated Divergence-aware EMA update (FedEMA) to further mitigate non-IID data challenges. Since FedBYOL contains all components that help improve performance, we adopt it as the baseline and optimize the model update operation.
Non-IID data causes the global model to diverge from centralized training (Zhuang et al., 2021a). Inspired by the insight that retaining local knowledge of non-IID data helps improve performance, we propose to update the online network via EMA of the global network. Compared with FedBYOL that replaces the online network with the global network, FedEMA fuses local and global knowledge effectively through EMA update, where the decay rate of EMA is dynamically measured by model divergences. Figure 4 depicts our proposed FedEMA method. The formulation is as followed:
W rk = µW r−1 k + (1− µ)W r g , (1)
W p,rk = µW p,r−1 k + (1− µ)W p,r g , (2)
µ = min(λ ∥∥W rg −W r−1k ∥∥ , 1), (3)
where W rk and W p,r k are the online encoder and predictor of client k at training round r; W r g and W p,rg are the global encoder and predictor; µ is the decay rate, measured by the divergence between global and online encoders; λ is a scaler to adjust the level of model divergence, which is measured by calculating the l2-norm of the global and online encoders. We summarize FedEMA in Algorithm 1. FedEMA can be regarded as a generalization of FedBYOL — they are the same when λ = 0.
Scaler λ plays a vital role to adapt FedEMA for different levels of divergence caused by the data. The divergence between global and online encoders varies when the settings of federated learning change. For example, different degrees of non-IID settings would result in different divergences. Since characteristics of data are unknown before training as they are unlabeled, we propose a practical autoscaler to calculate a personalized λk for each client k automatically. The formula is λk =
τ ||W r+1g −W rk || , where τ ∈ [0, 1] is the expected value of µ at round r. We calculate λk only once at the earliest round r that client k is sampled for training. When the same set of clients are sampled for training, λk = τ||W 1g−W 0k || is calculated at round r = 1.
The intuition of FedEMA is to retain more local knowledge when divergence is large and incorporate more global knowledge when divergence is small. When model divergence is large, keeping more local knowledge is more beneficial for the non-IID data. Since the global network is the aggregation of online networks, representing global knowledge from clients. When divergence is small, adapting more global knowledge help improve model generalization. Since model divergence is larger at the start of training (Figure 6), it is practical to choose larger τ ∈ [0.5, 1); τ = 1 is not considered because only local knowledge is used when τ = 1. We use τ = 0.7 by default in experiments.
Algorithm 1 Our proposed FedEMA 1: ServerExecution: 2: Init W 0g and W p,0 g , init λk to be null
3: for each round r = 0, 1, ..., R do 4: St ← (Selection of K clients) 5: for client k ∈ St in parallel do 6: W rk ,W p,r k ←Client(W rg ,W p,rg ,r,λk)
7: W r+1g ← ∑ k∈St nk n W r k
8: W p,r+1g ← ∑ k∈St nk n W p,r k
9: for client k ∈ St do 10: λk ← τ||W r+1g −W rk || if λk is null 11: Return WRg 12: Client (W rg , W p,rg , r, λk): 13: if λk is null or not selected in r − 1 then 14: Wk,W tk,W p k ←W rg ,W rg ,W p,rg 15: else 16: µ← min(λk
∥∥W rg −W r−1k ∥∥ , 1) 17: Wk ← µW r−1k + (1− µ)W rg 18: W pk ← µW p,r−1 k + (1− µ)W p,rg 19: for local epoch e = 0, 1, ..., E − 1 do 20: for b ∈ B data batches with size B do 21: W ok ←W ok − η▽LW ok ,W tk(W o k ; b) 22: W tk ← mW tk + (1−m)Wk 23: Return W rk , W p,r k
image
Global Encoder Global Predictor
Divergence-aware EMA Update
EMAAugmentation
Server
Client k
stop-gradient
Online Encoder Predictor
Target Encoder
Figure 4: Illustration of our proposed Federated Divergence-aware Exponential Moving Average update (FedEMA). Compared with FedBYOL that simply updates the online network W ok of client k with the global network W og , we propose to update them via EMA of the global network following Eqn 1 and 2, where the decay rate µ is dynamically measured the divergences between the online encoder Wk and the global encoder Wg (Eqn 3). The online network, W ok = (Wk,W p k ), is the concatenation of the online encoder Wk and the predictor W p k .

heading: 5 EVALUATION
text: This section follows the experimental setup in Section 3.3 to evaluate FedEMA in the linear evaluation and semi-supervised learning. We also provide ablation studies of important hyperparameters.

heading: 5.1 ALGORITHM COMPARISONS
text: To demonstrate the effectiveness of FedEMA, we compare it with the following methods: 1) standalone training, where a client learns independently using BYOL; 2) FedCA, which is proposed in Zhang et al. (2020a); 3) FedBYOL as the baseline; 4) FedU, which is proposed in (Zhuang et al., 2021a). Besides, we also present results of possible upper bounds that learn representations with centralized data using BYOL.
Linear Evaluation Table 3 shows that FedEMA outperforms other methods on different settings of CIFAR datasets. Specifically, the performance is more 3% higher than existing methods in most settings. Besides, our proposed autoscaler achieves similar results as λ = 0.8. More experiments on larger number of clients K and random selection of clients are provided in Table 6 in Appendix C.
Semi-supervised Learning We also assess the quality of representations following the semisupervised learning protocol (Zhai et al., 2019; Chen et al., 2020a) — we add a new two-layer MLP on top of the encoder and fine-tune the whole model with limited (1% and 10%) labeled data for 100 epochs. Table 4 indicates that FedEMA consistently outperforms other methods on non-IID settings of CIFAR datasets and our autoscaler outperforms manual-selected λ = 1.

heading: 5.2 ABLATION STUDIES
text: Ablation on FedEMA We analyze whether we need to update both online encoder (Eqn 1) and predictor (Eqn 2) in FedEMA. Figure 5 shows that updating only the encoder or predictor leads to better performance; only updating predictor also leads to faster convergence. Their combination re-
0 10 20 30 40 50 60 70 80 90 100 Rounds
0.0 0.2 0.4 0.6 0.8 1.0 D iv er ge nc e
Client 0 Client 1
Client 2 Client 3
Client 4
Figure 6: Changes of divergence throughout training.
0.0 0.2 0.4 0.6 0.8 1.0 1.2 1.4 1.6 1.8 2.0 Scaler (λ)
30 40 50 60 70 80 T op -1 A cc .
(% )
CIFAR-10 CIFAR-100
(a) Scaler λ
0.1 0.3 0.5 0.7 0.9 Decay Rate µ
76 78 80 82 84 T op -1 A cc .
(% )
Ours FedBYOL
Encoder only Predictor only Both
(b) Constant µ
2 4 6 8 10 (iid) Classes Per Client
76 78 80 82 84 86 88 T op -1 A cc . (% ) FedBYOL FedEMA
(c) Non-IID level
Figure 7: Ablation study on scaler λ, decay rate µ, and non-IID levels of the CIFAR-10 dataset: (a) analyzes the impact of scaler λ on performance; (b) compares using constant µ on encoder, predictor, or both; (c) studies the impact of different non-IID levels.
sults in the best performance. These results demonstrate the effectiveness of updating both predictor and encoder in FedEMA. More results on other settings are provided in Table 5 in Appendix C.
Changes of Divergence Figure 6 illustrates that the divergence between global encoder and online encoder (Eqn 3) decreases gradually as training proceeds. It validates our intuition that more local knowledge is used at the start of training when divergence is larger. Besides, clients can update at their own pace depending on the divergence caused by their local dataset.
Scaler λ We study the impact of λ with values in [0, 2] with interval of 0.2 in Figure 7a. λ > 1 leads to a significant performance drop because it results in µ = 1 at the start of training on CIFAR datasets, implying that no aggregated global network is used. When λ ∈ (0, 1), the performances are consistently better than FedBYOL (λ = 0) as both local and global knowledge are effectively aggregated. These analyses are mainly suitable for our experiment setting. The range values of λ depend on the characteristics of data and the hyper-parameters (e.g., local epoch) of FL settings. A practical way to tune λ manually is to understand the divergence by running the algorithm for several rounds and choose the λ that scales µ to (0.5, 1). Nevertheless, we recommend using autoscaler and provide ablation study of τ of autoscaler in Figure 10a in Appendix C.
Constant Values of µ We further demonstrate the necessity of dynamic EMA by comparing with using constant values of µ in Eqn 1 and 2. Figure 7b shows that a good choice of constant µ can outperform FedBYOL, but FedEMA outperforms using constant µ for the online encoder, predictor, or applying both. We also provide results that encoder and predictor use different µ in Appendix C.
Non-IID Level Figure 7c compares the performance of different non-IID levels, ranging from 2 to 10 classes per client on the CIFAR-10 dataset. We use autoscaler for these experiments. FedEMA consistently outperforms FedBYOL in these settings.

heading: 6 CONCLUSION
text: We uncover important insights of federated self-supervised learning (FedSSL) from in-depth empirical studies, using a newly introduced generalized FedSSL framework. Inspired by the insights, we propose a new method, Federated Divergence-aware Exponential Moving Average update (FedEMA), to further address the non-IID data challenge. Our experiments and ablations demonstrate that FedEMA outperforms existing methods in a wide range of settings. In the future, we plan to implement FedSSL and FedEMA on larger-scale datasets. We hope that this study will provide useful insights for future research.

heading: 7 REPRODUCIBILITY STATEMENT
text: To facilitate reproducibility of experiment results, we first provide basic experimental setups in Section 3.3, including datasets, implementation details, and evaluation protocols. Then, we describe more experimental details in Appendix B, including datasets, data transformation, network architecture, training details, and default settings. Also, we indicate the settings and hyper-parameters of experiments when their settings are different from the default. Moreover, we plan to open-source the codes in the future.

heading: ACKNOWLEDGMENTS
text: We would like to thank reviewers of ICLR 2022 for their constructive and helpful feedback. This study is in part supported by the RIE2020 Industry Alignment Fund – Industry Collaboration Projects (IAF-ICP) Funding Initiative, as well as cash and in-kind contribution from the industry partner(s); the National Research Foundation, Singapore under its Energy Programme (EP Award NRF2017EWT-EP003-023) administrated by the Energy Market Authority of Singapore, and its Energy Research Test-Bed and Industry Partnership Funding Initiative, part of the Energy Grid (EG) 2.0 programme, and its Central Gap Fund (“Central Gap” Award No. NRF2020NRF-CG001-027); Singapore MOE under its Tier 1 grant call, Reference number RG96/20.

heading: A DIFFERENCES OF SELF-SUPERVISED LEARNING METHODS
text: We study four SSL methods using the FedSSL framework in Section 3. These four SSL methods have two major differences that impact the executions of local training, model communication, and model aggregation. Figure 8 depicts these differences: 1) BYOL and SimSiam have predictors, whereas MoCo and SimCLR do not have them; 2) SimSiam and SimCLR share weights between two encoders, whereas BYOL and MoCo have different parameters for the online and target encoders.

heading: B EXPERIMENTAL DETAILS
text: In this section, we provide more details about the dataset, network architecture, and training and evaluation setups.
B.1 DATA
Datasets CIFAR-10 and CIFAR-100 are two popular image datasets (Krizhevsky et al., 2009). Both datasets consist of 50,000 training images and 10,000 testing images. CIFAR-10 contains 10 classes, where each class has 5,000 training images and 1,000 testing images. While CIFAR-100 contains 100 classes, where each class has 500 training images and 100 testing images. To simulate federated learning, we equally split the training set into K clients. We simulate non-IID data using label heterogeneity — data among clients is more skewed when each client contains less number of classes. Hence, we simulate different levels of non-IID data with l number of classes per client, where l = {2, 4, 6, 8, 10} for CIFAR-10 and l = {20, 40, 60, 80, 100} for CIFAR-100. For example, when simulating 5 clients with l = 4 classes per client in CIFAR-10, we need 5× 4 = 20 total sets of data over 10 classes. Thus, we split the training images of each class equally into two sets (2,500 images in each set) and assign random four sets without overlapping classes to a client. The setting is IID when each client contains all classes of a dataset. By default, we run experiments with K = 5 clients with non-IID setting l = 2 classes per client for CIFAR-10 dataset and l = 20 classes per client for CIFAR-100 dataset.
Transformation In local training of the FedSSL framework, we take two augmentations of each image as the inputs for online and target networks, respectively. We obtain the augmentations by transforming the images with a set of transformations: For SimCLR, BYOL, SimSiam, and MoCoV2, we adopt the transformations from Chen et al. (2020a); For MoCoV1, we use the transformation described in its paper (He et al., 2020).
B.2 NETWORK ARCHITECTURE
Predictor The network architecture of the predictor is a two-layer multilayer perceptron (MLP). The two-layer MLP starts from a fully connected layer with 4096 neurons. Followed by onedimension batch normalization and a ReLU activation function, it ends with another fully connected layer with 2048 neurons.
Encoder We use ResNet-18 He et al. (2016) as the default network architecture of the encoder in the majority of experiments. Besides, we also provide results of ResNet-50 in Table 5 and 7.
Our ResNet architecture differs from the implementation in PyTorch (Paszke et al., 2017) in three aspects: 1) We use kernel size 3 × 3 for the first convolution layer instead of 7 × 7; 2) We use an average pooling layer with kernel size 4 × 4 before the last linear layer instead of adaptive average pooling layer; 3) We replace the last linear layer with a two-layer MLP. The network architecture of the MLP is the same as the predictor.
B.3 TRAINING AND EVALUATION DETAILS
We implement FedSSL in Python using EasyFL (Zhuang et al., 2022), an easy-to-use federated learning platform based on PyTorch (Paszke et al., 2017). The following are the details of training and evaluation.
Training We use Stochastic Gradient Descent (SGD) as the optimizer in training. We use η = 0.032 as the initial learning rate and decay the learning with a cosine annealing (Loshchilov & Hutter, 2017), which is also used in SimSiam. By default, we train R = 100 rounds with local
epochs E = 5 and batch size B = 128 using K = 5 clients. We simulate training of K clients on K NVIDIA V100 GPUs and employ the PyTorch Paszke et al. (2017) communication backend (NCCL) for communications between clients and the server. If not specified, we use λ = 1 by default or autoscaler with τ = 0.7 for FedEMA. As for experiments of FedU, we follow the hyperparameters described in paper (Zhuang et al., 2021a).
Cross-silo FL vs Cross-device FL This paper primarily focuses on cross-silo FL where clients are stateful with high availability. Clients can cache local models and carry these local states from round to round. Extensive experiments demonstrate that FedEMA achieves the best performance under this setting. On the other hand, cross-device FL assumes there are millions of stateless clients that might participate in training just once. Due to the constraints of experimental settings, the majority of studies conduct experiments with at most hundreds of clients (Wang et al., 2020; Jeong et al., 2021). FedEMA can work under such experimental settings by caching the states of clients in the server. When the number of clients scales to millions, FedEMA degrades to FedBYOL that updates both encoders — without keeping any local states.
Evalaution We assess the quality of learned representations using linear evaluation (Kolesnikov et al., 2019; Grill et al., 2020) and semi-supervised learning (Zhai et al., 2019; Chen et al., 2020a) protocols. We first obtain a trained encoder (or learned representations) using full training set for linear evaluation and 99% or 90% of the training set for semi-supervised learning (excluding the 1% or 10% for fine-tuning). Then, we conduct evaluations based on the trained encoder. For linear evaluation, we train a new fully connected layer on top of the frozen trained encoder (fixed parameters) for 200 epochs, using batch size 512 and Adam optimizer with learning rate 3e-3. For semi-supervised learning, we add a new two-layer MLP on top of the trained encoder and fine-tune the whole model using 1% or 10% of data for 100 epochs, using batch size 128 and Adam optimizer with learning rate 1e-3. In both evaluation protocols, we remove the two-layer MLP of the encoder by replacing it with an identity function.

heading: C ADDITIONAL EXPERIMENTAL RESULTS AND ANALYSIS
text: In this section, we provide more experimental results of algorithm comparisons and further analyze FedEMA in different data amounts, training rounds R, and batch sizes B.
C.1 MORE EXPERIMENTAL RESULTS
Table 5 presents top-1 accuracy comparison under linear evaluation of a wide range of methods on CIFAR datasets using both ResNet-18 and ResNet-50. It supplements the algorithm comparisons in Section 3.4 and 5.1. Interestingly, FedMoCoV1 achieves good performances on IID settings of CIFAR-100 dataset. Since decentralized data are mostly non-IID, we focus more on the non-IID setting. FedEMA outperforms all the other methods in non-IID settings of CIFAR datasets. We use λ = 0.8 when using ResNet-18 and λ = 1 when using ResNet-50.
Table 6 shows results of scaling to larger numbers of clients K with subsampling clients in each training round. We run two sets of experiments: 1) randomly selecting 5 out of 20 clients in each round with local epoch E = 5 and total rounds R = 400; 2) randomly selecting 8 out of 80 clients in each round with local epoch E = 2 and total rounds R = 800. We run FedEMA with autoscaler. FedEMA consistently outperforms FedBYOL with both encoders updated. We conduct these experiments using ResNet-18.
Table 7 supplements the semi-supervised learning results on Table 4, providing additional results using ResNet-50 as the network architecture for the encoder. FedEMA consistently outperforms all the other methods.
Besides, Table 8 and 9 compare FedSimSiam, FedBYOL, and variances of FedBYOL to further demonstrate the insights from empirical studies. They supplement results in Table 2 and Figure 3.
C.2 FURTHER ANALYSIS
τ for Autoscaler We analyze the impact of τ on performances in Figure 10a. Generally, using autoscaler with τ ∈ (0, 1) is better than FedBYOL (τ = 0). The performance of τ = 1 yields worse results because only local knowledge are used in model update (the global knowledge is neglected) as discussed in Section 4. Besides, performances of τ ∈ [0.5, 1) are generally better other values, which verifies our intuition discussed in Section 4. These results also show that we can achieve even higher performance on the CIFAR-100 dataset on Table 3 by tunning τ . We run experiments on non-IID settings using ResNet-18.
Constant µ To further illustrate the effectiveness of our dynamic EMA, we provide results of using different combinations of constant µo on the online encoder and constant µp on the predictor in Figure 10b. The results of µo = 0.9 and µp = 0.9 is only 54.52%, which is far lower than the others. Among these combinations, µo = 0.5 and µp = 0.3 achieve the best performance. It suggests that better performances may be achieved if we can construct different dynamic µ for the encoder and the predictor, while we leave this interesting insight for future exploration. Although good choices of constant µo and µp achieve better performance than FedBYOL, FedEMA consistently outperforms all these methods. These results complement Figure 7b in the main manuscript.
Impact of Training Rounds R Figure 9 compares FedBYOL and FedEMA with increasing number of training (communication) rounds R. Performances of both FedBYOL and FedEMA increases as training proceeds and FedEMA consistently outperforms FedBYOL. We run these experiments with λ = 0.5 for FedEMA on the non-IID setting of CIFAR-10 dataset.
Impact of Batch Size B We investigate the impact of batch size in Table 10. The performances of batch size B = 128 and B = 256 are similar, outperforming the other batch sizes. It indicates that the batch size should not be either too small or too large. Besides, FedEMA outperforms FedBYOL in all batch sizes. We run the experiments with autoscaler (τ = 0.7) on the non-IID setting of the CIFAR-10 dataset.
Communication vs Computation Cost Table 11 shows the needed communication rounds and computation epochs to reach a target accuracy using different local epochs E with FedEMA. Increasing E reduces communication cost as the needed rounds decrease, but it generally requires a higher computation cost. For example, compared with E = 5 that needs 80 rounds to reach 80% with 400 epochs of computation, E = 20 only uses 40 rounds but needs 800 epochs computation cost. These results indicate the trade-off between communication cost and computation cost.
Data Amount Table 12 shows that increasing the data amount improves the performance significantly. By default, we split the CIFAR-10 dataset into 5 clients, each client contains 10,000 training images, denoting as 100% data amount. As a result, p% data amount means that each client contains 10, 000 ∗ p% images. For example, 25% data amount means that each client contains 2,500 images. With lesser data points in each client, we can construct more clients to conduct training as the total data amount is fixed. Table 12 shows that when the data amount is same in clients, increasing the number of clients in each training round do not improve performance. However, increasing the data amount in each client increases the performance significantly. These results indicate that it is important for clients to have sufficient data to participate in training in FedSSL.

Review summaries:
This paper investigates a generic federated SSL recipe that applies FedAvg to a range of existing SSL works including SimCLR, MoCo, BYOL and SimSiam. Each of these SSL blocks comprises two encoding networks, including an online net and a target net. The two nets were trained via optimizing a similarity loss such that the distance between encodings of different augmented versions (e.g. rotation) of the same input is small and vice versa. Depending on the specific SSL block, the parameterization of the two nets might be identical (SimCLR, SimSiam) or not (MoCo, BYOL). At each iteration, the online net of each client is uploaded to a server to be averaged. Then, the (global) averaged version is sent back to each client. Each client then reset its online net as a weighted average of the global and local version of the online net (to account for potential data heterogeneity), and continue updating the SSL block via gradient updates and so on.      Here, the main contribution of the paper is a series of ablation studies (Table 1, Figure 2, Table 2, Figure 3) that measure the isolated impact of several fundamental components of the aforementioned SSL blocks (e.g., SimCLR, SimSiam, MoCo and BYOL). This includes the predictor, stop-gradient, exponential moving average (EMA), non-identical online and target encoding net. The results indicate that BYOLD contains all the essential elements (predictor, EMA and non-identical online & target net & no stop-gradient) so the federated SSL recipe is rotated towards having BYOL as the model choice. Then, a customized, divergence-aware EMA is proposed for this recipe that is shown to improve substantially over existing FL-SSL baselines (Tables 3-4; Tables 6-7). There is also a bit of ablation study that measures the isolated impact of the new EMA which is also substantial as shown in Figure 5.

The authors propose a new approach called Federated Divergence-aware Exponential moving Average update (FedEMA) to avoid the IID assumption. FedEMA is built onto of FedSSL which is a framework for self-supervised learning in a Federated Learning context. The authors proposes a new approach to fuse the local and global knowledge effectively through a EMA update, where the decay rate of EMA is dynamically measured by a divergence.