IConE: Batch Independent Collapse Prevention for Self-Supervised Representation Learning

TL;DR

IConE introduces a dataset-level embedding approach for self-supervised learning that prevents collapse independently of batch size, enabling effective training on small or imbalanced datasets, especially in biomedical applications.

Contribution

IConE decouples collapse prevention from batch size by using global auxiliary embeddings, improving stability and robustness in small-batch and imbalanced data regimes.

Findings

Outperforms contrastive and non-contrastive baselines in small-batch regimes

Maintains high intrinsic dimensionality in learned representations

Demonstrates robustness to severe class imbalance

Abstract

Self-supervised learning (SSL) has revolutionized representation learning, with Joint-Embedding Architectures (JEAs) emerging as an effective approach for capturing semantic features. Existing JEAs rely on implicit or explicit batch interaction -- via negative sampling or statistical regularization -- to prevent representation collapse. This reliance becomes problematic in regimes where batch sizes must be small, such as high-dimensional scientific data, where memory constraints and class imbalance make large, well-balanced batches infeasible. We introduce IConE (Instance-Contrasted Embeddings), a framework that decouples collapse prevention from the training batch size. Rather than enforcing diversity through batch statistics, IConE maintains a global set of learnable auxiliary instance embeddings regularized by an explicit diversity objective. This transfers the anti-collapse mechanism from the transient batch to a dataset-level embedding space, allowing stable training even when batch statistics are unreliable, down to batch size 1. Across diverse 2D and 3D biomedical modalities, IConE outperforms strong contrastive and non-contrastive baselines throughout the small-batch regime (from B=1 to B=64) and demonstrates marked robustness to severe class imbalance. Geometric analysis shows that IConE preserves high intrinsic dimensionality in the learned representations, preventing the collapse observed in existing JEAs as batch sizes shrink.