The Details Matter: Preventing Class Collapse in Supervised Contrastive Learning

Supervised contrastive learning optimizes a loss that pushes together embeddings of points from the same class while pulling apart embeddings of points from different classes. Class collapse—when every point from the same class has the same embedding—minimizes this loss but loses critical information that is not encoded in the class labels. For instance, the “cat” label does not capture unlabeled categories such as breeds, poses, or backgrounds (which we call “strata”). As a result, class collapse produces embeddings that are less useful for downstream applications such as transfer learning and achieves sub-optimal generalization error when there are strata. We explore a simple modification to supervised contrastive loss that prevents class collapse by uniformly pulling apart individual points from the same class. More importantly, we introduce a theoretical framing to analyze this loss through a view of how it embeds strata of different sizes. We show that our loss maintains distinctions between strata in embedding space, even though it does not explicitly use strata labels. We empirically explore several downstream implications of this insight. Our loss produces embeddings that achieve lift on three downstream applications by distinguishing strata: 4.4 points on coarse-to-fine transfer learning, 2.5 points on worst-group robustness, and 1.0 points on minimal coreset construction. Our loss also produces more accurate models, with up to 4.0 points of lift across 9 tasks.