Year: 2021

Abstract:
Federated learning has been a popular approach to enable collaborative learning on multiple parties without exchanging raw data. However, the model performance of federated learning may degrade a lot due to non-IID data. While most existing studies focus on non-IID labels, federated learning on non-IID features has largely been overlooked. Different from typical federated learning approaches, the paper proposes a new learning concept called ADCOL (Adversarial Collaborative Learning) for non-IID features. Instead of adopting the widely used model-averaging scheme, ADCOL conducts training in an adversarial way: the server aims to train a discriminator to distinguish the representations of the parties, while the parties aim to generate a common representation distribution. Our experiments on three realworld datasets show that ADCOL achieves better accuracy and is much more communication-efficient than state-of-the-art federated learning algorithms on nonIID features. More importantly, ADCOL points out a promising research direction for collaborative learning.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 BACKGROUND AND RELATED WORK

heading: 2.1 NON-IID DATA
text: 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.

heading: 2.2 FEDERATED LEARNING ON NON-IID LABELS
text: 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.

heading: 2.3 FEDERATED LEARNING ON NON-IID FEATURES
text: 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.

heading: 2.4 PERSONALIZED FEDERATED LEARNING
text: 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.

heading: 3 METHOD

heading: 3.1 PROBLEM STATEMENT
text: 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.

heading: 3.2 MOTIVATION
text: 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.

heading: 3.3 MODEL ARCHITECTURE
text: 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

heading: 3.4 THE OVERALL FRAMEWORK
text: The overall framework is shown in Figure 1 and Algorithm 1. There are four steps in each round: (1) The server sends the discriminator to the parties. (2) The parties update their local models. (3) The parties send representations to the server. (4) The server updates the discriminator.
Step 1 In the first step, the server sends the discriminator to parties (line 4 of Algorithm 1).
Step 2 In the second step, the parties update their models using their local datasets (lines 10-17 of Algorithm 1). In addition to the objective which aims to minimize the cross-entropy loss (i.e., ℓCE) on the local dataset, ADCOL introduces an additional regularization term which aims to maximize the probability that the discriminator cannot distinguish the local representations. For each input x, ADCOL feeds the representation G(x) to the discriminator. ADCOL expects the discriminator to output probability vector [ 1N ]
N (i.e., the probability of each class is 1N ) such that it cannot distinguish which party that the representation comes from. Thus, ADCOL uses Kullback–Leibler (KL) divergence loss to measure the difference between the output of the discriminator D(G(x)) and the target [ 1N ] N . The final loss of an input (x, y) is computed as
ℓ = ℓCE(F (x), y) + µℓKL([ 1
N ]N || D(G(x))) (2)
where µ is a hyper-parameter to control the weight of KL divergence loss, ℓCE is the cross-entropy loss, and ℓKL is the KL divergence loss.
Step 3 After local training, the parties feed their data into the local models and transfer the representations to the server (line 5 of Algorithm 1).
Step 4 The server updates the discriminator using the received representations (lines 6-9 of Algorithm 1). Specifically, the server can build a training set, where the feature values are the representations and the labels are the party IDs that the representations come from. The training set is used to update the discriminator using stochastic gradient descent.
Algorithm 1: The ADCOL algorithm Input: number of communication rounds T , number of parties N , number of local epochs E, learning rate η, hyper-parameter µ Output: The local models Fi (i ∈ [N ])
1 Server executes: 2 for t = 1, 2, ..., T do 3 for i = 1, 2, ..., N in parallel do 4 send the discriminator D to party i 5 Ri ← PartyLocalTraining(i, D) 6 R← {(Ri, i)}Ni=1 7 for each batch b = {Ri, i} of R do 8 ℓ← CrossEntropyLoss(D(Ri), i) 9 D ← D − η∇ℓ 10 PartyLocalTraining(i, D): 11 for epoch e = 1, 2, ..., E do 12 for each batch b = {x, y} of Di do 13 ℓCE ← CrossEntropyLoss(Fi(x), y) 14 R← Gi(x) 15 ℓKL ← KLDiv([ 1N ] N ||D(R)) 16 ℓ← ℓCE + µℓKL 17 Fi ← Fi − η∇ℓ
18 return Gi(xi) to server

heading: 3.5 COMPARISONS WITH EXISTING FEDERATED LEARNING APPROACHES
text: ADCOL is quite different from typical federated learning algorithms. Instead of transferring and averaging models, ADCOL transfers the representations from the parties to the server. More importantly, the training is conducted in an adversarial way between the parties and the server. The server tries to distinguish the representations while the parties try to generate non-distinguishable representations. Here we compare ADCOL with FedAvg in terms of communication size and privacy.
Communication Size For simplicity, we do not consider party sampling during training (i.e., all parties participate in learning in each round). We use SL to denote the size of the local model. Then, the communication size of FedAvg is 2NSL, including the server sends the model to all parties and the parties send their local models to the server. We use n to denote the total number of examples (i.e., n = ∑N i=1 |Di|), d to denote the dimension of the representations, and SD to denote the size of the discriminator. Suppose each float value costs 4 bytes to store. In each round, the total communication size of ADCOL is (4nd+NSD), including the parties send the representations to the server and the server sends the discriminator to the parties. Thus, when n > (2SL − SD)N/4d, the communication size of ADCOL is larger than FedAvg in each round.
Note that the above threshold is usually large in practice. For example, in our experiments, when N = 6, d = 2048, SL = 142MB, and SD = 5MB, the threshold is equal to 204346. As we will show in the experiments, ADCOL is much more communication-efficient than FedAvg. Moreover, ADCOL requires a smaller bandwidth requirement as the size of each representation is much smaller than the local model.
Privacy While FedAvg transfers models and ADCOL transfers the representations, FedAvg and ADCOL both suffer from possible attacks on the exchanged information. There are many existing studies (Shokri et al., 2017; Nasr et al., 2019) that infer sensitive information from exchanged gradients/models. Also, there are studies (Salem et al., 2020) on reconstruction attacks wn the output of a model. While existing studies have shown that the mutual information between the input data and the final representation is small (Shwartz-Ziv and Tishby, 2017), it is not clear that which one is more private against inference attacks, which can be an interesting future direction. Local differential privacy (Kairouz et al., 2014) can be applied to protect the representation, which has been adopted in previous studies (Lyu et al., 2020). We have conducted experiments to show the feasibility of differential privacy to protect the communicated representations in Appendix B.13.

heading: 4 THEORETICAL ANALYSIS
text: As shown in Equation (2), the local loss has two parts: the cross-entropy loss part to update the whole network F and the KL divergence loss part to update the representation generator G. Ideally, to achieve minimum of ℓ, each part should achieve minimum. Since the cross-entropy loss part is same as FedAvg, we focus on the effect of the KL divergence loss. For simplicity, we ignore the
cross-entropy loss and study the KL divergence loss in our theoretical analysis1. The local objective of party i is:
min Gi
Ex∼DiℓKL([ 1
N ]N || D(Gi(x))) (3)
The objective of the discriminator is
max D N∑ i=1 Ex∼Di log(Di(Gi(x))), (4)
where Di(·) is the i-th output of the prediction vector D(·) (i.e., the probability of class i). Here we analyze the convergence property of the training process. All the proofs are available in Appendix A.
Theorem 1. We use PGi to denote the distribution of the representations generated in party i and PGi(z) is the probability of representation z in distribution PGi . Then, the optimal discriminator D ∗ of Equation (4) is
D∗k(z) = PGk(z)∑N i=1 PGi(z) . (5)
Theorem 2. Given the optimal discriminator D∗ from Equation (5), the global minimum of Equation (3) is achieved if and only if
PG1 = PG2 = · · · = PGN (6)
Theorem 1 and Theorem 2 show that to achieve the minimum of the objectives of the local parties and the discriminator, the parties will generate the same representation distribution, which matches the goal of ADCOL.
Theorem 3. Suppose P ∗G is the optimal solution shown in Theorem 2. If Gi (∀i ∈ [1, N ]) and D have enough capacity, and PGi is updated to minimize the local objective (i.e., Equation (3)), given the optimal discriminator D∗ from Equation (5), then PGi converges to P ∗ G.
The above theorem provides insights on the convergence of the training. In practice, we optimize the parameter θ of the local networks rather than PGi itself, which is reasonable due to the excellent performance as claimed in (Goodfellow et al., 2014).
Note that there are collapsing solutions for Equation (3) and Equation (4). The representations of each party can simply be constant vectors, which can achieve global minimum of Equation (3). Thus, the cross-entropy loss is necessary in Equation (2), which ensures that the generated representations are meaningful.

heading: 5 EXPERIMENTS

heading: 5.1 EXPERIMENTAL SETUP
text: Baselines We compare ADCOL with six baselines including SOLO (i.e., each party trains the model individually without collaborative learning), FedAvg (McMahan et al., 2016), FedBN (Li et al., 2021b), FedProx (Li et al., 2020), Per-FedAvg (Fallah et al., 2020), and FedRep (Collins et al., 2021). Here FedBN is the state-of-the-art FL approach on non-IID features. FedProx is a popular FL approach for non-IID data. Per-FedAvg and FedRep are two state-of-the-art personalized FL approaches. Like FedAvg (McMahan et al., 2016), we use weighted average according to the data volume of each party for all baselines.
Models All approaches use the same local model architecture for a fair comparison. The architecture of the local model is similar as SimSam (Chen and He, 2021), which has the following three components: (1) Base encoder: ResNet-50 (He et al., 2016). (2) Projection head: a 3-layer MLP with BN applied to each fully-connected layer. The input dimension is 4096. The dimension of the hidden
1Note that G is a part of F and the two losses are not independent of each other. For simplicity, we only analyze the KL divergence loss to study its effect.
layer and the output layer is 2048. (3) Predictor: a 2-layer MLP with BN applied to its hidden layer. The input dimension is 2048. The dimension of its hidden layer is 512. The discriminator is a 3-layer MLP. The input dimension is 2048. The dimension of the hidden layers is 512. The output dimension is equal to the number of parties.
Datasets We use the same datasets as in the study of FedBN. There are three real-world tasks in our experiments2: (1) Digits: The Digits task has the following five digit data sources from different domains: MNIST (LeCun et al., 1998), SVHN (Netzer et al., 2011), USPS (Hull, 1994), SynthDigits (Ganin and Lempitsky, 2015), and MNIST-M (Ganin and Lempitsky, 2015). (2) Office-Caltech-10 (Gong et al., 2012): The dataset has four data sources acquired using different camera devices or in different real environments with various backgrounds: Amazon, Caltech, DSLR, and WebCam. (3) DomainNet (Peng et al., 2019a): The dataset contains natural images coming from six different data sources with different image styles: Clipart, Infograph, Painting, Quickdraw, Real, and Sketch.
Setup By default, the number of parties is equal to the number of data sources, where each party has data from one of the data sources. For each dataset, we randomly split 1/5 of the original dataset as the test dataset, while the remained dataset is used as the training dataset. Same as (Li et al., 2021b), data of different parties have heterogeneous features but share the same labels and label distribution, which naturally forms the non-IID feature setting. The number of local epochs is set to 10 by default for all FL approaches. The number of epochs is set to 300 for SOLO. For ADCOL and FedProx, we tune µ ∈ {10, 1, 0.1, 0.01, 0.001} and report the best results. For FedRep, we tune β (i.e., step size for the second batch training) from 0.001, 0.01 and report the best results. We use the prediction layers as the shared representation in FedRep. We use PyTorch v1.81 (Paszke et al., 2019) to implement all approaches. We use the SGD optimizer for training with a learning rate of 0.01. The SGD weight decay is set to 10−5 and the SGD momentum is set to 0.9. The batch size is set to 64, 32, and 32 for Digits, Office-Caltech-10, and DomainNet, respectively. We run the experiments on a server with 8 * NVIDIA GeForce RTX 3090 and a cluster with 45 * NVIDIA GeForce RTX 2080 Ti.
Due to the page limit, we only present some representative experimental results in the main paper. For more experimental results and details, please refer to Appendix B.

heading: 5.2 OVERALL COMPARISON
text: Table 1 reports the test accuracy of different approaches on three tasks. We have the following observations. First, ADCOL is more effective than the other approaches. It can achieve the best test accuracy on most datasets. Moreover, ADCOL can outperform the other approaches by more than 2% accuracy on average. Second, while the parties may not benefit from FL approaches in some cases (e.g., Caltech-10), ADCOL always achieves better accuracy than SOLO, which demonstrates the robustness of ADCOL. Last, the personalized FL approaches (i.e., Per-FedAvg and FedRep) have a poor performance on the non-IID feature setting, which are even worse than SOLO.
2FedBN uses four tasks in their experiments. We only use three tasks because one task is not publicly available in FedBN’s GitHub website.

heading: 5.3 COMMUNICATION EFFICIENCY
text: To show the communication efficiency of ADCOL, like existing studies (Karimireddy et al., 2020; Lin et al., 2020), we compare the number of communication rounds and communication size of each approach to achieve the same target performance. The results on Digits are shown in Table 2. The results on the other two datasets are available in Appendix B.4. We can observe that no approach consistently outperforms the other approaches in terms of the number of communication rounds. However, the communication size of ADCOL is always much smaller than the other approaches. ADCOL can save at least 10 times the communication costs to achieve the same accuracy as FedAvg. The speedup can even be up to 34× on Digits. The results demonstrate that ADCOL is much more communication-efficient than the other FL approaches.

heading: 5.4 SENSITIVITY STUDIES
text: We study the effect of different factors on ADCOL and the other baselines using the Digit task.
Scalability and Heterogeneity We adopt the same approach as (Li et al., 2021b) to study the effect of number of parties and heterogeneity. We divide each dataset into ten parts randomly and equally and allocate each part into one party. The parties from the same dataset are treated as IID and the parties from different datasets are treated as non-IID. We add two parties from each dataset each time, which results in the number of parties N ∈ {10, 20, 30, 40, 50}. Moreover, the degree of heterogeneity decreases as the number of parties increases since the number of IID parties increases. The test accuracies are reported in Figure 2a. We can observe that the accuracy of all approaches can be slightly improved when increasing the number of parties due to the reduced heterogeneity and increased total amount of data. Given a different number of parties, ADCOL consistently outperforms the other baselines. Although the number of classes to distinguish increases for the discriminator when increasing N , ADCOL still shows a good and stable performance with a large number of parties.
Number of Local Epochs We vary the number of local epochs E ∈ {1, 2, 5, 10, 20} and report the results in Figure 2b. We run all approaches for 100 rounds. If the number of local epochs is too small, the local update is small in each round and the convergence speed is slow. Thus, the accuracy of all approaches is relatively low after running for 100 rounds with a small number of local epochs. ADCOL still consistently outperforms the other approaches with a different number of epochs.
Effect of Local Dataset Size We vary the percentage of the original local dataset used in each party from 20% to 100%. The results are shown in Figure 2c. The improvement of ADCOL is more
significant when the size of the local dataset is small. If the size of the local dataset is large, each party can already achieve satisfactory accuracy by SOLO. The accuracy of all approaches is close when the percentage is 100%. It is not necessary to conduct collaborative learning in such a case.
Effect of µ We vary µ ∈ {0, 0.1, 1, 10} and report the accuracy of ADCOL as shown in Figure 2d. We can observe that ADCOL can achieve the best accuracy when µ = 1. If µ is too small, the KL divergence loss of Equation (2) has little effect on the local training. Then, the goal of learning a common representation distribution may not achieve. If µ is too large, the cross-entropy loss of Equation (2) has little effect on the local training, and the representations may not be useful for classification at all (e.g., all representations are a constant vector). ADCOL with µ = 10 may even be worse than SOLO (i.e., µ = 0). Thus, an appropriate µ is important in ADCOL. Through our experimental studies, we find that setting µ = 1 is a good default choice.

heading: 6 CONCLUSION
text: In this paper, we propose ADCOL, a novel collaborative learning approach for non-IID features. ADCOL trains the models in an adversarial way between the parties and the server from a perspective of representation distributions. The parties aim to learn a common representation distribution, while the server aims to distinguish the representations by party IDs. Our experiments on three real-world tasks show that ADCOL achieves higher accuracy than the other state-of-the-art federated learning approaches on non-IID features. ADCOL provides new insights for collaborative learning.
Ethics Statement Not applicable.
Reproducibility Statement We have provided the experimental details in Section 5.1 and Section B.1 for reproducibility. Moreover, we will make the code publicly available.

heading: A THEORETICAL ANALYSIS
text: Theorem 1. We use PGi to denote the distribution of the representations generated in party i and PGi(z) is the probability of representation z in distribution PGi . Then, the optimal discriminator D ∗ of Equation (4)
D∗k(z) = PGk(z)∑N i=1 PGi(z) . (7)
Proof. From the view of the distribution of representations z, we can reformulate Equation (4) and the objective is to maximize:
N∑ i=1 ∫ z PGi(z) log(Di(z))dz (8)
Let V (D) = ∑N i=1 PGi(z) log(Di(z)). Note that ∑N
i=1 Di(z) = 1. Let F (D) = V (D) + λ(1 −∑N i=1 Di(z)). We have
∂F (D) ∂Di(z) = PGi(z) Di(z) − λ (9)
Let ∂F (D)∂Di(z) = 0 for i ∈ [1, N ], we have
PG1(z)
D1(z) =
PG2(z) D2(z) = · · · = PGN (z) DN (z) = λ (10)
Thus, V (D) can achieve maximum when
D∗k(z) = PGk(z)∑N i=1 PGi(z)
(11)
Theorem 2. Given the optimal discriminator D∗ from Equation (7), the global minimum of Equation (3) is achieved if and only if
PG1 = PG2 = · · · = PGN (12)
Proof. From Equation (3), the local objective of party k is to minimize
W (Gk) = −Ex∼Dk 1
N N∑ i=1 log(N ·Di(Gk(x)))
= −Ex∼Dk 1
N N∑ i=1 log( N · PGi(Gk(x))∑N j=1 PGj (Gk(x)) )
(13)
Note that ∑N i=1 PGi (Gk(x))∑N
j=1 PGj (Gk(x)) = 1. Similar to the proof in Theorem 1, Equation (13) can achieve
minimum when
PG1(Gk(x))∑N j=1 PGj (Gk(x)) = PG2(Gk(x))∑N j=1 PGj (Gk(x)) = · · · = PGN (Gk(x))∑N j=1 PGj (Gk(x)) = 1 N . (14)
For ∀k ∈ [1, N ] and ∀i ∈ [1, N ], we have PGi(Gk(x)) = ∑N j=1 PGj (Gk(x))
N . Given a representation z, we have PG1(z) = PG2(z) = · · · = PGN (z) = ∑N j=1 PGj (z)
N (15)
Thus, PG1 = PG2 = · · · = PGN .
Theorem 3. Suppose P ∗G is the optimal solution shown in Theorem 2. If Gi (∀i ∈ [1, N ]) and D have enough capacity, and PGi is updated to minimize the local objective (i.e., Equation (3)), given the optimal discriminator D∗ from Equation (7), then PGi converges to P ∗ G.
Proof. In Equation (13), consider W (Gk) = U(PGi) as a function of PGi . Then
∂U(PGi)
∂PGi = −
∑ k ̸=i PGk
NPGi(PGi + ∑ k ̸=i PGk) . (16)
We have ∂2U(PGi)
∂PGi 2 =
1 P 2Gi − 1 (PGi + ∑ k ̸=i PGk) 2 ≥ 0 (17)
Thus, U(PGi) is convex in PGi . Therefore, with sufficiently small updates of PGi , PGi converges to P ∗G, concluding the proof.

heading: B ADDITIONAL EXPERIMENTAL RESULTS

heading: B.1 ADDITIONAL EXPERIMENTAL DETAILS
text: In each experiment, like FedBN (Li et al., 2021b), to remove the effect of quantity skew, we truncate the size of all datasets to their smallest number with random sampling. For Digits, we resize all images to 28×28×3 and normalize them with mean 0.5 and standard derivation 0.5 for each channel. For Office-Caltech-10, we resize all images to 64× 64× 3 with random horizontal flip and random rotation. For DomainNet, we resize all images to 64×64×3 with random horizontal flip and random rotation. Like FedBN, we take Digits as the benchmark task for most studies.
The statistics of all the datasets are shown in Table 3. The label distributions are shown in Figure 3. The portion of samples with each class is close to 0.1. The label distribution is balanced among the parties. We train a ResNet-50 on all datasets (i.e., parties) from a task and extract the feature distributions of each dataset. Then, we use t-SNE to visualize the representation as shown in Figure 4. We can observe that the feature distribution of each party is different.

heading: B.2 CALTECH-10 AND DOMAINNET
text: Table 4 and 5 show the test accuracy of different approaches on Caltech-10 and DomainNet, respectively. We can observe that ADCOL still outperforms the other approaches in most cases.

heading: B.3 TRAINING CURVES
text: The training curves of different approaches on Digit are shown in Figure 5. We can observe that ADCOL is much more communication-efficient than the other approaches. ADCOL can convergence with a much smaller communication size than the other approaches.

heading: B.4 COMMUNICATION EFFICIENCY
text: Like Table 2, we show the communication efficiency of ADCOL on Caltech-10 and DomainNet in Table 6 and Table 7. We can observe that ADCOL is much more communication-efficient than the other approaches. The speedup can be even up to 300 times.

heading: B.5 PARTY SAMPLING
text: Party sampling is a technique usually used in the cross-device setting, where a fraction of parties is sampled to participate in federated learning in each round. Here we set the sample fraction to 0.4 in Digit and choose FedAvg and FedBN as the baselines. The training curves are shown in Figure 6. We can observe that all approaches have an unstable accuracy during training due to sampling. Moreover, FedAvg and FedBN have a very poor accuracy, which shows that existing federated learning approaches cannot well support party sampling on non-IID features. ADCOL significantly outperforms the other approaches.
We increase the number of parties to 100 (i.e., divide each dataset to 20 subsets) and vary the sampling rate from {0.1, 0.2, 0.5, 1}. We run all approaches for 200 rounds. The results are shown in Table 8. We can observe that when the sampling rate decreases, the performance of all approaches decreases. Moreover, the training is more unstable if the sampling rate is smaller. However, ADCOL still significantly outperforms FedAvg and FedBN. It is still a challenging task to develop effective algorithm on the cross-device setting with a low sampling rate.

heading: B.6 STUDY ON THE DISCRIMINATOR
text: One natural question is how to increase the information contained in the discriminator to improve the performance of ADCOL. We have tried two approaches.
Changing Model Architecture One approach is to increase the capacity of the discriminator. We change the model architecture to ResNet-50. The results are shown in Table 9. The performance of ADCOL cannot be improved by increasing the capacity of the discriminator.
Increasing the Number of Discriminators The other one approach is to increase the number of discriminators. Suppose the number of discriminators is Nd and the current round is t. Then, we use discriminators from round max(1, t−Nd) to round (t− 1) in the local training. The KL divergence loss is computed as
ℓ = 1
Nd Nd∑ i=1 ℓKL([ 1 N ]N || Di(G(x))), (18)
where Di is the discriminator trained in round max(1, t − i). The results are shown in Table 10. ADCOL cannot benefit from more discriminators. When the number of discriminators is larger, the accuracy of ADCOL is even worse. It is a future work to investigate how to integrate more useful information into the discriminator.

heading: B.7 DIMENSION OF REPRESENTATIONS
text: Same as SimSam (Chen and He, 2021), we set the dimension of representations (i.e., the output dimension of the projection head, the input dimension of the discriminator) to 2048 by default. As shown in Table 11, we report the performance of ADCOL varying the representation dimension. ADCOL can benefit from a larger representation dimension, where the representations are more informative. The mean accuracy can be improved by about 5% by increasing the dimension from 512 to 2048.

heading: B.8 NON-IID LABELS
text: We test the performance of ADCOL on non-IID label settings. Specifically, we sample pk ∼ DirN (0.5) and allocate a pk,j proportion of the instances of class k to party j, where Dir(0.5) is the Dirichlet distribution with a concentration parameter 0.5. The results are shown in Table 12. ADCOL cannot achieve a better performance than FedAvg. Intuitively, the task-specific representations of images from different classes should be very different. If the label distribution varies across parties, the representation distribution naturally also varies a lot. The intuition of ADCOL, which aims to learn a common representation distribution, is not appropriate on non-IID label settings.
B.9 IMBALANCED LOCAL DATA SIZES
In the main paper, we truncate the local data to the same size to control non-related factors. Here we show experiments with imbalanced number of samples across parties. We use 10% original data size for all parties without truncating the data. The results are shown in Table 13. ADCOL is still the best approach in most cases.

heading: B.10 COMPUTATION OVERHEAD
text: As shown in Table 14, the training time of ADCOL is larger than the other approaches. ADCOL requires the training of a discriminator in the server side, while the other approaches only need to average the models in the server side. However, in practice, the server usually has much powerful computation resources than the parties. Thus, the computation overhead in the server side is affordable.

heading: B.11 SHARING THE PREDICTOR LAYERS
text: The parties only send the representations to the server in ADCOL. While ADCOL aims to learn a common representation distribution z, an interesting extension is to share the predictor layers between
the parties and the server, which ideally helps in regularizing p(y|z). The results are shown in Table 15. We can observe that ADCOL without sharing the predictor layers is generally more effective than sharing the predictor layers. In practice, the distribution p(y|xi) is not exactly the same across parties. Thus, it is not necessary to regularize p(y|z) among the parties. Leaving the parties to fine tune their own predictor layer is more capable to learn the personalized local distribution.

heading: B.12 REPRESENTATION SAMPLING
text: An approach to reduce the communication cost is to sample a subset of representations and send them to the server. Here we change the sampling rate from {20%, 60%, 80%, 100%}. The final accuracy and the communication efficiency are shown in Table 16. We can observe that the communication cost of ADCOL can be significantly reduced with representation sampling technique. Moreover, there is little accuracy loss when the sampling rate is large than 60%.

heading: B.13 DIFFERENTIAL PRIVACY
text: We conduct experiments to apply differential privacy (Dwork, 2011) to protect the communicated representations. Specifically, we normalize the representations and add noises from Lap(0, 1/ϵ) before sending them to the server, where Lap(0, 1/ϵ) is the Laplace distribution with mean 0 and scale 1/ϵ. Then, in each round, the transferred representations satisfy ϵ−differential privacy (Lyu et al., 2020). We run 50 rounds for all approaches. For each round, we sample 150 of the representations of each party without replacement so that the privacy cost of each round is not accumulated. The results are shown in Table 17. We can observe that the accuracy of ADCOL is very close to the non-private version with a modest privacy budget (i.e., 10). It is promising to apply differential privacy if there are concerns about the information leakage of the representations.

heading: C DISCUSSION
text: Relation to Domain Adaptation Domain adaptation aims to train a model on a source domain (or multi-source domain), which has a good accuracy on a target domain. A classic and popular approach in domain adaptation is to perform adversarial training, i.e., training a discriminator to encourage
domain-invariant features (Ganin et al., 2016; Peng et al., 2019b). For the difference between our approach and domain adaptation, Peng et al. (2019b) proposed FADA, which extends domain adaptation in a federated setting. One connection between our approach and domain adaptation is that each party can be viewed as a source domain, and the target domain is the unknown oracle optimal. Then, our approach is to extract domain-invariant features from multiple source domain, which is used to regularize the training. To highlight the differences between our approach and the domain adaptation techniques, we compare our approach with the federated domain adaptation study (Peng et al., 2019b) (FADA) and show the main differences: (1) Setting: FADA aims to train a model on multiple source domain, which has a good accuracy on a target domain. Our study aims to train a personalized model for each party, which has good accuracy on its local data. (2) Discriminator: FADA uses multiple discriminators, where each discriminator is used for binary classification for one source-target domain pair. Our study uses a single discriminator for the multi-classification among all parties. Moreover, we have provided the theoretical analysis on the convergence properties. (3) Framework: FADA uses adversarial training to generate domain-invariant and domain-specific features. Our study uses adversarial training to regularize the local training in federated learning.
Intuitively, we cannot directly compare ADCOL and FADA in the experiments since the settings are different. In our experiments, there is no a single target domain for testing in FADA. One method is to treat each party as a target domain and applying FADA N times, where N is the number of parties. However, the computation and communication overhead is significantly large. Moreover, such an approach does not utilize the labels of the target dataset. We have compared ADCOL and FADA using the above method and the results are shown in Table 18. ADCOL significantly outperforms FADA. Moreover, the test accuracy of FADA is even smaller than local training in many cases since it does not exploit the labels of the target dataset.
Limitations ADCOL is a collaborative learning method for non-IID features. As shown in Appendix B.8, the performance of ADCOL is poor compared with federated learning approaches on non-IID label setting. Note that ADCOL aims to learn a common representation distribution. Intuitively, the task-specific representations of images from different classes should be very different, which can be easily classified by a small MLP. Thus, if the label distribution varies across parties, the representation distribution naturally also varies a lot across parties. The current objective of ADCOL does not fit into the non-IID label setting.
As shown in Section 3.5, the communication size of ADCOL is related to the number of examples. If the number of examples is very large and the size of the model is small, the communication cost of ADCOL will be larger than other federated learning approaches. However, local training can usually achieve satisfactory performance if the dataset size is very large. In such cases, besides ADCOL, existing federated learning approaches may also not help.
Insights and Future Work The key insights from ADCOL are (1) a GAN-style training scheme and (2) regularization from a view of representation distribution. While ADCOL does not have a requirement on the vanilla local training algorithm, it can also be extended to self-supervised federated learning, where the cross-entropy loss is replaced by the loss used in self-supervised learning (e.g., contrastive loss (Chen et al., 2020; Chen and He, 2021)). Moreover, while ADCOL only works on non-IID feature settings currently, the adversarial collaborative training scheme can potentially be
applied to address other data settings by modifying the objectives of local training and server training. There are many research opportunities based on the findings of this paper.

Review summaries:
This paper focuses on federated learning on non-IID features. This is a crucial problem when applying federated learning. The authors propose a new federated learning scheme, called ADCOL (Adversarial Collaborative Learning) for non-IID features. Specifically, the server is designed to train a discriminator to distinguish the local representations from local parties. While the local parties train the local models and expect the representations not to be distinguished by the discriminator. Authors conduct experiments on multiple parties where the data have heterogeneous features but share the same labels and label distribution and the results clearly show the effectiveness.

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.