Probabilistic Contrastive Learning for Long-Tailed Visual Recognition

TL;DR

This work tackles long-tailed visual recognition by moving beyond explicit sampling in supervised contrastive learning. It introduces Probabilistic Contrastive Learning (ProCo), which models normalized features on the unit sphere with a mixture of von Mises–Fisher distributions and derives a closed-form, infinite-sample loss that eliminates the need for large batches. The authors provide theoretical generalization and excess-risk bounds, extend ProCo to semi-supervised learning, and demonstrate strong empirical gains across LT and balanced classification, as well as LT object detection. The approach yields robust improvements in tail-class performance while remaining compatible with existing LT methods and scalable to large datasets. The key practical impact is a batch-efficient, distribution-aware contrastive framework that generalizes across tasks and data regimes.

Abstract

Long-tailed distributions frequently emerge in real-world data, where a large number of minority categories contain a limited number of samples. Such imbalance issue considerably impairs the performance of standard supervised learning algorithms, which are mainly designed for balanced training sets. Recent investigations have revealed that supervised contrastive learning exhibits promising potential in alleviating the data imbalance. However, the performance of supervised contrastive learning is plagued by an inherent challenge: it necessitates sufficiently large batches of training data to construct contrastive pairs that cover all categories, yet this requirement is difficult to meet in the context of class-imbalanced data. To overcome this obstacle, we propose a novel probabilistic contrastive (ProCo) learning algorithm that estimates the data distribution of the samples from each class in the feature space, and samples contrastive pairs accordingly. In fact, estimating the distributions of all classes using features in a small batch, particularly for imbalanced data, is not feasible. Our key idea is to introduce a reasonable and simple assumption that the normalized features in contrastive learning follow a mixture of von Mises-Fisher (vMF) distributions on unit space, which brings two-fold benefits. First, the distribution parameters can be estimated using only the first sample moment, which can be efficiently computed in an online manner across different batches. Second, based on the estimated distribution, the vMF distribution allows us to sample an infinite number of contrastive pairs and derive a closed form of the expected contrastive loss for efficient optimization. Our code is available at https://github.com/LeapLabTHU/ProCo.

Figures (3)

• Figure 1: Illustration of Probabilistic Contrastive Learning. ProCo estimates the distribution of samples based on the features from different batches and samples contrastive pairs from it. Moreover, a closed form of expected contrastive loss is derived by sampling an infinite number of contrastive pairs, which eliminates the inherent limitation of SCL on large batch sizes.
• Figure 2: Performance of the representation branch. We train the model for 200 epochs.
• Figure 3: Parameter analysis of the loss weight ratio $\alpha$. LA denotes the logit adjustment method.

Theorems & Definitions (11)

• Proposition 1
• proof
• Proposition 2: Generalization Error Bound
• Proposition 3: Excess Risk Bound
• Lemma 1: Asymptotic expansion
• proof
• Lemma 2: Bennett's inequality Maurer2009
• Lemma 3: Variance inequality
• proof
• proof