Distribution Alignment: A Unified Framework for Long-tail Visual Recognition
Songyang Zhang, Zeming Li, Shipeng Yan, Xuming He, Jian Sun
Introduction
While deep convolutional networks have achieved great successes in many vision tasks, it usually requires a large number of training examples for each visual category. More importantly, prior research mostly focuses on learning from a balanced dataset , where different object classes are approximately evenly distributed. However, for large-scale vision recognition tasks, partially due to the non-uniform distribution of natural object classes and varying annotation costs, we typically learn from datasets with a long-tail class label distribution. In such scenarios, the number of training instances per class varies significantly, from as few as one example for tail classes to hundreds or thousands for head classes .
The intrinsic long-tail property of our visual data introduces a multitude of challenges for recognition in the wild , as a deep network model has to simultaneously cope with imbalanced annotations among the head and medium-sized classes, and few-shot learning in the tail classes. A naively learned model would be largely dominated by those few head classes while its performance is much degraded for many other tail classes.
Early works on re-balancing data distribution focus on learning one-stage models, which achieve limited successes due to lack of principled design in their strategies . More recent efforts aim to improve the long-tail prediction by decoupling the representation learning and classifier head learning . However, such a two-stage strategy typically relies on heuristic design to adjust the decision boundary of the initially learned classifier head, which often requires tedious hyper-parameter tuning in practice. This severely limits its capacity to resolve the mismatch between imbalanced training data distribution and balanced evaluation metrics.
In this work, we first perform an ablative analysis on the two-stage learning strategy to shed light on its performance bottleneck. Specifically, our study estimates an ‘ideal’ classification accuracy using a balanced dataset to retrain the classifier head while keeping the first-stage representation fixed. Interestingly, as shown in Fig. 6, we find a substantial gap between this ideal performance and the baseline network, which indicates that the first-stage learning with unbalanced data provides a good representation, but there is a large room for improvement in the second stage due to the biased decision boundary (See Sec. 3.1 for details).
Based on those findings, we propose a simple and yet effective two-stage learning scheme for long-tail visual recognition problems. Our approach focuses on improving the second-stage training of the classifier after learning a feature representation in a standard manner. To this end, we develop a unified distribution alignment strategy to calibrate the classifier output via matching it to a reference distribution of classes that favors the balanced prediction. Such an alignment strategy enables us to exploit the class prior and data input in a principled manner for learning class decision boundary, which eliminates the needs for tedious hyper-parameter tuning and can be easily applied to various visual recognition tasks.
Specifically, we develop a light-weight distribution alignment module for calibrating classification scores, which consists of two main components. In the first component, we introduce an adaptive calibration function that equips the class scores with an input-dependent, learnable magnitude and margin. This allows us to achieve a flexible and confidence-aware distribution alignment for each data point. Our second component explicitly incorporates a balanced class prior by employing a generalized re-weight design for the reference class distribution, which provides a unified strategy to cope with diverse scenarios of label imbalance in different visual recognition tasks.
We extensively validate our model on four typical visual recognition tasks, including image classification on three benchmarks (ImageNet-LT , iNaturalist and Places365-LT ), semantic segmentation on ADE20k dataset , object detection and instance segmentation on LVIS dataset . The empirical results and ablative study show our method consistently outperforms the state-of-the-art approaches on all the benchmarks. To summarize, the main contributions of our works are three-folds:
We conduct an empirical study to investigate the performance bottleneck of long-tail recognition and reveal a critical gap caused by biased decision boundary.
We develop a simple and effective distribution alignment strategy with a generalized re-weight method, which can be easily optimized for various long-tail recognition tasks without whistles and bells.
Our models outperform previous work with a large margin and achieve state-of-the-art performance on long-tail image classification, semantic segmentation, object detection, and instance segmentation.
Related Works
To alleviate the adverse effect of the long-tail class distribution in visual recognition, prior work have extensively studied the one-stage methods, which either leverage the re-balancing ideas or explore knowledge transfer from head categories. The basic idea of resample-based methods is to over-sample the minority categories or to under-sample the frequent categories in the training process . Class-aware sampling proposes to choose samples of each category with equal probabilities, which is widely used in vision tasks . Repeat factor sampling is a smoothed sampling method conducting repeated sampling for tail categories, which demonstrates its efficacy in instance segmentation . In addition, proposes to increase the sampling rate for categories with low performance after each training epoch and balances the feature learning for under-privileged categories.
An alternative strategy is to re-weight the loss function in training. Class-level methods typically re-weight the standard loss with category-specific coefficients correlated with the sample distributions . Sample-level methods try to introduce a more fine-grained control of loss for imbalanced learning. Other work aim to enhance the representation or classifier head of tail categories by transferring knowledge from the head classes . Nevertheless, these methods require designing a task specific network module or structure, which is usually non-trivial to be generalized to different vision tasks.
More recent efforts aims to improve the long-tail prediction by decoupling the learning of representation and classifier head. Decouple proposes an instance-balanced sampling scheme, which generates more generalizable representations and achieves strong performance after properly re-balancing the classifier heads. The similar idea is adopted in , which develop effective strategies for long-tail object detection tasks. improve the two-stage ideas by introducing a post-process to adjust the prediction score. However, such a two-stage strategy typically relies on heuristic design in order to adjust the decision boundary of initially learned classifiers and requires tedious hyper-parameter tuning in practice.
Visual recognition community has witnessed significant progress with deep convolutional networks in recent years. In this study, we focus on four types of visual tasks, including image classification, object detection, semantic and instance segmentation, which have been actively studied in a large amount of prior work. For object detection, we consider the typical deep network architecture used in the R-CNN series method , which detects objects based on the region proposals. For instance segmentation, we take the Mask R-CNN as our example, which extends the Faster R-CNN by adding a branch for predicting the object masks in parallel with the existing branch for bounding box recognition. For the pixel-wise task, semantic segmentation, we use the FCN-based methods and the widely-adopted encoder-decoder structures . Despite those specific choices, we note that our strategy can be easily extended to other types of deep network methods for those visual recognition tasks.
Our Approach
Our goal is to address the problem of large-scale long-tail visual recognition, which typically has a large number of classes and severe class imbalance in its training data. To this end, we adopt a two-stage learning framework that first learns a feature representation and a classifier head from the unbalanced data, followed by a calibration stage that adjusts the classification scores. Inspired by our ablative study on existing two-stage methods, we propose a principled calibration method that aligns the model prediction with a reference class distribution favoring the balanced evaluation metrics. Our distribution alignment strategy is simple and yet effective, enabling us to tackle different types of large-scale long-tail visual recognition tasks in a unified framework.
Below we start with a brief introduction to the long-tail classification and an empirical study of two-stage methods in Sec.3.1. We then describe our proposed distribution alignment strategy in Sec.3.2. Finally, we present a comparison with previous methods from the distribution match perspective in Sec.3.3.
We now introduce the problem setting of long-tail classification and review the two-stage learning framework for deep networks. Subsequently, we perform an empirical ablative study on a large-scale image classification task, which motivates our proposed approach.
The task of long-tail recognition aims to learn a classification model from a training dataset with long-tail class distribution. Formally, we denote the input as , and the target label space as , where is the number of classes. The classification model defines a mapping from the input to the label space: , where and are its parameters. Our goal is to learn the model parameter from an imbalanced training dataset so that achieves optimal performance on an evaluation dataset with respect to certain balanced metrics (e.g., mean accuracy).
In the two-stage framework, we typically consider a deep network model with two main components: a feature extractor network and a classifier head . The feature extractor first extracts an input representation , which is then fed into the classifier head to compute class prediction scores as follows:
where and are the parameter of and , respectively. Here indicate the class prediction scores for classes and the model predicts the class label by taking .
In this work, we instantiate the classifier head as a linear classifier or a cosine similarity classifier as follows:
The two-stage learning method tackles the long-tail classification by decoupling the representation and the classifier head learning . Specifically, it first learns the feature extractor and classifier head jointly, and then with the representation fixed, re-learns the classifier head with a class balancing strategy. While such design achieves certain success, an interesting question to ask is which model component(s) impose a bottleneck on its balanced performance. In the following, we attempt to address the question by exploiting the full set of the ImageNet dataset. Particularly, we follow the decoupling idea to conduct a series of ablative studies on two model components under an ‘ideal’ balanced setting.
We first investigate whether the feature representation learned on the imbalanced dataset is restrictive for the balanced performance. To this end, we start from learning the feature extractor on the imbalanced ImageNet-LT training set with several re-balancing strategies (e.g. instance-balanced, class-balanced, or square-root sampling). We then keep the representation fixed and re-train the classifier head with the ideal balanced ImageNet train set (excluding ImageNet-LT val set). Our results are shown in the left panel of Fig. 2, which indicate that the first stage produces a strong feature representation that can potentially lead to large performance gain and the instance-based sampling achieves better overall results (cf. ).
Moreover, we conduct an empirical study on the effectiveness of the recent decoupling method (e.g. cRT ) compared with the above ’ideal’ classifier head learning. The right panel of Fig. 2 shows that there remains a large performance gap between the existing methods and the upper-bound. Those empirical results indicate that the biased decision boundary in the feature space seems to be the performance bottleneck of the existing long-tail methods. Consequently, a better strategy to address this problem would further improve the two-stage learning for the long-tail classification.
2 Distribution Alignment
To tackle the aforementioned issue, we now introduce a unified distribution alignment strategy to calibrate the classifier output via matching it to a reference distribution of classes that favors the balanced prediction. In this work, we adopt a two-stage learning scheme for all visual recognition tasks, which consists of a joint learning stage and a distribution calibration stage as follows.
1) Joint Learning Stage. The feature extractor and original classifier head (denoted as for clarity) are jointly learned on imbalanced with instance-balanced strategy in the first stage, where the original is severely biased due to the imbalanced data distribution.
2) Distribution Calibration Stage. For the second stage, the parameters of are frozen and we only focus on the classifier head to adjust the decision boundary. To this end, we introduce an adaptive calibration function (in Sec. 3.2.1) and a distribution alignment strategy with generalized re-weighting (in Sec. 3.2.2) to calibrate the class scores.
To learn the classifier head in the second stage, we propose an adaptive calibration strategy that fuses the original classifier head (parameters of are frozen) and a learned class prior in an input-dependent manner. As shown below, unlike previous work (e.g. cRT), our design does not require a re-training of the classifier head from scratch and has much fewer free parameters. This enables us to reduce the adverse impact from the limited training data of the tail categories. Moreover, we introduce a flexible fusion mechanism capable of controlling the magnitude of calibration based on input features.
Specifically, denote the class scores from as , we first introduce a class-specific linear transform to adjust the score as follows:
where and are the calibration parameters for each class, which will be learned from data. As mentioned above, we then define a confidence score function to adaptively combine the original and the transformed class scores:
where the confidence score has a form of , which is implemented as a linear layer followed by a non-linear activation function (e.g., sigmoid function) for all input . The confidence controls how much calibration is needed for a specific input . Given the calibrated class scores, we finally define a prediction distribution for our model with the Softmax function:
2.2 Alignment with Generalized Re-weighting
Given a train dataset , we introduce a calibration strategy based on distribution alignment between our model prediction and a reference distribution of classes that favors the balanced prediction.
Formally, denote the reference distribution as , we aim to minimize the expected KL-divergence between and the model prediction as follows:
where the expectation is approximated by an empirical average on and is a constant.
In this work, we adopt a re-weighting approach and introduce a generalized re-weight strategy for the alignment in order to exploit the class prior. Formally, we represent the reference distribution as a weighted empirical distribution on the training set,
where is the class weight, and is the Kronecker delta function(equals 1 if , otherwise equals 0). We then define the reference weight based on the empirical class frequencies on the training set:
where is a scale hyper-parameter to provide more flexibility in encoding class prior. Note that our scheme reduces to the instance-balance re-weight method with , and to the class-balanced re-weight method with . We illustrate the curve of re-weight coefficients based on ImageNet-LT dataset in Fig. 3.
3 Connection with Recent Work
Below we discuss the connections between our proposed distribution alignment strategy and recent two-stage methods. Detailed comparison is reported in Tab. 2. Notably, Logit Adjustment and Deconfound introduce a hand-craft margin to adjust the distribution while keep the magnitude as 1.0, and incorporate the class prior directly in or without re-training. LWS and -normalized try to achieve a similar goal by learning a magnitude scale and discarding the margin adjustment.
All these methods can be considered as the special cases of our DisAlign approach, which provides a unified and simple form to model the distribution mismatch in a learnable way. Moreover, the resample based strategy is not easy to be applied for the instance-level (object detection/instance segmentation) or pixel-level (semantic segmentation) tasks, our generalized re-weight provides an alternative solution to incorporate the class prior in a simple and effective manner. Experimental results in Sec. 4 also demonstrate the strength of our method compared with the aforementioned works.
Experiments
In this section, we conduct a series of experiments to validate the effectiveness of our method. Below we present our experimental analysis and ablation study on the image classification task in Sec. 4.1, followed by our results on semantic segmentation task in Sec. 4.2. In addition, we further evaluate our methods on object detection and instance segmentation tasks in Sec. 4.3.
To demonstrate our methods, we conduct experiments on three large-scale long-tail datasets, including ImageNet-LT , iNaturalist 2018 , and Places-LT . We follow the experimental setting and implementation of Detailed configuration and results are provided in the supplementary materials.. For the ImageNet-LT dataset, we report performance with ResNet/ResNeXt-{50,101,152} as backbone, and mainly use ResNet-50 for ablation study. For iNaturalist 2018 and Places-LT, our comparisons are performed under the settings of ResNet-{50,101,152}.
1) ImageNet-LT. We present the quantitative results for ImageNet-LT in Tab. 10. Our approach achieves 52.9% in per-class average accuracy based on ResNet-50 backbone and 53.4% based on ResNeXt-50, which outperform the state-of-the-art methods by a significant margin of 2.5% and 1.6%, respectively.
2) iNaturalist. In Tab. 13, our method DisAlign with cosine classifier achieves 69.5% per-class average accuracy using ResNet-50 backbone and 90 epochs of training, surpassing the prior art LDAM by a large margin at 1.5%. It also shows that our performance can be further improved with larger backbone and/or more training epochs. 3) Places-LT. In Tab. 14, we show the experimental results under the same setting as on Places-LT. Our method achieves 39.3% per-class average accuracy based on ResNet-152, with a notable performance gain at 1.4% over the prior methods. We also report the detailed performance of these three datasets with ResNet-{50,101,152} in the supplementary materials.
1) Different Backbone: We validate our method on different types of backbone networks, ranging from ResNet-{50,101,152} to ResNeXt-{50, 101, 152}, reported in Fig. 4. Our method achieves 54.9% with ResNet-152, and 55.0% with ResNeXt-152. It’s worth noting that even when adopting stronger backbones, the gain of DisAlign compared to the state-of-the-art methods is still significant. This demonstrates that our DisAlign is complementary to the capacity of backbone networks.
2) Model Components: We conduct a series of ablation studies to evaluate the importance of each component used in our DisAlign method. Tab. 5 summarizes the results of our ablation experiments, in which we compare our full model with several partial model settings. From the table, we find the learnable magnitude has a significant improvement compared with baseline and the learnable margin also achieves competitive results at 49.9%, which demonstrate the effectiveness of individual modules in our design. 3) Generalized Re-weight Scale We also investigate the influence of the generalized re-weight scale on the validation set of ImageNet-LT and plot the accuracy-scale curve in Fig. 5. It is evident that adjusting generalized reweight is able to achieve significant performance improvement. Moreover, we find the setting of is able to outperform the class-balanced re-weight (), which indicates that the generalized re-weight is more effective in coping with long-tail distributions.
2 Semantic Semgnetaion on ADE20k Dataset
To further validate our method, we apply DisAlign strategy to segmentation networks and report our performance on the semantic segmentation benchmark, ADE20k .
Follow a similar protocol as in image classification, we divide the 150 categories into 3 subsets according to the percentage of pixels in every category over the entire dataset. Specifically, we define three disjoint subsets as follows: head classes (each with more than 1.0% of total pixels), body classes (each with a percentage ranging from 0.1% to 1% of total pixels) and tail classes (each with less than 0.1% of total pixels). The complete list of the split is reported in supplementary material.
We evaluate our method using two widely-adopted segmentation models (FCN and DeepLabV3+ ) based on different backbone networks, ranging from ResNet-50, ResNet-101 to the latest ResNeSt-101, and report the performance in Tab. 6. Our method achieves 2.0 and 2.3 improvement in mIoU using FCN-8s with ResNet-50 and ResNet-101, respectively. The performance on the body and tail are improved significantly. Moreover, our method outperforms the baseline with large margin at 5.7 in mean accuracy with ResNet-101 backbone. Even with a stronger backbone: ResNeSt-101 , our method also achieves 0.7 mIoU and 2.8 improvement in mean accuracy, where the tail categories have a performance gain of 1.2 in mIoU and 3.7 in mean accuracy. We further validate our method using DeepLabV3+, which is a more powerful semantic segmentation model. Our DisAlign improves the performance of DeepLabV3+ by a margin of 0.5 based on ResNeSt-101, which achieves the new state-of-the-art (47.8 in mIoU) on the ADE20k dataset.
3 Object Detection and Instance Segmentation
We conduct experiments on LVIS dataset. For evaluation, we use a COCO-style average precision (AP) metric that averages over categories and different box/mask IoU threshold .
We first compare our method with recent work and report quantitative results in Tab. 8. We find our DisAlign with cosine classifier head achieves 25.6 in , and 26.3 in when applied to the Mask R-CNN+FPN with the ImageNet pre-trained ResNet-50 backbone. Moreover, our strategy can be further improved to achieve 27.6 in and 27.9 in based on the COCO pre-trained model. In both cases, our method is able to maintain the performance of the frequent (also called head) categories, and gain significant improvement on common (also called body) and rare (also called tail) categories. We also report performance with more power detection framework (e.g.Cascade R-CNN) and stronger backbones (e.g. ResNet-50/101, and ResNeXt-101) in Tab. 7 and Tab. 9. It is worth noting that even with the stronger backbones or frameworks, the performance gain of our DisAlign over the baseline is still significant.
Conclusion
In this paper, we have presented a unified two-stage learning strategy for the large-scale long-tail visual recognition tasks. To tackle the biased label prediction, we develop a confidence-aware distribution alignment method to calibrate initial classification predictions. In particular, we design a generalized re-weight scheme to leverage the category prior for the alignment process. Extensive experiments show that our method outperforms previous works with a large margin on a variety of visual recognition tasks(image classification, semantic segmentation, and object detection/segmentation).
References
Appendix A Experiments of Image Classification
In this section, we first introduce the dataset and evaluation metrics for image classification task in Sec.A.1. Then the training configuration will be detailed in Sec.A.2, followed by results on three benchmarks in Sec.A.3.
To demonstrate our methods, we conduct experiments on three large-scale long-tailed datasets, including Places-LT, ImageNet-LT, and iNaturalist 2018. Places-LT and ImageNet-LT are artificially generated by sampling a subset from their balanced versions (Places-365 and ImageNet-2012) following the Parento distribution. iNaturalist 2018 is a real-world, naturally long-tailed dataset, consisting of samples from 8,142 species.
We report the class-balanced average Top-1 accuracy on the corresponding validation/test set, and also calculate the accuracy of three disjoint subsets, ‘Many’, ‘Medium’ and ‘Few’, which are defined according to the amount of training data per class .
A.2 Training Configuration
Following , we use PyTorch framework for all experiments. For ImageNet-LT, we report performance with ResNet-{50,101,152} and ResNeXt-{50,101,152} and mainly use ResNet-50 for ablation study. For iNaturalist 2018, performance is reported with ResNet-{50,101,152}. For Places-LT, ResNet-152 is used as backbone and we pre-train it on the full ImageNet-2012 dataset.
We use the SGD optimizer with momentum 0.9, batch size 256, cosine learning rate schedule gradually decaying from 0.1 to 0, and image resolution 224224. For the joint learning stage, the backbone network and original classifier head are jointly trained with 90 epochs for ImageNet-LT, and 90/200 epochs for iNaturalist-2018. For the Places-LT dataset, the models are trained with 30 epochs with the all layers frozen expect the last ResNet block in the first stage.
In the second distribution alignment stage, we restart the learning rate and train it for 10/30 epochs as while keeping the backbone network and original classifier head fixed(10 epochs for ImageNet-LT and Places-LT, 30 epochs for iNaturalist-2018). For all three datasets, we set the generalized re-weight scale for dot-product classifier head, for cosine normalized classifier head. The and are initialized with 1.0 and 0.0, respectively.
A.3 Detailed Experimental Results
We present the detailed quantitative results for ImageNet-LT in Table 10.
To further demonstrate our method, we conduct experiments on two extra large-scale long-tail benchmarks and report the performance in Table 13 and Table 14.
A.4 Ablation Study
We report an ablation study of the two main components of our method with ResNeXt-50 in Tab. 11, which shows that both adaptive calibration and generalized re-weighting(G-RW) contribute to the performance improvement of our approach.
We plot the learned magnitude and margin according to the class sizes below. They share a similar trend, in which the tail/body classes have larger value than head. Thus our calibration alleviates the bias in the original prediction by boosting the tail scores.
We study confidence-based calibration in the table below, which shows that the input-aware calibration outperforms the input-agnostic counterpart and the baselines using only magnitude or margin. We also observe that the example whose biased prediction probability is low on its ground-truth class tends to be improved with higher confidence.
Appendix B Experiments of Semantic Segmentation
Similar to image classification, the large-scale semantic segmentation task still suffers from the long-tail data distribution. To further validate the effectiveness of our method, we also apply DisAlign on large-scale semantic segmentation benchmark: ADE-20k.
ADE20K dataset is a scene parsing benchmark, which contains 150 stuff/object categories. The dataset includes 20K/2K/3K images for training, validation, and testing. Compared with the image classification, the imbalance of ADE20K is more serve than the image classification, which has an imbalance ratio of 788(Max/Min). Follow the similar protocol in image classification, we divide the 150 categories into 3 groups according to the ratio of pixel number over the whole dataset. Specifically, three disjoint subsets are: head classes(classes each with a ratio over 1.0%), body classes(classes each with a ratio ranging from 0.1% to 1%) and tail classes(classes under a ratio of 0.1%), the complete list of the split is reported in Tab.16.
For the evaluation metric, we use the mean intersection of union(mIoU) and mean pixel accuracy(mAcc). We also report the mIoU and mAcc of each group(head, body and tail) for clarity.
B.2 Training Configuration
We implement our method based on MMSegmentation toolkit. In the joint learning training phase, we set the learning rate to 0.01 initially, which gradually decreases to 0 by following the ’poly’ strategy as . The images are cropped to and augmented with randomly scaling(from 0.5 to 2.0) and flipping. ResNet-50, ResNet-101 and ResNeSt-101 are used as the backbone. For the evaluation metric, we use the mean intersection of union(mIoU) and mean pixel accuracy(mAcc). All models are trained with 160k iterations with a batch size of 32 based on 8 V100 GPUs. In the DisAlign stage, we follow a similar protocol as stage-1 and only training the model with 8k iterations. We set for all experiments.
B.3 Quantitative Results
We evaluate our method with two state-of-the-art segmentation models(FCN and DeepLabV3+)based on different backbone networks, ranging from ResNet-50, ResNet-101 to the latest ResNeSt-101, and report the performance in Tab.15.
Appendix C Experiments on LVIS Dataset
LVIS v0.5 dataset is a benchmark dataset for research on large vocabulary object detection and instance segmentation, which contains 56K images over 1230 categories for training, 5K images for validation. This challenging dataset is an appropriate benchmark to study the large-scale long-tail problem, where the categories can be binned into three types similar with ImageNet-LT: rare(1-10 training images), common(11-100 training images), and frequent( training images).
C.2 Training Configuration
We train our models for object detection and instance segmentation based on Detecron2, which is implemented in PyTorch. Unless specified, we use the ResNet backbone(pre-trained on ImageNet) with FPN. Following the training procedure in , we resize the images so that the shorter side is 800 pixels. All baseline experiments are conducted on 8 GPUs with 2 images per GPU for 90K iterations, with a learning rate of 0.02 which is decreased by 10 at the 60K and 80K iteration. We use SGD with a weight decay of 0.0001 and momentum of 0.9. Scale jitter is applied for all experiments in default same with .
For the DisAlign, we freeze all network parameters and learn the magnitude and margin for extra 9K iterations with a learning rate of 0.02. Generalized re-weight is only used for fore-ground categories. Generalized re-weight scale is set to 0.8 for all experiments.
C.3 Quantitative Results
We report the detailed results in Table.17 and Tab.18.