SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zikai Xiao $^{1*}$ , Zihan Chen $^{2*}$  Songshang Liu $^{1}$  Hualiang Wang $^{3}$  Yang Feng $^{4}$  Jin Hao $^{5}$  Joel Tianyi Zhou $^{6}$  Jian Wu $^{1}$  Howard Hao Yang $^{1\dagger}$  Zuozhu Liu $^{1\dagger}$ $^{1}$ Zhejiang University,  
 $^{2}$ Singapore University of Technology and Design,  
 $^{3}$ The Hong Kong University of Science and Technology,  
 $^{4}$ Angelalign Technology Inc,  
 $^{5}$ State Key Laboratory of Oral Diseases, Sichuan University,  
 $^{6}$ Centre for Frontier AI Research (CFAR), A*STAR, Singapore  
zikai@zju.edu.cn
Data privacy and long-tailed distribution are the norms rather than the exceptions in many real-world tasks. This paper investigates a federated long-tailed learning (Fed-LT) task in which each client holds a locally heterogeneous dataset; if the datasets can be globally aggregated, they jointly exhibit a long-tailed distribution. Under such a setting, existing federated optimization and/or centralized long-tailed learning methods hardly apply due to challenges in (a) characterizing the global long-tailed distribution under privacy constraints and (b) adjusting the local learning strategy to cope with the head-tail imbalance. In response, we propose a method termed Fed-GraB, comprised of a Self-adjusting Gradient Balancer (SGB) module that re-weights clients' gradients in a closed-loop manner based on the feedback of global long-tailed prior derived from a Direct Prior Analyzer (DPA) module. Using Fed-GraB, clients can effectively alleviate the distribution drift caused by data heterogeneity during the model training process and obtain a global model with better performance on the minority classes while maintaining the performance of the majority classes. Extensive experiments demonstrate that Fed-GraB achieves state-of-the-art performance on representative datasets such as CIFAR-10-LT, CIFAR-100-LT, ImageNet-LT, and iNaturalist. Our codes are available at https://github.com/ZackZikaiXiao/FedGraB.
Federated learning (FL) is an approach for massively distributed clients to train a global machine learning model collaboratively without exposing their private data, which has garnered ever-increasing attention in academia and industry alike [1, 2]. Unlike conventional machine learning methods, FL brings the learning objective directly onto the end-user devices for local training, where only the intermediate parameters (e.g., gradients) need to be sent to a server for model aggregation and update. This approach not only substantially reduces the communication overheads but, more importantly, facilitates the clients in obtaining a generic global model without sharing their private data, thereby contributing to the development of trustworthy intelligent systems [3]. Despite its great potential,
\*Co-first author.
† Co-Corresponding author.
Figure 1: (a): An illustration of the Fed-GraB framework, in which SGB is mounted on all classes of each client based on the prior vector  $\mathcal{P}_c$  derived by the DPA module, ensuring federated balanced local training; (b): Comparison for the difference of positive and negative gradients.
realizing FL also requires one to overcome new hurdles in real-world implementations, among which a particular challenge stems from data heterogeneity, i.e., the non-IID distribution and imbalanced dataset sizes across clients [4-7]. These factors often give rise to a pernicious issue where the data set of different clients jointly follows a global long-tailed distribution, which is commonly known as the Federated Long-Tailed learning (Fed-LT) problem [8, 9]. For instance, patients' diagnosis varies substantially across medical centers but collaboratively form long-tailed distributions for certain diseases [10, 11]. In a broad range of applications, performance on the minority classes (i.e., tail classes) [12], such as rare diseases, dangerous behaviors in autonomous driving [13, 14], and abnormal breath or heart rates in wearable devices [15], play a pivotal role in developing reliable and robust solutions to an FL system.
The difficulty in addressing the Fed-LT problem primarily stems from the fact that although clients' datasets can be locally balanced, globally, i.e., if the datasets were aggregated, the distribution may be long-tailed. As clients would not disclose their data information due to privacy concerns, identifying the global long-tailedness in data distribution could be onerous. To that end, a natural question arises:
Question 1: How to leverage the global long-tailed statistics, especially for the minorities that need to be carefully tailored, without conflicting with privacy concerns?
On the other hand, even if the globally long-tailed data information is available, existing methods do not directly apply to address such a problem. Specifically, typical FL methods cope with data heterogeneity via generic techniques such as dynamic regularizations, client selections, data augmentation, distillation, and personalization training [16-19, 4, 20, 21]. These approaches ignored the diverged/imbalanced class levels that can result in the global long-tailed issue [22, 23]. Hence, they cannot guarantee decent performance for the tail classes. A key ingredient to improve long-tailed learning is to boost the performance of tail
Figure 2: An example of the global and local distributions of CIFAR-10-LT (Left) and CIFAR100-LT (Right) under non-IID setting. The bar on the top demonstrates the sorted global long-tailed distribution. Each row represents the diverse statistics of class imbalance of each client.
classes without degrading the performance of the head classes [24, 25]. While several learning methods dedicated to long-tailed data exist, they usually assume global class priors for re-balancing, which are not available in the context of FL due to privacy concerns [8]. Moreover, these methods may not be applicable in local training as local distributions could present diverse long-tailed characteristics and/or are not even long-tailed [9]. As such, it begs another question as follows:
Question 2: How to perform local training to synergistically aggregate a global model that excels on both majority and minority classes under the FL setting?
A possible choice is to re-balance the clients' local gradients via approaches such as Seasaw loss [26] or Equalization Loss v2 [27]. However one needs to customize these methods on different clients with heuristically selected hyper-parameters, resulting in complicated models with limited generalization capability. Moreover, the trained local models may have huge variance, as illustrated in Fig. 4 (a), and hardly formulate a synergistic model to handle tail class.
In light of the above challenges, this paper proposes Fed-GraB, comprised of a long-tailed statistic analyzer and a gradient rebalancer, for the Fed-LT problem, as illustrated in Fig. 1. More specifically, we develop a federated weight-norm-based method, coined as the Direct Prior Analyzer (DPA), to derive a prior vector for global data statistics, utilizing the weight parameters from the global classifier. To overcome the degradation caused by the distribution discrepancy between the server and clients (or amongst the clients) in Fed-LT, we establish a self-adjusting gradient balancer (SGB) based on the estimated global head-tail characteristics, which is integrated at every client to rebalance the positive and negative gradients in a class-wise and privacy-preserving manner. A marked advantage of Fed-GraB is that it trains all heterogeneous clients in a feedback-based closed-loop manner, encouraging all clients to contribute to the global long-tailed recognition task collaboratively. We also provide comprehensive theoretical and empirical analyses to verify the effectiveness of Fed-GraB. We conduct excessive experiments on multiple benchmark datasets (including the CIFAR-10/100-LT, ImageNet-LT, and iNaturalist) with both long-tailed and non-IID data distributions. The results demonstrate that Fed-GraB outperforms state-of-the-art (SOTA) baselines, including federated optimization for heterogeneous data, FL methods for LT classification, and centralized LT learning methods in FL settings.
2.1 Federated learning with data heterogeneity
Numerous FL approaches have been developed to cope with the heterogeneous or imbalanced data distribution. For instance, FedProx [16], FedDyn [17], and MOON [28] modify the local loss function by adding regularizers at the local side. SCAFFOLD [18] proposes a control variates-based method to reduce the clients' distribution drift brought by the discrepancies. FedIR [4] dynamically adjusts the sampling ratio of each client based on each client's contribution and data distribution situation. And FedNova [20] is proposed to tackle the induced objective inconsistency problem. On the server side, FedAvgM [19] and FedAdaGrad [29] are proposed to mitigate the performance degradation. Moreover, FL-oriented client selection [30, 31] and data augmentation strategies [32] have been investigated to enhance the performance. To address the poor generalization of a single generic model, personalization training methods have been recently explored to maintain multiple personalized local models [33, 34, 21]. However, most of these methods pay more attention to the discrepancies in inter-client distributions while ignoring the inconsistency among different classes, failing to achieve satisfactory performance on the tail classes. In lieu of this, an alternate paradigm [35] solely addresses the personalized long-tail conundrum, disregarding the variances within inter-client distributions. Notable exceptions could be found in [8, 36, 6] to tackle class imbalance. Still, they require the exchange of local features [8], private local data distribution [36], and auxiliary data [6], which would lead to potential privacy issues.
2.2 Long-tailed learning
Long-tailed data widely exists in real-world machine learning tasks, where head classes dominate the training [10, 37]. To tackle the poor performance of tail classes in such scenarios, many approaches have been explored to boost the performance of tail classes via class re-balancing, information augmentation, and module improvements [38-41, 25]. Re-balancing strategies, as the mainstream approaches in long-tailed learning, can be categorized into class re-sampling, re-weighting, and logit adjustments [10, 42]. Re-sampling approaches [43-46] address the class imbalance via over-sampling tail classes or under-sampling head classes. On the other hand, the re-weighting approaches aim to balance the loss( [47-50]) or gradients (Seesaw loss [26], Eqlv2 [27]). Specifically, Focal loss [39] increases the prediction probabilities to achieve better prediction performance for tail classes. And a label-distribution-aware margin loss LDAM is introduced by [49] using label frequency. Moreover, a
multi-expert framework (TADE [51]), prototype learning (OLTR [24]), and head-to-tail knowledge transfer (LEAP [52]) have been proposed for improved long-tailed data performance. However, such methods' effectiveness may be constrained in Fed-LT due to discrepancies in local and global statistics.
Federated Learning: Consider an FL system with  $N$  clients and an  $M$ -class visual recognition dataset  $\mathcal{D} = \{\mathcal{D}_k\}_{k=1}^N$  for classification tasks.  $\mathcal{D}_k$  denotes the local dataset of client  $k$  with size  $|\mathcal{D}_k|$ . Usually, the FL training process iterates multiple communication rounds until convergence. In a typical communication round  $l$ , the central server randomly selects a subset of clients  $\mathcal{S}_l$  and distributes the latest global model  $w^l$  to them. For client  $k \in \mathcal{S}_l$ , it performs local update to  $w^l$  based on  $\mathcal{D}_k$ , where the locally computed model is denoted as  $w_k^l$ . At the end of the round, the global model could be updated by the aggregation criterion (e.g., FedAvg) for the next round computation.
Long-tailed distribution in FL: To characterize the long-tailed data distribution in Fed-LT, we measure the imbalance factors in global perspectives, which is denoted by  $\mathrm{IF_G}$  and computed over the global dataset  $\mathcal{D}$ . We could employ a predefined imbalance  $\mathrm{IF_G}$  to sample from balanced datasets, constructing long-tailed datasets; see Fig. 2 for visualization.
Positive and Negative Gradients: Notably, positive and negative cumulative gradients have an important effect on long-tailed tasks such as visual recognition [53], object detection [27], and instance segmentation [26]. Suppose we have a neural network-based classification model with Cross Entropy loss as  $\mathcal{L}(z) = -\sum_{i=1}^{M} y_i \log(\sigma_i)$ , with  $\sigma_i = \frac{e^{z_i}}{\sum_{j=1}^{M} e^{z_j}}$ , where  $y_i$  is the ground truth label,  $z = [z_1, z_2, \ldots, z_M]$  are the logits, and  $\sigma = [\sigma_1, \sigma_2, \ldots, \sigma_M]$  denotes the probabilities of the classifier. Given a sample with label  $j$ , the positive gradients are defined as the derivative of the loss with respect to  $z_j$ , while the rest shall be the negative gradients, i.e.,
$$
\nabla_ {z _ {j}} ^ {\text {p o s}} (\mathcal {L}) = \sigma_ {j}, \quad \nabla_ {z _ {i \neq j}} ^ {\text {n e g}} (\mathcal {L}) = \sigma_ {i} - 1. \tag {1}
$$
3.2 Direct Prior Analyzer (DPA)
This section details the proposed DPA module which calculates a prior vector for balanced training. It utilizes the same information required by FedAvg, without the need for extra data or distribution details. Essentially, DPA derives a prior probability vector that discerns potential tail classes by examining the L2-norm of the weight parameters within the classifiers. The reason for this design arises from existing studies [25, 54], revealing that a model recognizes head or tail classes via classifier and tends to give higher scores for head classes when trained with imbalanced data. Thus, the weight norm serves as an effective indicator of the imbalance degree since the norm of the active weights could empirically reflect the behavior and frequency of different classes. In particular, let  $w_{g}^{j}$  denote the weights of classifier  $j$  in the global model classifier with weights  $w_{g}$ . The estimated global distribution would be given
Figure 3: The closed-loop in SGB.  $\beta_j^{pos}(t),\beta_j^{neg}(t)$  are updated according to the controller, and the re-weighted gradients are stored in the collector for future cumulative computation of  $\Delta_j(t)$ .
by  $\mathcal{P}_c = \{p_j\mid \bar{h} (\| w_g^j\| _2)\}$ , where  $h(\cdot)$  is a linear transformation for scaling the norm to share of the total. This estimation can be executed after global model aggregation with full client participation.
In the Fed-GraB implementation, the prior vector  $\mathcal{P}_c$  derives from DPA at the onset of local training. It is also worth mentioning that  $\mathcal{P}_c$  does not need to be perfect and only requires information that reflects the model prediction inclinations at the class level, rather than a complete correspondence with the distribution itself. See more detailed analysis in experiments.
3.3 Self-adjusting Gradient Balancer for Fed-LT
This part establishes the Self-adjusting Gradient Balancer (SGB), an effective tool that can be integrated into biased classifiers to re-balance local training and improve the overall performance. Particularly, we consider the difference between cumulative positive and negative gradients  $\Delta(t)$  with  $\Delta_j(t)$  for class  $j$  defined as  $\Delta_j(t) = g_j^{pos}(t) - g_j^{neg}(t)$ , where  $g_j(t) = \sum_{i=0}^t \nabla_{z_j}(\mathcal{L}^t)$  is the cumulative gradients of class  $j$ , capturing the overall imbalance degree throughout the training process. In a toy balanced scenario with  $M$  samples each with equal probability to the  $M$  classes, the expected target  $z_j(t)$  of  $\Delta_j(t)$  is:
$$
E \left(z _ {j} (t)\right) = \sum_ {i = 1} ^ {M - 1} \left(E \left(\sigma_ {i} (t)\right) - 1\right) - E \left(\sigma_ {j} (t)\right) = 0, \tag {2}
$$
indicating that  $\Delta (t)$  approaches zero when the distributions become identical.
The core idea of SGB is to consistently align  $\Delta_j(t)$  to 0 for tail classes in a closed-loop and self-adjusting manner. Specifically, we carry out a proportional-integral-derivative controller [55] to continually update the re-weighting coefficients of gradients with regard to the logits in a class-wise manner, as demonstrated in Fig. 3. The re-weighting process is based on adaptive self-adjustment with error feedback until  $\Delta(t)$  reaches the pre-determined balanced target, rather than following heuristic methods as [26, 27]. The error feedback  $e_j(t) = \Delta_j(t) - z_j(t)$  for class  $j$  represents the distance between the current status  $\Delta_j(t)$  and a target  $z_j(t)$  during training. Given  $e(t)$ , the output of the controller in SGB is
$$
u _ {j} (t) = K _ {P} e _ {j} (t) + K _ {I} \sum_ {j = 0} ^ {t} e _ {j} (t) + K _ {D} \left(e _ {j} (t) - e _ {j} (t - 1)\right), \tag {3}
$$
where  $K_{P}, K_{I}$ , and  $K_{D}$  are controlling factors. The proportional item with  $K_{P}$  shows the controller would give bigger re-weighting coefficients  $\beta_{j}$  to compensate the imbalance for current larger errors. Meanwhile, the integral item with  $K_{I}$  could look back on the past errors to reduce the steady-state offset which could not be handled by the proportional item. Furthermore, as computing  $\beta_{j}$  only with the proportional and integral loop could hardly achieve steady and smooth re-weighting, we introduce the third differential item with  $K_{D}$  to control the fluctuations based on future expectations. For instance, when  $\Delta_{j}(t)$  is close to 0, the differential item would prohibit  $u(t)$  from diverging re-weighting and vise versa. Overall, the three items work together to achieve quick, precise, and stable adjustments.
Algorithm 1 Local training process of Fed-GraB
Input:  $w_{k}^{l}$ , local model at round  $l$  for client  $k$
Output: Updated local model  $w_{k}^{l + 1}$  functionCLIENTUPDATE( $w_{k}^{l}$ )
1: Update prior vector  $\mathcal{P}_c$  by DPA with  $w_k^l$
2: for batch  $b \in \mathcal{B}$  do
3: for each classifier  $z_{j}\in z$  do #Initialize  $\Delta_j(t) = 0$
4: Compute  $\nabla_{z_j}(\mathcal{L};b) = \left[\nabla_{z_j}^{pos}(\mathcal{L};b),\nabla_{z_j}^{neg}(\mathcal{L};b)\right]$
5: Sample a random value  $r \sim \text{Uniform}(0, 1)$  #Gradient Balancer
6: Input of SGB  $e(t) \gets \Delta_j(t)$
7: Compute output of SGB  $u(t)$  by Eq. (3)
8: Compute re-weighting coefficients  $\beta_{j}$  by Eq. (4) #Gradient Collector
9: Update  $\Delta_j(t)$  with the re-weighted gradients by Eq. (5)
0:  $w_{k}^{l + 1}\gets w_{k}^{l} - \eta \sum_{j = 0}^{M - 1}\beta_{j}\nabla_{z_{j}}(\mathcal{L};b)$
1: return  $w_{k}^{l + 1}$
We first map  $u_{j}(t)$  through a simple activation function  $\phi(\cdot)$  to calculate  $\beta_{j}^{neg} = \phi(-u_{j}(t))$ ,  $\beta_{j}^{pos} = \phi(u_{j}(t))$ , where  $\phi(x) = \frac{\gamma}{1 + \delta e^{-\zeta x}}$ . Then, we combine the global information obtained by DPA to calculate the final coefficients for positive and negative gradients per iteration as:
$$
\beta_ {j} = \mathbb {I} _ {r > \mathcal {P} _ {c} [ j ]} \cdot \left[ \beta_ {j} ^ {\text {p o s}}, \beta_ {j} ^ {\text {n e g}} \right] + \mathbb {I} _ {r \leq \mathcal {P} _ {c} [ j ]} \cdot [ 1, 1 ], \tag {4}
$$
where  $\mathbb{I}(\cdot)$  is the indicator function. Consequently, the re-weighted gradients are computed as:
$$
\left\{ \begin{array}{l} \nabla_ {z _ {j}} ^ {\text {p o s} ^ {\prime}} (\mathcal {L} ^ {t}) = \beta_ {j} ^ {\text {p o s}} (t) \nabla_ {z _ {j}} ^ {\text {p o s}} (\mathcal {L} ^ {t}), \\ \nabla_ {z _ {j}} ^ {\text {n e g} ^ {\prime}} (\mathcal {L} ^ {t}) = \beta_ {j} ^ {\text {n e g}} (t) \nabla_ {z _ {j}} ^ {\text {n e g}} (\mathcal {L} ^ {t}). \end{array} \right. \tag {5}
$$
The working mechanism of SGB at a typical client is summarized in Fig. 3. As depicted in Fig. 1 (a), SGB is implemented across all the classes in each client, with parameters adjusted in accordance with the prior vector derived by DPA. By utilizing global information, this approach encourages all clients to contribute coherently to the global model while simultaneously rectifying prediction biases arising from long-tailed datasets. As shown in Fig. 1 (b), SGB improves the balance between cumulative positive and negative gradients. See experiments and Supplementary for further analysis.
3.4 Algorithm Summary for Fed-Grad
The training pipeline of our Fed-GraB framework is presented in Algorithm 1, which consists of two stages. First, DPA calculates the prior vector  $\mathcal{P}_c$  after receiving aggregated weights  $w_k^l$ . Then, SGB is applied to all the classes. The local classifier would be dynamically adjusted with weighted positive and negative gradients via a  $\Delta(t)$ -based closed-loop controller in accordance with prior vector  $\mathcal{P}_c$  derived from DPA. The combination with SGB and DPA is detailed in Eq. (4).
3.5 Privacy discussions of DPA
As the local prior computation is performed at the client side, there would be a few privacy concerns with the DPA method. It is, however, noteworthy that the potential privacy issue exists in the general FL frameworks rather than specific to our proposed DPA method. For instance, gradient inversion [56] can pose a threat to almost all gradient transmission-based FL methods without any privacy-preservation techniques. As the privacy issue of the FL framework is beyond the scope of this work, we briefly include the discussion in this subsection.
Baselines: We consider two types of SOTA baselines: (1) FL-oriented methods to tackle data heterogeneity (FedProx [16], FedNova [20]) or federated long-tailed data-oriented (FedIR [4] and CReFF [8]), and FedAvg is also included for reference; (2) Long-tailed learning (LT)-oriented methods ( $\tau$ -norm [25], Eq1v2 [27], LDAM [49], Focal-loss [39] and GCL-loss [42]) applied at local training of each client. We also provide the results of the long-tailed methods in centralized learning (CL) settings as the oracle performance upper bound, including SGB.
Datasets: We conduct the experiments in three benchmark datasets for long-tailed classification, i.e., CIFAR-10/100-LT [57], ImageNet-LT [58]. CIFAR-10/100-LT are sampled into long-tailed distribution by exponential distribution controlled by IF [49], and we use the same configurations as in [24] for ImageNet-LT with the number of images per class ranging from 5 to 1280. To evaluate the performance on real-world data, we also conduct experiments on iNaturalist-User-160k, with 160k examples of 1,023 species classes and partitioned on the basis of iNaturalist-2017 [59].
Federated settings: We use non-IID data partitions for all experiments, implemented via symmetric Dirichlet distributions with concentration parameter  $\alpha$  to control the identicalness of local data distributions among all the clients. We train a ResNet-18 over  $N = 40$  clients on CIFAR-10-LT. ResNet-34 and ResNet-50 are used on CIFAR-100-LT and ImageNet-LT respectively with  $N = 20$  clients. For iNaturalist-160k, we use the same settings as ImageNet-LT.
Implementation of baselines: For the sake of fairness, we keep consistent settings external for experiments. We conduct experiments using a starting model from FedAvg. For CReFF, the number of federated features is 100, we use 0.1, 0.01 as federated feature learning rate and main net learning rate respectively on CIFAR-10/100-LT. We report the official result [8] on ImageNet-LT. As  $\tau$ -norm is a one-shot method, we provide a pre-trained model by FedAvg and adjust the classifier weights on the server as post-process. Focal-Loss and EQL v2 do not require distribution prior. We directly replace the local cross-entropy loss with their proposed re-balancing methods. For LDAM, we use the local data distribution for its local training.
4.2 Comparison with State-of-the-art Methods
Evaluation on CIFAR-10-LT: