Engineering the Neural Collapse Geometry of Supervised-Contrastive Loss
This work addresses the challenge of engineering embedding geometries for classification tasks, offering incremental improvements in representation learning for vision datasets.
The paper tackled the problem of controlling the geometry of feature embeddings learned by supervised-contrastive loss by introducing prototypes during training, resulting in embeddings that align with prototype geometry and establishing a connection to cross-entropy loss with a fixed classifier.
Supervised-contrastive loss (SCL) is an alternative to cross-entropy (CE) for classification tasks that makes use of similarities in the embedding space to allow for richer representations. In this work, we propose methods to engineer the geometry of these learnt feature embeddings by modifying the contrastive loss. In pursuit of adjusting the geometry we explore the impact of prototypes, fixed embeddings included during training to alter the final feature geometry. Specifically, through empirical findings, we demonstrate that the inclusion of prototypes in every batch induces the geometry of the learnt embeddings to align with that of the prototypes. We gain further insights by considering a limiting scenario where the number of prototypes far outnumber the original batch size. Through this, we establish a connection to cross-entropy (CE) loss with a fixed classifier and normalized embeddings. We validate our findings by conducting a series of experiments with deep neural networks on benchmark vision datasets.