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

Konstantinos Almpanakis; Anna Kreshuk

cs.CVOct 14, 2026

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

Authors: Konstantinos Almpanakis, Anna Kreshuk

Organizations: European Molecular Biology Laboratory (EMBL), Heidelberg, Germany

Abstract

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.

Figures & tables

Figure 1: Minibatch and OOB regularization. (a) In-batch methods estimate the anti-collapse term from active samples at each step. (b) OOB regularization periodically encodes the dataset, constructs detached per-instance anchors, and retrieves the anchors indexed by the online minibatch.

	Refresh phase
1	Run inference in minibatches over the entire dataset, collecting one view representation per sample.
2	After the full inference pass, construct $a_{i}^{(r)}$ jointly from all $N$ representations.
3	Detach $A^{(r)}$ and store it by sample index in CPU memory.
	Online phase, repeated until the next refresh
4	Sample $\mathcal{B}$ and generate $V$ augmented views per sample.
5	Retrieve $\{a_{i}^{(r)}:i\in\mathcal{B}\}$ and evaluate eq. 4 .

Table 1: OOB training separates global target construction from online parameter updates.

Figure 2: Minibatch and cached gradient agreement with the full-dataset reference.

Figure 3: Image classification results using linear probes and 5-NN.

Benchmark	5-NN	Linear	5-NN	Linear	5-NN	Linear	5-NN	Linear
	SIGReg		SIGReg-OOB (Ours)		SimCLR		Contrastive-OOB (Ours)
Natural images	34.63	45.55	48.80 ( +14.17 )	58.44 ( +12.89 )	50.21	59.35	52.50 ( +2.29 )	59.49 ( +0.14 )
2D biomedical	57.45	67.44	62.21 ( +4.76 )	72.84 ( +5.40 )	50.00	63.55	60.83 ( +10.83 )	70.13 ( +6.58 )
3D biomedical	51.39	55.89	55.61 ( +4.22 )	63.95 ( +8.06 )	51.76	57.52	55.47 ( +3.71 )	61.78 ( +4.26 )

Table 2: Classification balanced accuracy (%). Parentheses show OOB gains in percentage points.

Figure 4: Empirical and class-balanced gradient agreement under class imbalance.

Figure 5: Minority-class recall under controlled class imbalance.

Figure 6: Instance-retrieval mAP and example nearest neighbors.

			SIGReg-OOB (Ours)			SIGReg			Contrastive-OOB (Ours)			SimCLR
Dataset	Encoder	Original	$B=2$	$B=8$	$B=64$	$B=2$	$B=8$	$B=64$	$B=2$	$B=8$	$B=64$	$B=2$	$B=8$	$B=64$
BloodMNIST	S (22M)	92.0	+6.6	+6.7	+6.8	-6.5	+6.7	+6.0	+6.1	+6.2	+5.8	+1.6	+6.7	+6.5
	B (86M)	90.1	+8.8	+8.6	+8.6	-5.7	+8.7	+8.4	+8.0	+8.4	+8.1	+4.1	+8.7	+8.7
	L (304M)	89.4	+9.3	+9.4	+9.3	-4.6	+9.4	+9.0	+9.1	+8.9	+9.5	+2.5	+9.4	+9.2
	G (1.1B)	93.0	+5.3	+5.1	+5.8	-6.8	+5.8	+5.3	+2.8	-5.6	+3.7	+3.6	+5.5	+5.7
OrganAMNIST	S (22M)	85.7	+7.1	+6.1	+5.9	+0.9	+8.0	+4.1	+8.7	+8.2	+8.1	-15.9	+4.6	+4.7

Table 3: Continued pretraining of DINOv2 and 5-NN of adapted models relative to original.

Figure 7: Best classification performance under GPU-memory budgets.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Collection	Domains
Natural images	CIFAR-10, CIFAR-100, DTD, ImageNet-100, Imagenette, Imagewoof, STL-10
2D Biomedical	BloodMNIST, BreastMNIST, DermaMNIST, OCTMNIST, OrganAMNIST, OrganCMNIST, OrganSMNIST, PathMNIST, PneumoniaMNIST, TissueMNIST, BBBC010, plankton(private)
3D Biomedical	AdrenalMNIST, FractureMNIST, NoduleMNIST, OrganMNIST, SynapseMNIST, VesselMNIST

Appendix

Table 4: Datasets used in the classification and oracle experiments.

Modality	Spatial transforms	Appearance and normalization
2D	Random resized crop to $224$ (scale $0.2$ – $1.0$ ); horizontal flip	Color jitter $(0.4,0.4,0.4,0.1)$ , $p=0.8$ ; grayscale, $p=0.2$ ; Gaussian blur (kernel 23, $\sigma\in[0.1,2.0]$ ); ImageNet normalization
3D	Random resized crop to $64^{3}$ (scale $0.5$ – $1.0$ ); independent axis flips, $p=0.5$ ; random $90^{\circ}$ rotation in one plane, $p=0.5$	Intensity shift in $[-0.1,0.1]$ ; contrast in $[0.8,1.2]$ ; Gaussian noise ( $\sigma=0.1$ ), $p=0.3$ ; $3^{3}$ average-pool blur, $p=0.3$

Appendix

Table 5: Augmentation pipelines used for self-supervised training.

Method	Configuration
SIGReg-OOB	Three-layer MLP projector with 2048-dimensional hidden and output layers, LayerNorm and ReLU after each non-final layer; $M=256$ projections; $(\lambda_{\mathrm{align}},\lambda_{\mathrm{ac}})=(5,1)$ ; projection directions resampled at every refresh.
SIGReg	The same 2048-dimensional projector
Contrastive-OOB	Two-layer ReLU projector with one backbone-width hidden layer and 128-dimensional output anchors; $(\lambda_{\mathrm{align}},\lambda_{\mathrm{ac}})=(1,1)$ ; 20 anchor updates with Adam at learning rate 0.05; sampled anchor subset 256; $\tau_{u}=2$ , $\lambda_{f}=0.1$ , and previous-cache mixture weight 0.5.
SimCLR	The same two-layer 128-dimensional projector, normalized output, four augmented views in the classification experiments, and temperature 0.5.
MoCoV2	Two-layer 128-dimensional query/key projectors, four augmented views, queue size 4096, encoder momentum 0.999, and temperature 0.2.

Appendix

Table 6: Method-specific settings used in the principal experiments.

Figure 8: Controlled ablations of SIGReg-OOB. Each factor varies within a matched slice, and bars report mean linear-probe balanced accuracy over the dataset–batch cells shared by all values in that group.

Figure 9: Robustness of the oracle comparison to dataset size. The MedMNIST sweep uses 500, 1k, 2k, and 5k training examples. Panels A and C evaluate the minibatch estimator, while B and D follow global-cache lifetime at epochs 50 and 100, respectively.

Figure 10: Frozen classification across MedMNIST training-set sizes. Shading denotes one standard deviation across training seeds.

Dataset	SIGReg	SIGReg- [-1pt]OOB	SimCLR	MoCoV2	Contrastive [-1pt]-OOB	VICReg	IConE	IConE- [-1pt]SIGReg	BYOL	SwAV	DINO
BloodMNIST	82.4	84.4	68.9	68.1	78.2	76.3	81.1	80.8	72.9	75.7	72.6
BreastMNIST	71.3	72.5	66.6	65.2	72.0	70.9	68.8	68.5	67.8	68.1	65.5
DermaMNIST	54.7	60.0	50.9	48.6	59.5	52.4	45.2	45.2	49.1	49.3	46.3
OCTMNIST	48.2	52.6	41.2	39.3	41.8	48.9	42.9	44.2	41.5	43.8	38.2
OrganAMNIST	61.0	70.7	54.3	52.8	66.9	61.0	73.1	72.6	55.2	56.8	52.5
OrganCMNIST	58.9	65.6	51.2	47.8	63.5	64.3	70.9	71.5	50.2	51.6	53.0

Appendix

Table 7: Per-dataset 2D MedMNIST balanced accuracy (%) across training-set sizes, averaged over batch sizes. Entries are linear probe results. Within each method family, the best result is bold; the best result overall is underlined.

Dataset	SIGReg	SIGReg-OOB	SimCLR	MoCoV2	Contrastive [-1pt]-OOB
BBBC010	90.9 / 93.8	90.4 / 91.5	89.9 / 90.9	90.8 / 90.9	90.0 / 91.2
CIFAR-10	49.2 / 60.0	79.7 / 82.2	76.2 / 81.1	68.2 / 73.8	74.9 / 79.4
CIFAR-100	20.7 / 33.5	42.2 / 53.9	39.1 / 53.1	34.6 / 46.9	33.4 / 49.4
DTD	17.5 / 25.2	26.0 / 32.7	26.5 / 35.8	22.3 / 29.1	32.4 / 37.0
ImageNet-100 (full)	44.3 / 53.9	52.4 / 60.6	– / –	– / –	– / –
ImageNet-100-5k	39.3 / 54.4	55.9 / 67.8	53.8 / 65.3	49.2 / 63.2	63.9 / 69.5

Appendix

Table 8: Per-dataset balanced accuracy (%) on natural/scientific 2D and 3D MedMNIST datasets, averaged over batch sizes. Entries are 5-NN / linear probe. Within each method family, the best result is bold; the best result overall is underlined.

Figure 11: Effective rank and RankMe of frozen representations. Lines average over the common datasets and bands show standard errors across datasets.

Figure 12: Minority-class recall for the controlled-imbalance 2D experiments at batch size 64: 5-NN (left) and linear probe (right).

Figure 13: Worst-class recall and minority–majority recall gap for Contrastive-OOB, SimCLR, and MoCoV2. Lines average dataset–batch experiments and bands show standard errors across those results.

Figure 14: Transductive and held-out retrieval for mAP and Recall@

1/5/10

. For each seed, scores are first averaged over the three datasets; lines then show the three-seed mean and bands show one standard deviation.

Protocol	Metric	SIGReg	SIGReg-OOB	SimCLR	Contrastive-OOB	MoCoV2
Transductive	mAP	49.1	62.2	57.3	71.0	63.4
	R@1	59.5	70.8	66.1	76.5	71.0
	R@5	70.1	79.0	76.8	83.5	80.6
	R@10	74.2	82.2	80.8	86.1	84.3
Held-out	mAP	60.6	67.4	66.7	79.7	72.8
	R@1	67.1	73.1	71.7	83.9	76.7

Appendix

Table 9: Retrieval performance (%), averaged over the three datasets, batches 2–64, and three training seeds within each protocol. Bold denotes an OOB result that exceeds its corresponding in-batch comparator: SIGReg-OOB versus SIGReg, or Contrastive-OOB versus both SimCLR and MoCoV2.

Figure 15: Per-dataset DINOv2 continued-pretraining performance on biomedical datasets for SIGReg and SIGReg-OOB. Solid lines show linear-probe accuracy and dashed lines show 5-NN accuracy. Each point averages DINOv2-S/B/L/G encoders; bands show standard errors across encoder sizes.

Figure 16: Detailed resource measurements across dataset sizes and batches. Lines average the datasets evaluated for each method, and bands show standard errors across datasets.

Figure 17: OOB improvement over the corresponding in-batch method during pretraining. Solid lines show linear-probe accuracy and dashed lines show 5-NN accuracy. Lines average matched dataset–batch cells and bands show standard errors across those cells.

Comparison	Probe	Epoch 25	Epoch 50	Epoch 100	Epoch 200
SIGReg-OOB $-$ SIGReg	5-NN	+9.09	+11.94	+13.88	+14.43
SIGReg-OOB $-$ SIGReg	Linear	+6.91	+9.14	+10.80	+11.62
Contrastive-OOB $-$ SimCLR	5-NN	+1.58	+3.68	+3.71	+2.43
Contrastive-OOB $-$ SimCLR	Linear	+0.89	+1.73	+1.71	+0.47

Appendix

Table 10: Mean OOB gain in percentage points at selected epochs.

Explore similar work

$λ$ -JEPA Spectral Anti-Collapse Regularization for Self-Supervised Learning
Sep 28, 2026Berker Demirel, Clémentine Dominé, Valentino Maiorca +3Spectral Regularization Joint-Embedding Predictive Architecture
Institute of Science and Technology Austria (ISTA) Am Campus 1, 3400 Klosterneuburg, Austria
VISReg: Variance-Invariance-Sketching Regularization for JEPA training
Jun 1, 2026Haiyu Wu, Randall Balestriero, Morgan LevineRepresentation Collapse Joint-Embedding Predictive Architecture
Altos Labs · Brown University
IConE: Batch Independent Collapse Prevention for Self-Supervised Representation Learning
Mar 16, 2026Konstantinos Almpanakis, Anna KreshukSelf-Supervised Learning Representation Collapse
European Molecular Biology Laboratory, Heidelberg, Germany

Sep 28, 2026J/K move · Enter open · S save

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

Berker Demirel, Clémentine Dominé, Valentino Maiorca, Marco Fumero+2

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

VISReg: Variance-Invariance-Sketching Regularization for JEPA training

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

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