Representation Collapse — Visual Intuition¶
Research note: research/02-representation-collapse.md
Drop the decoder and train an encoder so two views of the same input have similar latents —
min ||f(x) - f(x')||^2 — and there is a trivial escape: map everything to one point
(f = const, loss = 0). That is representation collapse, the default failure of predicting
in latent space. This notebook shows it happen, distinguishes its two faces, demonstrates the
mechanism that prevents it, and builds the diagnostic Layer A needs.
What we will see:
- Invariance-only training collapses — embedding variance decays to zero.
- Two faces: complete vs dimensional collapse — and why one metric is not enough.
- The variance/covariance mechanism (VICReg) that makes a non-zero solution optimal.
- Effective dimension as a continuous health metric for collapse.
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
np.random.seed(42)
plt.rcParams.update({
'figure.dpi': 110,
'font.size': 11,
'axes.titlesize': 12,
'axes.labelsize': 11,
'figure.facecolor': 'white',
'axes.facecolor': '#f8f9fa',
'axes.grid': True,
'grid.alpha': 0.35,
})
def effective_dim(Z):
# exp(entropy of the covariance spectrum): = dim when flat, -> 1 when one dim dominates.
Zc = Z - Z.mean(axis=0, keepdims=True)
cov = (Zc.T @ Zc) / (len(Z) - 1)
sv = np.linalg.svdvals(cov)
sv = np.clip(sv, 1e-12, None)
p = sv / sv.sum()
return float(np.exp(-(p * np.log(p)).sum()))
def total_variance(Z):
return float(Z.var(axis=0).sum())
print('Setup complete.')
Setup complete.
Experiment 1: Invariance-only training collapses¶
We train a linear encoder W to make two augmented views agree, with nothing else in the
loss: L = mean ||W x - W x'||^2. The two views differ by isotropic augmentation noise. The global
optimum is W = 0 (loss exactly 0), so gradient descent should shrink every embedding toward a
single point. We watch the embedding's total variance decay.
d, n = 16, 512
rng = np.random.RandomState(0)
# Structured (anisotropic) data so the embedding starts genuinely spread out.
A = rng.randn(d, d)
X = rng.randn(n, d) @ A
X = (X - X.mean(0)) / X.std(0)
# Two views differ by ISOTROPIC augmentation noise; the difference cov is therefore c * I,
# so the invariance gradient (2/n) W (Diff^T Diff) shrinks every direction of W uniformly.
noise = 0.5
Da = X + noise * rng.randn(n, d)
Db = X + noise * rng.randn(n, d)
Diff = Da - Db # what invariance penalises
iso = float((Diff ** 2).mean()) # isotropic per-element variance of the augmentation difference
W = np.eye(d)
eta = 0.05
epochs = 120
hist_var, hist_eff, hist_inv = [], [], []
for ep in range(epochs + 1):
Z = X @ W.T
hist_var.append(total_variance(Z))
hist_eff.append(effective_dim(Z))
hist_inv.append(np.mean(np.sum((Da @ W.T - Db @ W.T) ** 2, axis=1)))
# Population invariance gradient under isotropic augmentation: uniform shrink of W.
W = W * (1.0 - 2 * eta * iso)
Z0 = X @ np.eye(d).T
Zf = X @ W.T
print(f"total variance: start {hist_var[0]:.2f} -> end {hist_var[-1]:.4f}")
print(f"effective dim: start {hist_eff[0]:.2f} -> end {hist_eff[-1]:.2f} (scale-blind!)")
print(f"invariance loss: start {hist_inv[0]:.2f} -> end {hist_inv[-1]:.4f}")
total variance: start 16.00 -> end 0.0001 effective dim: start 9.82 -> end 9.82 (scale-blind!) invariance loss: start 7.79 -> end 0.0000
fig, axes = plt.subplots(1, 3, figsize=(15, 4.4))
# PCA to 2D (same projection for both) to show the cloud shrinking.
u, s, vt = np.linalg.svd(Z0 - Z0.mean(0), full_matrices=False)
P = vt[:2].T
axes[0].scatter((Z0 @ P)[:, 0], (Z0 @ P)[:, 1], s=8, alpha=0.5, color='#2980b9', label='start')
axes[0].scatter((Zf @ P)[:, 0], (Zf @ P)[:, 1], s=8, alpha=0.7, color='#c0392b', label='after training')
axes[0].set_title('Embedding cloud collapses to a point')
axes[0].set_xlabel('PC 1'); axes[0].set_ylabel('PC 2'); axes[0].legend()
axes[1].semilogy(hist_var, color='#c0392b', label='total variance')
axes[1].semilogy(hist_inv, color='#2980b9', label='invariance loss')
axes[1].set_title('Variance and loss both -> 0')
axes[1].set_xlabel('epoch'); axes[1].set_ylabel('value (log)'); axes[1].legend()
axes[2].plot(hist_eff, color='#27ae60')
axes[2].axhline(d, ls=':', color='gray', label=f'dim = {d}')
axes[2].set_title('Effective dim is scale-blind')
axes[2].set_xlabel('epoch'); axes[2].set_ylabel('effective dimension')
axes[2].set_ylim(0, d + 1); axes[2].legend()
plt.tight_layout(); plt.show()
What you see¶
The invariance-only objective drives the encoder straight into complete collapse: the embedding cloud (left, red) contracts onto a single point, and the total variance falls toward zero in lockstep with the loss (middle). The encoder has learned to forget the input — the cheapest way to make two views agree.
The right panel is the subtle lesson: effective dimension barely moves. Because the augmentation
noise is isotropic, W shrinks uniformly, so the shape of the spectrum (which effective dimension
measures) is preserved even as its scale vanishes. Complete collapse is visible in absolute
variance but invisible to the scale-invariant effective-dimension metric — you need both.
Experiment 2: The two faces of collapse¶
Collapse comes in two forms. Complete collapse: all points pile onto one spot (variance ~ 0). Dimensional collapse: points still spread, but only within a low-dimensional subspace — the covariance becomes low-rank. The second is dangerous because it hides: a 2D PCA scatter can look perfectly healthy while most dimensions are dead. We construct all three cases and compare.
rng = np.random.RandomState(3)
n, d = 800, 16
# (a) Healthy: full-rank, fairly flat spectrum.
Z_health = rng.randn(n, d) * np.linspace(1.0, 0.6, d)
# (b) Dimensional collapse: signal lives in only 2 directions, rest is tiny noise.
basis = np.linalg.qr(rng.randn(d, d))[0]
latent2 = rng.randn(n, 2) @ basis[:2]
Z_dim = latent2 + 0.02 * rng.randn(n, d)
# (c) Complete collapse: everything near one point.
Z_complete = np.full((n, d), 3.0) + 0.01 * rng.randn(n, d)
for name, Z in [('healthy', Z_health), ('dimensional', Z_dim), ('complete', Z_complete)]:
print(f"{name:12s} total_var={total_variance(Z):8.3f} eff_dim={effective_dim(Z):5.2f} "
f"min_std={Z.std(0).min():.4f}")
healthy total_var= 10.363 eff_dim=15.10 min_std=0.5811 dimensional total_var= 1.958 eff_dim= 2.04 min_std=0.0537 complete total_var= 0.002 eff_dim=15.82 min_std=0.0097
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
for name, Z, col in [('healthy', Z_health, '#27ae60'),
('dimensional collapse', Z_dim, '#e67e22'),
('complete collapse', Z_complete, '#c0392b')]:
Zc = Z - Z.mean(0)
sv = np.linalg.svdvals((Zc.T @ Zc) / (n - 1))
axes[0].semilogy(range(1, d + 1), np.clip(sv, 1e-6, None), 'o-', color=col, label=name)
axes[0].set_title('Covariance spectra (the tell-tale)')
axes[0].set_xlabel('singular-value index'); axes[0].set_ylabel('variance (log)')
axes[0].legend()
# Top-2 PCA of the dimensional-collapse case: looks perfectly spread.
Zc = Z_dim - Z_dim.mean(0)
vt = np.linalg.svd(Zc, full_matrices=False)[2]
proj = Zc @ vt[:2].T
axes[1].scatter(proj[:, 0], proj[:, 1], s=10, alpha=0.5, color='#e67e22')
axes[1].set_title('Dimensional collapse HIDES in top-2 PCA\n(eff_dim ~ 2, but the scatter looks fine)')
axes[1].set_xlabel('PC 1'); axes[1].set_ylabel('PC 2')
plt.tight_layout(); plt.show()
What you see¶
The spectra (left) separate the three regimes cleanly. Healthy (green) keeps all 16 singular values high. Complete collapse (red) pushes every value to the floor — its total variance ~ 0.002 is the giveaway, even though its spectrum is flat so effective dimension reads a misleadingly high ~15.8. Dimensional collapse (orange) is the treacherous middle: two directions stay high and the other fourteen fall off a cliff, so total variance looks unremarkable (~2) but effective dimension ~ 2 exposes it.
That contrast is the whole point: the two metrics catch different failures. Total variance catches complete collapse but is blind to dimensional collapse; effective dimension catches dimensional collapse but is blind (scale-invariant) to complete collapse. The right panel shows the trap — a top-2 PCA scatter of the dimensionally-collapsed embeddings looks beautifully spread, because the two surviving directions are the two PCA axes and the 14 dead dimensions are invisible. Only the full spectrum, read through both metrics, tells the truth.
Experiment 3: The variance/covariance mechanism (VICReg)¶
How do you stop collapse without a decoder? VICReg adds two terms: a variance hinge that forces each embedding dimension's std above a threshold gamma, and a covariance term that decorrelates dimensions. We visualise their effect on the loss landscape: along a global embedding scale alpha, invariance-only is minimised at alpha = 0 (collapse), but invariance + variance has its minimum at a non-zero alpha — collapse is no longer optimal.
# Base embedding directions with per-dim std sigma_j; W = alpha * W0 scales them.
rng = np.random.RandomState(5)
d = 16
sigma = np.linspace(1.0, 0.5, d) # per-dim std at alpha = 1
gamma = 0.7 # variance target (std threshold)
alphas = np.linspace(0.0, 2.0, 200)
# Invariance loss grows with scale (two views disagree more when embeddings are larger): ~ alpha^2.
inv = alphas ** 2
# Variance hinge summed over dims: sum_j max(0, gamma - alpha * sigma_j).
var_term = np.array([np.sum(np.maximum(0.0, gamma - a * sigma)) for a in alphas])
L_inv_only = inv
L_vicreg = inv + 4.0 * var_term # mu = 4
a_inv = alphas[np.argmin(L_inv_only)]
a_vic = alphas[np.argmin(L_vicreg)]
print(f"invariance-only minimised at alpha = {a_inv:.3f} (collapse)")
print(f"invariance+var minimised at alpha = {a_vic:.3f} (non-zero -> no collapse)")
# Covariance term as a function of correlation between two dims (unit std).
rho = np.linspace(-1, 1, 200)
cov_penalty = rho ** 2 # off-diagonal covariance squared
invariance-only minimised at alpha = 0.000 (collapse) invariance+var minimised at alpha = 1.317 (non-zero -> no collapse)
fig, axes = plt.subplots(1, 2, figsize=(14, 4.8))
axes[0].plot(alphas, L_inv_only, color='#c0392b', label='invariance only')
axes[0].plot(alphas, L_vicreg, color='#27ae60', label='invariance + variance')
axes[0].scatter([a_inv], [L_inv_only.min()], color='#c0392b', zorder=5)
axes[0].scatter([a_vic], [L_vicreg.min()], color='#27ae60', zorder=5)
axes[0].axvline(0, ls=':', color='gray')
axes[0].set_title('Variance term moves the optimum off zero')
axes[0].set_xlabel('embedding scale alpha'); axes[0].set_ylabel('loss')
axes[0].annotate('collapse', (a_inv, L_inv_only.min()), textcoords='offset points', xytext=(20, 20),
arrowprops=dict(arrowstyle='->'), color='#c0392b')
axes[0].annotate('healthy', (a_vic, L_vicreg.min()), textcoords='offset points', xytext=(10, 30),
arrowprops=dict(arrowstyle='->'), color='#27ae60')
axes[0].legend()
axes[1].plot(rho, cov_penalty, color='#8e44ad')
axes[1].axvline(0, ls=':', color='gray')
axes[1].set_title('Covariance term decorrelates dimensions')
axes[1].set_xlabel('correlation between two dims (rho)'); axes[1].set_ylabel('covariance penalty')
axes[1].annotate('minimised at\nrho = 0 (full rank)', (0, 0), textcoords='offset points',
xytext=(20, 40), arrowprops=dict(arrowstyle='->'))
plt.tight_layout(); plt.show()
What you see¶
Left: the invariance-only loss (red) slides monotonically to its minimum at alpha = 0 — exactly the
collapse of Experiment 1. Adding the variance hinge (green) inserts a steep penalty whenever any
dimension's std drops below the threshold gamma, lifting the loss near zero and relocating the minimum
to a non-zero scale. Collapse stops being the optimal solution; the encoder is forced to keep the
embeddings spread.
Right: the covariance term penalises off-diagonal correlation as rho^2, minimised at rho = 0.
Where the variance term fights complete collapse (keeps each dimension alive), the covariance term
fights dimensional collapse (keeps dimensions independent, hence full-rank). Together they are the
explicit version of what BYOL/SimSiam/DINO achieve implicitly through architecture — the subject of
the next notes.
Experiment 4: Effective dimension as a continuous health metric¶
Collapse is rarely all-or-nothing during training; dimensions die gradually. A good monitor should degrade smoothly as health is lost, giving early warning. We sweep a controlled collapse knob — progressively suppressing trailing covariance directions — and confirm effective dimension tracks it continuously, the diagnostic Layer A would expose for any embedding.
rng = np.random.RandomState(11)
n, d = 1000, 20
Z_full = rng.randn(n, d) * np.linspace(1.0, 0.8, d) # healthy, near-flat spectrum
# Knob 1 (discrete): keep only k live directions, suppress the rest.
ks = np.arange(1, d + 1)
eff_by_k = []
for k in ks:
scale = np.ones(d); scale[k:] = 0.03
eff_by_k.append(effective_dim(Z_full * scale))
# Knob 2 (continuous): smoothly decay the tail of the spectrum by factor t in [0,1].
ts = np.linspace(0, 1, 60)
eff_by_t = []
ramp = np.linspace(1.0, 0.0, d) # later dims suppressed more as t grows
for t in ts:
scale = (1 - t) + t * (1 - ramp) # t=0 -> all ones; t=1 -> tail -> 0
eff_by_t.append(effective_dim(Z_full * scale))
print(f"eff_dim with 20 live dims: {eff_by_k[-1]:.1f} with 1 live dim: {eff_by_k[0]:.2f}")
eff_dim with 20 live dims: 19.6 with 1 live dim: 1.12
fig, axes = plt.subplots(1, 2, figsize=(14, 4.6))
axes[0].plot(ks, eff_by_k, 'o-', color='#2980b9')
axes[0].plot([1, d], [1, d], ls=':', color='gray', label='y = x (ideal)')
axes[0].set_title('Effective dim tracks the number of live directions')
axes[0].set_xlabel('k = number of live dimensions'); axes[0].set_ylabel('effective dimension')
axes[0].legend()
axes[1].plot(ts, eff_by_t, color='#c0392b')
axes[1].set_title('Smooth early warning as the spectrum decays')
axes[1].set_xlabel('collapse knob t (tail suppression)'); axes[1].set_ylabel('effective dimension')
axes[1].annotate('healthy', (0.02, eff_by_t[0]), textcoords='offset points', xytext=(20, -5))
axes[1].annotate('collapsing', (0.95, eff_by_t[-1]), textcoords='offset points', xytext=(-70, 15))
plt.tight_layout(); plt.show()
What you see¶
Effective dimension behaves exactly as a health metric should. Against the discrete knob (left) it
rises almost along the ideal y = x line: when k directions are alive it reports ~k, so a reading
far below the ambient dimension is an immediate red flag for dimensional collapse. Against the
continuous knob (right) it falls smoothly as the spectral tail is suppressed — no cliff, no
threshold-hunting — giving early warning while collapse is still partial and reversible.
This is the concrete tool the research note hands to Layer A: a single scalar, computable on any embedding matrix, that catches the silent failure Experiment 2 warned about. With it, every later method in the tier (stop-gradient, EMA, JEPA) can be audited for the one thing they all must achieve — not collapsing.
Summary / Key Takeaways¶
- Collapse is the default (Exp 1). Invariance-only training maps everything to a point; the loss and the embedding variance fall to zero together. Removing the decoder removes the information anchor.
- Two faces, two metrics (Exp 2). Complete collapse shows in absolute variance; dimensional collapse hides behind a healthy-looking PCA scatter and is only visible in the full spectrum / effective dimension. Monitor both.
- The fix is a non-zero optimum (Exp 3). Variance + covariance regularisation (VICReg) makes collapse sub-optimal: the variance term keeps each dimension alive, the covariance term keeps them independent. BYOL/SimSiam/DINO do this implicitly.
- Effective dimension is the diagnostic (Exp 4). It tracks live directions and degrades smoothly, giving Layer A an early-warning health metric for any latent space.
- Bridge. Every remaining Tier 8 method is a different anti-collapse mechanism. Next:
stop-gradient and asymmetric architecture — collapse prevention with no negatives and no explicit
regulariser, just one
stop_gradient.