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