JEPA Collapse Lab

Why do joint-embedding models collapse — and how do Barlow Twins and VICReg prevent it?

A small, self-contained experimental lab that reproduces representation collapse in a joint-embedding architecture, diagnoses it quantitatively, and shows how two classic non-contrastive objectives — Barlow Twins and VICReg — prevent it without negative pairs, without stop-gradient, and without a predictor.

This is not a full JEPA (there is no predictor). It is the experiment that demonstrates you understand the fundamental problem that motivated JEPA in the first place.

---

The question

Given an image $x$, sample two augmented views $x_A = t_A(x)$ and $x_B = t_B(x)$, encode them with the same encoder, and pull the embeddings together:

                 ┌──────────────────┐
 x ── t_A(x) ──► │                  ├──► h_A ──► Projector ──► z_A ──┐
                 │   Encoder E      │                                ├──► Loss
 x ── t_B(x) ──► │   (shared)       ├──► h_B ──► Projector ──► z_B ──┘
                 └──────────────────┘

There is a trivial shortcut: output the same vector for every input. The invariance loss goes to zero, and the representation carries no information whatsoever. This is representation collapse.

The interesting part is not that collapse happens — it is how we measure it, and which mechanisms provably prevent it.

The three experiments

All three runs share the same backbone, same projector, same augmentations, same budget. Only the loss changes.

Experiment A — Naive invariance (the collapse)

$$\mathcal{L}_{\text{inv}} = |z_A - z_B|_2^2$$

The network progressively discovers it can emit a constant vector. Expected signature:

• per-dimension variance of embeddings $\to 0$,
• all embeddings nearly identical,
• effective rank of the embedding matrix $\to 1$ (in practice: noise-dominated, see Findings),
• linear-probe accuracy $\to$ chance level (in practice: random-feature level, see Findings).

Key lesson: a low loss does not mean a good representation.

Experiment B — Barlow Twins

Compute the cross-correlation matrix of the two embedding batches:

$$\mathcal{C}_{ij} = \frac{\sum_b z^A_{b,i}, z^B_{b,j}}{\sqrt{\sum_b (z^A_{b,i})^2},\sqrt{\sum_b (z^B_{b,j})^2}}$$

and push it toward the identity:

$$\mathcal{L}_{BT} = \sum_i (1 - \mathcal{C}_{ii})^2 + \lambda \sum_i \sum_{j \neq i} \mathcal{C}_{ij}^2$$

• Diagonal terms enforce invariance between the two views.
• Off-diagonal terms reduce redundancy between embedding dimensions.

A constant output cannot satisfy the unit diagonal without exploding the off-diagonal penalty, so collapse is ruled out by construction.

Experiment C — VICReg

Three explicit terms:

$$\mathcal{L} = \lambda, \underbrace{s(Z_A, Z_B)}_{\text{invariance}} + \mu, \underbrace{\big[v(Z_A) + v(Z_B)\big]}_{\text{variance}} + \nu, \underbrace{\big[c(Z_A) + c(Z_B)\big]}_{\text{covariance}}$$

• Invariance $s$: MSE between the two views (pulls them together).
• Variance $v$: hinge loss that keeps the per-dimension standard deviation above a target $\gamma$ — each dimension is forbidden to become constant.
• Covariance $c$: penalizes off-diagonal covariance — dimensions are forbidden to copy the same information.

VICReg makes the fight against collapse fully explicit: invariance, variance, decorrelation.

Setup

Component	Choice
Dataset	STL-10: 100k unlabeled images for pretraining; labeled split used only for linear evaluation. CIFAR-10 as a cheap debug config.
Backbone	ResNet-18 (small-image stem: 3×3 conv, no maxpool)
Projector	MLP 512→1024→1024→128, BatchNorm in hidden layers (h: probe space, z: loss space)
Augmentations	RandomResizedCrop, horizontal flip, color jitter, grayscale, blur
Protocol	Same backbone / augmentations / epochs / optimizer for A, B, C; fixed seeds

Diagnostics — how we prove collapse (or its absence)

Diagnostic	What it shows
Paired-view cross-correlation	Alignment (diagonal) and redundancy (off-diagonal) between the two augmented views
Per-dimension variance	The collapse curve: std of each embedding dimension over training
Covariance heatmap	Redundant vs. decorrelated dimensions
Singular value spectrum	How many directions the representation actually spans
Effective rank	$\mathrm{erank}(Z) = \exp(-\sum_i p_i \log p_i)$, $p_i = \sigma_i / \sum_j \sigma_j$
UMAP / PCA projection	Visual cluster structure (colored by true label)
Linear probe accuracy	Frozen encoder + logistic regression on the labeled split

Results

STL-10 pretraining on the 100k unlabeled split — identical budget for the three runs (10 epochs, batch 256, Adam lr 3e-4, seed 0, ResNet-18 + 512→1024→1024→128 projector, single RTX 3090). Linear probe: multinomial logistic regression on frozen features, labeled train/test splits only. Full machine-readable table: results/results_table.md.

Method	Loss (last)	z_std (last)	Mean per-dim std	Effective rank	Probe acc (h)	Probe acc (z)
A — Naive	2.6e-5 (degenerate)	0.004	0.006	n/a (collapsed)	42.9 %	33.1 %
B — Barlow Twins	2.81	0.86	0.91	61.6	71.8 %	66.1 %
C — VICReg	9.08	1.00 (= γ)	1.03	87.6	73.5 %	67.8 %

The naive model collapses, the other two don't — visible in every figure, no explanation needed. Per-dimension std over training (the collapse curve) and the singular spectrum of the embedding matrix:

	A — Naive (collapse)	C — VICReg (healthy)
Training curves
Per-dim std
Singular spectrum
UMAP, colored by label

All per-run figures (incl. covariance heatmaps and PCA projections): results/figures/. Reproduce with scripts/diagnose.py --all and scripts/probe.py --all (see Quickstart).

Findings and surprises

1. Collapse lives in the projector, not the backbone. With the naive loss, z_std drops below 5e-3 within four epochs — yet the probe on backbone features h still reads 42.9 %, roughly the random-feature level (a random-init encoder scores ~34–43 % in our smoke checks). The encoder barely moves from its random init; all the "learning" is the projector collapsing to a near-constant map. A low loss does not mean a good representation — but a collapsed loss does not mean a destroyed backbone either.
2. A standardized probe can half-see through collapse. Probe(z) on the collapsed run reads 33.1 %, far above the ~10 % chance level: the pipeline standardizes features, and dividing by a ~1e-4 per-dim std re-amplifies the residual variation to O(1), where a linear model partially reads labels again. Interpret probe(z) on collapsed runs with care — collapse metrics (std, spectrum) are the ground truth.
3. Effective rank is noise-dominated under collapse. On a near-constant embedding matrix the SVD sees only numerical noise, whose almost-flat spectrum yields erank ≈ 37–110 instead of ~1 — the opposite of the naive expectation. We therefore report erank only when the mean per-dim std is above a noise floor (COLLAPSE_STD_FLOOR = 1e-2 in diagnostics/metrics.py); collapsed runs show "n/a (collapsed)".
4. Practical note — VRAM. In fp32, the small-image ResNet-18 (no maxpool) keeps 64×96×96 activations through layer 1; two views + backward ≈ 20.5 GiB at batch 256, so batch 512 does not fit in 24 GiB. All three runs use batch 256 — identical budget matters more than large batches.

Paired-view cross-correlation

Each heatmap compares the first 64 projector dimensions from the two augmented views of the same images. A healthy representation has a bright diagonal (alignment) and a muted off-diagonal (low redundancy). The nearly uniform naive heatmap must be read alongside its near-zero per-dimension variance: correlation is no longer meaningful once the representation has collapsed.

Naive invariance — collapse	Barlow Twins	VICReg

Repository structure

├── configs/                     # stl10.yaml, cifar10_debug.yaml
├── src/jepa_collapse_lab/
│   ├── config.py                # YAML config loading
│   ├── utils.py                 # seeds, checkpoint save/load
│   ├── data/                    # two-view augmentations, dataset & loader builders
│   ├── models/                  # ResNet-18 small-image backbone, MLP projector, SSLModel
│   ├── losses/                  # naive, Barlow Twins, VICReg
│   ├── diagnostics/             # variance, covariance, spectrum, rank, UMAP
│   └── eval/                    # frozen linear probe
├── scripts/
│   ├── visualize_pairs.py       # augmentation sanity check
│   ├── train.py                 # train naive / barlow_twins / vicreg
│   ├── diagnose.py              # collapse figures from any checkpoint
│   └── probe.py                 # linear probe accuracy
├── tests/
└── results/
    ├── figures/                 # generated figures (pairs + per-run diagnostics)
    └── checkpoints/             # run dirs: {dataset}_{experiment}/

Quickstart

uv sync
uv run pytest

# Phase 1 sanity check: visualize augmented pairs (downloads the dataset on first run)
uv run scripts/visualize_pairs.py --config configs/cifar10_debug.yaml   # fast debug
uv run scripts/visualize_pairs.py --config configs/stl10.yaml           # main dataset

# Phase 3: train the three variants (same trainer, only the loss changes)
uv run scripts/train.py --config configs/cifar10_debug.yaml --experiment naive
uv run scripts/train.py --config configs/cifar10_debug.yaml --experiment barlow_twins
uv run scripts/train.py --config configs/cifar10_debug.yaml --experiment vicreg
# STL-10 full budget:
# uv run scripts/train.py --config configs/stl10.yaml --experiment barlow_twins

# Phase 4: collapse diagnostics on a checkpoint (or every run under results/checkpoints)
uv run scripts/diagnose.py --checkpoint results/checkpoints/cifar10_naive/last.pt
uv run scripts/diagnose.py --all --checkpoints-root results/checkpoints

# Phase 5: frozen linear probe (logistic regression on backbone features h)
uv run scripts/probe.py --checkpoint results/checkpoints/cifar10_vicreg/last.pt
uv run scripts/probe.py --all --checkpoints-root results/checkpoints

The augmentation pipeline in action — two independent views of the same image (random crop, flip, color jitter, grayscale, blur):

CIFAR-10 (debug)	STL-10 (main dataset)

Roadmap

See ROADMAP.md for the full phased plan and the current status.

What this project demonstrates

• what representation collapse is, and how to diagnose it (not just observe it);
• why a low SSL loss is not evidence of a good representation;
• the role of variance and decorrelation in non-contrastive learning;
• the conceptual bridge toward JEPA: predicting in representation space only works if the representation space hasn't collapsed.

References

• Zbontar et al., Barlow Twins: Self-Supervised Learning via Redundancy Reduction, ICML 2021.
• Bardes et al., VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning, ICLR 2022.
• LeCun, A Path Towards Autonomous Machine Intelligence, 2022.
• Assran et al., Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture (I-JEPA), CVPR 2023.