[Paper Review] Distributional Robustness Loss for Long-tail Learning

Dvir Samuel, Gal Chechik|arXiv (Cornell University)|Apr 7, 2021

Domain Adaptation and Few-Shot Learning42 references4 citations

TL;DR

This paper proposes DRO-LT Loss, a distributionally robust optimization-based loss that improves representation learning in long-tailed data by minimizing worst-case risk within uncertainty regions around class centroids. By learning robust representations that reduce bias toward head classes while maintaining head accuracy, the method achieves new state-of-the-art performance on CIFAR100-LT, ImageNet-LT, and iNaturalist benchmarks, with significant gains in tail-class recognition.

ABSTRACT

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.

Motivation & Objective

• To address the under-explored issue of representation bias in deep models trained on long-tailed data, particularly in the feature extractor (backbone) layers.
• To investigate whether improving representation quality—beyond classifier reweighting—can significantly enhance recognition performance on long-tailed benchmarks.
• To develop a robust optimization-based loss that accounts for uncertainty in tail-class representations due to limited samples.
• To derive an efficient, differentiable upper bound of the distributionally robust loss that can be jointly optimized with standard classification loss.
• To demonstrate that robust representation learning can improve tail-class accuracy without degrading head-class performance.

Proposed method

• Formulates long-tailed learning as a distributionally robust optimization (DRO) problem, where the model is trained to be robust against worst-case distributions within uncertainty neighborhoods of class centroids.
• Defines the uncertainty region around each class centroid as a ball of radius ε, with larger radii for tail classes due to higher variance in few-shot samples.
• Derives an upper bound of the DRO loss that is analytically tractable and differentiable, enabling efficient optimization via gradient descent.
• Introduces a learnable robustness safety margin ε for each class, jointly optimized during training to avoid manual hyperparameter tuning.
• Combines the DRO-LT loss with standard cross-entropy loss, allowing joint optimization of representation learning and classifier head tuning.
• Applies the method to representations at multiple layers of the network, with ablation showing strongest gains at the penultimate layer.

Figure 1: Our distributional robustness loss is designed for learning a representation where samples are kept close to the centroid of their class. Here, the empirical centroid $\widehat{\mu_{2}}$ (framed pink triangle) is estimated based only on few samples (four pink triangles) and as a result it

Experimental results

Research questions

• RQ1 Does representation learning in deep networks trained on long-tailed data suffer from bias toward head classes, even in early layers?
• RQ2 Can distributionally robust optimization improve representation quality for both head and tail classes in long-tailed learning?
• RQ3 Is it possible to derive an efficient, differentiable upper bound of the DRO loss that enables end-to-end training?
• RQ4 Can the robustness margin ε be jointly learned per class without additional hyperparameter tuning?
• RQ5 Does improving representation robustness lead to improved tail-class recognition while preserving head-class accuracy?

Key findings

• The feature extractor of deep networks exhibits significant bias toward head classes even in early layers, with accuracy gaps between head and tail classes persisting across all layers.
• The proposed DRO-LT loss reduces the generalization gap between training and test error for tail classes, as shown in Figure 4, while maintaining head-class performance.
• t-SNE visualization in Figure 5 confirms that DRO-LT produces more compact and well-separated feature clusters for both head and tail classes.
• On CIFAR100-LT with an imbalance ratio of 50, DRO-LT achieves new state-of-the-art performance, significantly improving tail-class accuracy without sacrificing head accuracy.
• The model learns larger robustness margins (ε) for tail classes compared to head classes, reflecting higher uncertainty in their representations.
• The method achieves consistent SOTA results across three benchmarks: CIFAR100-LT, ImageNet-LT, and iNaturalist, demonstrating broad applicability.

Figure 2: Learned $\varepsilon$ values for each class of a ResNet-32 trained on CIFAR100-LT with imbalance ratio 50. Blue: Polynomial fit.