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: This paper introduces the self-supervised learning (SSL) framework for FL. Different SSL methods are investigated to study their feasibility under the FL setting. With the popular SimSiam framework, personalized federated SSL is proposed. The performance comparison indicates that representation regularization-based personalization method is able to outperform other variants.