Adversarial Collaborative Learning on Non-IID Features

1 INTRODUCTION . Deep learning is data hungry . While data are always dispersed in multiple parties ( e.g. , mobile devices , hospitals ) in reality , data are not allowed to transfer to a central server for training due to privacy concerns and data regulations . Collaborative learning among multiple parties without the exchange of raw data has been an important machine learning topic . Federated learning ( McMahan et al. , 2016 ; Kairouz et al. , 2019 ; Li et al. , 2019b ; a ) has been a popular form of collaborative learning without exchanging raw data . A basic federated learning framework is FedAvg ( McMahan et al. , 2016 ) , which uses a model-averaging scheme . In each round , the parties update their local models and send them to the server . The server averages all local models to update the global model , which is sent back to the parties as the new local model in the next round . FedAvg has been widely used due to its effectiveness and simpleness . Most existing federated learning approaches are designed based on FedAvg . However , as shown in many existing studies ( Hsu et al. , 2019 ; Li et al. , 2020 ; Wang et al. , 2020b ; Karimireddy et al. , 2020 ; Acar et al. , 2021 ; Li et al. , 2021b ) , the performance of FedAvg and its alike algorithms may be significantly degraded in non-IID data among parties . While many studies try to improve FedAvg on non-IID data , most of them ( Li et al. , 2020 ; Wang et al. , 2020b ; Karimireddy et al. , 2020 ; Acar et al. , 2021 ; Li et al. , 2021a ; Wang et al. , 2020a ) focus on the label imbalance setting , where the parties have different label distributions . In their experiments , they usually simulate the federated setting by unbalanced partitioning the dataset into multiple subsets according to the labels . As summarized in ( Hsieh et al. , 2020 ; Kairouz et al. , 2019 ) , besides the label distribution skew , feature imbalance is also an important case of non-IID data . In the feature imbalance setting , the feature distribution Pi ( x ) varies across parties . This setting widely exists in reality , e.g. , people have different stroke width and slant when writing the same word . Another example in practice is that medical images of different hospitals have different intensity and contrast . However , compared with non-IID labels , federated learning on non-IID features has been less explored . Since the feature distributions are different across parties , there exists bias in the local training . It is challenging to mitigate the bias while exploiting the heterogeneous feature distributions in the training . In this paper , we think out of the model-averaging scheme used in federated learning , and propose a novel learning concept called adversarial collaborative learning . While the feature distribution of each party is different , to mitigate the local bias , we aim to learn a common representation distribution among parties . Specifically , the server aims to train a discriminator to distinguish the local representations by the party IDs , while the parties train the local models such that the generated representations can not be distinguished by the discriminator . Our experiments show that ADCOL outperforms state-of-the-art federated learning algorithms ( Li et al. , 2021b ; 2020 ; Fallah et al. , 2020 ; Collins et al. , 2021 ) on three real-world tasks . More importantly , ADCOL points out a promising research direction on collaborative learning . Inspired by our idea of conducting collaborative learning in an adversarial way , we believe there will be more interesting future works in this direction . 2 BACKGROUND AND RELATED WORK . 2.1 NON-IID DATA . We use Pi ( x , y ) to denote the data distribution of party i , where x is the features and y is the label . According to existing studies ( Kairouz et al. , 2019 ; Hsieh et al. , 2020 ) , we can categorize non-IID data in federated learning into the following four classes : ( 1 ) non-IID labels : the marginal distribution Pi ( y ) varies across parties . ( 2 ) non-IID features : the marginal distribution Pi ( x ) varies across parties . ( 3 ) concept drift : The conditional distributions Pi ( y|x ) or Pi ( x|y ) varies across parties . ( 4 ) quantity skew : the amount of data varies across parties . In this paper , we focus on non-IID features , which widely exist in reality . For example , while the pneumonia distribution of patients may be close among different hospitals , the distributions of CT images ( such as resolution , color , length and width ) may vary across hospitals , which are collected from different equipment and environments . 2.2 FEDERATED LEARNING ON NON-IID LABELS . Non-IID data is a key challenge in federated learning . There have been many studies trying to improve the performance of federated learning under non-IID data . However , most existing studies ( Li et al. , 2020 ; Wang et al. , 2020a ; Hsu et al. , 2019 ; Li et al. , 2021a ; Acar et al. , 2021 ; Karimireddy et al. , 2020 ) focus on label imbalance , where they usually simulate the federated setting with heterogeneous label distributions in the experiments . For example , FedProx ( Li et al. , 2020 ) introduces a proximal term in the objective of local training , which limits the update of the local model by the distance between the local model and the global model . FedMA ( Wang et al. , 2020a ) applies Bayesian non-parametric methods to update the global model in a layer-wise manner in the model aggregation phase . FedDyn ( Acar et al. , 2021 ) introduces a dynamic regularization term in the local training based on the global model and the local model from the previous round . More recently , MOON ( Li et al. , 2021a ) proposes model-contrastive federated learning , which corrects the local training by a contrastive loss which maximizes the agreement between the representation learned by the global model and the representation learned by the local model . As we will show in the experiments , existing algorithms have severe performance degradation on parties with non-IID features . More recently , a contemporary study ( Zhang et al. , 2021a ) proposed FedUFO , where each party trains a discriminator to apply feature and objective consistency constrains to address the non-IID data issue . However , FedUFO needs to transfer each local model to all the other parties , which causes massive communication overhead . Moreover , FedUFO is still based on FedAvg , while we propose a novel new collaborative training scheme based on adversarial training in this paper . 2.3 FEDERATED LEARNING ON NON-IID FEATURES . To the best of our knowledge , there is only one study named FedBN ( Li et al. , 2021b ) that focuses on non-IID features . FedBN introduces lightweight modifications on FedAvg . Instead of transferring and averaging all model parameters in FedAvg , FedBN keeps all the local batch normalization parameters and does not synchronize them with the global model . The operations for non-batch normalization parameters are the same as FedAvg . Our proposed ADCOL is based on a novel approach : adversarial learning for a common representation among multiple parties . ADCOL outperforms FedBN on three real-world datasets in both accuracy and communication-efficiency . 2.4 PERSONALIZED FEDERATED LEARNING . There are basically two kinds of objective in federated learning : learning a single global model or learning a personalized local model for each party ( i.e. , personalized federated learning ( Fallah et al. , 2020 ; Dinh et al. , 2020 ; Hanzely et al. , 2020 ; Zhang et al. , 2021b ; Huang et al. , 2021 ; Collins et al. , 2021 ) ) . When the training data is non-IID across parties , personalized federated learning is reasonable since a single global model may not be good for every local distribution . In this paper , same as FedBN , ADCOL is designed under the personalized federated learning setting , i.e. , each party aims to learn a local model for its own local task . 3 METHOD . 3.1 PROBLEM STATEMENT . Suppose there are N parties , where party i has a local datasetDi = { xi , yi } . The feature distributions P ( xi ) are different among parties while the label distributions P ( yi ) are similar among parties . The parties conduct collaborative learning over D ≜ ⋃ i∈ [ N ] Di with the help of a central server without exchanging the raw data . The goal of each party is to train a machine learning model which has good accuracy on its local test dataset . 3.2 MOTIVATION . The key idea of ADCOL is intuitive . While the feature distributions are different among parties , to mitigate the local bias , ADCOL aims to learn a common task-specific representation distribution . Inspired by GANs ( Goodfellow et al. , 2014 ) and adversarial learning for distribution matching ( Tzeng et al. , 2017 ; Peng et al. , 2019b ) , ADCOL adopts an adversarial way to restrict the representation distribution . Next we explain the motivation from a distribution perspective . The ultimate goal of each party in typical federated learning is to learn the oracle optimal distribution p∗ ( y|x ) . Here we introduce the representation z of the data x , which is supposed to be sufficient for the prediction of labels y ( Federici et al. , 2020 ) . Then , x , y , and z forms a Markov chain x→ z→ y . Thus , we decomposite the oracle optimal into two parts : p∗ ( z|x ) and p∗ ( y|z ) . The first part is to learn a good representation and the second part is to predict the label by the representation with simple linear layers . To learn the optimal representation distribution p∗ ( z|x ) , the objective of each party is min θi Ex∼DiℓKL ( p ( z|x ; θi ) || p∗ ( z|x ) ) , ( 1 ) where ℓKL is the KL divergence loss . We adopt KL divergence due to its popularity for evaluating the distance between probability distributions . However , the optimal representation distribution is unknown . Thus , we approximate the objective by two aspects : 1 ) By training in a novel adversarial way , we use the similarity between the representations of different parties to regularize the representation , i.e. , the parties learn a common representation z . 2 ) We ensure that the generated representation contains necessary information for the prediction of y by training on the labeled local data . Next , we describe our model architecture and training algorithm in details . 3.3 MODEL ARCHITECTURE . There are two kinds of models in ADCOL : the local models trained in the parties and the discriminator trained in the server . As ADCOL works from the perspective of representation , the architecture of the local model is similar as existing studies ( Chen et al. , 2020 ; He et al. , 2020 ; Chen and He , 2021 ) on self-supervised representation learning . The local model has three components : a base encoder , a projection head , and a predictor . The base encoder ( e.g. , ResNet-50 ) extracts representation vectors from inputs . Like SimCLR ( Chen et al. , 2020 ) and SimSam ( Chen and He , 2021 ) , an additional projection head is introduced to map the representation to a space with a fixed dimension . The final predictor is used to output predicted probabilities for each class . For ease of presentation , we use F ( · ) to denote the whole model and G ( · ) to denote the model before the predictor ( i.e. , G ( x ) is the mapped representation vector of input x ) . For the discriminator , we simply use a multi-layer perceptron in our experiments . … discriminator server ( a ) Step 1 : The server sends the discriminator to the parties . Representations discriminator targetsoutputs ℓ ! '' 1 𝑁ℓ # $ Base encoder Projection head Predictor ( b ) Step 2 : The parties update their local models Representations Representations … server ( c ) Step 3 : The parties send representations to the server . Representations Representations Representations … discriminator 1 2 N targetsoutputs ℓ ! '' … … ( d ) Step 4 : The server updates the discriminator . Figure 1 : The ADCOL framework

Summary: This paper studies federated learning when different devices have non-iid features. To address the heterogeneity problem,  it proposes a federated learning scheme called ADCOL based on adversarial learning.  In ADCOL, the devices transfer local representations to the server while sending the discriminator to the devices. The server aims to distinguish the devices' local representations, while the devices aim to train local models that generate non-distinguishable representations. To make the representations non-distinguishable, ADCOL adds an additional regularization term to the devices' loss functions. This practice aims to maximize the probability that the discriminator cannot distinguish the local representations. The experimental results show that the proposed method has some advantages over several baselines.