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

Konstantinos Almpanakis; Anna Kreshuk

cs.CVMar 16, 2026

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

Authors: Konstantinos Almpanakis, Anna Kreshuk

Organizations: European Molecular Biology Laboratory, Heidelberg, Germany

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.

Figures & tables

Figure 1: Decoupling invariance and anti-collapse gradients with IConE. Left: Standard JEAs apply both the multi-view invariance objective ( $\mathcal{L}_{\text{invariance}}$ ) and the anti-collapse mechanism ( $\mathcal{L}_{\text{regularization}}$ ) directly within the encoder’s dynamic representation space $\vec{Z}$ . This couples the anti-collapse constraint to the current batch. Right: IConE transfers the anti-collapse mechanism to an explicitly different parameter space. A persistent embedding table yields instance-specific anchors $\vec{E}$ , which are independently structured via $\mathcal{L}_{\text{regularization}}$ . The encoder $f$ is optimized purely through attractive forces: local view consistency ( $\mathcal{L}_{\text{invariance}}$ ) and alignment to its regularized target ( $\mathcal{L}_{\text{alignment}}$ ). This structural decoupling removes the need for batch-dependent negative sampling or variance statistics, applying the repulsive force to the entire dataset at the same time.

Figure 2: Batch-size stability. Mean and standard deviation linear-probe top-1 balanced accuracy aggregated over (A) 2D datasets and (B) 3D datasets.

	$B{=}4$			$B{=}16$			$B{=}64$
Method	kNN-1	kNN-5	kNN-20	kNN-1	kNN-5	kNN-20	kNN-1	kNN-5	kNN-20
BYOL	54.6	53.7	49.6	58.1	56.3	49.7	58.1	55.3	53.0
DIET	55.6	55.6	54.7	55.3	55.3	55.3	57.1	57.0	54.4
DINO	55.8	54.6	50.5	57.7	54.9	51.9	57.4	56.6	54.6
MoCo-v2	50.8	51.1	49.7	56.7	56.2	55.4	57.2	56.3	53.3
SimCLR	49.9	48.4	47.7	52.0	51.6	51.8	53.1	54.5	53.9

Table 1: Self-supervised k-NN performance. Mean k-NN balanced accuracy (%) across 6 3D datasets. Best results in bold .

Figure 3: Batch sensitivity analysis. Left: Performance drop from largest to smallest batch size (lower is better). Right: Pearson Correlation between batch size and performance versus absolute top-1 balanced accuracy.

Figure 4: UMAP visualization of learned representations. Embeddings for OrganMNIST3D across batch sizes (rows) and methods (columns). IConE (leftmost column, green border) maintains well-separated class clusters at all batch sizes, while batch-dependent methods show progressive collapse as batch size decreases.

Figure 5: Representation geometry analysis. (A) Alignment vs. uniformity loss for 3D datasets, colored by downstream accuracy. (B) Accuracy vs. effective rank (3D). (C) RankMe vs. batch size (2D). (D) LiDAR vs. batch size (2D). IConE achieves superior geometry metrics across all measures.

Figure 6: Robustness to class imbalance. (A) Overall top-1 balanced accuracy vs. imbalance ratio. (B) Minority class accuracy vs. imbalance ratio. IConE maintains superior performance on both metrics, with particularly strong advantages on minority classes at high imbalance ratios. Results are averaged over all 2D and 3D datasets.

	2D						3D
Method	$B{=}1$	$B{=}4$	$B{=}8$	$B{=}16$	$B{=}32$	$B{=}64$	$B{=}1$	$B{=}4$	$B{=}8$	$B{=}16$	$B{=}32$	$B{=}64$
Cross-Entropy	76.20	78.17	79.72	79.87	80.67	80.32	66.61	72.70	71.81	69.65	73.03	72.85
SupCon	72.20	74.08	76.32	79.70	79.93	79.73	59.49	66.50	72.97	72.59	67.99	71.55
IConE-Instance	76.81	78.43	78.56	79.53	79.58	79.63	68.77	75.16	73.14	76.56	75.10	75.31
IConE-Class	78.71	82.42	81.97	82.16	82.28	82.31	69.59	74.25	72.13	77.14	74.34	76.21

Table 2: Supervised learning performance. Aggregated top-1 balanced accuracy (%) across batch sizes for 2D and 3D datasets. Best in bold , second best underlined .

	2D						3D
Regularizer	$B{=}1$	$B{=}2$	$B{=}4$	$B{=}8$	$B{=}16$	$B{=}32$	$B{=}1$	$B{=}2$	$B{=}4$	$B{=}8$	$B{=}16$	$B{=}32$
IConE (default)	59.6	59.5	60.6	60.2	62.0	60.1	62.1	67.9	69.7	66.7	68.8	65.4
IConE-VCReg	59.3	60.2	60.8	60.9	60.4	61.2	63.0	66.1	68.8	69.2	69.1	67.9
IConE-SIGReg	57.3	58.4	58.3	58.5	58.5	58.4	62.9	67.5	68.4	68.9	67.5	68.0

Table 3: Regularizer ablation. Top-1 accuracy (%) for IConE with default orthogonality, VCReg, SIGReg regularization on 2D and 3D datasets across batch sizes.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	2D	3D
Optimizer	AdamW	AdamW
Base learning rate	$1\times 10^{-4}$	$1\times 10^{-4}$
Weight decay	0.05	0.05
LR schedule	Cosine annealing	Cosine annealing
Training epochs	100	100
Number of views $V$	2	2

Appendix

Table 4: Optimization hyperparameters. IConE uses identical settings across 2D and 3D, with the only difference being embedding dimension (determined by the backbone).

Figure 7: Per-dataset 2D results. Linear probe balanced accuracy vs. batch size for each 2D MedMNIST dataset. In most datasets, IConE maintains consistent performance across all batch sizes, while other baselines show characteristic degradation patterns at small batch sizes.

Figure 8: Per-dataset 3D results. Linear probe balanced accuracy vs. batch size for each 3D MedMNIST dataset. IConE’s batch-agnostic design provides consistent advantages across all volumetric imaging tasks.

Figure 9: Joint batch size and dataset size analysis. Balanced accuracy heatmaps for dataset sizes $N\in\{500,1\text{k},2\text{k},5\text{k}\}$ . IConE maintains stable high performance (dark blue) across all conditions, while baselines degrade significantly at small batch sizes, particularly in data-scarce regimes.

Views ( $V$ )	$1$	$2$	$4$	$8$	$16$	$32$	Average
	Batch Size ( $B$ )
$2$ (default)	73.4	74.7	76.1	75.8	75.5	75.3	75.1
$4$	72.6	74.8	76.1	76.1	76.0	75.4	75.1
$8$	72.8	74.6	76.1	75.9	76.1	75.4	75.2
$12$	72.1	74.5	75.9	75.8	75.2	75.2	74.8

Appendix

Table 5: Number of views ablation. Top-1 Balanced accuracy (%) on 2D datasets across various batch sizes.

V=2

provides strong average performance with minimal memory overhead.

Init scale ( $\sigma$ )	$1$	$2$	$4$	$8$	$16$	$32$	Average
	Batch Size ( $B$ )
$0.02$ (default)	74.3	74.2	75.8	75.5	75.7	76.0	75.2
$0.05$	72.3	75.5	76.0	76.2	75.1	76.1	75.2
$0.10$	71.3	75.4	75.6	75.9	75.7	75.9	75.0
$0.20$	70.9	74.9	76.8	75.3	76.0	75.5	74.9
$0.50$	72.2	75.6	76.8	76.3	75.0	75.0	75.1

Appendix

Table 6: Initialization scale ablation. Balanced accuracy (%) on 2D datasets across various batch sizes. Performance remains robust across initialization scales spanning an order of magnitude.

Figure 10: Peak GPU memory (GB) across batch sizes and dataset sizes for all evaluated SSL methods on BloodMNIST. Each panel corresponds to a different batch size. IConE-SIGreg and IConE-VCReg match the memory footprint of competitive baselines across all settings, whereas the default IConE variant incurs increasing overhead as

N

grows.

Figure 11: Wall-clock training time (hours) across batch sizes and dataset sizes for all evaluated SSL methods on BloodMNIST. The default IConE variant exhibits super-linear growth in training time with dataset size due to its

\mathcal{O}(N^{2})

diversity regularisation, while IConE-SIGreg and IConE-VCReg scale comparably to the remaining baselines.

	$N{=}5$ k			$N{=}10$ k			$N{=}50$ k			$N{=}100$ k
Method	$B{=}4$	$B{=}16$	$B{=}64$	$B{=}4$	$B{=}16$	$B{=}64$	$B{=}4$	$B{=}16$	$B{=}64$	$B{=}4$	$B{=}16$	$B{=}64$
IConE	2.2	3.7	10.0	6.6	8.3	15.0	158.2	164.8	192.5	637.2	660.5	759.2
IConE-VCReg	1.0	2.5	8.9	1.0	2.6	8.9	1.3	2.9	9.2	1.7	3.2	9.6
IConE-SIGReg	1.3	2.9	9.2	1.6	3.2	9.5	3.9	5.5	11.8	6.9	8.4	14.7

Appendix

Table 7: Peak GPU memory scaling. Algorithmic peak GPU memory (GB) for

224^{2}

inputs at dataset size

N

and embedding dimension

d{=}384

. Lowest per column in bold .

Figure 12: Representation quality metrics on 2D MedMNIST datasets. Each row shows a different dataset; each column shows a different metric as a function of batch size. Higher is better for RankMe, LiDAR, and Effective Rank (indicating better dimensional utilization). Lower is better for

\mathcal{L}_{\text{align}}

(tighter same-class clusters) and

\mathcal{L}_{\text{uniform}}

(more uniform coverage). IConE (green) consistently achieves the best collapse-related metrics while maintaining competitive alignment and uniformity, especially at small batch sizes where baseline methods exhibit severe collapse.

Figure 13: Representation quality metrics on 3D MedMNIST datasets. Same format as Fig. 12 . IConE’s advantages are even more pronounced in 3D, where the increased computational cost of volumetric data makes small-batch training essential. At

B=1

and

B=2

, baseline methods show near-complete dimensional collapse (low RankMe, LiDAR, Effective Rank), while IConE maintains high-quality representations with good dimensional utilization.

Figure 14: Natural image evaluation. Linear-probe top-1 balanced accuracy across batch sizes on a 20-class ImageNet subset (250 images per class). IConE retains its batch-size stability outside the biomedical domain.

Figure 15: Final representations for each ablation variant. (a) Original 2D Gaussian mixture data with 5 overlapping classes. (b) Full IConE achieves well-separated clusters on the unit circle (88% kNN, 88% linear). (c) Without

\mathcal{L}_{div}

, representations collapse into a small arc (57% kNN). (d) Without

\mathcal{L}_{vi}

, no learning occurs (39% kNN, near random). (e) Without

\mathcal{L}_{vv}

, separation is degraded (72% kNN).

Variant	$\mathcal{L}_{vv}$	$\mathcal{L}_{vi}$	$\mathcal{L}_{div}$
Full IConE	✓	✓	✓
No $\mathcal{L}_{div}$	✓	✓	$\times$
No $\mathcal{L}_{vi}$	✓	$\times$	✓
No $\mathcal{L}_{vv}$	$\times$	✓	✓

Appendix

Table 8: Ablation variants configuration.

Figure 16: Training evolution across all ablation variants. Each row shows a different variant; each column shows representations at a specific epoch (0, 5, 25, 100, 200, 299). Points are colored by class label.

Figure 17: Classification accuracy during training. (a) kNN-5 accuracy. (b) Linear probe accuracy. Full IConE maintains stable high performance, while ablated variants show degraded or unstable learning.

Figure 18: Alignment and uniformity metrics during training. (a) Alignment loss (lower = tighter same-class clusters). (b) Uniformity loss (lower = more spread). Full IConE achieves the best trade-off with both low alignment and low uniformity. Without

\mathcal{L}_{vv}

, alignment suffers; without

\mathcal{L}_{div}

, uniformity suffers.

Variant	kNN-5 (%) $\uparrow$	Linear (%) $\uparrow$	$\mathcal{L}_{\text{align}}$ $\downarrow$	$\mathcal{L}_{\text{uniform}}$ $\downarrow$	Silhouette $\uparrow$
Full IConE	87.9	87.9	0.281	-1.389	0.475
No $\mathcal{L}_{div}$	57.4	59.6	0.234	-0.999	0.071
No $\mathcal{L}_{vi}$	39.3	53.1	0.000 †	-0.000 †	0.015
No $\mathcal{L}_{vv}$	72.3	72.8	0.599	-1.384	0.255

Appendix

Table 9: Ablation study results on synthetic 2D data. Best values are shown in bold , second best are underlined .

\dagger

indicates degenerate values due to representation collapse.

Explore similar work

Oct 14, 2026

Dataset-Level Regularization for Joint-Embedding Self-Supervised Learning

Self-supervised joint embedding learning aligns representations of augmented views while regularizing their distribution to prevent collapse. Regularization aims to shape the overall representation distribution, but most joint embedding methods restrict its computation to the active minibatch. We propose out-of-batch regularization, which maintains a dataset-level state in CPU memory and uses globally constructed per-instance targets during minibatch training. Evaluated on natural images and biomedical datasets, these anchors approximate the full-dataset regularization gradient more faithfully than conventional minibatch estimation at small batch sizes. Applications to contrastive (SimCLR) and non-contrastive (SIGReg) approaches show consistent improvement over the corresponding in-batch objectives for image classification, minority-class performance under controlled imbalance, instance retrieval, and continued self-supervised pretraining, with the largest gains at small batch sizes. These results suggest that dataset-level regularization can be decoupled from minibatch optimization, making batch size an optimization choice that does not affect the geometry of the regularizer.

Konstantinos Almpanakis, Anna Kreshuk

European Molecular Biology Laboratory (EMBL), Heidelberg, Germany

Self-Supervised Visual Representation Learning Representation Learning Contrastive Learning

Ask PDF

Apr 11, 2026

A Minimal Model of Representation Collapse: Frustration, Stop-Gradient, and Dynamics

Self-supervised representation learning is central to modern machine learning because it extracts structured latent features from unlabeled data and enables robust transfer across tasks and domains. However, it can suffer from representation collapse, a widely observed failure mode in which embeddings lose discriminative structure and distinct inputs become indistinguishable. To understand the mechanisms that drive collapse and the ingredients that prevent it, we introduce a minimal embedding-only model whose gradient-flow dynamics and fixed points can be analyzed in closed form, using a classification-representation setting as a concrete playground where collapse is directly quantified through the contraction of label-embedding geometry. We illustrate that the model does not collapse when the data are perfectly classifiable, while a small fraction of frustrated samples that cannot be classified consistently induces collapse through an additional slow time scale that follows the early performance gain. Within the same framework, we examine collapse prevention by adding a shared projection head and applying stop-gradient at the level of the training dynamics. We analyze the resulting fixed points and develop a dynamical mean-field style self-consistency description, showing that stop-gradient enables non-collapsed solutions and stabilizes finite class separation under frustration. We further verify empirically that the same qualitative dynamics and collapse-prevention effects appear in a linear teacher-student model, indicating that the minimal theory captures features that persist beyond the pure embedding setting.

Louie Hong Yao, Yuhao Li, Shengchao Liu

Center for Soft Matter and Biological Physics, Virginia Tech, Blacksburg, VA 24061, USA · Department of Computer Science and Engineering, The Chinese University of Hong Kong, Hong Kong, China

Representation Collapse Self-Supervised Learning

Ask PDF

Sep 28, 2026

$λ$ -JEPA Spectral Anti-Collapse Regularization for Self-Supervised Learning

Joint-embedding self-supervised learning typically combines an invariance objective across augmented views with additional mechanisms to prevent representational collapse. These objectives are often applied after a projection head, while downstream tasks use the backbone representation before the projector. We find that this mismatch does not necessarily prevent dimensional collapse in the backbone, which can retain low effective rank and potentially limit downstream transfer. To address this, we introduce SACReg, a spectral anti-collapse regularizer motivated by an analysis of

λ

-balance, which captures the relative scale of weight matrices across layers. In a two-layer linear network, we show that (i)

λ

-balance prevents collapse, and (ii) our regularizer applied to the backbone induces

λ

-balance. In the nonlinear case, this regularizer leads to anti-collapse as well and, in realistic architectures on ImageNet100, it empirically increases the representations' ranks. We apply SACReg to JEPA and propose

λ

-JEPA, which improves over LeJEPA and VISReg on ImageNet-1k classification and in average linear-probe transfer performance across eight downstream image datasets. On video self-supervised learning,

λ

-JEPA improves over LeVJEPA and V-JEPA 2 on the Something-Something-v2 and Kinetics-400 benchmarks. Code is available at https://github.com/berkerdemirel/lambda-jepa.

Berker Demirel, Clémentine Dominé, Valentino Maiorca +3

Institute of Science and Technology Austria (ISTA) Am Campus 1, 3400 Klosterneuburg, Austria

Spectral Regularization Joint-Embedding Predictive Architecture Video Representation Learning Self-Supervised Visual Representation Learning

Ask PDF

Abstract

Explore similar work

Dataset-Level Regularization for Joint-Embedding Self-Supervised Learning

A Minimal Model of Representation Collapse: Frustration, Stop-Gradient, and Dynamics

λλλ-JEPA Spectral Anti-Collapse Regularization for Self-Supervised Learning

$λ$ -JEPA Spectral Anti-Collapse Regularization for Self-Supervised Learning