SlowGuess/NeurIPS_2023_OCR huggingface dataset:
ravid.shwartz.ziv@nyu.edu
Micah Goldblum*
Yucen Lily Li
yucenli@nyu.edu
C. Bayan Bruss
bayan.bruss@capitalone.com
andrewgw@cims.nyu.edu
Real-world datasets are often highly class-imbalanced, which can adversely impact the performance of deep learning models. The majority of research on training neural networks under class imbalance has focused on specialized loss functions, sampling techniques, or two-stage training procedures. Notably, we demonstrate that simply tuning existing components of standard deep learning pipelines, such as the batch size, data augmentation, optimizer, and label smoothing, can achieve state-of-the-art performance without any such specialized class imbalance methods. We also provide key prescriptions and considerations for training under class imbalance, and an understanding of why imbalance methods succeed or fail.
Only a minuscule proportion of credit card transactions are fraudulent, and most cancer screenings come back negative. In reality, some events are common while others are exceedingly rare. As a result, machine learning systems, often developed in class-balanced settings [e.g., 48, 46, 16], are routinely trained and deployed on class-imbalanced data where relatively few samples are associated with certain minority classes, while majority classes dominate the datasets. Class-imbalanced training data can negatively impact performance. Consequently, a wide body of literature focuses on specially tailored loss functions and sampling methods for counteracting the negative effects of imbalance [10, 14, 25, 27, 50, 61]. In striking contrast to such approaches, we instead show that simply tuning existing components of standard neural network training routines can achieve state-of-the-art performance on class-imbalanced image and tabular benchmarks at little implementation overhead and without requiring any specialized loss functions or samplers designed specifically for imbalance. Like Wightman et al. [69], who found that modern training routines allow ResNets to achieve performance competitive with that of later architectures, we show that modern training techniques cause the benefits of specialized class-imbalance methods to nearly vanish.
Moreover, our carefully tuned training routine can be combined with existing class imbalance methods for additional performance boosts. Conducting evaluations on real-world datasets, we find that existing methods, which performed well on web-scraped natural image benchmarks on which they were designed, underperform in the real-world setting, whereas our approach is robust.
Our investigation provides key prescriptions and considerations for training under class imbalance:
1. The impact of batch size on performance is much more pronounced in class-imbalanced settings, where small batch sizes shine.
2. Data augmentations have an amplified impact on performance under class imbalance, especially on minority-class accuracy. The augmentation strategies in our experiments which achieve the best performance on class-balanced benchmarks yield inferior performance on imbalanced problems.
3. Large architectures, which do not overfit on class-balanced training sets, strongly overfit on imbalanced training sets of the same size. Moreover, newer architectures which work well on class-balanced benchmarks do not always perform well under class imbalance.
4. Adding a self-supervised loss during training can improve feature representations, leading to performance boosts on class-imbalanced problems.
5. A small modification of Sharpness-Aware Minimization (SAM) [19] pulls decision boundaries away from minority samples and significantly improves minority-group accuracy.
6. Label smoothing [57], especially on minority class examples, helps prevent overfitting.
To understand why exactly such training routine improvements confer significant benefits, we investigate the role of overfitting in class-imbalanced training. Our analysis shows that naive training routines overfit on minority samples, causing neural collapse [60], whereby features extracted from the penultimate layer concentrate around their class-mean. Combining this analysis with decision boundary visualizations, we demonstrate that unsuccessful methods for class-imbalanced training overfit strongly, whereas successful methods regularize.
A long line of research has been conducted on class-imbalanced classification. There are several archetypal approaches specially designed to address imbalance:
Resampling the data. In early ensemble learning studies, boosting and bagging algorithms were adjusted to take account of imbalanced data by resampling. Traditionally, resampling involves oversampling minority class samples by simply copying them [25, 10, 27], or undersampling majority classes by removing samples [17, 32, 3, 8], so that minority and majority class samples appear equally frequently in the training process.
Loss reweighting. Reweighting methods assign different weights to majority and minority class loss functions, increasing the influence of minority samples which would otherwise play little role in the loss function [14, 34]. For instance, one may scale the loss by inverse class frequency [28] or reweight it using the effective number of samples [14]. As an alternative approach, one may focus on hard examples by down-weighing the loss of well-classified examples [50] or dynamically rescaling the cross-entropy loss based on the difficulty of classifying a sample [61]. Bertsimas et al. [6] encourage larger margins for rare classes, while Goh and Sim [21] learn robust features for minority classes using class-uncertainty information which approximates Bayesian methods.
Two-stage fine-tuning and meta-learning. Two-stage methods separate the training process into representation learning and classifier learning [54, 59, 38, 4]. In the first stage, the data is unmodified, and no resampling or reweighting is used to train good representations. In the second stage, the classifier is balanced by freezing the backbone and fine-tuning the last layers with resampling or by learning to debias the class confidences. These methods assume that the bias towards majority classes exists only in the classifier layer or that tweaking the classifier layer can correct the underlying biases.
Several works have also inspected representations learned under class imbalance. Kang et al. [38] find that representations learned on class-imbalanced training data via supervised learning perform better when the linear head is fine-tuned on balanced samples. Yang and Xu [71] instead examine the effect of self- and semi-supervised training on imbalanced data and conclude that imbalanced labels are significantly more useful when accompanied by auxiliary data for semi-supervised learning. Kotar et al. [44], Yang and Xu [71] make the observation that self-supervised pre-training is insensitive to imbalance in the upstream training data. These works study SSL pre-training for the purpose of transfer learning, sometimes using linear probes to evaluate the quality of representations. Inspired by their observations, we find that the addition of an SSL loss function on the same class-imbalanced dataset, even when no upstream data is available, can significantly improve generalization.
In summary, existing works propose countless approaches to address class imbalance during training. In contrast, we show that strong performance can be achieved on class-imbalanced datasets simply by tuning the components of standard neural networks training routines, without specialized loss functions or sampling methods designed specifically for imbalance. Our tuned routines require little to no additional implementation compared to standard neural network training pipelines and can be combined with existing specialized approaches for class imbalance. We additionally provide novel prescriptions and considerations for training under class imbalance, as well as an understanding of how regularization can contribute to the success of training under imbalance.
3 Optimizing Training Routines for Imbalanced Data
While previous research on training class imbalance mainly concentrated on developing new loss functions and sampling methods tailored to imbalanced data, little attention has been paid to how the components of traditional training routines interact with such data.
In this section, we delve into various methods that can be optimized or modified specifically for imbalanced training scenarios. For clarity and depth of discussion, we divide these methods into two distinct groups, each discussed in its own subsection.
Figure 1: Imbalanced data prefers small batch sizes. We plot the percent improvement in accuracy over the baseline batch size of 128 for different train ratios as a function of batch size. Positive values indicate higher accuracy than the baseline. Balanced training sets yield flatter lines, indicating insensitivity to batch size. Experiments conducted with ResNet-50 on CIFAR-100.
Subsection 3.1 deals with what we identify as the "fundamental building blocks" of conventional balanced training. This includes elements like batch size, data augmentation, pre-training, model architecture, and optimizers. We scrutinize the effects of altering these elements' hyperparameters, emphasizing their influence on the performance of models under imbalanced training. Our discussion is fortified by experi
ments conducted on both vision datasets (CIFAR-10, CIFAR-100, CINIC-10 and Tiny-ImageNet) and tabular datasets across different network architectures.
Subsection 3.2 introduces a set of optimization methods—Joint-SSL, SAM, and label smoothing—that we have reformed to cater more appropriately to imbalanced training. These methods, initially designed for balanced datasets, are subjected to modifications making them more amenable to imbalanced scenarios. The effectiveness of these methods, including comparisons to other state-of-the-art techniques across various domains and datasets (vision and tabular), is thoroughly evaluated in Section 3.1.1.
This structure allows a separate discourse on the fundamental aspects of deep learning training and the particular adjustments needed for imbalanced training. Our aim is to provide a comprehensive overview of both traditional and innovative techniques, giving readers a broad spectrum of strategies for tackling imbalanced data.
3.1 Tuning the Building Blocks of Neural Network Training
The success of deep learning models hinges on the precise orchestration of their training routine. The success of models is especially sensitive to this orchestration under imbalanced training data. We now explore the building blocks of neural network training routines and their importance in optimizing models for imbalanced data. This study equips researchers and practitioners with practical strategies to improve the performance of their models under class imbalance. In this subsection, we focus on batch size, data augmentation, pre-training, model architecture, and optimizers.
3.1.1 Experimental Setup
Datasets. To conduct our investigation, we leverage three benchmark image datasets: CIFAR-10 [46], CIFAR-100, and CINIC-10 [15], along with three tabular datasets: Otto Group Product Classification [37], Covertype [7], and Adult datasets [43]. For naturally balanced datasets, we adopt the imbalanced setup proposed by Liu et al. [54], which employs an exponential distribution to imbalance classes, closely mirroring real-world long-tailed class distributions. The Class-imbalance ratio  $(r)$  represents the ratio of samples in the rarest to the most frequent class. A dataset with  $r = 1$  is fully balanced, while  $r = 0.1$  indicates that the majority class has ten times more samples than the minority class. We investigate varying imbalance ratios in both the training and test sets.
Models. For image datasets, we utilize ResNets [29] and WideResNet [73] with different depths (8, 32, 50, and 152). In addition, to evaluating the role of architecture in imbalanced training, we also employ DenseNet [35], MobileNetV2 [62], Inception v3 [67], EfficientNet [68], and VGG [64]. Tabular datasets are processed using XGBoost [11], SVM, and MLP. We follow the supervised pre-training protocol of Kang et al. [39], while for self-supervised pre-training, we employ SimCLR [12] and VICReg [5], where the fine-tuning is as described in Kotar et al. [44]. Further details can be found in Appendix B.
Metrics. In addition to overall test accuracy, we provide minority and majority class accuracy (representing the  $20\%$  of classes with the smallest and largest number of samples) in the appendix, if not in the main body of the paper. We conduct five runs with different seeds for each evaluation in our experiments and report the mean along with one standard error.
Batch size. Studies conducted on balanced training data suggest that small batch sizes may exhibit superior convergence behavior, while large batch sizes can reach optimal minima unattainable by smaller sizes [41]. In class-imbalanced settings, one intuition is that larger batch sizes may be necessary to obtain enough samples from the minority class and counteract forgetting. On the other hand, large batches can also increase the risk of overfitting.
To examine the influence of batch size in class-imbalanced contexts, we train networks with varying batch sizes across multiple training ratios. In Figure 1, we plot the percent improvement in accuracy over the baseline (which is set as best batch size 128 as a function of batch size for several training ratios. Positive values represent higher accuracy compared to the baseline, whereas negative values denote lower accuracy.
Our analysis reveals that batch size has a much greater impact in highly imbalanced settings and very little impact in balanced settings. Notably, data with a high degree of class imbalance tends to benefit from smaller batch sizes, even though small batches often do not contain minority samples, possibly due to the regularization effects that help mitigate overfitting to the majority classes. See Appendix A.1 for additional details and experiments.
Data augmentation. Data augmentation is a feature of virtually all modern image classification pipelines. We now investigate the impacts of various augmentation policies across varying levels of class imbalance. Our experiments show that the effects of data augmentation are greatly amplified on imbalanced data, especially for minority classes (Figure 2 - left). This finding suggests that augmentation serves as a regularizer, supporting recent studies on the role of data augmentation in preventing overfitting during class-balanced training [20]. Moreover, we find that the optimal augmentation policy can depend on the level of imbalance.
We assess our findings using a variety of augmentation methods including horizontal flips, random crops, AugMix [31], TrivialAugmentWide [58], AutoAugment [13], Mixup, CutOut, and Random Erasing.
To identify the most potent data augmentation strategy, we gauge the improvement in accuracy, represented as the percentage increase compared to training without augmentation, in Appendix A.2. While TrivialAugment outperforms other methods on balanced training data, AutoAugment emerges as the most effective for imbalanced data
Furthermore, we examine the sensitivity of performance to the specific augmentation method used. By assessing the variance in performance across different augmentations for balanced and imbalanced
Figure 2: Left: Augmentations yield far bigger improvements on minority classes. We compare the percent improvement in test accuracy of TrivialAugment compared to training without any augmentation as a function of the training ratio. Right: The type of augmentation matters more on imbalanced data. Augmentation variance measures the variance in percent improvement over training without augmentation across different augmentation types. Variance across augmentation types on imbalanced data is much greater for minority classes, indicating the importance of choosing an appropriate augmentation policy. Experiments conducted with ResNet-50 on CIFAR-100. Error bars represent one standard error over 5 trials.
Figure 3: Left: Deeper architectures overfit on class-imbalanced data. While deeper ResNet models perform better on balanced data, they can overfit and underperform on imbalanced data. Right: Performance on balanced and imbalanced datasets is virtually uncorrelated across a wide variety of architectures (Pearson correlation coefficient 0.14). Experiments conducted on CIFAR-100 with an imbalanced train ratio of 0.001. Error bars represent one standard error over 5 trials.
situations (see Figure 2), we find that minority class performance is particularly sensitive to the chosen augmentation policy. See Appendix A.2 for additional details and experiments.
Model architecture. While larger networks often enhance performance on class-balanced datasets without overfitting, their efficacy on imbalanced data remains unexplored. To probe this behavior, we train ResNets of various sizes on both balanced and imbalanced data with a training ratio of 0.01. We plot the test accuracy on minority classes as a function of the network's size in Figure 3 (left). While the network's performance on balanced data monotonically improves with increasing size, its performance on imbalanced data peaks at a size with 20 million parameters, and declines thereafter. This dip in performance suggests that, unlike their behavior on balanced data, larger networks may be susceptible to overfitting minority classes in the face of severe class imbalance. We then plot the test accuracy of a wide variety of architectures on the balanced and imbalanced CIFAR-100 variants in Figure 3 (right), and find that accuracies in these two settings are virtually uncorrelated! Computer vision architectures have largely been developed on well-known class-balanced benchmarks like ImageNet [16]. These developments may not generalize to imbalanced settings which are common in the real world. See Appendix A.3 for additional details and experiments.
Pre-training. Fine-tuning models pre-trained on expansive upstream datasets can markedly enhance performance across a variety of domains by equipping the model with an informative initial parameter
vector learned from pre-training data [12, 65]. Self-Supervised Learning (SSL) has emerged as a highly effective strategy for representation learning in fields such as computer vision, natural language processing (NLP), and tabular data. Networks pre-trained using SSL often produce more transferable representations than those pre-trained through supervised learning. Furthermore, self-supervised pre-training strategies for transfer learning display greater robustness to upstream imbalance compared to supervised pre-training [51]. We observe in our experiments that pre-trained backbones provide considerably greater benefits under severe downstream class imbalance than under balanced downstream training sets. The ability for SSL pre-training to mitigate overfitting could therefore be particularly valuable in the class-imbalanced setting.
To gauge the effectiveness of pre-training, we fine-tune various pre-trained models on downstream datasets with a range of class imbalance ratios. We make use of pre-trained ResNet-50 weights learned on ImageNet-1K [16], ImageNet-21K, and two self-supervised learning methods, SimCLR and VICReg. Figure 4 depicts the relative improvement in test accuracy compared to training from random initialization across different pre-training models. A positive value signifies a performance improvement. All pre-training methods notably outperform random initialization. However, we observe a considerably larger improvement under class imbalanced scenarios, where models pretrained on larger datasets yield greater boosts in accuracy. Moreover, SSL methods conducted on ImageNet-1K surpass supervised pre-training on the same upstream training set. See Appendix A.4 for additional details and experiments.
3.2 Improved Optimization Methods for Class Imbalance
While the fundamental building blocks above provide a strong foundation for training perform-. mant models, we can also customize modern optimization techniques specifically for imbalanced training and achieve further gains. This subsection showcases adapted variants of SSL, SAM, and label smoothing.
Self-Supervision. Self-Supervised Learning (SSL) has received substantial attention in representation learning, particularly in computer vision, NLP, and tabular data [12, 40, 65, 49], especially for training on massive volumes of unlabeled data. Networks pre-trained via SSL often showcase more transferable representations compared to those pre-trained through supervised methods [24]. Moreover, self-supervised pre-training strategies for transfer learning exhibit more robustness to upstream imbalance compared to supervised pre-training [51]. Despite these advantages, the availability of massive pre-training datasets can often be a limiting factor in many use cases.
Figure 4: Pre-training is more impactful on imbalanced downstream data. % Improvement refers to the improvement in test accuracy compared to training from random initialization. The benefits of SSL over supervised pre-training are also amplified under class imbalance. Experiments conducted with ResNet-50 on CIFAR-100.
Traditionally, pre-training involves a two-step process: learning on an upstream task, followed by fine-tuning on a downstream task. In our approach, we instead integrate supervised learning with an additional self-supervised loss function during from-scratch training, bypassing the need for pre-training. The simple integration of an SSL loss function, which does not depend on class-imbalanced labels and is insensitive to imbalance, results in improved feature representations and better generalization. We refer to this combined training procedure as Joint-SSL. It is important to note that our method differs from those used in Kotar et al. [44], Yang and Xu [71], Liu et al. [51], which investigate SSL pre-training on larger datasets for transfer learning, rather than directly training with an SSL objective on downstream data. See Appendix A.5 for additional details.
Sharpness-Aware Minimization (SAM). SAM [19] is an optimization technique for finding "flat" minima of the loss function, often improving generalization. The technique involves taking an inner ascent step followed by a descent step to find parameters that minimize the increase in loss from the ascent step. Huang et al. [36] shows that flat minima correspond to wide-margin decision boundaries.
We thus adapt SAM for class-imbalanced cases by increasing the flatness especially for minority class loss terms. To do so, we increase the ascent step size in SAM's inner loop for minority classes, denoting this method SAM-Asymmetric (SAM-A). Plotting the decision boundaries of a small multi-layer perceptron on a toy 2D dataset (Appendix Figure 18), we observe that classifiers naturally form small margins surrounding minority class samples. SAM-A expands these margins, preventing the model from overfitting to the minority samples. See Appendix A.6 for additional details.
Label smoothing. Conventionally, classifiers are trained with hard targets, minimizing the cross-entropy between true targets  $y_{k}$  and network outputs  $p_k$  as in  $H(y;p) = \sum_{k=1}^{K} -y_k\log(p_k)$ , with  $y_{k}$  equal to 1 for the correct class and 0 otherwise. Label smoothing uses a smoothing parameter,  $\epsilon$  to instead minimize the cross-entropy between smoothed targets  $y_{LS_k}$  and network outputs  $p_k$ , where  $y_{LS_k} = y_k(1 - \epsilon) + \epsilon / K$  [57].
We adapt label smoothing for the class-imbalanced setting by applying more smoothing to minority-class examples than to majority-class examples. This procedure prevents overfitting on minority samples. See Appendix A.7 for additional details and experiments.