Source-linked AI summary

Distributional Robustness Loss for Long-tail Learning

Dvir Samuel, Gal Chechik

arXiv:2104.03066v2cs.LGcs.CV

TL;DR

Long-tailed data can bias deep models toward frequent classes, and the effects of imbalance on learned representations have received less attention. The paper develops a distributionally robust loss with an efficiently minimized upper bound for learning representations. It reports state-of-the-art results on three long-tail benchmarks, improving tail recognition while maintaining high head-class accuracy.

  • Problem

    The impact of long-tailed imbalance on learned representations is less understood, even though deep models tend to favor head classes and fail to recognize tail classes.

    Source: §1 · pp. 1–2
  • Method

    The paper develops a distributionally robust loss for learning balanced representations and derives an upper bound that can be minimized efficiently.

    Source: §9 · p. 8
  • Results

    DRO-LT with a standard classifier sets state-of-the-art results on CIFAR100-LT, ImageNet-LT, and iNaturalist, improving tail-class performance while maintaining high head-class accuracy.

    Source: §9 · p. 8
  • Takeaways & Limitations

    The results suggest that training representations appropriately for unbalanced data can substantially affect downstream accuracy.

    Source: §9 · p. 8

    The approach models each class’s samples as distributed around a centroid, with likelihood decaying exponentially with feature-space distance.

    Source: §5.2 · p. 4

Abstract

from arXiv · show

Real-world data is often unbalanced and long-tailed, but deep models struggle to recognize rare classes in the presence of frequent classes. To address unbalanced data, most studies try balancing the data, the loss, or the classifier to reduce classification bias towards head classes. Far less attention has been given to the latent representations learned with unbalanced data. We show that the feature extractor part of deep networks suffers greatly from this bias. We propose a new loss based on robustness theory, which encourages the model to learn high-quality representations for both head and tail classes. While the general form of the robustness loss may be hard to compute, we further derive an easy-to-compute upper bound that can be minimized efficiently. This procedure reduces representation bias towards head classes in the feature space and achieves new SOTA results on CIFAR100-LT, ImageNet-LT, and iNaturalist long-tail benchmarks. We find that training with robustness increases recognition accuracy of tail classes while largely maintaining the accuracy of head classes. The new robustness loss can be combined with various classifier balancing techniques and can be applied to representations at several layers of the deep model.

1. Introduction

The introduction argues that long-tail learning research has focused mainly on balancing data, losses, or classifiers, leaving the effects of imbalance on learned representations poorly understood. It presents a distributionally robust representation loss and an efficiently minimized upper bound as a way to address that gap.

1. Introduction

  • Prior approaches balance training data, losses, or classifiers, while the impact of imbalance on learned representations remains poorly understood.The cited work [21] suggested that high-quality representations can be learned despite imbalance and that classifier adjustment may suffice; whether better representations improve recognition remained an open question.Source: §1 · pp. 1–2
  • The loss accounts for uncertain tail-class centroids by pulling same-class samples and pushing away other-class samples within an uncertainty region.The illustrated uncertainty regions are typically larger for tail classes than for head classes with more samples.Source: Figure 1 · p. 1
  • DRO-LT consistently outperforms previous models on CIFAR100-LT, CIFAR10-LT, ImageNet-LT, and iNaturalist benchmarks.The paper formulates unbalanced learning as robust optimization and highlights high variance as a hindrance to tail accuracy.Source: §1 · p. 2
  • DRO-LT uses robustness to learn a balanced feature extractor for head and tail classes, with a computable upper bound that can be minimized efficiently.The paper also describes learning a class-specific robustness safety margin jointly with the task, avoiding additional hyperparameter tuning.Source: §1 · p. 2

2. Related Work

Prior long-tail methods manipulate data, losses, training stages, or model ensembles, while distributionally robust optimization addresses shifts between training and test distributions. The paper situates its approach among these methods and notes that DRO had not, to the authors’ knowledge, been applied to long-tail learning.

2.1. Long-tail recognition

  • Long-tail recognition methods include resampling, loss reweighting, two-stage fine-tuning, and ensemble models.Resampling can overfit tail classes, reduce generalization, or require expensive augmentation; two-stage methods typically freeze the backbone while balancing the classifier.Source: §2.1 · pp. 2–3

2.2. Distributionally robust optimization

  • DRO defines an uncertainty set around the training distribution and minimizes risk against its worst-case test distribution; the authors knew of no prior application to long-tail learning.DRO also has reported connections to standard robust optimization, regularization, risk-aversion methods, and game theory.Source: §2.2 · p. 3
  • Adversarially robust representation learning is outside this paper’s scope.Source: §2.2 · p. 3

3. Overview of our approach

The approach targets last-layer representations whose class distributions are separated enough for linear classifiers, addressing the unreliable training samples of rare classes. Its loss accounts for uncertain class centroids and uses an efficiently computed upper bound as a surrogate.

3. Overview of our approach

  • When rare classes have few samples, their training distributions may poorly represent the true data distributions, producing representations that hinder classification.Source: §3 · p. 3
  • The representation loss extends contrastive losses by pulling samples toward their class centroid and pushing other-class samples away, while accounting for noisy centroid estimates.It optimizes against worst-case centroids within a safety hyperball around the empirical centroid.Source: §3 · p. 3
  • An easily computed upper bound serves as a surrogate for the potentially hard-to-compute worst-case loss and yields better representations.The loss uses feature-space distance to an estimated class centroid and robustness margins Δ and Δ′.Source: §3 · p. 3

4. Distributional Robust Optimization

This section contrasts empirical risk minimization, which uses the empirical training distribution, with distributionally robust optimization over a family of possible distributions. It motivates the latter when training data poorly approximates the test distribution.

4. Distributional Robust Optimization

  • ERM minimizes expected loss over the empirical training distribution, using observed training samples to estimate the unknown data distribution.The empirical distribution Pδ assigns probability mass to observed sample-label pairs, with δ denoting the Kronecker delta.Source: §4 · p. 3
  • When training data are limited, the empirical distribution may be far from the test distribution, so alternative estimates may better reduce test loss.Source: §4 · p. 3
  • DRO minimizes expected loss against the worst-case distribution in an uncertainty set around the empirical training distribution.The set is typically a hyperball of radius ε under a distribution discrepancy measure, often KL-divergence or Wasserstein distance.Source: §4 · p. 3

5. Our approach

DRO-LT learns class-separated latent representations while accounting for uncertainty in class centroids estimated from imbalanced samples. It bounds a worst-case negative log-likelihood with a computable surrogate and combines that representation loss with classification loss.

5.1. Preliminaries

  • The feature extractor maps each input to a latent vector, where each class is represented by an empirical centroid distinct from its true-distribution centroid.The empirical centroid is the mean of that class’s feature vectors.Source: §5.1 · pp. 3–4

5.2. A loss for representation learning

  • The representation loss brings same-class features toward their class centroid and separates them from other classes, with equal class weighting preventing head classes from dominating.The loss is formed as a weighted average of per-class negative log-likelihoods.Source: §5.2 · p. 4

5.3. A robust loss

  • DRO-LT bounds the negative log-likelihood with high probability by evaluating worst-case candidate centroids, addressing error from noisy empirical centroids.The surrogate increases the sample-to-centroid distance according to the class uncertainty-ball radius.Source: §5.3 · pp. 4–5
  • Training applies cross-entropy to the classifier output and the robust loss to the penultimate-layer representation, combining them with a tunable weight.The combined objective is λ times cross-entropy plus (1−λ) times the robust loss.Source: §5.3 · p. 5

6. Training

Training selects an uncertainty radius that balances coverage of the true centroid against bound tightness, then estimates centroids and adds DRO loss after initial representation learning.

6. Training

  • A radius that is too small lowers the probability that the true centroid lies within the uncertainty region, while a larger radius makes the bound more likely to hold but less tight.Tail classes have noisier centroid estimates because they contain fewer samples, motivating a larger radius.Source: §6 · p. 5
  • The radius can be shared across classes, scaled by 1/√n to give tail classes a larger safety radius, or learned separately for each class.The shared and sample-count-scaled radii are tuned on a validation set; the learned radius is optimized during training.Source: §6 · p. 5
  • Training first uses cross-entropy to obtain initial features and centroids, then adds DRO loss; finally, the feature extractor is frozen while the classifier is fine-tuned with balanced sampling.Centroids are computed from all sample features at each epoch’s start and held fixed for that epoch, rather than estimated per batch.Source: §6 · p. 5

7. Experiments

The experiments cover three long-tailed recognition benchmarks and compare DRO-LT variants with established balancing approaches. Evaluation uses overall accuracy and, on two benchmarks, class-frequency splits.

7.1. Datasets

  • Experiments use CIFAR100-LT at imbalance factors 100, 50, and 10, ImageNet-LT with 115.8K images across 1,000 classes, and iNaturalist with 437.5K images across 8,142 classes.ImageNet-LT has 1,280 to 5 images per class; iNaturalist is long-tailed by nature.Source: §7.1 · pp. 5–6

7.2. Compared methods

  • Comparisons include resampling, loss-manipulation, two-stage fine-tuning, representation-learning methods, and DRO-LT variants using four uncertainty-radius settings.The radius settings are ε=0 (ERM), shared ε, ε/√n, and a learned value per class.Source: §7.2 · p. 6

7.3. Evaluation Protocol

  • Accuracy is measured on class-balanced test sets; CIFAR-100 and ImageNet-LT also report Many, Med, and Few class splits.The splits are defined by training counts: more than 100, 20–100, and fewer than 20 samples, respectively.Source: §7.3 · p. 6

7.4. Implementation details

  • The experiments use SGD with momentum 0.9, ResNet-32 for CIFAR-100 and ResNet-50 for ImageNet-LT and iNaturalist, with schedules differing by benchmark.Hyperparameters and early stopping are selected on validation data; the selected λ is 0.5, and accuracy is reported as stable across a range of λ values.Source: §7.4 · p. 6

8. Results

DRO-LT performs well across long-tail benchmarks, improving tail-class recognition while maintaining head-class accuracy. Results also point to gains from balancing the learned representations, not only the classifier.

Benchmark results

  • DRO-LT variants achieve the best results across CIFAR100-LT imbalance factors, perform well on ImageNet-LT tail and head classes, and slightly improve on iNaturalist baselines.On ImageNet-LT, the reported comparison contrasts with methods that sacrifice head accuracy for better tail classification.Source: §8 · p. 6, Table 2 · p. 7, Table 3 · p. 7

Classifier and representation

  • DRO-LT representations enable CE, resampling, and balanced classifiers to reach good performance compared with other representation-learning approaches.The comparison evaluates representation learning with CE, balanced resampling, or DRO-LT, then freezes the feature extractor and trains classifiers using CE, resampling, or LWS.Source: §8 · p. 7, Table 4 · p. 8

Adaptive robustness

  • Training on CIFAR100-LT with imbalance ratio 50 learns slightly larger uncertainty radii for tail classes than for head classes.The learned per-class radii are denoted ε.Source: §8 · p. 7

Feature-space visualization

  • DRO-LT produces a more compact feature space with margins around head and tail classes; tail classes have larger margins because their features are estimated less accurately.The comparison uses a t-SNE projection of CIFAR100-LT representations trained with cross-entropy and DRO-LT.Source: §8 · p. 8, Figure 5 · p. 8

Representation imbalance

  • Nearest-centroid accuracy shows a substantial head–tail gap at every tested ResNet-32 layer on CIFAR100-LT, including lower layers.With balanced CIFAR-100 training, validation accuracy grows at higher layers; with long-tailed training, accuracy also grows higher up but the class gap remains.Source: Figure 3 · p. 7, §8 · p. 8

Robustness

  • DRO-LT narrows the head–tail accuracy gap, and its train–test error gap is much lower for tail classes while remaining the same for head classes.The figure compares vanilla cross-entropy with DRO-LT.Source: §8 · pp. 7–8, Figure 4 · p. 8

9. Discussion

The paper argues that long-tailed training biases deep feature extractors toward head classes and that this harms recognition. It proposes an efficiently minimized DRO-based loss and reports state-of-the-art results on three long-tail benchmarks while maintaining head-class accuracy.

Contribution

  • The paper identifies head-class bias in feature extractors as a recognition problem, contrasting with prior studies that suggested classifier rebalancing alone was sufficient.It develops a Distributionally Robust Optimization-based loss and an efficiently minimized upper bound to learn balanced representations.Source: §9 · p. 8
  • The robustness safety margin can be learned during training without additional hyperparameter tuning.This margin-learning procedure is presented as part of the paper’s contribution.Source: §9 · p. 8

Conclusion

  • Combining DRO-LT with a standard classifier achieves state-of-the-art results on CIFAR100-LT, ImageNet-LT, and iNaturalist, improving tail performance while maintaining high head accuracy.The authors conclude that training representations appropriately for unbalanced data can substantially affect downstream accuracy.Source: §9 · p. 8

Supplemental Material

The supplemental analyses define the coverage probability used for the uncertainty radius and examine bound tightness and key DRO-LT settings. They find that combining robustness with cross-entropy works best, while accuracy is maintained across a broad range of shared uncertainty radii.

A. Additional information about p(ϵ)

B. Tightness of the bound

  • For CIFAR100-LT, the average relative magnitude of the bound gap was 0.07, suggesting the bounds are quite tight.The gap is estimated between Eq. 17 and Eq. 20, averaged over samples.Source: B. Tightness of the bound · p. 11

C. DRO-LT Sample count ϵ/√n

  • Sample count ϵ/√n assumes a shared sample-distribution standard deviation across classes, with centroid uncertainty varying only by sample count.The shared standard deviation is tuned as a hyperparameter.Source: C. DRO-LT Sample count ϵ/√n · p. 11

D. Loss trade-off parameter λ

E. Robustness vs Performance

  • With shared ε on CIFAR100-LT, accuracy is maintained across a broad range of radii, but very small or large ε reduces it.Very large ε makes the worst-case centroid too distant from the class centroid, loosening the bound.Source: Figure 7 · p. 11, E. Robustness vs Performance · p. 11

F. More analysis on iNaturalist:

The supplemental results report DRO-LT comparisons on iNaturalist and CIFAR-10-LT, and probe how it affects representations across ResNet-32 layers. The reported representation gap narrows mainly at the last convolutional layer.

H. Balancing latent representations

  • With a nearest-centroid classifier, DRO-LT narrows the head–tail accuracy gap mostly at conv30 in a ResNet-32.The comparison uses a model trained with cross-entropy and one trained with DRO-LT applied to the last convolutional layer.Source: Figure 8 · p. 12, H. Balancing latent representations · p. 12

F. More analysis on iNaturalist:

G. Results for CIFAR-10-LT:

Loading 2104.03066v2…