Divergence-aware Federated Self-Supervised Learning

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 . 3 AN EMPIRICAL STUDY OF FEDERATED SELF-SUPERVISED LEARNING . 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 . 3.1 PROBLEM DEFINITION . 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 can not 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 . 3.2 GENERALIZED FRAMEWORK . 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 .

Summary: 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.