[Paper Review] Long-Tailed Classification by Keeping the Good and Removing the Bad Momentum Causal Effect

Kaihua Tang, Jianqiang Huang+1 more|arXiv (Cornell University)|Sep 27, 2020

Fault Detection and Control Systems231 citations

TL;DR

The paper introduces a causal framework that treats SGD momentum as a confounder in long-tailed classification, learns via de-confounded training, and infers the direct effect to achieve state-of-the-art results on multiple benchmarks.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for principled understanding of long-tailed classification beyond re-weighting/re-sampling heuristics.
• Develop a causal model that explains how SGD momentum acts as a confounder and mediator in long-tailed settings.
• Propose a one-stage, re-training-free solution that disentangles direct causal effects from mediation to improve tail accuracy.
• Show empirical gains on benchmarks including Long-tailed CIFAR-10/-100, ImageNet-LT, and LVIS to validate the theory.

Proposed method

• Construct a causal graph with momentum M, feature X, head-projection D, and prediction Y to model confounding and mediation.
• Apply backdoor adjustment to derive a de-confounded training objective that estimates P(Y|do(X)) via inverse probability weighting.
• Formulate a logit for P(Y=i|do(X=x)) as an energy-based, multi-head, normalized classifier (Eq. 7).
• Compute Total Direct Effect (TDE) for Y given X by counterfactual inference that subtracts the indirect mediation effect (Eq. 8).
• Use Background-Exempted Inference for tasks with a background class to retain head bias while evaluating TDE-based predictions.
• Provide theoretical connections to prior two-stage and normalization-based methods (Table 1) and explain when de-confound-TDE outperforms alternatives.

Experimental results

Research questions

• RQ1 How does SGD momentum causally influence feature representations and predictions in long-tailed data?
• RQ2 Can we remove the bad (confounding) effect of momentum while keeping the good (mediation) effect to improve tail accuracy?
• RQ3 Does de-confounded training plus direct-effect inference (TDE) provide a one-stage, retraining-free solution with robust gains across datasets?
• RQ4 How does the proposed approach relate to and explain existing re-balancing and normalization-based methods?

Key findings

• Achieves new state-of-the-arts on three long-tailed visual recognition benchmarks (Long-tailed CIFAR-10/-100, ImageNet-LT) and LVIS for object detection/segmentation.
• Demonstrates significant performance gains: 3.5% absolute improvement on mask AP and 3.1% absolute improvement on box AP on LVIS using the same Cascade Mask R-CNN backbone.
• Shows that de-confounded training plus TDE inference consistently outperforms prior re-balancing and one-stage methods across many-shot, medium-shot, and few-shot regimes.
• Provides a principled explanation for why two-stage training methods work and why one-stage de-confounded TDE can be more effective and training-efficient.
• Visualizes that de-confound-TDE focuses on discriminative regions rather than broad context, aligning with a direct-effect emphasis.