Enhancing the Performance of Global Model by Improving the Adaptability of Local Models in Federated Learning

1 Introduction

Owing to the accessibility of large-scale datasets, deep neural networks have achieved great success over the years [He et al., 2016]. However, a vast amount of training data may be distributed across plenty of clients, posing challenges on how to effectively utilize the client data [Li et al., 2019; Kairouz et al., 2019]. To handle these challenges, federated learning [McMahan et al., 2017] has emerged as a distributed learning paradigm with privacy-preserving property. Federated learning enables the clients to collaboratively learn a well-performed global model by sharing local updates. It demonstrates the strength in fully utilizing data from the clients without requiring the upload of private data to the server. Due to its communication efficiency and privacy preservation, federated learning has been widely used in multiple domains [Hard et al., 2020; Kang et al., 2020; Jiang et al., 2021; Zheng et al., 2020; Khan et al., 2021; Adnan et al., 2022].

∗Wei Wang is the corresponding author.

However, the generalization performance of federated learning relies on the assumption that the client data are independent and identically distributed (IID) [McMahan et al., 2017]. In real scenarios, federated learning often encounters data heterogeneity, where the clients hold Non-IID data. As mentioned in earlier studies, data heterogeneity affects the effectiveness of federated learning [Zhao et al., 2018; Li et al., 2020]. Recent studies have proposed numerous methods to address the issue of degradation of model generalization performance in Non-IID settings. A line of work relies on clients uploading extra auxiliary variables during the model upload phase, such as gradients [Karimireddy et al., 2020; Dai et al., 2023], statistical distribution information [Duan et al., 2021], etc., to modify server-side model updates. However, these methods introduce additional communication costs and the reliability of the uploaded information may impact the effectiveness of federated learning [Li et al., 2022].

Another line of work focuses on aligning the local models with the global model. [Li et al., 2020; Acar et al., 2021] propose adding penalty terms to local objectives to prevent local models from deviating from the global model during local training. [Chen and Chao, 2022; Zhang et al., 2022] propose replacing the local loss with a balanced loss to improve the local model's performance on classes with few samples. These methods aim to make local models converge near the global optimal stationary point. However, since each local model is trained with its own data and does not know the data distributions over other clients, these methods still struggle to achieve effective alignment between the global and local models.

In this paper, we introduce the adaptability of local models, i.e., the average performance of local models on data distributions over clients, and enhance the performance of the global model by improving the adaptability of local models. Since each client does not know the data distributions of other clients, the adaptability of the local model cannot be directly optimized. First, we provide the property of an appropriate local model which has good adaptability on the data distributions over clients. Then, we formalize the property into the local training objective with a constraint and propose a feasible solution to train the local model. During the model aggregation phase, we further propose a model aggregation method that allows local models with good adaptability to have large aggregation weights. We call this method Federated Learning with Adaptability over Client Distributions (FedACD). Extensive experiments on federated learning benchmarks demonstrate that our method significantly improves the adaptability of local models and leads to a well-performed global model that consistently outperforms the baseline methods.

2 RelatedWorks

Federated learning [Li et al., 2019; Kairouz et al., 2019] enables clients to collaboratively learn a global model by sharing local updates conducted on the local data. The general federated learning method, FedAvg[McMahan et al., 2017], introduces an iterative model averaging approach. However, previous studies demonstrate that the robust generalization performance of federated learning relies on the assumption that the client data are independent and identically distributed (IID) [McMahan et al., 2017]. In Non-IID scenarios, the learned model may experience degradation in the generalization performance [Zhao et al., 2018].

Many methods based on FedAvg have been proposed to tackle data heterogeneity. One line of work relies on clients to upload extra auxiliary variables during the model upload phase to modify the server-side model updates. Scaffold [Karimireddy et al., 2020] requires clients to upload extra gradient information besides the updated model, utilizing control variates (variance reduction) to correct client drift in local updates, thereby improving the convergence rate of FL. CReFF [Shang et al., 2022] uses the gradient information uploaded by the clients to train the virtual features [Luo et al., 2021] on the server to retrain the classifier of the global model. FedNH [Dai et al., 2023] initializes uniform class prototypes on the server and then sends the fixed prototypes to clients to guide training. However, these methods introduce additional communication costs, and the model's performance heavily relies on the reliability of the auxiliary variables. In scenarios with severe data heterogeneity or partial client sampling, the reliability and update frequency of auxiliary variables decrease, leading to undecent generalization performance [Li et al., 2022].

Another line of work focuses on that data heterogeneity leads to a misalignment between local models and the global model and attempts to adjust local training loss to align local and global models. FedProx [Li et al., 2020] proposes adding a proximal term into local objectives and penalizing gradient updates far from the global model. FedNova [Wang et al., 2020] employs a normalized averaging approach to eliminate objective inconsistency while preserving rapid error convergence. FedROD [Chen and Chao, 2022] proposes to view each client's local training as an independent class imbalance problem and utilizes the balanced loss to replace the crossentropy loss to adjust local objectives. FedLC [Zhang et al., 2022] introduces a deviation bound to measure the gradient deviation after local updates and then calibrates the logit of each class before softmax cross-entropy based on local label distribution to alleviate the deviation. However, the alignment of local and global models remains an ongoing exploration.

3 Methods

We consider the following Federated Learning (FL) scenario. There are M clients with distribution D1 ,D2 , ...,DM and each client m has data S m = {(xmk , ym k )}N m k=1 drawn from distribution Dm, where xmk ∈ Rd is the d dimension sample , ym k ∈ [1, 2, ...,C] is the label of xmk , and N m = |S m| is the local sample size, m ∈ [M]. The goal of FL is to learn a global model ϕw with parameters w over all training data S ≜ ⋃mS m without data transmission. For a sample x, fw(x) is the logit vector of the global model ϕw on x. The probability output, denoted p(x) = (p1, . . . , pC), is derived from fw(x), where pi is the i-th element of p(x) and pi = ef iw (x) ∑Ck=1 efkw (x) , where f kw (x) is the k-th element of fw(x). The predicted label is given by ϕw(x) = argmax(p1, . . . , pC). In federated learning, the global model ϕw is to minimize the following risk over client distributions:

Let P(ϕw(x) ̸= y | y = i) represent the error rate of the global model ϕw on class i and P(y = i) represent the prior probability of class i, i ∈ [C]. We can get that P(ϕw(x) ̸= y) = ∑Ci=1 P(ϕw(x) ̸= y | y = i)P(y = i). For simplicity of notation, let ϵϕw = (ϵϕw 1 , ..., ϵϕw C ) denote the error rate of the global model ϕw, where ϵϕw i ≥ 0 represents the error rate of the global model on class i, i.e., ϵϕw i = P(ϕw(x) ̸= y|y = i). Let πm = (πm 1 , . . . , πm C ) denote the prior distribution of client m, where πm i represents the prior distribution of class i on client m. From Equation (1), we get that

In federated learning, the global model is aggregated from local models. Prior works attempt to enhance global model performance by aligning local models with the global model, e.g., introducing regularization terms [Li et al., 2020; Acar et al., 2021] or using balanced losses [Chen and Chao, 2022; Zhang et al., 2022]. However, since clients have heterogeneous data distributions, and each client does not know the data distributions of other clients, these methods have struggled to achieve effective alignment. The global model ϕw is to minimize the risk on the data distributions π1 , . . . , πM over clients in Equation (2). Let ϕwm denote the local model of client m with parameters wm, m ∈ [M]. This motivates us that the local model ϕwm should also achieve small risk on the data distributions π1 , . . . , πM over clients:

R(ϕwm ) implies the adaptability of local model ϕwm . If R(ϕwm ) is small, i.e., ϕwm has good average performance on the data distributions π1 , . . . , πM , we say that ϕwm has good adaptability. Unfortunately, since client m does not know the data distributions π1 , . . . , πm-1 , πm+1 , . . . , πM , it is impossible to directly minimize Equation (3) to guide the training process of ϕwm . Here, we discuss how to find a surrogate loss for Equation (3). For client m, we consider the specific model ϕw∗ whose error rate ϵϕw∗ satisfies ϵϕw∗ i = ϵϕw∗ j for i ̸= j and i, j ∈ [C]. The risk R(ϕw∗) over the client distributions for ϕw∗ is:

From Equation (4), we can find that R(ϕw∗) = ∥ϵϕw∗ ∥1 C for any client distributions π1 , . . . , πM . This implies that ϕw∗ may be a good choice for serving as the local model. Now, we provide the following theorem:

Theorem 1. For client m, let πm denote the data distribution of client m, and ϕwm is the local model on client m with error rate ϵϕwm. Suppose ∥ϵϕwm∥1 = ∥ϵϕw∗ ∥1 = ϵ∗ . If ϵϕwm does not satisfy the condition that ϵϕwm i = ϵϕwm j for i ̸= j and i, j ∈ [C], there exist sets of client distributions ̂π1 , . . . , ̂πm-1 ,πm, ̂πm+1 , . . . , ̂πM on which the risk R(ϕwm ) > R(ϕw∗ ).

Proof. Without loss of generality, we assume that πm = ( 1C + δm 1 , . . . , 1C + δm C ) where - 1C ≤ δm 1 , . . . , δm C ≤ 1 - 1C and ∑Ck=1 δm k = 0. Let r = argmaxk∈[C] ϵϕwm k be the class index with maximal error rate, s = argmink∈[C] ϵϕwm k is the class index with minimal error rate. Since ϵϕwm does not satisfy the condition that ϵϕwm i = ϵϕwm j for i ̸= j and i, j ∈ [C], we have ϵϕwm r > ϵϕwm s . Nowwe construct distribution ̂πn = (̂ πn 1 , . . . , ̂ πnC ) where ̂ πn k = 1C +δn k for k ∈ [C], n ∈ [M]\{m} with - 1C ≤ δn k ≤ 1 - 1C and ∑Ck=1 δn k = 0. Let 0 < θ < 1C be a constant. If δn k , k ∈ [C], n ∈ [M]\{m} satisfy the following condition 1 :

then the risk R(ϕwm ) over client distributions is

1 Here, we consider the non-trivial case that the number of clients M is larger than the number of class C.

It is easy to find that the condition in Equation (5) can be satisfied, e.g.,

for k ∈ [C] and t ∈ [M]\{m}. Since ∥ϵϕw∗ ∥1 = ϵ∗ , with Equation (4) we get R(ϕw∗) = ϵ∗ C . Thus, we have R(ϕwm ) > R(ϕw∗ ).

Theorem 1 indicates that for the model ϕwm with the same generalization ability as ϕw∗, i.e., ∥ϵϕwm∥1 = ∥ϵϕw∗ ∥1 = ϵ∗ , if the model ϕwm does not satisfy the condition that ϵϕwm i = ϵϕwm j for i ̸= j and i, j ∈ [C], ϕw∗ is a better choice than ϕwm . The reason is that there exist sets of client distributions ̂π1 , . . . , ̂πm-1 ,πm, ̂πm+1 , . . . , ̂πM on which R(ϕwm ) > ϵ∗ C , while R(ϕw∗) is always ϵ∗ C for any distributions. In federated learning, the clients have heterogeneous data distributions, and each client does not know the distributions of other clients. In this way, we can train the local model according to the optimization in Equation (6) for each client m, m ∈ [M]:min ∥ϵϕwm ∥1,

Minimizing ∥ϵϕwm∥1 can be achieved by minimizing the cross-entropy (CE) loss over samples of local data S m. However, ensuring that the error rate ϵϕwm satisfies the constraints ϵϕwm i = ϵϕwm j for i ̸= j is challenging. In the following section, we propose a feasible solution for this. As defined, ϵϕwm i = P(ϕwm (x) ̸= y | y = i) = ∑j̸=i P(ϕwm (x) = j | y = i), where P(ϕwm (x) = j | y = i) represents the probability that local model ϕwm predicts a sample of class i as class j and can be computed by averaging the j-th element of the probability outputs of the samples of class i. Therefore, it is possible to optimize the probability outputs of local model ϕwm to force the error rate ϵϕwm to meet the constraints in Equation (6). First, we introduce the probability matrix P m for client m:

where P m ij = P(ϕwm (x) = j | y = i). P m ij can be evaluated as: P m ij = 1 N m i ∑(x,y)∼Dm i pj, i, j ∈ [C], where pj is the j-th element of the probability output p(x) of sample x. Specifically, pj = ef jwm (x) ∑Ck=1 efkwm (x) , where f kwm (x) is the k-th element of the logit vector fwm (x) of sample x. It is easy to find that ϵϕwm i = 1-P m ii , i ∈ [C]. In this way, the constraints ϵϕwm i = ϵϕwm j , i ̸= j, in Equation (6) become the probability constraints: P m ii = P m jj , i ̸= j. We make the constraints hold during the local training process in the following way:

first, we force the probability matrix P m to have the form of P m ij = 1-P m ii C-1 , i ̸= j; then, we make P m ii = P m jj for i ̸= j by adjusting P m ij = P m ji , i.e., 1-P m ii C-1 = 1-P m jj C-1 .

The local probability matrix P m is evaluated by averaging the probability outputs of each class's samples of the local data. Therefore, to force the local probability matrix P m to have the form of P m ij = 1-P m ii C-1 , i ̸= j, we flatten the misclassification probability for each sample . For a sample x with label y, its probability output p(x) = (p1, ..., pC), where py is the correct classification probability, and pk, k ̸= y, is the misclassification probability. Based on the probability output p(x), we construct the target vector q(x) = (q1, ..., qC), where the misclassification probability qk, k ̸= y, is flattened. Specifically, qy = py and qk = 1-py C-1 for k ̸= y. Thus, we minimize the Kullback-Leibler (KL) divergence between p(x) and q(x) shown in Equation (7) to flatten the misclassification probability for the sample x.

In this way, by adjusting the probability output p(x) of each sample x with label y to ( 1-py C-1 , ..., py, ..., 1-py C-1 ), the probability matrix P m has the form of P m ij = 1-P m ii C-1 for i ̸= j.

To ensure that the local probability matrix P m meets the probability constraints P m ii = P m jj for i ̸= j, for the pair of classes (i, j), the local probability matrix P m should further satisfy P m ij = P m ji . With the condition P m ij = 1-P m ii C-1 for i ̸= j, P m ij = P m ji is equivalent to 1-P m ii C-1 = 1-P m jj C-1 , i.e., P m ii = P m jj , i ̸= j. To make P m ij = P m ji , if P m ij > P m ji , P m ij should be decreased; otherwise, P m ij should be increased. Directly adjusting P m ij and P m ji during the local training process is difficult. Since P m is evaluated based on the sample's logit vector, we can implicitly adjust P m ij and P m ji by modifying the logit margin between f iwm (x) and f jwm (x) for sample x. Based on previous work [Menon et al., 2021], we dynamically adjust the cross-entropy (CE) loss by modifying the logit margin term {f iwm (x) - f ywm (x)} for each sample x with label y with the following loss function:

For simplicity, we define ∆yi = P m yi P m iy . If ∆yi < 1, the logit margin term {f iwm (x) - f ywm (x)} will be reduced to suppress class y's relative margin towards class i, causing an increase of ∆yi. Conversely, if ∆yi > 1, the logit margin term {f iwm (x) - f ywm (x)} will be increased to relax class y's relative margin towards class i, leading to a decrease of ∆yi. Finally, ∆yi = 1 implies that P m ii = P m jj for i ̸= j. Thus, the loss function to guide the local training process of ϕwm is:

where the parameter λ controls the contribution of L1 and L2.

When confronted with severe data heterogeneity, certain classes of some clients may hold limited samples or even no samples. If class k is a missing class with N m k = 0, ∆yk in Equation (8) is unknown due to the unknown value of P m ky . To deal with missing classes, previous work [Zhang et al., 2022] suggests that the gradient updates of missing classes should be constrained, which inspires us to set ∆yk as a small value. The evaluation of the local probability matrix relies on averaging the probability outputs of samples, while the means of the probability outputs of the samples belonging to the classes with limited samples may deviate from the true means. To address this issue, we incorporate Input Mixup [Zhang et al., 2018] into the training process. Specifically, within each local training epoch's minibatch, the data mixup technique is employed to increase the occurrence frequency of samples for the classes with limited samples.

During the model aggregation phase, local models with good adaptability should be assigned with large aggregation weights. For client m, m ∈ [M], if the local model ϕwm has good adaptability, the constraints in Equation (6) should hold, implying that the probability matrix P m satisfies that P m ii = P m jj for i ̸= j. In this way, we can assign the aggregation weights by measuring the degree to which the constraints are satisfied. We set a template matrix Q, where Qii = Qjj for i ̸= j, and measure the degree to which the constraints are satisfied by calculating the KL divergence between the local probability matrix P m and the template matrix Q. Intuitively, the ideal template matrix is the identity matrix I, but we can not calculate the KL divergence between P m and I. Hence, we define a template matrix Q as that shown in Equation (10). For client m, we compute its aggregation score Vm by calculating the KL divergence between the local probability matrix P m and the template matrix Q as follows:

where 0 < τ < 1 is a parameter. Our Federated Learning with Adaptability over Client Distributions (FedACD) follows the rules below:

Local: w⋆m ← argminw L1 + λL2, initialized with w;

After local training, client m uploads the aggregation score Vm to the server. It is noteworthy that Vm is a scalar, and the server cannot infer the local probability matrix or any privacysensitive information from Vm, which effectively protects the privacy of the client data.

4 Experiments

4.1 Settings

Datasets and Models. We perform extensive experiments on three benchmark datasets: CIFAR-10, CIFAR-100 [Krizhevsky, 2009], and Tiny-ImageNet [Deng et al., 2009; Le and Yang, 2015]. Tiny-ImageNet is a subset of ImageNet with 100k samples of 200 classes. Following [Chen and Chao, 2022], we adopt a simple Convolutional neural network for CIFAR-10 and CIFAR-100, while using Resnet18 [He et al., 2016] for the Tiny-ImageNet. We implement all compared federated learning methods with the same model for a fair comparison.

Client settings. We employ Dirichlet sampling to generate Non-IID data for each client. Dirichlet sampling is a common technique used in FL for creating Non-IID data [Chen and Chao, 2022; Zhang et al., 2022; Dai et al., 2023]. It yields distinct label distributions for each client, with the degree of data heterogeneity controlled by β. Smaller β refers to severer heterogeneity and when β < 1, some clients may lack samples of certain classes. In our experiments, we adopt β ∈ {0.3, 0.1, 0.05}.

Hyper-parameters. For the local training process, SGD optimizer is used with a 0.01 initial learning and 0.9 momentum. We employ a weight decay of 10-5 for CIFAR-10 and CIFAR-100, and 10-3 for TinyImageNet to mitigate overfitting. The batch size is set as 64. The number of clients is set as 20 and the participation ratio is set as 40%. The local training epoch is set as 5 and the total communication round is set as 200. λ is set as 1 and τ is set as 1 - 10-5 .

Baselines. We select three types of FL methods as the baselines. 1) Generic FL: FedAvg [McMahan et al., 2017]; 2) Classical FL with Non-IID data: FedProx [Li et al., 2020], FedNova [Wang et al., 2020]; 3) FL methods most related to us: FedROD [Chen and Chao, 2022], CReFF[Shang et al., 2022], FedNTD[Lee et al., 2022], FedLC [Zhang et al., 2022], FedDecorr[Shi et al., 2023], FedNH [Dai et al., 2023].

4.2 Main results

We evaluate the performance of the global model on the test set with uniform data distribution (shown in Table 1) and the test sets constructed based on the data distribution of each client (shown in Table 2). Our method achieves the best performance across three datasets with diverse data heterogeneity. To further demonstrate that our method improves the adaptability of local models, we also calculate the average risk on the distributions over clients of local models (shown in Section 4.1). The results show that the risk over client distributions of the local models in our method is consistently lower than that in the baseline methods with diverse data heterogeneity. This confirms that our method enhances the performance of the global model by improving the adaptability of local models. Classical FL methods with Non-IID data, such as FedProx and FedNova, show similar results to FedAvg, which indicates that it is hard to align the local models that are trained on data with different distributions. The classifier re-training method, CReFF, fails to demonstrate effectiveness under various settings, which could be attributed to the unreliability of synthesized features. The prototype-based method, FedNH, exhibits only marginal improvement in settings with severe data heterogeneity, which indicates that prototype learning is also impacted by data heterogeneity. FedDecorr focuses on representation learning with data heterogeneity, but also fails to achieve significant performance improvements. FedROD and FedLC demonstrate an improvement in accuracy compared to FedAvg on CIFAR-10, which indicates that the balanced loss helps alleviate the impact of data heterogeneity to some extent. However, the balanced loss fails to exhibit effectiveness when confronted with challenging scenarios, such as CIFAR-100 and TinyImageNet.

4.3 Ablation Study

Different number of clients K with various participation ratio γ. We select the best baselines in Table 1 and conduct experiments on skewed CIFAR-100 under different number of clients K with various participation ratio γ. As shown in Table 4, to validate the impact of client number, we set γ = 0.4 and K ∈ {20, 40, 60, 80, 100}. As the number of clients gradually increases, achieving convergence in FL becomes harder. Most methods experience a decline in accuracy, while our method consistently outperforms the baseline methods. To validate the impact of the participation ratio, we set K = 20 and γ ∈ {0.2, 0.4, 0.6}. When γ is small, the client data distributions among different rounds vary significantly, leading to divergent gradient directions. However, our method consistently achieves the best performance. All experiments demonstrate that our method performs well under different numbers of clients with partial participation ratios.

Input Mixup. We incorporate Input Mixup into local training. For a fair comparison, we select several methods compatible with Input Mixup in Table 1 and incorporate Input Mixup into the training process. As shown in Table 5, our method without Input Mixup outperforms the baseline methods without Input Mixup under all settings. After incorporating Input Mixup, our method consistently outperforms the baseline methods with Input Mixup by a large margin under various data heterogeneity. Unlike other methods where Input Mixup mainly augments data diversity, we leverage it to improve the reliability of the evaluation process of the local probability matrix by ensuring that the means of the probability outputs of samples belonging to classes with limited samples are closer to the true means, thus incorporating Input Mixup can improve the performance of our method.

Different local epoch E. We select the best baselines in Table 1 and vary the number of local training epoch E ∈ {3, 5, 7, 9, 11, 13, 15} for each client in every round. As shown in Table 6, our method consistently achieves the best performance across different settings. Meanwhile, it is noteworthy that as the local training epoch E increases, the performance of most methods exhibits a declining trend. However, our method maintains relatively stable performance for different local epochs, especially with a large E ∈ {9, 11, 13, 15}. The reason is that with the increase of local training epochs, the local models tend to overfit the local data distributions. This increases the risk of local models on heterogeneous data distributions over clients, i.e., the adaptability of local models becomes worse.

Parameter λ. Parameter λ in Equation (9) controls the contribution of L1 and L2. We conduct experiments with λ ∈ {0.5, 1, 1, 5, 2.0}. As shown in Table 7, our method performs best with λ = 1 and remains robust to variations of λ. Therefore, we set λ = 1 in all experiments.

Aggregation method. We propose an aggregation method that allows local models with good adaptability over client distributions to have large aggregation weights. We conduct experiments by comparing our aggregation method with the aggregation method with uniform weights (denoted as FedACDbase). As shown in Table 8, our aggregation method consistently outperforms FedACDbase under various data heterogeneity, demonstrating its effectiveness.

5 Conclusion

In this paper, we introduce the adaptability of local models, i.e., the average performance of local models on data distributions over clients, and focus on improving the adaptability of local models to enhance the performance of the global model. Extensive experiments on federated learning benchmarks demonstrate that our method achieves the wellperformed global model that outperforms the baseline methods.

Acknowledgements

This work is supported by the National Science Foundation of China (62276125) and the Collaborative Innovation Center of Novel Software Technology and Industrialization.