A Refined Generalization Analysis for Extreme Multi-class Supervised Contrastive Representation Learning

Contrastive Representation Learning (CRL) has achieved strong empirical success in multiple machine learning disciplines, yet its theoretical sample complexity remains poorly understood. Existing analyses usually assume that input tuples are identically and independently distributed, an assumption violated in most practical settings where contrastive tuples are constructed from a finite pool of labeled data, inducing dependencies among tuples. While one recent work analyzed this learning setting using U-Statistics to estimate the population risk, the techniques used therein require the risk of each class to concentrate uniformly, making excess risk bounds scale in the order of $\rho_{\min}^{-{1}/{2}}$ where $\rho_{\min}$ denotes the probability of the rarest class. Such a dependency can be overly pessimistic in the extreme multiclass settings where there are many tail classes which contribute minimally to the overall population risk. Our contributions are two-fold. Firstly, we improve upon the previous work and prove a bound with a sample complexity of the same order as the number of classes $R$, regardless of the distribution over classes. Furthermore, we formulate a different estimator that captures the concentration of the risk \textit{across classes}, enabling sharper bounds in extreme multi-class learning scenarios, especially where class distributions are long-tailed. Under mild assumptions on the class distributions, the resulting sample complexity is $\mathcal{{O}}(k)$ where $k$ is the number of samples per tuple.

Supervised contrastive representation learning is a machine learning framework in which a model is trained on labeled data to learn an embedding space where examples are organized according to class membership: 1. Representations of examples from the same class are encouraged to be close together. 2. Representations of examples from different classes are encouraged to be far apart. Rather than operating on individual examples, these models are typically trained using tuple-wise comparisons. In each training step, one constructs a tuple consisting of an anchor, a positive example from the same class, and k negative examples drawn from different classes. The learning objective then encourages the anchor to be closer to the positive than to the negatives in the representation space. The core research question is "how many labeled examples are needed in order for representation models to achieve reasonable generalization performance". We study this question in a multi-class contrastive learning problem with $R$ classes and $k$ negative samples per tuple and ask whether there exists an empirical learning objective that estimates the population contrastive risk with sample complexity scaling with $O(k)$ rather than $O(R)$. The answer we provide in this work is "yes, as long as you have enough small classes". We formulated this learning objective using tools in U-Statistics literature and derive sample complexity guarantees that scale with the tuple size rather than the total number of classes. This makes the theoretical bounds in this work especially competitive compared to prior works in extreme multi-class scenario (when $R$ is large) or in certain long-tailed class distributions.