🤖 machine learning

Optimal Representations for Generalized Contrastive Learning with Imbalanced Datasets

This paper characterizes the geometry of optimal representations in contrastive learning under class imbalance, proving that while balanced classes exhibit Neural Collapse, imbalanced classes follow a proportion-dependent angular symmetry that can lead to "Minority Collapse" when imbalance exceeds a specific threshold.

Original authors: Thuan Nguyen, Shuchin Aeron, D. Richard Brown III, Prakash Ishwar

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC0 1.0

Original authors: Thuan Nguyen, Shuchin Aeron, D. Richard Brown III, Prakash Ishwar

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching a Robot to Sort a Messy Closet

Imagine you are teaching a robot to sort a huge pile of clothes. The robot's job is to learn what makes a "shirt" look like a shirt and a "pants" look like pants.

In Contrastive Learning (CL), the robot learns by playing a game:

The Anchor: The robot picks up one item (e.g., a blue shirt).
The Positive: It finds another item that is similar (another blue shirt). It tries to pull these two close together in its mind.
The Negatives: It looks at items that are different (a pair of jeans, a hat, a sock). It tries to push these far away from the blue shirt.

The goal is for the robot to create a mental map where similar things are clustered together and different things are far apart.

The Problem: The "Imbalanced" Closet

Most research assumes the closet is perfectly balanced: 100 shirts, 100 pants, 100 hats. But in the real world, closets are messy. You might have 1,000 shirts, but only 5 pairs of pants and 3 hats. These are imbalanced datasets.

The authors of this paper asked: What happens to the robot's mental map when the closet is heavily skewed? Does it still sort things correctly?

The Main Discovery: "Neural Collapse" and "Minority Collapse"

The paper proves two main things about how the robot organizes its mind when it tries to do the best possible job.

1. The "Perfect Clump" (Intra-class Variance Collapse)

When the robot learns perfectly, it stops treating every single shirt as unique. Instead, it realizes, "Hey, all these blue shirts are basically the same."

The Analogy: Imagine a flock of birds. Before learning, they fly around in a scattered cloud. After learning, every single bird in the "blue shirt" group flies to the exact same spot in the sky. They all collapse into a single point.
The Result: The paper proves that for any class (even the rare ones), the best way for the robot to learn is to make every item in that class identical to the "average" of that class.

2. The "Geometric Dance" (The Shape of the Clumps)

Once the robot has collapsed all the shirts into one point, all the pants into another, and all the hats into a third, how are these points arranged?

Balanced Closet: If you have equal numbers of shirts, pants, and hats, the points arrange themselves in a perfect, symmetrical shape (like an equilateral triangle or a perfect pyramid). This is called an Equiangular Tight Frame (ETF).
Imbalanced Closet: If you have way more shirts than pants, the symmetry breaks. The "shirt" point moves, and the "pants" and "hat" points shift to accommodate the weight of the shirts.
The Discovery: The paper provides a mathematical recipe (a convex optimization problem) to calculate exactly where these points should sit to be the most efficient, even when the numbers are uneven.

3. The "Minority Collapse" (The Dangerous Threshold)

This is the most dramatic finding. The authors looked at what happens when the imbalance is extreme.

The Scenario: Imagine you have 95% shirts, 2.5% pants, and 2.5% hats.
The Phenomenon: The paper proves that if the "minority" groups (pants and hats) become too small, they stop being distinct. The robot gets confused and decides, "Pants and hats are so rare, I'll just treat them as the exact same thing."
The Analogy: Imagine two small islands in a vast ocean. If the islands get too tiny, the water between them disappears, and they merge into a single, tiny speck. The robot collapses the "pants" point and the "hat" point into a single vector, often pointing in the exact opposite direction of the "shirt" point.
The Threshold: The authors calculated a specific "tipping point" (a threshold). If the majority class takes up more than about 93% of the data (for a specific type of learning algorithm), the minority classes will inevitably collapse into one.

How They Proved It

The authors didn't just guess; they used heavy math to prove this happens in the "ideal" scenario where the robot has infinite brainpower (an "Unconstrained Features Model").

Lower Bound: They created a "floor" for how bad the robot's performance could be. They proved that to hit this perfect floor, the robot must collapse the items into their class averages.
Convex Optimization: They showed that finding the perfect arrangement of these collapsed points is like solving a puzzle where there is only one correct solution, and you can find it using standard math tools.
Simulation: They ran computer experiments using real image data (from the CIFAR-10 dataset). They artificially created imbalanced datasets and watched the robot learn.
- Result: The computer results matched the math perfectly. The "minority" images physically collapsed into a single dot in the visualization, just as the theory predicted.

Summary of Key Takeaways

Imbalance changes the geometry: When data is unbalanced, the perfect arrangement of learned features is no longer a symmetrical shape; it's a distorted shape determined by the proportions of the data.
Everything collapses: To learn perfectly, the robot must make every item in a class look exactly like the average of that class.
Extreme imbalance is dangerous: If one class dominates too much (e.g., >93%), the rare classes lose their individuality and merge into a single point, making it impossible for the robot to distinguish between them.
It's solvable: Even though the geometry is complex, the paper provides a way to calculate the optimal arrangement using convex optimization.

The paper essentially maps out the "physics" of how machine learning models organize information when the world they are learning from is messy and unbalanced.

Technical Summary: Optimal Representations for Generalized Contrastive Learning with Imbalanced Datasets

Problem Statement

Contrastive Learning (CL) aims to learn representation maps that pull "similar" (positive) samples closer and push "different" (negative) samples apart. While the geometry of optimal representations is well-understood for balanced datasets—exhibiting Neural Collapse (NC) where class means form an Equiangular Tight Frame (ETF)—the behavior of CL under imbalanced class distributions remains largely uncharacterized. Existing theoretical works have primarily focused on:

Empirical versions of the InfoNCE loss with specific, near-maximal numbers of negative samples ( $k \approx$ batch size).
Balanced class settings where all classes have equal probability.
Classifier networks with explicit label-dependent losses (e.g., MSE, Cross-Entropy), rather than the general CL risk.

This paper addresses the gap in understanding the optimal representation geometry for generalized CL losses under arbitrary class imbalances, specifically investigating whether the ETF structure persists and identifying new phenomena like Minority Collapse (MC).

Methodology

The authors adopt an Unconstrained Features Model (UFM) framework, treating the final-layer feature vectors as free optimization variables rather than network weights. This allows for a tractable analysis of the global optimum.

Key Assumptions and Setup

Generalized Loss Function: The study considers a broad family of CL losses defined as $\ell_k = \psi(\dots)$ , where $\psi: \mathbb{R}^k \to \mathbb{R}$ is strictly convex and argument-wise strictly increasing. This subsumes popular losses like InfoNCE.
Sampling Distributions: The analysis covers both Unsupervised CL (UCL), where negatives are drawn from the entire dataset (allowing "class collisions"), and Supervised CL (SCL), where negatives are drawn from classes different from the anchor.
Class Imbalance: The paper models datasets where class probabilities $\lambda_i$ are not necessarily equal, including scenarios with extreme imbalance (one majority class and multiple equiprobable minority classes).
Optimization Strategy:
1. Lower Bound Construction: A tight lower bound $G(M)$ for the CL risk is derived, expressed solely in terms of the class mean vectors $M = [\mu_1, \dots, \mu_C]$ .
2. Convex Relaxation: The problem of minimizing the risk is transformed into minimizing a function $S(A)$ over the Gram matrix $A = M^\top M$ . The authors prove that $S(A)$ is strictly convex over the set of valid correlation matrices.
3. Analytical Characterization: The properties of the unique minimizer $A^*$ are analyzed to determine the geometric structure of the optimal class means.

Key Contributions

1. Generalized Lower Bound and Variance Collapse

The paper establishes a novel lower bound (Lemma 1) for the CL risk that holds for any strictly convex, argument-wise increasing loss function and any number of negative samples $k$ .

Intra-class Variance Collapse: The authors prove that to achieve this lower bound, the representation map must collapse all samples within a class to their class mean (i.e., $f(x) = \mu_{y(x)}$ ).
Unit Norm: The optimal class means must have unit Euclidean norm.

2. Geometry of Optimal Class Means

For representation dimensions $d \ge C-1$ :

Rank Deficiency: The unique optimal Gram matrix $A^*$ is proven to be rank-deficient with rank at most $C-1$ , even without explicit rank constraints.
Convex Computation: The optimal class means can be computed by solving a convex optimization problem over the Gram matrix, followed by spectral decomposition.
Symmetry Structure: Unlike the balanced case where means form an ETF, the optimal geometry for imbalanced classes exhibits an equiangular symmetry structure determined by relative class proportions. Specifically, if two classes have equal probability, their means maintain identical inner products with all other class means (Theorem 3).

3. Minority Collapse (MC) Phenomenon

The paper provides a rigorous theoretical proof of Minority Collapse in CL.

Definition: When class imbalance is extreme (specifically, when the majority class probability $\lambda_1$ exceeds a non-asymptotic threshold $\tau$ ), the optimal representations of all samples from the minority classes collapse into a single vector.
Threshold Derivation: A sufficient condition for MC is derived in terms of the number of classes $C$ , the number of negatives $k$ , and the Lipschitz properties of the loss function $\psi$ .
InfoNCE Specifics: For the InfoNCE loss, the paper derives parameter-free thresholds:
- $\tau \approx 0.9292$ for the UCL setting.
- $\tau \approx 0.9438$ for the SCL setting.
  These thresholds guarantee MC regardless of $C$ or $k$ .

4. Unification of Settings

The results are shown to hold for both UCL and SCL settings, with minor adjustments to the probability distributions, thereby unifying the theoretical understanding of contrastive learning across different sampling mechanisms.

Results and Validation

Theoretical Consistency: The derived lower bounds and geometric properties are mathematically proven to be necessary and sufficient for optimality under the UFM.
Numerical Experiments:
- Intra-class Collapse: Experiments on imbalanced CIFAR-10 subsets (3 classes) using ResNet-50 confirm that trained networks collapse intra-class variance to class means. The resulting Gram matrices match those computed via convex optimization (CVX).
- Minority Collapse: Experiments with extreme imbalance ( $\lambda_1 = 0.9$ ) demonstrate that minority class means collapse into a single vector diametrically opposite to the majority class mean, validating the theoretical MC phenomenon.
- Geometry: The experiments confirm that imbalanced classes do not form an equilateral triangle (ETF) but rather an isosceles triangle, consistent with the derived symmetry properties.

Significance and Claims

The paper claims to provide the first computable characterization of optimal CL representations under general class imbalance. Its significance lies in:

Generalizing Neural Collapse: Extending the NC phenomenon from balanced settings to imbalanced ones, showing that while intra-class collapse persists, the inter-class geometry (ETF) does not.
Explaining Minority Collapse: Providing a theoretical foundation for the empirically observed "minority collapse" in CL, identifying it as an optimal (though potentially undesirable) solution to the risk minimization problem under extreme imbalance.
Computational Efficiency: Demonstrating that optimal class means can be found via convex optimization, suggesting that high-dimensional representation spaces may be inefficient for learning optimal features in CL.
Robustness of Theory: Establishing that these phenomena are inherent to the geometry of the loss function and sampling distribution, rather than artifacts of specific network architectures or training dynamics.

The authors conclude by identifying open problems, including the characterization of optimal geometry when $d < C-1$ (relevant for Large Language Models) and the analysis of hard-negative sampling effects.