SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Enhancing Minority Classes by Mixing: An Adaptive Optimal Transport Approach for Long-tailed Classification
Jintong Gao $^{1}$ , He Zhao $^{2}$ , Zhuo Li $^{3,4}$ , Dandan Guo $^{1*}$
$^{1}$ School of Artificial Intelligence, Jilin University  $^{2}$ CSIRO's Data61
$^{3}$ Shenzhen Research Institute of Big Data
4The Chinese University of Hong Kong, Shenzhen
gaojt20@mails.jlu.edu.cn
221019088@link.cuhk.edu.cn
guodandan@jlu.edu.cn
Real-world data usually confronts severe class-imbalance problems, where several majority classes have a significantly larger presence in the training set than minority classes. One effective solution is using mixup-based methods to generate synthetic samples to enhance the presence of minority classes. Previous approaches mix the background images from the majority classes and foreground images from the minority classes in a random manner, which ignores the sample-level semantic similarity, possibly resulting in less reasonable or less useful images. In this work, we propose an adaptive image-mixing method based on optimal transport (OT) to incorporate both class-level and sample-level information, which is able to generate semantically reasonable and meaningful mixed images for minority classes. Due to its flexibility, our method can be combined with existing long-tailed classification methods to enhance their performance and it can also serve as a general data augmentation method for balanced datasets. Extensive experiments indicate that our method achieves effective performance for long-tailed classification tasks. The code is available at https://github.com/JintongGao/Enhancing-Minority-Classes-by-Mixing.
Large-scale balanced datasets play a vital role in the remarkable success of deep neural networks for various tasks. However, datasets in many real-world applications often exhibit unexpected long-tailed distributions where most of the data belongs to several majority classes while the rest spreads across lots of minority classes. The model trained on such a long-tailed dataset will be biased toward majority classes, leading to poor generalizations about minority classes.
Re-weighting [1, 2, 3, 4, 5, 6, 7], over-sampling [8, 9, 10, 11, 12], under-sampling [13, 14, 15, 16], data augmentation [17, 18, 19, 20, 21, 22, 23], two-stage methods [24, 25, 26, 27], and other methods [28, 29] are common solutions to the long-tailed problem. Among them, over-sampling and data augmentation aim to balance the data distribution by oversampling or generating closely related minority classes. As the representative data augmentation techniques, mixup and its variants [30, 31, 32] have performed satisfactorily in the computer vision fields, whose key idea is constructing mixed samples by performing linear interpolations between data/features and corresponding labels. Considering mixup-based methods are designed for balanced data, applying them to long-tailed classification without any adjustments may ignore the specificity of long-tailed data distribution. Recently, some novel mixup-based methods for the long-tailed problem have been explored [22, 33,
*Corresponding author. This work was supported by NSFC (62306125).
Figure 1: Visualization of mixed samples created by CMO and ours(OTmix), respectively.
34, 21]. As a representative example, Context-rich Minority Oversampling (CMO) [20] generates minority-centric images by mixing the background and foreground images, where background images are mainly from the majority classes while foregrounds are mainly from minority classes. The critical concept of CMO is to paste an image from a minority class (foreground) onto rich-context images from a majority class (background). Despite the effectiveness of CMO, it cannot be guaranteed that all the generated images are realistic and reasonable by randomly mixing the background and foreground images. For example, we visualize mixed samples created by ours and CMO in Fig. 1, where we take "Panda" and "Grey Whale" as two foreground images. CMO generates unreasonable mixed samples since pandas are hardly seen in the sky, the beaches, or the streets, and grey whales rarely appear in the desert, the flowers, or the sky. Therefore, instead of arbitrarily pairing the foreground and background images, it is essential to consider the semantic distances between the foreground and background images to generate more meaningful and advantageous samples.
In this work, we propose a more adaptive image-mixing method based on optimal transport (OT), which is a powerful tool for computing the distance between two distributions under the form of multiple sampling [35]. We consider an empirical distribution  $P_{\mathrm{back}}$  over the background images and an empirical distribution  $Q_{\mathrm{fore}}$  over the foreground images, where the former are more likely from the majority and the latter from minority classes. To learn the semantic similarity between background and foreground images, we formulate the learning of the similarity as the OT problem between  $P_{\mathrm{back}}$  and  $Q_{\mathrm{fore}}$ . In this way, we can view the learned transport plan as the similarity between foreground and background samples, which can be naturally used as guidance to select the most relevant background image for each foreground image. Due to the importance of cost function in learning the transport plan for OT, we further design the cost function using the sample-level representation and class-level label information. Hence, the rich context of the majority classes as background images can be better transferred to the minority classes as foreground images for providing an adaptive way to learn the similarity between them and generate more semantically meaningful mixed images. Interestingly, our proposed OTmix can not only be combined with existing long-tailed methods but also be used as a new image-mixing strategy for balanced datasets.
Our main contributions are summarized as follows: (1) We propose a novel image-mixing method for generating reasonably mixed samples for long-tailed classification, where we introduce a distribution over the background images mainly from majority classes and another distribution over the foreground images mainly from minority classes. (2) By minimizing the OT distance between these two distributions, we view the learned transport plan as similar in guiding the image-mixing process. (3) We design the cost function based on the feature information and confusion matrix to learn the desired transport plan. Extensive experiments demonstrate the effectiveness of our method for long-tailed classification and also balanced classification.
Data Processing Learning relatively balanced classes from the data perspective is an effective solution to the long-tailed problem [36], which can be roughly divided into over-sampling, undersampling, and data augmentation. Since our proposed method is related to the over-sampling and data augmentation methods, we describe them in detail below. Over-sampling aims to emphasize the minority classes and increase the instance number of the minority classes [20, 8, 9]. For example, the classical Synthetic Minority Over-sampling Technique (SMOTE) [12] uses the interpolation between a given minority sample and its nearest minority neighbors to create new samples. Data augmentation is another way of data processing to compensate the minority classes by generating and synthesizing new samples [18, 23]. For example, Major-to-minor Translation (M2m) [17] defines an optimization phase to augment minority classes via translating samples from majority classes. The Meta Semantic Augmentation (MetaSAug) approach [24] is proposed to perform effective
semantic data augmentation by learning more meaningful class-wise covariance. Sample-Adaptive Feature Augmentation (SAFA) [19] aims to extract diverse and transferable semantic directions from majority classes, and adaptively translate minority-class features along extracted semantic directions for augmentation.
In addition, mixup-based methods [32, 37, 38, 30, 39, 31, 40, 41, 42], the effective data augmentation methods in balanced datasets, have been developed to solve the long-tailed problem recently and are closely related to our work. For example, Mixup Shifted Label-aware Smoothing Model (MiSLAS) [27] uses mixup in its Stage-1 training without any adjustments. Rebalanced Mixup (Remix) [21] provides a disproportionately higher weight to the minority class when mixing two samples. Uniform Mixup (UniMix) [22] adopts a minority-favored mixing factor to encourage more majority-minority pairs occurrence and is further combined with the Bayes Bias (Bayias) caused by the inconsistency of prior. Targeted copy-paste augmentation [43] aligned with the novel domain enhances out-of-domain robustness on long-tailed camera trap dataset. [44] improves minority class performance from synthetic data by contrasting night and day backgrounds. To mitigate the overfitting issue in subpopulation shift, Uncertainty-aware Mixup (UMIX) [33] adopts the sampling uncertainty to reweight the mixed samples. CMO [20] constructs mixed samples between the background images more likely from the majority classes and foreground images more likely from the minority classes. However, it ignores the semantic similarity between images by randomly mixing them. Different from CMO, we formulate the image-mixing problem as the OT problem between two distributions and use the learned transport plan as guidance to mix images, providing a general and adaptive image-mixing method for long-tailed learning.
Loss Function Engineering Designing effective training objectives is another solution to fight against class imbalance [45, 7, 5]. Label-Distribution-Aware Margin Loss (LDAM) [1] is proposed from the view of the generalization error bound, and LDAM-DRW adopts a deferred class-level re-weighting method. Balanced Softmax (BALMS) [2] accommodates the label distribution shift between training and testing and proposes a Meta Sampler that learns to re-sample training set by meta-learning. Label Distribution Disentangling (LAD) [4] loss disentangles the source label distribution from the model prediction based on the optimal bound of the Donsker-Varadhan representation. Besides, Focal Loss [46] determines the weights for samples with the sample difficulty in the object detection task, and Influence-balanced Loss (IB) [3] assigns different weights to samples according to their influence on a decision boundary. Our method can be naturally combined with many loss functions.
Other Methods Recently, two-stage algorithms have been proposed [1, 5, 24], such as MiSLAS [27]. Meanwhile, a Bilateral Branch Network (BBN) [25] unifies the representation and classifier learning stages to form a cumulative learning strategy. RoutIng Diverse Experts (RIDE) [28] uses multiple experts to reduce the variance and bias of the long-tailed classifier. In addition, some approaches employ meta-learning, such as MetaSAug [24] and BALMS [2]. Causal Norm (CN) [29] pinpoints the causal effect of momentum and extracts the unbiased direct impact of each instance.
Optimal Transport As a powerful tool, OT has been applied to generative models [47, 48, 49, 50, 51], computer vision [52, 53, 7, 54, 55, 56, 57], text analysis [58, 59, 60, 61], and etc. To the best of our knowledge, there are still very limited works for the imbalanced classification by means of OT. For example, [7] proposes a re-weighting method based on OT, and [62] introduces Optimal Transport via Linear Mapping (OTLM) to perform the post-hoc correction. Different from them, ours falls into the data augmentation group by mixing images. Another recent work is Optimal Transport for OverSampling (OTOS) [63], which moves random points from a prior uniform distribution into that of minority class samples based on OT. However, our proposed OTmix introduces OT to guide the mix between majority and minority images, which leverages the rich context of the majority classes guided by OT. Alignmixup [38] adopts OT to align two images and interpolate between two sets of features in a standard classification task, which is distinct from ours in terms of task and technical detail.
Long-tailed Classification Given a training set  $\mathcal{D} = \{(x_i, y_i)\}_{i=1}^N$  for a multi-class problem with  $K$  classes, if each class  $k$  contains  $n_k$  samples, we have that  $\sum_{k=1}^{K} n_k = N$ . Without loss of generality, we can always assume  $n_1 \geq n_2 \geq \dots \geq n_K$  for the long-tailed problem. Denote
Figure 2: Adaptive image-mixing method with optimal transport.
the model parameterized with  $\theta$  as  $g(\theta)$ , which is trained on  $\mathcal{D}$  with the well-known Empirical Risk Minimization (ERM) algorithm [32] ignores such class imbalance and performs poorly on the minority classes [33].
CutMix Here, we use the CutMix [37] to mix samples because of its simplicity and effectiveness, which replaces the image region with a patch from another training image. Denote  $x \in \mathbb{R}^{W \times H \times C}$  and  $y$  denote a training image and its label, respectively. CutMix combines two training samples  $(x_i, y_i)$  and  $(x_j, y_j)$  and generates the new sample  $(\tilde{x}_{ij}, \tilde{y}_{ij})$  as follows:
$$
\tilde {x} _ {i j} = \mathbf {M} \odot x _ {i} + (\mathbf {1} - \mathbf {M}) \odot x _ {j}, \quad \tilde {y} _ {i j} = \lambda y _ {i} + (1 - \lambda) y _ {j}, \tag {1}
$$
where the combination ratio  $\lambda$  is sampled from the beta distribution  $\mathrm{Beta}(\alpha, \alpha)$ ,  $\odot$  is element-wise multiplication, and the binary mask  $\mathbf{M} \in \{0,1\}^{W \times H}$  indicates where to drop out and fill in from two images,  $\mathbf{1}$  is a binary mask filled with ones. We provide more details about  $\mathbf{M}$  in Appendix 7.1.
Optimal Transport OT measures the minimal cost to transport between two probability distributions [35, 64, 65, 66, 67]. We only provide a brief introduction to OT for discrete distributions and refer the readers to [35] for more details. Denote two discrete probability distributions  $p = \sum_{i=1}^{n} a_i \delta_{x_i}$  and  $q = \sum_{j=1}^{m} b_j \delta_{y_j}$ , where both  $\mathbf{a}$  and  $\mathbf{b}$  are discrete probability vectors summing to 1,  $x_i$  and  $y_j$  are the supports of the two distributions respectively, and  $\delta$  is a Dirac function. Then the OT distance is formulated as follows:  $\mathrm{OT}(p,q) = \min_{\mathbf{T} \in \Pi(p,q)} \langle \mathbf{T}, \mathbf{C} \rangle$ , where  $\mathbf{C} \in \mathbb{R}_{\geq 0}^{n \times m}$  is the cost matrix with element  $C_{ij} = C(x_i, y_j)$  which reflects the cost between  $x_i$  and  $y_j$  and the transport probability matrix  $\mathbf{T} \in \mathbb{R}_{\geq 0}^{n \times m}$  is subject to  $\Pi(p,q) := \left\{\mathbf{T} \mid \sum_{i=1}^{n} T_{ij} = b_j, \sum_{j=1}^{m} T_{ij} = a_i\right\}$ . The optimization problem above is often adapted to include a popular entropic regularization term  $H = -\sum_{ij} T_{ij} \ln T_{ij}$  for reducing the computational cost, denoted as Sinkhorn algorithm [68].
In this work, we propose a novel adaptive image-mixing data augmentation method based on OT for long-tailed classification, where the overall framework is shown in Fig. 2. We consider a distribution  $P_{\mathrm{back}}$  over the background images mainly from majority classes and another distribution  $Q_{\mathrm{fore}}$  over the foreground images mainly from minority classes. By minimizing the OT distance between these two distributions, we design the sample-level and class-level cost functions and use the learned transport plan as guidance for the image-mixing process to generate more semantically meaningful mixed samples between the two distributions.
Main Objective Since majority classes usually have sufficient data and rich information, it is natural to leverage the majority classes to enhance minority classes. In this work, we aim to generate minority-centric images with majority contexts, combining the background and foreground images based on CutMix. The critical question is how to pair a background image and a foreground image. The first factor of the pairing strategy is that we would expect a background image to be more likely from the majority classes and a foreground image to be more likely from the minority classes, which has been modeled in CMO [20]. However, in addition to the first one, the second factor is the semantic similarity between the background and foreground images, which has not been studied in the literature and is the main focus of our paper. Specifically, suppose a set of candidate images
already satisfy the first factor, i.e., background and foreground images are from majority and minority classes, respectively. Pairing these images arbitrarily may not guarantee that the combined images are semantically meaningful and helpful, e.g., in Fig. 1.
In this paper, we aim to generate semantically meaningful images by proposing a pairing process that satisfies the first and second factors at the same time. Our proposal is derived from a distribution-matching perspective. First, we consider the distribution  $P_{\mathrm{back}}$  of the background images as the empirical distribution of all the images in the training dataset, which is a (discrete) uniform distribution:
$$
P _ {\text {b a c k}} = \sum_ {i = 1} ^ {N} \frac {1}{N} \delta \left(x _ {i}, y _ {i}\right). \tag {2}
$$
Although  $P_{\mathrm{back}}$  is uniformly distributed, because the majority classes have significantly more images than the minority classes, it is natural that the samples from  $P_{\mathrm{back}}$  are more likely to be images from the majority classes, which is what we expect. Secondly, we introduce the distribution over the foreground images, which is a discrete distribution over all the training images:
$$
Q _ {\text {f o r e}} = \sum_ {j = 1} ^ {N} w _ {j} \delta \left(x _ {j}, y _ {j}\right), \tag {3}
$$
where  $w_{j}$  is the weight of the  $j$ -th sample. In this case, we expect that a sample from  $Q_{\mathrm{fore}}$  is more likely to be an image of the minority classes. Thus,  $Q_{\mathrm{fore}}$  cannot be a uniform distribution. Accordingly, we can specify  $w_{j}$  in several ways, such as the (smoothed) inverse class frequency [69, 70, 71], or the adequate number of samples [72]. Taking the inverse class frequency as the example, we can sample the  $k$ -th class with the following probability:  $q_{k} = \frac{1 / n_{k}^{r}}{\sum_{k=1}^{K} 1 / n_{k}^{r}}$ , where  $r > 0$  controls the smoothness of  $q_{k}$ ,  $r = 1$  indicates the inverse class frequency and  $r = 1 / 2$  means the smoothed version, i.e., if  $r$  increases, the weight of the minority class becomes increasingly larger than that of the majority class. Finally, we can use  $w_{j} = q_{k} \frac{1}{n_{k}}$  if  $y_{j} = k$ .
For now, with our construction of  $P_{\mathrm{back}}$  and  $Q_{\mathrm{fore}}$ , the first factor is satisfied. For the second factor that requires the semantic similarity between a pair of background and foreground images, we formulate it as an optimization problem of the entropic OT:
$$
\mathrm {O T} _ {\epsilon} \left(P _ {\text {b a c k}}, Q _ {\text {f o r e}}\right) \stackrel {\text {d e f}} {=} \min  _ {T _ {i j} \in \Pi \left(P _ {\text {b a c k}}, Q _ {\text {f o r e}}\right)} \sum_ {i, j} ^ {N, N} C _ {i j} T _ {i j} - \epsilon \left[ \sum_ {i, j} ^ {N, N} - T _ {i j} \ln T _ {i j} \right], \tag {4}
$$
where  $\epsilon > 0$  is the hyper-parameter for the entropic constraint,  $C_{ij}$  is the transport cost function, and the transport probability  $T_{ij}$  satisfies  $\Pi(P_{\mathrm{back}}, Q_{\mathrm{fore}}) \coloneqq \left\{ T_{ij} \mid \sum_{i=1}^{N} T_{ij} = w_j, \sum_{j=1}^{N} T_{ij} = \frac{1}{N} \right\}$ . Notably, as an upper-bounded positive metric,  $T_{ij}$  indicates the transport probability between the  $i$ -th background image and the  $j$ -th foreground image, which can be naturally used to measure the importance of each background image for the foreground image when performing image-mixing.
Cost Function  $C_{ij}$  measures the distance between  $(x_i, y_i)$  in  $P_{\mathrm{back}}$  and  $(x_j, y_j)$  in  $Q_{\mathrm{fore}}$ . As the main parameter for defining the transport distance between probability distributions,  $C_{ij}$  plays an important role in learning the optimal transport plan, which can be flexibly defined in different ways. For clarity, we can reformulate the concerned model  $g(\theta)$  as  $g_2(g_1(\theta_1); \theta_2)$ , where  $g_1(\theta_1)$  denotes the feature extractor parameterized with  $\theta_1$  and  $g_2(\theta_2)$  is the classifier parameterized with  $\theta_2$ . We explore a few conceptually intuitive options of  $C_{ij}$ . A simple but straightforward way is defining the sample-level  $C_{ij}$  with the features of  $x_i$  and  $x_j$ :
$$
C _ {i j} = 1 - \cos \left(z _ {i}, z _ {j}\right), \tag {5}
$$
where  $\cos (\cdot ,\cdot)$  is the cosine similarity,  $z_{i} = g_{1}(x_{i};\theta_{1})\in \mathbb{R}^{e}$  and  $z_{j} = g_{1}(x_{j};\theta_{1})\in \mathbb{R}^{e}$  denote the  $e$  -dimensional representation of  $x_{i}$  and  $x_{j}$ , respectively. Now  $C_{ij}$  will be small if two samples have similar or close features, where the cost is influenced by the feature extractor  $\theta_{1}$ . Besides, we can also define the class-level  $C_{ij}$  only using the label information. Specifically, a normalized confusion matrix  $\mathbf{F}\in \mathbb{R}_{>0}^{K*K}$  can be estimated with an unbiased validation set or a small balanced subset sampled from  $\bar{\mathcal{D}}$ . Summarizing the prediction results of the target model  $g$ , the confusion matrix  $\mathbf{F}$  shows how the model  $g$  is confused when it makes predictions. Since we focus on combining the foreground images mainly from minority classes and background images mainly from minority classes, we can set the diagonal element in  $\mathbf{F}$  to 0, denoted as  $\hat{\mathbf{F}}$ , to avoid mixing the images from the same class. After that, we can define  $C_{ij}$  with the ground-truth labels of two samples:
$$
C _ {i j} = 1 - \hat {F} _ {y _ {i}, y _ {j}}, \tag {6}
$$
where  $\hat{F}_{y_i,y_j}$  indicates the element of row  $y_{i}$ , column  $y_{j}$ , of the matrix  $\hat{\mathbf{F}}$ . If  $y_{i}$  and  $y_{j}$  are from the same class,  $\hat{F}_{y_i,y_j} = 0$  and  $C_{ij} = 1$ ; if  $y_{i}$  and  $y_{j}$  are not from the same class but are easy get to confused by the model  $g$ ,  $\hat{F}_{y_i,y_j}$  would be large and  $C_{ij}$  would be small. Now,  $C_{ij}$  is affected by the feature extractor  $\theta_{1}$  and classifier  $\theta_{2}$ . Intuitively, based on the features and labels of samples, the cost function can also be defined as:
$$
C _ {i j} = \omega \left(1 - \cos \left(z _ {i}, z _ {j}\right)\right) + (1 - \omega) \left(1 - \hat {F} _ {y _ {i}, y _ {j}}\right), \tag {7}
$$
where  $\omega$  is a hyper-parameter for balancing the feature and label information and  $C_{ij}$  will be small if two samples have similar features and are from two confused classes but  $y_i \neq y_j$ . Once the cost function  $C_{ij}$  is defined, it can be fed into (4) for learning the transport plan  $T_{ij}$  between foreground images and background images, where a small  $C_{ij}$  tends to produce a large  $T_{ij}$ .
Image-mixing Process By minimizing the OT problem, the resultant transport plan  $\mathbf{T}$  provides an adaptive way to weigh the similarity between background and foreground images. Therefore, we can select the most suitable background image  $x_{j'}$  for the foreground image  $x_i$  as follows:
$$
\underset {j ^ {\prime} \in N} {\arg \max } j ^ {\prime} = \{i \in N: j ^ {\prime} = \underset {j \in N} {\max } T _ {i j} \}. \tag {8}
$$
We aim to generate new mixed samples by combining the foreground image  $x_{i}$  and the most relevant background image  $x_{j'}$  and their corresponding labels. As their name implies,  $x_{j'}$  is used as a background image while  $x_{i}$  provides the foreground patch, which can be pasted onto  $x_{j'}$ . Recalling CutMix in (1), the detailed mixing process is expressed as: