Federated Learning on Heterogeneous and Long-Tailed Data via Classifier Re-Training with Federated Features

1 Introduction

The emergence of federated learning (FL) enables multiple clients to collaboratively learn a powerful global model without transmitting local private data to the server. Serving as a communication-efficient and privacy-preserving learning framework, FL has shown its potential to facilitate real-world applications in multiple domains, e.g., natural language processing [Jiang et al., 2021] and fraudulent credit card detection [Zheng et al., 2020]. However, during FL model training, one major practical challenge is data heterogeneity [Li et al., 2020a]. In the FL environment, data generated or collected by different clients may come from various sources, which results in distribution discrepancy between clients. Moreover, real-world data often exhibits long-tail distribution with heavy class imbalance, where the number of samples in some classes (called head classes) severely outnumbers that in some other classes (called tail classes). Building an unbiased classification model on data with this kind of distribution is termed long-tail learning, which has been extensively studied in recent years [Zhang et al., 2021]. In the FL scenario, if the training data across clients is long-tailed and heterogeneous at the same time, the joint problem becomes complicated and challenging because each client may hold different tail classes. For example, medical institutions aim to build a diagnosis model on private patient records, which are held by each institution locally. Each institution has different disease distributions, and different diseases make up long-tail distribution because some diseases are common and others are rare. Adopting FL in this case, local models will poorly perform on different tail classes due to data heterogeneity and long-tail distribution, and the aggregated global model will also be affected. One straightforward way is to simply adopt existing long-tail learning methods to the joint problem in FL. However, most of these methods require the information of local or global class distribution as prior knowledge for optimization, which may expose potential privacy issues.

∗Corresponding author: Yang Lu, luyang@xmu.edu.cn

One of the recent advances in long-tail learning is decoupling the training process into representation learning and classifier re-training [Kang et al., 2020]. That is, with the trained feature extractor fixed, a biased classifier is re-trained using a set of balanced data, which is also called two-stage learning. In the FL framework, an FL model can also be decoupled into two-stage training, where the feature extractor is obtained by aggregating client models and the classifier is retrained on the server using a set of balanced data. We conduct a quick experiment on CIFAR-10-LT (the long-tailed version of CIFAR-10 [Cao et al., 2019]) to show that two-stage learning performs well on heterogeneous and long-tailed data in FL. We use imbalance factors (IF) and the Dirichlet distribution coefficient α to control the degree of long-tail and heterogeneity, respectively. Higher IF means a higher degree of imbalance degree, and lower α means a higher degree of heterogeneity. Given a global model pre-trained on heterogeneous and long-tailed data, we fix its feature extractor and retrain its classifier on the server with a small balanced dataset (100 per class) collected from clients. It can be observed in Figure 1 that the strategy is particularly useful and the classification performance is significantly improved. In particular, re-training classifier achieves the highest performance gain of 18% when IF = 100 and α = 0.1 (the rightmost grey bar in the left subfigure of Figure 1), which is an extremely heterogeneous and long-tailed setting. However, this prerequisite of using a shared balanced dataset on the server is infeasible for most FL practical scenarios due to privacy concerns.

Motivated by the idea of two-stage learning and the privacy concerns of FL, we propose a novel and privacy-preserving FL method called Classifier Re-training with Federated Features (CReFF) to deal with the joint problem of data heterogeneity and long-tail distribution. We learn a set of balanced features called federated features because only the classifier needs to be re-trained on the server. It is based on an intuitive idea that the classifier re-trained on federated features should produce comparable performance as the one re-trained on the real data, which can be achieved by making two classifiers similar. Therefore, the classifier re-training optimization on federated features should follow the similar path as that of the real data. Specifically, we optimize the federated features to make their gradients close to the gradients of real data. As shown in Figure 1, the classifier re-trained on federated features (orange bars) achieves comparable performance as the one re-trained on real data (grey bars). Extensive experiments show that CReFF significantly outperforms the state-of-theart federated learning methods on image classification tasks with heterogeneous and long-tailed data distribution. The contributions of this paper can be summarized as follows:

• We study the joint problem of FL with heterogeneous and long-tailed data distribution, where the local and global class distribution is unknown to the server.
• We reveal an intriguing fact that the biased classifier is the primary factor leading to the poor performance of an FL global model on heterogeneous and long-tailed data.
• We propose CReFF, a novel FL algorithm to deal with heterogeneous and long-tailed data by re-training the classifier with learnable federated features on the server. CReFF has no privacy concerns because no real data or information of class distribution is required.

2 RelatedWork

2.1 Federated Learning with Heterogeneous Data

A variety of solutions have been proposed to tackle data heterogeneity in FL, which are mainly from two perspectives:

One focuses on optimization strategies to make the diversity between client models and global model limited [Li et al., 2020b; Huang et al., 2021]; the other adopt mechanisms on the server to alleviate the negative influence of data heterogeneity [Lin et al., 2020; Chen and Chao, 2021]. For example, CCVR [Luo et al., 2021] deals with data heterogeneity via classifier re-training using virtual features sampled from an approximated Gaussian mixture model. Although the abovementioned methods solve the data heterogeneity to some extent, they poorly perform on the tail classes due to the lack of consideration of the global long-tail distribution.

Some FL methods are designed for imbalanced data (not specifically long-tailed). One strategy is to adopt client selection to match clients with complementary class distribution [Duan et al., 2020; Yang et al., 2020]. In this case, all clients are required to upload their local class distributions to the server, which violates the principle of privacy protection of FL. Ratio loss [Wang et al., 2021a] utilizes balanced auxiliary data on the server to estimate the global class distribution for better local optimization. However, as we mentioned above, the auxiliary data is not available on the server in real applications. Compared with these related methods, the proposed CReFF has no privacy concerns because no real data or information of class distribution is required.

2.2 Long-tail Learning

Recently, long-tail learning has drawn much interest in deep learning [Zhang et al., 2021]. Some methods follow the ideas of imbalance learning to augment the feature space for rare classes [Kim et al., 2020; Zang et al., 2021] or re-weight different classes according to their frequencies [Cao et al., 2019]. Some recent methods decouple the training phase into representation learning and classifier re-training [Kang et al., 2020], which generate more generalizable representations and achieve strong performance after re-balancing the classifier. Based on decoupling, ensemble-based methods [Xiang et al., 2020; Wang et al., 2021b] adopt a multi-expert framework to learn diverse classifiers in parallel, which reduces the model bias towards the tail classes. However, most of them require the global class distribution, as we mentioned above. During FL model training, it is infeasible to gather the information of class distribution of each client to obtain the global class distribution, which makes the vast majority of long-tail learning methods not applicable to FL scenarios.

3 ProposedMethod

3.1 Preliminaries

Settings and Notations. We discuss a typical FL setting with K clients holding potentially heterogeneous data partition D1 ,D2 , ...,DK , respectively. The goal is to learn a global model over the whole training data D ≜ ⋃k Dk with the help of a central server without data transmitting. In this paper, we consider the setting when D is drawn from a long-tail distribution X = {xi, yi}|Ni=1 , yi ∈ {1, ...,C}. Let nkc be the number of samples of class c on client k, and nc = ∑Kk=1 nkc . Without loss of generality, a common assumption in long-tail learning is that the classes are sorted by cardinality in non-increasing order, i.e., if c1 < c2, we have nc1 ≥ nc2 and n1 ≫ nC. For an FL model, we typically consider a neural network φw with parameters w = {u, v}. It has two main components: 1) a feature extractor fu with parameters u, mapping each input sample x to a d-dim feature vector; 2) a classifier hv with parameters v, typically being a fully-connected layer which outputs logits to denote class confidence scores. The parameters of the kth client local model are denoted as wk.

Basic Algorithmof FL. FedAvg [McMahan et al., 2017] is the fundamental algorithm for FL. In round t, the server first sends a global model wt to clients. The clients then update the received model on their local data Dk , k = 1, ...,K:

After local updating, some clients in the set At are selected to upload their updated models to the server. Finally, the server performs weighted average to update the global model for round t + 1:

3.2 Motivation

Classifier re-training has been shown effective for heterogeneous data [Luo et al., 2021] and long-tailed data [Kang et al., 2020], separately. Now, we empirically show that classifier re-training also works for the joint problem. In other words, the biased classifier is the primary factor leading to the poor performance of the global model trained on heterogeneous and long-tailed data.

First, we divide a ResNet-8 network into five components, including four blocks 1 and one classifier. Then, we use a small set of balanced data (100 per class) collected from clients to re-train each component of a FedAvg model while keeping other components fixed. The FedAvg model is pretrained on CIFAR-10-LT with IF = 100 and α = 0.5. As shown in Figure 2, we can observe that re-training any single component can achieve certain improvement (blue line). Especially, re-training the classifier achieves the highest performance gain (around 15.35%). Therefore, it can be shown that the biased classifier is the primary factor leading to the poor performance of the global model because other components are less affected.

However, requiring external data on the server is impractical in real FL applications, although the amount of the data is small. Transmitting data from clients violates the key privacy-preserving principle of FL, and the data collected by the server itself may not follow the same distribution as the training data in clients.

3.3 Framework of the Proposed CReFF

Based on the above assumption and corresponding observation, we propose CReFF to deal with the joint problem of data heterogeneity and long-tail distribution. It is designed as a simple and effective approach based on FedAvg by only re-training a new classifier using a small set of learnable features on the server, which is called federated features. Our goal is to re-train the classifier on federated features such that its performance is comparable with the classifier re-trained on real features, which can be achieved by making two classifiers similar. Therefore, the classifier re-training optimization on federated features should follow the similar path as that of the real data. Specifically, we optimize the federated features to make their gradients close to the gradients of real data such that the classifiers trained on real and federated features converge to a similar solution in the parameter space. In this way, the classifier re-trained on the federated features can be approximated by the one re-trained on the real data in a privacy-preserving manner. As shown in Figure 2, CReFF produces comparable performance as on real data.

1 We regard the first convolutional layer with BN layer as the first block.

CReFF consists of two core components: local training on clients and federated features optimization on the server. In round t, each client receives two models from the server: 1) a global model wt ; 2) a re-trained model ̂wt . The former is used to update local models for global model aggregation, and the latter is used to calculate the real feature gradients for federated features optimization. After local training, some clients are selected to upload their local models and real feature gradients to the server. Finally, the server updates the global model and the re-trained model for round t + 1. The architecture of CReFF is illustrated in Figure 3.

Local Training. In round t, the server sends two models to clients. The global model wt is composed of a feature extractor ut and a classifier vt , and the re-trained model ̂wt has a re-trained classifier ̂vt with the same feature extractor ut . For each received model, the local training consists of a corresponding part. The first part is a typical FedAvg local model update as shown in Equation (1). The second part is calculating the real feature gradients for federated features optimization on the server. Client k produces d-dim real features Z k c = {zkc,i }|nkc i=1 for class c through the feature extractor ut . Then, real feature gradients gk c ∈ RC×d of class c can be computed using ̂vt :

Finally, client k uploads two parts to the server: 1) the local model wt+1 k for global model aggregation; 2) the real feature gradients {gk c |c ∈ C k } for federated features optimization. Herein, C k denotes the set of classes on client k because the real feature gradients on client k may only contain partial classes due to data heterogeneity. It is worth noting that clients only upload the average gradients of each class to the server, which erases the private information of local data to some context because the averaging operation is irreversible.

Federated Features Optimization. Once the server receives two parts from the selected clients in At , there are also two corresponding parts to aggregate the client models and optimize federated features for classifier re-training, respectively. The first part is to aggregate the local models by Equation (2) to obtain the aggregated global model wt+1 . The second part aims to learn a set of balanced d-dim federated features S t c = {stc,i }|mi=1 for each class c in round t, where m is the number of federated features of each class. First, the server aggregates real feature gradients of each class c by averaging over all selected clients in At :

where Atc is the subset of clients that hold class c. Then, the server computes the federated feature gradients gfed c over federated features S t of class c by using the same re-trained classifier ̂vt as calculating gagg c :

Both gfed c and gagg c are gradients of the same re-trained classifier ̂v, whose dimensions are C × d. We use the gradient matching loss [Zhao et al., 2021] to measure the difference between them by averaging the cosine dissimilarity of each part:

where g[j] denotes the jth row of the gradient. After optimizing the federated features S t on this loss, we obtain an updated S t+1 for classifier re-training.

Compared with sampling virtual features from the approximated Gaussian mixture for classifier re-training [Luo et al., 2021], CReFF learns federated features by optimizing the gradient matching loss. The advantages are twofold. On the one hand, CReFF can produce nearly the same gradients as the real features to update the classifier, which can not be achieved by CCVR. On the other hand, CCVR requires each client to upload its local class distribution to the server for computing global mean and covariance. In contrast, CReFF does not require the information of local class distribution in a more privacy-preserving manner.

The last step of CReFF is to re-train the classifier of the updated global model wt+1 to obtain the re-trained model ̂wt+1 . Initialized as ˜v, a new classifier ̂vt+1 is re-trained on the optimized balanced federated features S t+1 :

Finally, the updated global model wt+1 = {ut+1 , vt+1 } and the re-trained model ̂wt+1 = {ut+1 , ̂vt+1 } are sent to clients for round t + 1. The whole training process of CReFF is shown in Algorithm 1.

4 Experimental Results

In this section, we first show the ability of CReFF to deal with data heterogeneity and long-tail distribution on several

Algorithm 1: Training process of CReFF

Input: Initialized global model w0 , number of steps I for federated features optimization, number of steps J for classifier re-training, number of federated features per class m, number of training rounds T Output: Re-trained model ̂wT on round T 1 for t = 1 to T do 2 Randomly select a set of active clients At ; // Clients execute: 3 for k ∈ At do 4 Update local model wt+1 k by Equation (1); 5 Compute real feature gradients {gk c |c ∈ C k } by Equation (3); 6 Send wt+1 k and {gk c |c ∈ C k } to the server; 7 end // Server executes: 8 Aggregate local models to wt+1 by Equation (2); 9 Aggregate real feature gradients to gagg c by Equation (4); 10 Compute federated feature gradients gfed c by Equation (5); 11 Optimize federated features to S t+1 with the loss in Equation (6) for I epochs; 12 Re-train the classifier of wt+1 to ̂wt+1 on S t+1 by Equation (7) for J epochs; 13 Send wt+1 and ̂wt+1 to clients; 14 end

benchmark datasets, compared with the state-of-the-art FL methods. Then, we specifically study the reason of federated features being effective for classifier re-training on the server.

4.1 Experimental Setup

We conduct experiments on the following datasets:

• CIFAR-10/100-LT [Krizhevsky and Hinton, 2009]. We follow [Cao et al., 2019] to shape the original balanced CIFAR-10/100 into long-tail distribution with IF = 100, 50 and 10, respectively. Like previous studies [Lin et al., 2020], we use Dirichlet distribution to generate the heterogeneous data partition among clients. We set the value of α at 0.5 on CIFAR-10/100-LT.
• ImageNet-LT [Russakovsky et al., 2015]. It contains 115.8K images from 1,000 categories, with the largest and smallest categories containing 1,280 and 5 images, respectively. We set the value of α at 0.1 on ImageNet-LT.

We use ResNet-8 for CIFAR-10/100-LT, and ResNet-50 for ImageNet-LT as the base model. We implement all compared FL methods with the same model for a fair comparison. All experiments are run by PyTorch on two NVIDIA GeForce RTX 3080 GPUs. By default, we use standard cross-entropy loss and run 200 communication rounds. We set the number of total clients at 20 and an active client ratio 40% in each round. For local training, the batch size is set at 32. For server-side training, we initialize the federated features as random noise and set the number of them per class at 100, the optimization steps I on federated features at 100, the classifier re-training steps J at 300. We use SGD with a learning rate 0.1 as the optimizer for all optimization process.

4.2 Comparison with the State-of-the-art Methods

We compare CReFF with heterogeneity-oriented FL methods, including FedAvg [McMahan et al., 2017], FedAvgM [Hsu et al., 2019], FedProx [Li et al., 2020b], FedDF [Lin et al., 2020], FedBE [Chen and Chao, 2021], CCVR [Luo et al., 2021] and FedNova [Wang et al., 2020]. Moreover, we also compare with imbalance-oriented FL methods, including Fed-Focal Loss [Sarkar et al., 2020], Ratio Loss [Wang et al., 2021a] and FedAvg with τ-norm [Kang et al., 2020]. Note that Ratio Loss depends on a balanced auxiliary data, so we fetch the data from clients, although it is infeasible.

Results onCIFAR-10/100-LT. The results are summarized in Table 1. CReFF achieves the highest test accuracy on both datasets with different IFs. Compared with the baseline FedAvg, the performance gain of CReFF is the highest when IF = 100 (around 14.4% improvement for CIFAR-10-LT and 4.3% improvement for CIFAR-100-LT). It shows the generalization ability of CReFF when the global class distribution is highly long-tailed. For heterogeneity-oriented methods, most of them (e.g., FedAvgM and FedDF) perform similarly to FedAvg because they only deal with data heterogeneity without taking global class imbalance into account. For imbalanceoriented methods, some of them, e.g., Ratio Loss, perform well in some cases compared with FedAvg. However, there is still a performance gap compared with CReFF because they aim at dealing with the global imbalance problem but generally ignore the data heterogeneity problem across clients.

Results on ImageNet-LT. We further evaluate CReFF on ImageNet-LT as reported in Table 2. To better examine the performance of classes with different number of training samples, we report the accuracy on three sets of classes: Manyshot (more than 100 samples), medium-shot (20-100 samples), and few-shot (less than 20 samples). Compared with other FL methods, CReFF achieves the best results on all cases, which shows that the improved overall accuracy by CReFF does not sacrifice the accuracy of many-shot classes.

4.3 Model Validation

To further validate the effectiveness of CReFF, we study two key research questions.

Why does a classifier re-trained on federated features performsimilarly to the one re-trained on real features? We show the dissimilarity between the federated and the real feature gradients by the gradient matching loss of each class. We group the classes into three sets for CIFAR-10-LT: Many (more than 1500 samples), medium (200-1500 samples) and few (less than 200 samples), and then average the dissimilarities in each group shown by the solid lines in Figure 4. It can be observed that the dissimilarities decrease to around 0 in a few rounds, which means that optimizing on the gradient matching loss successfully makes the federated feature gradients close to the real feature gradients. Thus, the optimization of classifier re-training on federated features follows a similar path as that of the real features. Moreover, we also show the dissimilarity between the federated and the real features by calculating their cosine dissimilarity in Figure 4 (dotted lines). As training goes on, the dissimilarities between the federated and real features are also decreasing but they are far from 0. It means that the federated features do not have to be similar to the real features although their corresponding gradients are very close. Therefore, the privacy of the real features can be further guaranteed.

How many federated features per class is enough to re-train a good classifier? One important hyperparameter in CReFF is the number of federated features per class. We evaluate its influence by tuning it in the set {0, 1, 5, 10, 50, 100, 150, 200} on CIFAR-10-LT with different IFs. It can be observed from Figure 5 that more federated features produce higher accuracy in general. However, as the number of federated features increases, the improvement gradually becomes stable. This result is anticipated because the information provided by real feature gradients is limited. Meanwhile, a surprising result is that even if only one federated feature is learned for each class, CReFF can still improve 4.21%, 3.74%, and 0.92% for CIFAR-10-LT with IF = 100, 50 and 10, respectively.

5 Conclusion

In this paper, we first show that classifier re-training is a straightforward and effective way to deal with the joint problem of data heterogeneity and long-tail distribution in the federated learning framework. Based on this, CReFF is proposed to learn a small set of balanced features called federated features on the server for classifier re-training in a privacypreserving manner. The classifier re-trained on the federated features can produce comparable performance as the one retrained on the real data. Experiments have shown that CReFF outperforms the state-of-the-art FL methods in the setting of heterogeneity and long-tail distribution, and the effectiveness CReFF is validated empirically.

Acknowledgments

We thank Dr. Fei Chen for his helpful discussion and reviewers of IJCAI-ECAI 2022 for their constructive and helpful feedback. This work was supported in part by the National Natural Science Foundation of China (No. 62002302, U21A20514, 61872307 and 52007173), the Open Research Projects of Zhejiang Lab (No. 2021KB0AB03), the Natural Science Foundation of Fujian Province (No. 2020J01005), and Zhejiang Provincial Natural Science Foundation of China (No. LQ20E070002).