Symmetric Neural-Collapse Representations with Supervised Contrastive Loss: The Impact of ReLU and Batching

TL;DR

This paper demonstrates that adding a ReLU activation at the final layer restores symmetry in supervised contrastive learning representations under class imbalance, with theoretical and experimental validation across various datasets and architectures.

Contribution

It provides the first geometric characterization of supervised contrastive loss under class imbalance and introduces batch-binding as an effective batch selection strategy.

Findings

ReLU activation restores symmetry in SCL representations under imbalance.

Batch-binding ensures invariant symmetric representations.

ReLU does not reduce test accuracy.

Abstract

Supervised contrastive loss (SCL) is a competitive and often superior alternative to the cross-entropy loss for classification. While prior studies have demonstrated that both losses yield symmetric training representations under balanced data, this symmetry breaks under class imbalances. This paper presents an intriguing discovery: the introduction of a ReLU activation at the final layer effectively restores the symmetry in SCL-learned representations. We arrive at this finding analytically, by establishing that the global minimizers of an unconstrained features model with SCL loss and entry-wise non-negativity constraints form an orthogonal frame. Extensive experiments conducted across various datasets, architectures, and imbalance scenarios corroborate our finding. Importantly, our experiments reveal that the inclusion of the ReLU activation restores symmetry without compromising test accuracy. This constitutes the first geometry characterization of SCL under imbalances. Additionally, our analysis and experiments underscore the pivotal role of batch selection strategies in representation geometry. By proving necessary and sufficient conditions for mini-batch choices that ensure invariant symmetric representations, we introduce batch-binding as an efficient strategy that guarantees these conditions hold.