SSFL: Tackling Label Deficiency in Federated Learning via Personalized Self-Supervision

Federated Learning ( FL ) is transforming the ML training ecosystem from a centralized over-the-cloud setting to distributed training over edge devices in order to strengthen data privacy , reduce data migration costs , and break regulatory restrictions . An essential , but rarely studied , challenge in FL is label deficiency at the edge . This problem is even more pronounced in FL , compared to centralized training , due to the fact that FL users are often reluctant to label their private data and edge devices do not provide an ideal interface to assist with annotation . Addressing label deficiency is also further complicated in FL , due to the heterogeneous nature of the data at edge devices and the need for developing personalized models for each user . We propose a self-supervised and personalized federated learning framework , named ( SSFL ) , and a series of algorithms under this framework which work towards addressing these challenges . First , under the SSFL framework , we analyze the compatibility of various centralized self-supervised learning methods in FL setting and demonstrate that SimSiam networks performs the best with the standard FedAvg algorithm . Moreover , to address the data heterogeneity at the edge devices in this framework , we have innovated a series of algorithms that broaden existing supervised personalization algorithms into the setting of self-supervised learning including perFedAvg , Ditto , and local fine-tuning , among others . We further propose a novel personalized federated self-supervised learning algorithm , Per-SSFL , which balances personalization and consensus by carefully regulating the distance between the local and global representations of data . To provide a comprehensive comparative analysis of all proposed algorithms , we also develop a distributed training system and related evaluation protocol for SSFL . Using this training system , we conduct experiments on a synthetic non-I.I.D . dataset based on CIFAR-10 , and an intrinsically non-I.I.D . dataset GLD-23K . Our findings show that the gap of evaluation accuracy between supervised learning and unsupervised learning in FL is both small and reasonable . The performance comparison indicates that representation regularization-based personalization method is able to outperform other variants . Ablation studies on SSFL are also conducted to understand the role of batch size , non-I.I.D.ness , and the evaluation protocol . 1 INTRODUCTION . Federated Learning ( FL ) is a contemporary distributed machine learning paradigm that aims at strengthening data privacy , reducing data migration costs , and breaking regulatory restrictions ( Kairouz et al. , 2021 ; Wang et al. , 2021 ) . It has been widely applied to computer vision , natural language processing , and data mining . However , there are two main challenges impeding its wider adoption in machine learning . One is data heterogeneity , which is a natural property of FL in which diverse clients may generate datasets with different distributions due to behavior preferences ( e.g. , the most common cause of heterogeneity is skewed label distribution which might result from instances where some smartphone users take more landscape pictures , while others take more photos of daily life ) . The second challenge is label deficiency at the edge , which is relatively less studied . This issue is more severe at the edge than in a centralized setting because users are reluctant to annotate their private and sensitive data , and/or smartphones and IoT devices do not have a user-friendly interface to assist with annotation . To mitigate the data heterogeneity issue among clients , researchers have proposed algorithms for training a global model FedAvg ( McMahan et al. , 2017 ) , FedProx ( Li et al. , 2018 ) , FedNova ( Wang et al. , 2020 ) , FedOPT ( Reddi et al. , 2020 ) , as well as personalized FL frameworks ( e.g. , pFedMe , Ditto , Per-FedAvg ) . These algorithms all depend on the strong assumption that the data at the edge has sufficient labels . To address the label deficiency issue in FL , recent works ( Liu et al. , 2020 ; Long et al. , 2020 ; Itahara et al. , 2020 ; Jeong et al. , 2020 ; Liang et al. , 2021 ; Zhao et al. , 2020 ; Zhang et al. , 2020a ; b ) assume that the server or client has a fraction of labeled data and use semi-supervised methods such as consistency loss ( Miyato et al. , 2018 ) or pseudo labeling ( Lee , 2013 ) to train a global model . A more realistic but challenging setting is fully unsupervised training . Although a recent work in FL ( Saeed et al. , 2020 ) attempts to address this challenge through Siamese networks proposed around thirty years ago ( Bromley et al. , 1993 ) , its design does not tackle data heterogeneity for learning personalized models , and it only trains on small-scale sensor data in IoT devices . Moreover , these existing works in FL have not examined recent progress in the Self-Supervised Learning ( SSL ) community where methods such as SimCLR ( Chen et al. , 2020 ) , SwAV ( Caron et al. , 2021 ) , BYOL ( Grill et al. , 2020 ) , and SimSiam ( Chen & He , 2020 ) have shown tremendous improvement in reducing the amount of labeled data required to achieve state-of-the-art performance . As such , it remains still unclear how these SSL methods can be incorporated into FL and how well they would perform , especially when intertwined with the data heterogeneity challenge that does not exist in centralized training . In this paper , we propose Self-Supervised Federated Learning ( SSFL ) , a unified self-supervised and personalized federated learning framework , and a series of algorithms under this framework to address these challenges . As shown in Figure 1 , this framework brings state-of-the-art self-supervised learning algorithms to the realm of FL in order to enable training without using any supervision , while also integrating model personalization to deal with data heterogeneity ( Section 3.1 ) . More specifically , under the SSFL framework , we analyze the compatibility of various centralized selfsupervised learning methods in the FL setting and demonstrate that SimSiam networks performs the best with the standard FedAvg algorithm ( Section 3.2 ) . Moreover , to address the data heterogeneity at edge devices , we have innovated a series of algorithms that broaden the reach of existing supervised personalization algorithms into the setting of self-supervised learning , including perFedAvg ( Fallah et al. , 2020 ) , Ditto ( Li et al. , 2021 ) , and local fine-tuning , among others . We further propose a novel personalized federated self-supervised learning algorithm , per-SSFL ( Section 3.3 ) , which balances personalization and consensus by carefully regulating the distance between the local and global representations of data ( shown as the yellow block in Figure 1 ) . To provide a comprehensive and comparative analysis of the proposed algorithms , we also develop a distributed training system and evaluation protocol for SSFL . Using this training system , we conduct experiments on a synthetic non-I.I.D . dataset based on CIFAR-10 and a natural non-I.I.D . dataset GLD-23K . Our experimental results demonstrate that all algorithms in our framework work reliably . In FL , the gap of evaluation accuracy between supervised learning and unsupervised learning is small . Personalized SSFL performs better than FedAvg-based SSFL . We also conduct ablation studies to fully understand the SSFL framework , namely the role of batch size , different degrees of non-I.I.D.ness , and performance in more datasets . Finally , our unified API design can serve as a suitable platform and baseline , enabling further developments of more advanced SSFL algorithms . 2 PRELIMINARIES . SSFL builds upon two fundamental areas in machine learning : federated optimization and selfsupervised learning . Thus , we first introduce some basics and formulations in these areas . 2.1 FEDERATED OPTIMIZATION . Federated optimization refers to the distributed optimization paradigm that a network of K devices collaboratively solve a machine learning task . In general , it can be formulated as a distributed optimization problem with the form ( McMahan et al. , 2017 ) : minθ ∑K k=1 |Dk| |D| L ( θ , Dk ) . Here , each device k has a local dataset Dk drawn from a local distribution Xk . The combined dataset D = ∪Kk=1Dk is the union of all local datasets Dk . θ represents the model weight of a client model . L is the client ’ s local loss function that measures the local empirical risk over the heterogeneous dataset Dk . Under this formulation , to mitigate the non-I.I.D . issue , researchers have proposed algorithms such as FedProx ( Li et al. , 2018 ) , FedNova ( Wang et al. , 2020 ) , and FedOPT ( Reddi et al. , 2020 ) for training a global model , as well as personalized FL frameworks such as Ditto ( Li et al. , 2021 ) , and Per-FedAvg ( Fallah et al. ) . All of these algorithms have a strong assumption that data at the edge have sufficient labels , meaning that their analysis and experimental settings are based on a supervised loss function , such as the cross-entropy loss for image classification . 2.2 SELF-SUPERVISED LEARNING . Self-supervised learning ( SSL ) aims to learn meaningful representations of samples without human supervision . Formally , it aims to learn an encoder function fθ : X 7→ Rd where θ is the parameter of the function , X is the unlabeled sample space ( e.g . image , text ) , and the output is a d dimensional vector containing enough information for downstream tasks such as image classification and segmentation . The key to SSL ’ s recent success is the inductive bias that ensures a good representation encoder remains consistent under different perturbations of the input ( i.e . consistency regularization ) . One prominent example among recent advances in modern SSL frameworks is the Siamese architecture ( Bromley et al. , 1993 ) and its improved variants SimCLR ( Chen et al. , 2020 ) , SwAV ( Caron et al. , 2021 ) , BYOL ( Grill et al. , 2020 ) , and SimSiam ( Chen & He , 2020 ) . Here we review the most elegant architecture , SimSiam , and defer the description and comparison of the other three to Appendix A. SimSiam proposes a two-head architecture in which two different views ( augmentations ) of the same image are encoded by the same network fθ . Subsequently , a predictor Multi Layer Perceptron ( MLP ) hθ and a stop-gradient operation denoted by ·̂ are applied to both heads . In the SSL context , “ stop gradient ” means that the optimizer stops at a specific neural network layer during the back propagation and views the parameters in preceding layers as constants . Here , θ is the concatenation of the parameters of the encoder network and the predictor MLP . The algorithm aims to minimize the negative cosine similarity D ( · , · ) between two heads . More concretely , the loss is defined as LSS ( θ , D ) = 1 |D| ∑ x∈D D ( fθ ( T ( x ) ) , ¤ hθ ( fθ ( T ( x ) ) ) ) , ( 1 ) where T represents stochastic data augmentation and D is the data set . 3 SSFL : SELF-SUPERVISED FEDERATED LEARNING . In this section , we propose SSFL , a unified framework for self-supervised federated learning . Specifically , we introduce the method by which SSFL works for collaborative training of a global model and personalized models , respectively .

Summary: The submission is not anonymized, the author's name chaoyanghe appears frequently in the attached SSFL_Sumpplementary/SSFL-Source-Code. Therefore, it should be directly desk rejected.  --- Logging error --- Traceback (most recent call last):   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 1025, in emit     msg = self.format(record)   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 869, in format     return fmt.format(record)   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 608, in format     record.message = record.getMessage()   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 369, in getMessage     msg = msg % self.args