Stop-gradient & Asymmetric Architecture — Visual Intuition¶
Research note: research/03-stop-gradient-asymmetric.md
BYOL and SimSiam avoid collapse with no negatives and no explicit variance/covariance term — only two architectural details: a predictor on one branch and a stop-gradient on the other. This notebook trains a tiny linear cosine-SimSiam and shows, empirically, that both ingredients are needed, that the healthy run preserves the representation's effective dimension while every ablation loses it, that an EMA target (BYOL) works too, and that the mechanism is fragile.
What we will see:
- The 2x2 ablation: predictor x stop-gradient — only the corner with both stays healthy.
- Training trajectories: effective dimension decays for the ablations, stays flat for the winner.
- An EMA target encoder (BYOL-style) also prevents collapse, more smoothly.
- Fragility: starve the predictor and even the "good" configuration collapses.
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 covariance spectrum): ~ number of directions actually used.
Zc = Z - Z.mean(0)
sv = np.linalg.svdvals((Zc.T @ Zc) / (len(Z) - 1))
sv = np.clip(sv, 1e-12, None); p = sv / sv.sum()
return float(np.exp(-(p * np.log(p)).sum()))
def neg_cos_grad(P, T):
# gradient of -cos(P, T) w.r.t. P (T held fixed); batched, returns (n, d)
Pn = np.linalg.norm(P, axis=1, keepdims=True) + 1e-9
Tn = np.linalg.norm(T, axis=1, keepdims=True) + 1e-9
Ph, Th = P / Pn, T / Tn
dot = np.sum(Ph * Th, axis=1, keepdims=True)
return -(Th - dot * Ph) / Pn
print('Setup complete.')
Setup complete.
We use a deliberately minimal model so the only thing on trial is the architecture: a linear
encoder W and a linear predictor H, trained with the symmetrised negative-cosine loss on
augmented pairs. stop-gradient means the target branch is detached; EMA means the target comes
from a slowly-updated copy of the encoder (BYOL). We track the effective dimension of the
representation Z = X W^T — collapse shows up as effective dimension falling well below the ambient
16.
def train(use_pred, use_sg, ema=False, d=16, epochs=2000, lr=0.5, noise=0.3, wd=0.002,
pred_lr_mult=1.0, ema_tau=0.99, seed=0, record_every=40):
rng = np.random.RandomState(seed)
A = rng.randn(d, d)
Xall = rng.randn(4096, d) @ A
Xall = (Xall - Xall.mean(0)) / Xall.std(0)
n = 1024
W = np.eye(d) + 0.01 * rng.randn(d, d)
H = np.eye(d) + 0.01 * rng.randn(d, d)
W_tgt = W.copy()
traj = []
for ep in range(epochs):
idx = rng.choice(len(Xall), n, replace=False); X = Xall[idx]
X1 = X + noise * rng.randn(n, d); X2 = X + noise * rng.randn(n, d)
Z1, Z2 = X1 @ W.T, X2 @ W.T
gW = np.zeros_like(W); gH = np.zeros_like(H)
for (Za, Xa, Zb, Xb) in [(Z1, X1, Z2, X2), (Z2, X2, Z1, X1)]:
# target: EMA encoder if BYOL, else the (possibly detached) online encoder
T = (Xb @ W_tgt.T) if ema else Zb
P = Za @ H.T if use_pred else Za
gP = neg_cos_grad(P, T) / n
if use_pred:
gH += gP.T @ Za
gZa = gP @ H
else:
gZa = gP
gW += gZa.T @ Xa
if (not use_sg) and (not ema): # target branch also trainable
gW += (neg_cos_grad(T, P) / n).T @ Xb
gW /= 2; gH /= 2
W = W - lr * (gW + wd * W)
if use_pred:
H = H - lr * pred_lr_mult * (gH + wd * H)
if ema:
W_tgt = ema_tau * W_tgt + (1 - ema_tau) * W
if ep % record_every == 0:
traj.append(effective_dim(Xall @ W.T))
return effective_dim(Xall @ W.T), np.array(traj)
print('trainer ready')
trainer ready
Experiment 1: The 2x2 ablation¶
The famous SimSiam result: you need both the predictor and the stop-gradient. We run all four combinations and report the final effective dimension of the representation (ambient = 16).
configs = [(False, False, 'plain Siamese\n(no pred, no sg)'),
(False, True, 'stop-grad only\n(no predictor)'),
(True, False, 'predictor only\n(no stop-grad)'),
(True, True, 'predictor + stop-grad\n(SimSiam)')]
results = {}
for up, sg, name in configs:
final, traj = train(use_pred=up, use_sg=sg)
results[(up, sg)] = (name, final, traj)
print(f"{name.replace(chr(10),' '):42s} final eff_dim = {final:5.2f}")
plain Siamese (no pred, no sg) final eff_dim = 4.07
stop-grad only (no predictor) final eff_dim = 4.07
predictor only (no stop-grad) final eff_dim = 4.13
predictor + stop-grad (SimSiam) final eff_dim = 9.74
fig, ax = plt.subplots(figsize=(9, 4.8))
names = [results[(up, sg)][0] for up, sg, _ in configs]
finals = [results[(up, sg)][1] for up, sg, _ in configs]
colors = ['#c0392b', '#c0392b', '#c0392b', '#27ae60']
bars = ax.bar(range(4), finals, color=colors, alpha=0.85)
ax.axhline(16, ls=':', color='gray', label='ambient dim = 16')
ax.set_xticks(range(4)); ax.set_xticklabels(names, fontsize=9)
ax.set_ylabel('final effective dimension'); ax.set_ylim(0, 17)
ax.set_title('Only predictor + stop-gradient preserves the representation')
for b, v in zip(bars, finals):
ax.text(b.get_x() + b.get_width() / 2, v + 0.3, f'{v:.1f}', ha='center', fontweight='bold')
ax.legend(); plt.tight_layout(); plt.show()
What you see¶
Three of the four bars sit around an effective dimension of ~4 — barely a quarter of the ambient 16. Removing the predictor collapses (with or without stop-gradient); removing the stop-gradient collapses (with or without the predictor). Only the corner with both ingredients (green) keeps the representation near its full effective dimension (~9.7, essentially its initial value).
This reproduces Chen & He's central ablation: stop-gradient and predictor are each necessary, and only together are they sufficient. Neither is a regulariser in the loss — the loss is pure invariance (negative cosine) in every bar. The difference is entirely in how gradients are routed. That is the surprising claim of this note made concrete.
Experiment 2: Collapse happens during training¶
A final number hides the dynamics. Here we plot effective dimension over training for the four configurations. Collapse is not instantaneous — it is a slow drift down a basin that the winning configuration simply does not fall into.
fig, ax = plt.subplots(figsize=(9.5, 5))
styles = {(False, False): ('#c0392b', '-', 'plain Siamese'),
(False, True): ('#e67e22', '--', 'stop-grad only'),
(True, False): ('#8e44ad', '-.', 'predictor only'),
(True, True): ('#27ae60', '-', 'predictor + stop-grad (SimSiam)')}
for (up, sg), (col, ls, lab) in styles.items():
traj = results[(up, sg)][2]
xs = np.arange(len(traj)) * 40
ax.plot(xs, traj, ls, color=col, lw=2.2 if (up and sg) else 1.6, label=lab)
ax.axhline(16, ls=':', color='gray', alpha=0.7)
ax.set_xlabel('training epoch'); ax.set_ylabel('effective dimension')
ax.set_title('Effective dimension over training: who falls into the collapse basin')
ax.legend(); ax.set_ylim(0, 17); plt.tight_layout(); plt.show()
What you see¶
Every curve starts near the initial effective dimension (~9.8), then the three ablations slide down and plateau around 4 — they lose more than half their dimensions in the first few hundred epochs and never recover. The SimSiam curve (green, bold) is the lone exception: it dips slightly and then holds flat near its starting value for the entire run.
The shape matters for practice. Collapse is a gradual loss of rank, not a sudden switch, which is exactly why the effective-dimension monitor from the previous note is useful as an early warning: you can watch the ablations diverging from the healthy run long before the final epoch. The mechanism does not add a force that fights collapse — it removes the slope that leads there.
Experiment 3: An EMA target also works (BYOL)¶
SimSiam shares one encoder and detaches the target. BYOL instead builds the target from a separate EMA (momentum) copy of the encoder. SimSiam's contribution was showing the EMA is not required — but it does help. We compare the SimSiam target (shared + stop-grad) to an EMA target, both with a predictor.
final_simsiam, traj_simsiam = train(use_pred=True, use_sg=True)
final_byol, traj_byol = train(use_pred=True, use_sg=False, ema=True, ema_tau=0.99)
print(f"SimSiam (shared + stop-grad) final eff_dim = {final_simsiam:.2f}")
print(f"BYOL (EMA target) final eff_dim = {final_byol:.2f}")
print(f"std of last-half trajectory SimSiam={traj_simsiam[len(traj_simsiam)//2:].std():.3f}"
f" BYOL={traj_byol[len(traj_byol)//2:].std():.3f}")
SimSiam (shared + stop-grad) final eff_dim = 9.74 BYOL (EMA target) final eff_dim = 9.28 std of last-half trajectory SimSiam=0.126 BYOL=0.048
fig, ax = plt.subplots(figsize=(9.5, 5))
xs = np.arange(len(traj_simsiam)) * 40
ax.plot(xs, traj_simsiam, color='#27ae60', lw=2, label=f'SimSiam (stop-grad), final {final_simsiam:.1f}')
ax.plot(np.arange(len(traj_byol)) * 40, traj_byol, color='#2980b9', lw=2,
label=f'BYOL (EMA target), final {final_byol:.1f}')
ax.axhline(16, ls=':', color='gray', alpha=0.7)
ax.set_xlabel('training epoch'); ax.set_ylabel('effective dimension')
ax.set_title('Stop-gradient vs. EMA target: both avoid collapse')
ax.legend(); ax.set_ylim(0, 17); plt.tight_layout(); plt.show()
What you see¶
Both the shared stop-gradient target (green) and the EMA target (blue) keep the representation healthy — neither collapses. This is the point SimSiam made: the momentum encoder of BYOL is a stabiliser, not the source of collapse-avoidance; a plain stop-gradient on a shared encoder already does the job.
The EMA variant trades a tiny bit of final effective dimension for a smoother, more stable target (its target changes slowly by construction). That smoothness is why momentum/EMA targets became the default in larger systems — the subject of the next note, where the same EMA idea reappears in DINO and JEPA as the way to build a reliable prediction target without a decoder.
Experiment 4: The mechanism is fragile¶
Collapse-avoidance here is an implicit property of the architecture, not a guaranteed floor like VICReg's variance term. So it can break. We take the winning configuration and starve the predictor — scale down its learning rate — so it can no longer adapt. As the predictor freezes toward identity, the safeguard disappears and the representation collapses again.
mults = [0.0, 0.1, 0.25, 0.5, 1.0, 1.5]
finals_mult = []
for m in mults:
f, _ = train(use_pred=True, use_sg=True, pred_lr_mult=m)
finals_mult.append(f)
print(f"predictor lr x{m:<4} -> final eff_dim = {f:.2f}")
predictor lr x0.0 -> final eff_dim = 3.41
predictor lr x0.1 -> final eff_dim = 4.51
predictor lr x0.25 -> final eff_dim = 6.26
predictor lr x0.5 -> final eff_dim = 6.30
predictor lr x1.0 -> final eff_dim = 9.74
predictor lr x1.5 -> final eff_dim = 10.99
fig, ax = plt.subplots(figsize=(9, 4.8))
ax.plot(mults, finals_mult, 'o-', color='#2980b9', lw=2)
ax.axhline(16, ls=':', color='gray', alpha=0.7, label='ambient dim = 16')
ax.axhspan(0, 5, color='#c0392b', alpha=0.08)
ax.text(0.05, 4.4, 'collapse regime', color='#c0392b', fontsize=10)
ax.set_xlabel('predictor learning-rate multiplier'); ax.set_ylabel('final effective dimension')
ax.set_title('Starve the predictor and the safeguard fails')
ax.set_ylim(0, 17); ax.legend(); plt.tight_layout(); plt.show()
What you see¶
With the predictor learning at full speed (multiplier 1.0+) the representation stays healthy, but as its learning rate is scaled toward zero the effective dimension falls back into the collapse regime — a frozen predictor is effectively no predictor, the exact ablation from Experiment 1. The safeguard depends on the predictor being able to track the feature statistics; cripple that and collapse returns.
This is the honest caveat of the whole approach: it works through an implicit bias of the training dynamics, with no explicit term in the loss reporting health. That is why a framework treating latent space as first-class must monitor collapse directly (effective dimension) rather than trust that an architecture is safe — the architecture is only safe inside a regime you have to stay within.
Summary / Key Takeaways¶
- Both ingredients are necessary (Exp 1). Predictor and stop-gradient — remove either and the representation collapses to ~1/4 of its dimensions, even though the loss is pure invariance.
- Collapse is a slow drift (Exp 2). The ablations slide into a low-rank basin over training; the winning config simply never enters it. Effective dimension gives early warning.
- EMA targets work too (Exp 3). A BYOL-style momentum target also avoids collapse and is smoother — momentum is a stabiliser, not the cause of collapse-avoidance. This bridges to the next note.
- The mechanism is fragile (Exp 4). Starving the predictor reopens the collapse regime. There is no explicit health term, so Layer A must monitor effective dimension directly.
- Bridge. Next: contrastive learning — the other anti-collapse route, which adds an explicit repulsive force (negatives) and comes with a mutual-information bound the architectural tricks lack.