Source-linked AI summary
Long-Tailed Classification by Keeping the Good and Removing the Bad Momentum Causal Effect
Kaihua Tang, Jianqiang Huang, Hanwang Zhang
TL;DR
Long-tailed datasets are difficult to balance at scale, while re-weighting and re-sampling methods lack a fundamental theory and can mishandle the competing effects of imbalance. The paper treats SGD momentum as a causal confounder, removes its harmful effect through intervention and counterfactual inference, and reports state-of-the-art results across classification and instance segmentation benchmarks.
Problem
Long-tailed classification needs principled foundations because balanced data are difficult to collect at scale and existing re-balancing methods leave key effects theoretically unexplained.
Method
The framework uses de-confounded training to remove momentum's backdoor effect while retaining mediation, then uses the direct causal effect for prediction logits.
Results
The method achieves new state-of-the-art results across Long-tailed CIFAR-10/-100, ImageNet-LT, and LVIS, including 3.5% mask AP and 3.1% box AP improvements on LVIS.
Takeaways & Limitations
The proposed one-stage solution extracts an unbiased direct effect and is described as adaptive to class-distribution priors across long-tailed recognition tasks.
Takeaways & Limitations
The authors identify broader-impact risks, including possible malicious use to identify minority groups, and call for validation in wider application domains.
Abstract
from arXiv · showhide
As the class size grows, maintaining a balanced dataset across many classes is challenging because the data are long-tailed in nature; it is even impossible when the sample-of-interest co-exists with each other in one collectable unit, e.g., multiple visual instances in one image. Therefore, long-tailed classification is the key to deep learning at scale. However, existing methods are mainly based on re-weighting/re-sampling heuristics that lack a fundamental theory. In this paper, we establish a causal inference framework, which not only unravels the whys of previous methods, but also derives a new principled solution. Specifically, our theory shows that the SGD momentum is essentially a confounder in long-tailed classification. On one hand, it has a harmful causal effect that misleads the tail prediction biased towards the head. On the other hand, its induced mediation also benefits the representation learning and head prediction. Our framework elegantly disentangles the paradoxical effects of the momentum, by pursuing the direct causal effect caused by an input sample. In particular, we use causal intervention in training, and counterfactual reasoning in inference, to remove the "bad" while keep the "good". We achieve new state-of-the-arts on three long-tailed visual recognition benchmarks: Long-tailed CIFAR-10/-100, ImageNet-LT for image classification and LVIS for instance segmentation.
1 Introduction
Long-tailed datasets are difficult to balance at scale, and existing re-balancing methods leave the paradoxical effects of imbalance theoretically unexplained. The paper models these effects causally through SGD momentum and removes the harmful effect while preserving beneficial mediation, achieving strong benchmark results.
- Long-tailed data arise because expanding class vocabularies makes balanced collection disproportionately costly, especially when multiple classes co-occur in one image.
- Long-tailed classification is necessary for scaling deep models, but its fundamental theory remains missing despite progress on established benchmarks.
- Re-balancing exposes a paradox: it can improve classifier training while causing under-fitting or over-fitting in feature learning, and two-stage methods outperform end-to-end training without explaining why.
- SGD momentum acts as a confounder linking sample features and prediction logits, creating a backdoor path that can bias tail samples toward head classes while mediation benefits representation learning.
- The proposed framework uses de-confounded training to block the momentum backdoor path while retaining mediation, then uses direct causal effects as inference logits.
- 3.5% and 3.1% absolute improvements on mask AP and box AP, respectively, are achieved on LVIS using the same Cascade Mask R-CNN with R101-FPN backbone.
2 Related Work
Prior long-tailed recognition methods re-balance classes, mine difficult examples, transfer knowledge, or use causal inference. These approaches improve recognition but can introduce fitting, distribution-access, parameter, or training-complexity constraints.
- Re-Balanced Training: Re-sampling and re-weighting re-balance class contributions but inevitably risk under-fitting head classes and over-fitting tail classes.
- Re-Balanced Training: Re-balancing methods may also require access to the data distribution, limiting use in online and streaming settings.
- Hard Example Mining: Instance-level re-weighting addresses imbalance by focusing on hard samples rather than directly modifying class-prior distributions.
- Transfer Learning/Two-Stage Approach: Transfer-learning and two-stage approaches can be effective, but may increase parameters substantially or require complicated training strategies.
- Causal Inference: Causal inference is used as both an interpretation framework and a way to pursue desired causal effects, including removal of dataset bias in computer vision.
3 A Causal View on Momentum Effect
The paper represents momentum, features, head-direction projections, and predictions in a causal graph. In long-tailed training, momentum is dominated by head samples, influencing feature direction and prediction through confounding and mediation paths.
- The causal graph contains momentum M, object feature X, head-direction projection D, and prediction Y; M is a confounder and D is a mediator.
- SGD momentum is an exponential moving average of past gradients, with decay ratio µ and learning rate lr, and dampens oscillations from individual samples.
- In long-tailed data, momentum is dominated by head samples, producing the causal link M → X because backbone parameters generating features are trained under momentum.
- Assumption 1 defines the head direction as the normalized exponential moving average of features, using the same decay rate µ as momentum.
- The assumption also holds for balanced datasets, where the causal link X → Y is naturally unaffected by momentum M.
- Feature effects decompose into indirect mediation through D and direct effects through X → Y; the discriminative component and head-direction projection jointly determine the total effect.
4 The Proposed Solution
The proposed solution pursues the total direct causal effect of an input after de-confounded training, using intervention to remove confounding and counterfactual inference to retain mediation while reducing its harmful bias. It also connects this framework to prior re-balancing methods and extends inference for tasks with a background class.
- TDE inference: TDE measures the direct effect of input X on class-logit prediction by subtracting the null-input counterfactual from the intervened input prediction.The mediator D is held at d in both terms, and x0 is the null input 0.
- De-confounded training: De-confounded training applies backdoor adjustment to cut the momentum confounding path while retaining the mediation path through D.The resulting model is intended to preserve beneficial feature learning while removing harmful confounding bias.
- De-confounded training: Inverse probability weighting approximates the backdoor adjustment, while a multi-head strategy provides finer-grained sampling across feature and weight channels.The weighted probability is modeled with an energy-based formulation whose denominator normalizes variable magnitudes using class-specific and class-agnostic energies.
- TDE inference: After training, counterfactual consistency lets Eq. (7) compute the first TDE term, while replacing the direct feature component with zero computes the null-input term.The scale parameter α controls the trade-off between indirect and direct effects in the final calculation.
- Background-exempted inference: Background-exempted inference uses original total-effect inference for the background class while applying the proposed calculation to other classes.This preserves the effect of a predominantly observed background class in detection and instance-segmentation settings.
- Revisiting prior methods: The causal framework interprets one-stage re-weighting as controlled direct effect and two-stage re-balancing as natural direct effect, explaining their distinct treatment of mediation.The paper also states that normalized classifiers without causal guidance can underperform the proposed de-confounded model.
5 Experiments
The method is evaluated across long-tailed image-classification and instance-segmentation benchmarks, with controlled protocols, visualizations, ablations, and comparisons against prior methods. It consistently improves performance and adapts across imbalance settings and tasks.
- Datasets and Protocols: The method is evaluated on Long-tailed CIFAR-10/-100, ImageNet-LT, and LVIS across image classification, object detection, and instance segmentation.ImageNet-LT contains 1,000 classes over 186,000 images, while LVIS provides large-vocabulary instance-segmentation categories.
- Comparisons with State-of-The-Art Methods: Grad-CAM activation maps compare the linear classifier baseline, Decouple-LWS, and the proposed method on ImageNet-LT.The experiment visualizes classifier behavior across these three approaches.
- Datasets and Protocols: Evaluation uses Top-1 accuracy for CIFAR, four frequency-based accuracy splits for ImageNet-LT, and mask AP across IoU thresholds 0.5 to 0.95 for LVIS.LVIS results are also categorized into rare, common, and frequent classes.
- Ablation studies: De-confounded TDE achieves the best performance across ablation settings, while TDE inference improves all three normalized classifier models.Changing α produces smooth performance trade-offs between tail and head classes.
- Comparisons with State-of-The-Art Methods: The proposed method outperforms previous state-of-the-arts on all Long-tailed CIFAR-10/-100 imbalance ratios and ImageNet-LT splits and beats the LVIS 2019 winner under the same framework.The LVIS comparison uses Cascade Mask R-CNN with an R101-FPN backbone; the reported comparison also notes lower GPU and batch-size requirements than EQL.
6 Conclusions
The paper concludes that its causal framework explains momentum's role in long-tailed classification and enables one-stage extraction of unbiased instance effects. It reports state-of-the-art results on ImageNet-LT and LVIS while identifying broader validation and disentanglement as future work.
- 6 Conclusions: The causal framework theoretically explains previous methods and provides a one-stage solution for extracting each instance’s unbiased direct effect.Implementation combines de-confounded training with total direct effect inference and is described as simple, adaptive, and agnostic to class-distribution statistics.
- 6 Conclusions: New state-of-the-art results are achieved on ImageNet-LT and LVIS benchmarks across multiple tasks.The conclusion presents these results as evidence for the framework’s effectiveness.
- 6 Conclusions: Future work will validate the theory across more application domains and seek better feature-disentanglement algorithms for more precise counterfactual effects.
Broader Impact
The work is presented as improving classifier fairness and enabling larger-vocabulary datasets without compulsory class-balancing preprocessing.
- The method is described as improving classifier fairness, helping prevent potential discrimination from models that cater blindly to majority classes.
- It also allows larger-vocabulary datasets to be collected without compulsory class-balancing preprocessing.
A Additional Explanations of Assumption 1
The appendix explains how long-tailed sample frequencies bias momentum toward head-class gradients, producing head-like features, while balanced data retain the mediation structure without class preference.
- A Additional Explanations of Assumption 1: With 99 class-A samples and 1 class-B sample, the mean gradient becomes (0.98, 1), approximating the head-class gradient.With balanced 50-sample classes, the mean gradient is (0, 1), a direction shared by both classes.
- A Additional Explanations of Assumption 1: Momentum accelerates along the imbalanced mean-gradient direction, encouraging backbone parameters to generate head-like feature vectors.The appendix characterizes this as an unfair deviation toward the head class.
- A Additional Explanations of Assumption 1: Even balanced datasets can contain the projection D, but its similarity with features is then approximately equal across classes.Consequently, the momentum-mediated path does not create class preference in the balanced setting.
- A Additional Explanations of Assumption 1: Figure 6 ranks classifier-weight magnitudes by descending training-sample count after training with momentum μ = 0.9.
B Revisiting Previous Methods in Long-Tailed Classification
This section revisits normalized classifiers and re-balancing strategies through the proposed causal framework. It presents normalized classifiers as approximations of de-confounded training and describes several classifier and sampling designs.
- Normalized Classifiers: Normalized classifiers are widely used in long-tailed classification and, when correctly applied, approximate the proposed de-confounded training.
- Re-balancing Strategies: Table 5 reports the proposed TDE with and without Background-Exempted Inference on the LVIS V0.5 validation set using Cascade Mask R-CNN with an R101-FPN backbone.
- Normalized Classifiers: The general normalized classifier is introduced as a framework for analyzing prior normalized-classifier methods.
- Normalized Classifiers: The notation omits the classifier-head superscript because most previous methods set K to 1.
- Normalized Classifiers: Cosine, capsule, τ-norm, and Learnable Weight Scaling classifiers instantiate different normalization choices for visual features or classifier weights.
- Re-balancing Strategies: OLTR and Decouple use class-aware sampling in second-stage training, while BBN combines imbalanced and re-balanced sampling through shared bilateral branches.
C Background-Exempted Inference
Background-Exempted Inference preserves the beneficial background-class effect while applying TDE selectively in detection and segmentation settings. The strategy prevents TDE from harming foreground-background selection.
- Background-Exempted Inference excludes the legitimately biased background class from TDE calculation in object detection and instance segmentation.
- The strategy successfully prevents TDE from hurting foreground-background selection and enables TDE on a selected subset of categories.
D The Difference Between Re-balancing NDE and The Proposed TDE
The section distinguishes re-balancing NDE from TDE using causal intervention and a one-dimensional binary-classification illustration. Freezing the backbone implements the intervention underlying re-balanced classifier training.
- Re-balancing NDE is expressed as a difference between intervened predictions using the observed input and a dummy input.
- The dummy-input term is a constant offset because the balanced moving average does not point toward any specific class.
- Freezing the backbone breaks M → X, thereby implementing do(X = x); balanced re-sampling also produces a fair d′.
- Fine-tuning the backbone in OLTR violates this intervention, explaining its worse performance than Decouple-OLTR, which freezes the backbone.
- Figure 7 compares conventional, one-stage, and two-stage re-balancing classifiers with TDE in a one-dimensional binary classification example.
E Additional Ablation Studies
Additional ablations test hyper-parameters, classifier heads, backbones, and evaluation settings. They report benefits across classifier normalizations, backbones, and LVIS test-server evaluation, including improvements under stronger backbones.
- Hyper-parameters and Classifier Heads: Table 6 selects hyper-parameters using ImageNet-LT validation performance with a ResNeXt-50-32x4d backbone.
- Hyper-parameters and Classifier Heads: K = 2 experiments with cosine and capsule classifiers show that de-confounded-model gains do not come from increasing the number of heads.
- Hyper-parameters and Classifier Heads: Multi-head fine-grained sampling generally improves de-confounded training across both cosine and capsule normalization functions.
- Backbones: ResNeXt-101-32x4d and ResNeXt-101-64x4d backbones provide additional improvements on ImageNet-LT and LVIS V0.5, respectively.
- Evaluation Settings: The LVIS V0.5 test-server results use a single model with ResNeXt-101-64x4d, original hyper-parameters, no external dataset, and no model-enhancement tricks.