PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Konstantinos Almpanakis, Anna Kreshuk

arXiv:2610.20158v1cs.CV

TL;DR

Most joint-embedding methods compute anti-collapse regularization on the active minibatch, which may be too small to represent the dataset. This paper constructs dataset-level per-instance targets for online training; across evaluated tasks, OOB variants improve over matched in-batch objectives, especially at small batch sizes.

  • Problem

    Most joint-embedding methods compute anti-collapse regularization from the active minibatch, which may be too small to represent the dataset.

    Source: §1 · p. 1, §3.1 · p. 3

  • Method

    OOB maintains a detached dataset-indexed state and periodically constructs offline per-instance targets that online minibatches retrieve during training.

    Source: §1 · pp. 1–2

  • Results

    Across natural images and biomedical data, OOB variants yield positive mean gains over matched in-batch objectives, largest at small batch sizes, and extend to minority-class recall, retrieval, and continued pretraining.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    The active minibatch need not define the population over which representation geometry is regularized.

    Source: §5 · p. 10

    OOB requires storing and periodically refreshing dataset-level state as the encoder evolves, and the trade-off at larger scales remains to be understood.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Most joint-embedding regularizers use the active minibatch as a proxy for the dataset, although small batches may poorly represent the population and larger batches can be constrained or suboptimal for optimization. The paper introduces dataset-level out-of-batch targets and evaluates their gradient estimation and performance against matched in-batch objectives.

1 Introduction

  • Small batches may poorly represent the dataset, while larger batches can be hardware-constrained or leave few parameter updates per epoch on small datasets.The batch size needed to estimate a population regularizer can therefore differ from the best batch size for optimization.Source: §1 · p. 1
  • OOB regularization moves the statistical context from the active batch to a dataset-level state, so anti-collapse support is no longer limited by batch size or GPU capacity.The state is periodically refreshed to construct globally optimized targets, and online steps retrieve targets for current samples.Source: §1 · pp. 1–2
  • The analysis measures minibatch error and cache staleness against a finite-population gradient, while separating gradient-estimation error from population reweighting under class imbalance.It covers 19 real datasets.Source: §1 · p. 2
  • Across natural images and 2D/3D biomedical data, OOB variants yield positive mean gains over matched in-batch objectives, with the largest gains at small batch sizes.Improvements extend to minority-class recall, instance retrieval, and continued pretraining.Source: §1 · p. 2

2 Related Work

Prior work prevents collapse through several regularization mechanisms, often assuming that the active batch represents the dataset. Other approaches broaden context with stored embeddings or assigned targets; OOB instead periodically constructs detached, per-instance targets from a population-level regularizer.

2 Related Work

3 Method

OOB separates dataset-wide regularization targets from minibatch encoder updates: global anchors provide population context through per-instance losses. The method instantiates this design with Gaussian and spherical targets and compares cached and minibatch gradients against a full-dataset reference.

3 Method

  • OOB regularization constructs dataset-level targets as per-instance anchors, separating population-level regularization from minibatch encoder updates.The framework instantiates Gaussian and spherical anchors and compares minibatch and cached estimates with a full-dataset gradient reference.Source: §3 · p. 3, §3.1 · p. 3, §3.5 · p. 5

3.1 Alignment and population regularization

  • In-batch regularization uses the active minibatch distribution, while OOB training periodically constructs dataset-level anchors for separable sample-to-anchor losses between refreshes.The empirical distributions give each included representation equal probability.Source: §3.1 · p. 3

3.2 A refresh-and-train interface

  • At refresh, the encoder processes one view of every training sample, constructs detached dataset-indexed anchors, and stores them in CPU memory; online steps fetch only minibatch rows.Gradients pass through the online representations, not the anchors, and the anti-collapse term requires no relationships among samples in the minibatch.Source: §3.2 · p. 4, Table 1 · p. 4

3.3 SIGREG-OOB: projected Gaussian anchors

  • SIGREG-OOB assigns each sample’s projected rank in the full dataset to a fixed Gaussian quantile grid, caching the resulting coordinates as targets.Projection directions are resampled at each refresh; the SIGREG baseline uses the same sorted-quantile targets computed within its batch.Source: §3.3 · p. 4

3.4 CONTRASTIVE-OOB: spherical anchors

  • CONTRASTIVE-OOB constructs spherical anchors by balancing global population spreading with preservation of anchor–sample correspondence.It uses sample-to-anchor attraction online, without in-batch negatives or a momentum encoder; anchor spreading is performed during construction.Source: §3.4 · pp. 4–5

3.5 Comparing minibatch and cached gradients

  • The gradient study compares minibatch and cached estimates with the same full-dataset representation-space gradient, measuring directional agreement by cosine similarity and scale-sensitive error by rRMSE.For SIGReg, the oracle ranks the current full dataset, minibatches rank within sampled batches, and cached assignments are compared with the current oracle as cache age grows. A label-aware sampled population separates estimator error from changes in population weighting.Source: §3.5 · p. 5

4 Results

Across natural and scientific/biomedical image datasets, the experiments evaluate OOB regularization on gradient agreement, classification, imbalance, retrieval, continued pretraining, and memory use. OOB methods generally improve on matched in-batch objectives, with the clearest classification gains at small batch sizes.

4.1 Datasets

  • The experiments span natural images, scientific and biomedical images, including 2D and 3D datasets.The dataset set includes CIFAR-10/100, DTD, STL-10, ImageNet-100, Imagenette, Imagewoof, BBBC010, MedMNIST, and a private plankton dataset.Source: §4.1 · p. 6

4.2 Oracle gradient analysis

  • A fresh cache exactly recovers the full-dataset gradient, while minibatch estimates can oppose it and cached agreement deteriorates as the encoder changes.Minibatch agreement improves with batch size; cache deterioration is slower at later checkpoints, so the refresh schedule should depend on training stage.Source: §4.2 · p. 6, Figure 2 · p. 6

4.3 Image classification across batch sizes

  • Averaged over small batch sizes, SIGREG-OOB improves 5-NN/linear probe performance over SIGREG by 14.17/12.89pp on natural images, while CONTRASTIVE-OOB improves over SimCLR by 10.83/6.58pp on biomedical data.On 3D volumes, both OOB methods retain positive mean gains.Source: §4.3 · p. 6, Figure 3 · p. 7

4.4 Population weighting and minority classes under imbalance

  • A cached empirical estimator tracks the empirical objective, but OOB state does not correct the divergence between empirical and class-balanced gradients as minority classes are downsampled.When class information is available, population weighting can be used as a hyperparameter.Source: §4.4 · p. 7, Figure 4 · p. 7
  • SIGREG-OOB gains 6.24 points in 5-NN minority recall and 8.6 points in linear minority recall over SIGREG; CONTRASTIVE-OOB gains 10.69 and 10.28 points over SimCLR.Gains remain positive across the downsampling sweep.Source: §4.4 · p. 8, Figure 5 · p. 8

4.5 Instance retrieval

  • SIGREG-OOB improves retrieval mAP over SIGREG by 13.1 points, while CONTRASTIVE-OOB improves over SimCLR/MoCoV2 by 13.7/7.6 points at a given batch size.Retrieval uses cosine similarity between frozen features and reports mean average precision over three seeds.Source: §4.5 · p. 8, Figure 6 · p. 8

4.6 Continued self-supervised pretraining

  • OOB methods outperform corresponding in-batch variants across most continued-pretraining configurations, with the clearest advantage at small batches.At small batches, OOB regularization improves on pretrained initialization while in-batch training can degrade representation quality.Source: §4.6 · p. 9, Table 3 · p. 9

4.7 Performance under a memory budget

  • Within the evaluated GPU-memory range, OOB methods can match or exceed in-batch performance while using substantially less GPU memory.At the same batch size, OOB methods use the same GPU memory as their in-batch counterparts.Source: §4.7 · p. 9, Figure 7 · p. 9

5 Conclusion and Future Work

OOB regularization decouples population-level representation regularization from minibatch optimization using periodically constructed dataset-level targets, improving the match to the full-dataset signal, especially when minibatches are small. The section also documents the study’s reproducibility practices, evaluation protocols, training setup, and cache costs.

5 Conclusion and Future Work

  • Across contrastive and non-contrastive objectives, OOB more faithfully recovers the full-dataset regularization signal and improves representation quality, particularly when minibatches poorly approximate the population.The active minibatch need not define the population over which representation geometry is regularized.Source: §5 · pp. 9–10
  • Dataset-level state must be stored and periodically refreshed as the encoder evolves, and understanding this trade-off at larger scales remains an important direction.The paper also motivates separating globally coupled representation objectives from minibatch encoder updates, with regularizer-specific constructions as a possible route beyond the objectives studied.Source: §5 · p. 10

AI use statement

  • The authors reviewed AI-assisted text and references and independently inspected and executed AI-assisted code to verify its correctness.They report using generative AI to polish the manuscript, search literature, and help write code, but not to generate research ideas or propose experiments.Source: AI use statement · p. 10

Reproducibility statement

  • The appendix documents architectures, training configurations, augmentations, hyperparameters, evaluation protocols, hyperparameter selection, and computational costs to support independent reproduction.The authors plan to release code to facilitate reproduction and comparison.Source: Reproducibility statement · p. 10

A Reproducibility details

  • Matched OOB/in-batch classification pairs preserve the data split, augmentations, views, encoder, projector, optimizer, schedule, and downstream evaluator.Classification and retrieval results are aggregated over independent training seeds.Source: A Reproducibility details · p. 14

A.1 Datasets and evaluation protocols

  • Frozen classification uses deterministic normalized backbone features and a fixed shuffled 70/30 split, with balanced-accuracy evaluation by cosine k-NN and standardized logistic-regression probes.Retrieval uses normalized frozen features and two protocols: pretraining and retrieval on the same unlabeled collection, or an 80/20 instance-identity split; methods share the encoder, transformations, optimizer, and evaluator.Source: §A.1 · p. 14

A.2 Architectures, optimization, and augmentation

  • The 2D experiments use a ViT-Small/16 encoder and 224 × 224 inputs, while 3D experiments use a 3D ResNet-18 and 643 inputs; all methods train for 100 epochs.Training uses AdamW and a cosine learning-rate schedule; from-scratch runs use fp32 on one NVIDIA L40S GPU.Source: §A.2 · p. 14

A.3 Method-specific hyperparameters and cost

  • Both OOB methods refresh every 0.5 epoch, keep their caches in CPU memory, and transfer only active-minibatch rows to the GPU.At N = 13,000, the SIGReg-OOB and CONTRASTIVE-OOB caches occupy 12.7 MiB and 6.3 MiB, respectively; cache construction adds wall time but not online activation memory.Source: §A.3 · pp. 14–15

A.4 Ablations of the SIGReg-OOB configuration

Controlled SIGReg-OOB ablations find that projection resampling gives the largest improvement, while a 0.5-epoch refresh period performs best and balancing loss terms matters more than increasing projection count.

A.4 Ablations of the SIGReg-OOB configuration

  • The ablations comprise 1,792 evaluations across 39 configurations and 48 matched dataset–batch experiments, with one factor varied within each controlled slice.The sweep covers projection resampling, refresh period, alignment and anti-collapse weights, and projection count.Source: §A.4 · p. 15
  • The ablation slices use distinct fixed settings for the anti-collapse-weight and projection-count comparisons, with projection resampling enabled in both.The anti-collapse-weight slice fixes (p, λalign, M) = (2, 25, 512), while the projection-count slice fixes (p, λalign, λac) = (2, 25, 5).Source: §A.4 · p. 15
  • Projection resampling produces the largest improvement; a 0.5-epoch refresh performs best, while one epoch is close and overly frequent refreshes are not beneficial.An intermediate alignment weight is preferable to either extreme, and increasing the anti-collapse weight or projection count does not improve the matched slices.Source: §A.4 · p. 15

B Oracle-gradient protocol

The oracle study compares representation-space gradients from minibatch and cached estimates with a recomputed full-dataset reference. Agreement improves with larger batches, while cache agreement depends on how rapidly the encoder changes.

B Oracle-gradient protocol

  • The study differentiates the projected-Gaussian anti-collapse loss with respect to probe representations; its gradients are not scalar losses or full network-parameter gradients.The minibatch estimator uses a sampled batch, while the cached estimator reuses assignments made at an earlier encoder state; the reference is recomputed from current representations.Source: B Oracle-gradient protocol · pp. 15–16
  • The canonical experiment spans all 2D and 3D datasets, checkpoints at epochs 0, 50, and 100, batch sizes 1–64, and cache ages up to 128 optimizer updates.It uses three seeds, 16 probe samples, 16 sampled contexts, and 256 projection directions; agreement is measured by cosine similarity and relative root-mean-square error.Source: B Oracle-gradient protocol · p. 16
  • At B = 1, mean cosine agreement at epochs 0, 50, and 100 was −0.047, −0.417, and −0.362; at B = 64 it was 0.972, 0.903, and 0.939.A fresh cache agrees exactly by construction.Source: B Oracle-gradient protocol · p. 16
  • At cache age 128, mean agreement was 0.453, 0.720, and 0.912 at the default learning rate, versus 0.692, 0.909, and 0.906 at the matched 10−6 low-rate control.The largest low-rate improvements occur at epochs 0 and 50; cache lifetime depends on encoder-update speed, not update count alone.Source: B Oracle-gradient protocol · p. 16
  • The Figure 9 comparison indicates that the ordering between limited batch context and fresh dataset-level context persists across the tested dataset sizes.The oracle is a reference for estimating this chosen finite-population regularizer.Source: B Oracle-gradient protocol · p. 16

C Detailed frozen-classification results

OOB advantages in frozen classification are strongest at low batch sizes across the controlled training-set-size sweep and are distributed across domains. Representation-rank diagnostics qualify the interpretation: rank does not by itself explain semantic performance.

C Detailed frozen-classification results

C.1 Representation-rank diagnostics

  • In the smallest-batch regime, CONTRASTIVE-OOB has much higher representation rank than SimCLR and MoCoV2.Figure 11 reports effective rank and RankMe averaged over common datasets, with standard-error bands across datasets.Source: §C.1 · p. 16, Figure 11 · p. 19
  • On 2D biomedical data, SIGReg-OOB can have lower rank than SIGREG while matching or exceeding its downstream accuracy; on natural images, its rank is comparable or higher.Rank diagnoses one aspect of collapse but does not by itself explain semantic performance.Source: §C.1 · pp. 16, 19, Figure 11 · p. 19

D Detailed class-imbalance results

Out-of-batch gains extend beyond average accuracy to underrepresented classes and often the worst class, though the empirical cache preserves the population’s frequency weighting. The controlled-imbalance experiments report minority, worst-class, and minority–majority recall metrics.

D Detailed class-imbalance results

  • The experiments retain a 1/k fraction of designated classes, keep SSL label-free, and use labels afterward to report macro, worst-class, and minority–majority recall.The retained fractions use k ∈ {2, 5, 8, 10, 15, 20, 25}.Source: D Detailed class-imbalance results · p. 20
  • At batch size 64, Figure 12 reports 2D minority-class recall, while Figure 13 summarizes worst-class recall and the minority–majority gap for contrastive methods.Figure 13 averages dataset–batch experiments and shows standard-error bands.Source: Figure 12 · p. 20, Figure 13 · p. 20, D Detailed class-imbalance results · p. 20
  • OOB gains are visible for underrepresented classes and, in many settings, the worst class; class-balanced objectives still require an explicit weighting decision.The empirical cache preserves the empirical population’s frequency weighting.Source: D Detailed class-imbalance results · p. 20

E Instance-retrieval results

Contrastive-OOB has the highest average mAP and recall in both transductive and held-out retrieval. SIGReg-OOB also improves over SIGReg on every aggregate retrieval metric.

E Instance-retrieval results

F Continued pretraining with DINOv2

In DINOv2 continued pretraining, OOB regularization is more likely to retain or improve pretrained representations at small batches, but outcomes vary across methods, encoders, and datasets. OOB adds modest wall time at matched memory, and its advantage develops during training.

F Continued pretraining with DINOv2

  • At small batches, OOB is more likely to retain or improve pretrained representations than in-batch regularization, but outcomes vary by method, encoder, and target dataset.Runs start from public DINOv2 ViT-S/B/L/G checkpoints and use two augmented views for 100 epochs.Source: F Continued pretraining with DINOv2 · pp. 21–22, Figure 15 · p. 22

G Computational measurements

H.1 Convergence over pretraining time

  • The Gaussian OOB advantage grows throughout training, while the contrastive advantage emerges early and remains positive on average through epoch 200.Source: §H.1 · p. 22, Figure 17 · p. 23

Related papers

Loading 2610.20158v1…