Chapter 28 · Section 28.3 · 90 min

Non-contrastive methods and collapse

A non-contrastive method learns image features by asking two views of the same photograph to agree, and it asks nothing else. It uses no negatives, the other images that a contrastive method pushes away. That makes it cheaper and simpler, but it opens a trap. A network that outputs the same vector for every photograph makes every pair of views agree perfectly. Its loss is the best possible, and it has learned nothing. This failure is called collapse. This section explains how methods without negatives avoid collapse, and how you can tell when a training run has fallen into it.

Here is the setting, recalled from §28.2. Take one unlabeled photo and make two views of it. Each view is a random crop of the photo with a random color distortion. An encoder network turns each view into a vector. A contrastive loss pulls the two vectors of one photo together. It also pushes them away from the vectors of every other photo in the batch. Those other photos are the negatives, and they are what stops the encoder from mapping everything to one point. They are also expensive. SimCLR used batches of 4,096 images spread over 128 TPU cores, and MoCo kept a queue of 65,536 stored vectors.

Delete the negatives and only the pull is left. A loss that only pulls is minimized by a constant output. The book has shown this twice. In §1.10, a linear encoder trained only to match two views shrank to the zero map. In §7.2, a retrieval encoder trained only to pull pairs together retrieved no better than a blind guess.

And yet methods with no negatives work, and some beat SimCLR. The yardstick is the linear-probe accuracy of §28.1. You freeze the trained encoder, fit one linear classifier on its features with the ImageNet labels, and report that classifier's accuracy. With a ResNet-50, SimCLR reached 69.3 percent. BYOL, which has no negatives, reached 74.3 (Grill et al. 2020). SimSiam, simpler still, reached 67.7 percent after 100 epochs with an ordinary batch and plain SGD (Chen and He 2021). Remove one line of SimSiam's code, a stop-gradient defined below, and it scores 0.1 percent, which is chance on 1,000 classes. Nothing in either loss forbids the constant solution. So what keeps these networks away from it?

There are two kinds of answer. Either the loss makes the constant expensive, or training is arranged so that it never gets there. Methods of the first kind add a term to the loss. The term keeps every output dimension spread out, or decorrelates the dimensions, or forces each batch to use many clusters. Methods of the second kind make the two branches of the network different. One branch gets a stop-gradient, the other gets an extra small network, or the target changes slowly.

This section builds both kinds. We first make collapse precise and measurable. Then we take apart BYOL and SimSiam, which are of the second kind. In a linear network small enough to solve by hand, we prove that the stop-gradient is what saves them. Then we derive three methods of the first kind, Barlow Twins, VICReg, and SwAV. We end with dimensional collapse, a quieter failure that strikes even methods with negatives.

The symbols are those of the whole chapter, and here they are in one place.

$x$ is an unlabeled image. $t$ and $t'$ are two random augmentations drawn from the augmentation policy $\mathcal{T}$ of §28.2, so the two views are $x^{(1)} = t(x)$ and $x^{(2)} = t'(x)$ .
The encoder $f_\theta$ is a ResNet-50, or a ViT followed by pooling. It produces the representation $\mathbf{h} = f_\theta(x) \in \R^{d}$ , the vector every evaluation reads. For a ResNet-50, $d = 2048$ .
The projection head $g_\theta$ is a small MLP. It maps the representation to the embedding $\mathbf{z} = g_\theta(\mathbf{h}) \in \R^{d_z}$ , the vector the loss sees. §28.2 showed why the head is thrown away after training.
$\theta$ holds the trainable parameters. A batch holds $B$ images, so $2B$ views.
$\hat{\mathbf{z}} = \mathbf{z}/\lVert\mathbf{z}\rVert$ is the unit vector along $\mathbf{z}$ .
In this section $\lambda$ always means weight decay. Other weights get subscripts, such as $\lambda_{\text{BT}}$ .

What collapse looks like, and how to measure it

You cannot read collapse off the loss, because a collapsed network usually has an excellent loss. SimSiam without its stop-gradient quickly reaches $-1$ , the lowest value its loss can take. You have to look at the embeddings themselves. There are three ways for them to fail.

📐 Three kinds of collapse

Let $\mathbf{Z} \in \R^{n\times d_z}$ stack the embeddings of $n$ inputs, one per row.

Complete collapse. The embedding does not depend on the input, $\mathbf{z} = \mathbf{c}$ for every $x$ , so every row of $\mathbf{Z}$ is the same vector.
Dimensional collapse. The embeddings fill only a subspace of dimension $r \ll d_z$ . Their covariance matrix has $d_z - r$ eigenvalues that are zero or negligible. Equivalently, the centered matrix $\mathbf{Z} - \mathbf{1}\bar{\mathbf{z}}^\top$ , with the mean row $\bar{\mathbf{z}}$ subtracted from every row, has $d_z - r$ singular values that are zero or negligible.
Uniform collapse. Some methods end in a softmax, so each image gets a probability distribution over many outputs. SwAV below is one, and DINO in §28.4 is another. Uniform collapse gives every image the same uniform distribution.

Complete collapse is easy to see once you look. Dimensional collapse is the dangerous one. The network still tells inputs apart, but only along a few directions, so the loss and even a quick probe can look fine. Three measurements catch these failures, and we build them one at a time.

Measure one, the spread of each coordinate

Normalize every embedding to unit length. Take the standard deviation of each coordinate over a batch, and average over the coordinates. A completely collapsed network scores 0, because every unit vector is the same. A healthy network, whose unit vectors point evenly in every direction, scores $1/\sqrt{d_z}$ , up to the sampling noise of the batch. For $d_z = 2048$ that is $1/\sqrt{2048} = 0.0221$ . SimSiam's authors used this reference value. Their Figure 2 shows two curves, one sitting near 0.0221 with the stop-gradient and one falling to 0 without it.

∂ Proof by symmetry

Flipping the sign of coordinate $i$ is an orthogonal map, so $\hat z_i$ and $-\hat z_i$ have the same distribution, and $\mathbb{E}[\hat z_i] = 0$ . Swapping coordinates $i$ and $j$ is orthogonal too, so $\mathbb{E}[\hat z_i^2] = \mathbb{E}[\hat z_j^2]$ for all $i, j$ . These $d_z$ equal second moments add up to $\mathbb{E}\lVert\hat{\mathbf{z}}\rVert^2 = 1$ , so each is $1/d_z$ . The standard deviation is then $\sqrt{1/d_z - 0^2} = 1/\sqrt{d_z}$ . ∎

This test has a blind spot, which the cell below shows. Embeddings confined to an 8-dimensional subspace of $\R^{128}$ score almost as well as healthy ones. Their unit vectors still point in many different directions, so every coordinate still varies. To see dimensional collapse you need the spectrum.

Measure two, the spectrum

Stack the embeddings of a batch into the matrix $\mathbf{Z}$ , subtract the mean of each column, and compute the singular values (§0.10). Each nonzero singular value stands for one direction along which the embeddings vary, and its size says how much they vary along it. A completely collapsed batch has every singular value equal to zero. A batch confined to $r$ directions has $r$ nonzero singular values and the rest at zero. A healthy batch has $d_z$ singular values of similar size.

Measure three, the effective rank

A spectrum is a list of $d_z$ numbers. When you compare a hundred runs, you want one number per run. The effective rank counts how many directions share the spread. Three tiny cases show what it should do. Singular values $(2, 2, 0)$ mean two equal directions, so the count should be 2. Singular values $(1, 1, 1, 1)$ should count 4. Singular values $(10, 0.1, 0.1)$ should count a little more than 1, because one direction holds almost everything. The exponential of an entropy does exactly this.

📐 Effective rank (RankMe)

For $\mathbf{Z} \in \R^{n\times d_z}$ with singular values $\sigma_1 \ge \dots \ge \sigma_{r}$ , $r = \min(n, d_z)$ , not all zero, define

p_k = \frac{\sigma_k}{\sum_{j=1}^{r}\sigma_j} + \epsilon, \qquad \operatorname{erank}(\mathbf{Z}) = \exp\Big(-\sum_{k=1}^{r} p_k\log p_k\Big),

the exponential of the entropy of the normalized singular values (Roy and Vetterli 2007). Garrido et al. (2023) named it RankMe. Their formula has no centering step, so it reads the embedding matrix as it is. The small $\epsilon$ keeps zero singular values from producing $0\log 0$ . The paper gives no value for it, and this section uses $10^{-7}$ . Ignoring $\epsilon$ , $1 \le \operatorname{erank}(\mathbf{Z}) \le r$ . The lower bound holds exactly when one singular value is nonzero, and the upper bound exactly when all $r$ are equal.

∂ Proof of the bounds

Without $\epsilon$ , the $p_k$ form a probability distribution on $r$ outcomes. The entropy of such a distribution lies between 0, reached only by a point mass, and $\log r$ , reached only by the uniform distribution (§0.5). Exponentiating gives $1 \le \operatorname{erank} \le r$ . ∎

In plain words, the effective rank counts the directions that carry the embeddings' spread, each weighted by its share. Check it on the third tiny case. With $(10, 0.1, 0.1)$ the shares are $(0.980, 0.0098, 0.0098)$ , the entropy is $0.110$ , and the effective rank is $e^{0.110} = 1.12$ .

The effective rank is a cousin of the stable rank of §0.10. The stable rank divides the squared Frobenius norm by the largest squared singular value, so it compares the spectrum with its top value only. The effective rank depends on the shape of the whole spectrum.

Two details change the number in practice. RankMe uses the singular values themselves, not their squares (the eigenvalues of $\mathbf{Z}^\top\mathbf{Z}$ ). And RankMe's formula does not center, so a nonzero mean adds a direction. When you compare runs, use one convention and say which.

✏️ Collapse measures for three 4 by 3 matrices, by hand

Take four embeddings in $\R^3$ .

All four rows equal to $(1, 2, 2)$ . The matrix is $\mathbf{1}(1,2,2)$ , rank 1, with $\sigma_1 = \lVert(1,2,2)\rVert\cdot\lVert\mathbf{1}\rVert = 3\cdot 2 = 6$ , so $\operatorname{erank} = 1$ . Centered, it is the zero matrix. This is complete collapse.
Rows $(1,1,0)$ , $(1,1,0)$ , $(1,-1,0)$ , $(1,-1,0)$ . Then $\mathbf{Z}^\top\mathbf{Z} = 2(1,1,0)^\top(1,1,0) + 2(1,-1,0)^\top(1,-1,0) = \operatorname{diag}(4, 4, 0)$ , so $\sigma = (2, 2, 0)$ and $\operatorname{erank} = \exp(\log 2) = 2$ . Centering subtracts the mean row $(1, 0, 0)$ and leaves rows $(0, \pm 1, 0)$ , with singular values $(2, 0, 0)$ . The shared direction $(1,0,0)$ counts for RankMe and vanishes after centering.
Rows $(2,0,0)$ , $(0,1,0)$ , $(0,0,1)$ , $(0,0,0)$ . Then $\sigma = (2, 1, 1)$ , $p = (\tfrac12, \tfrac14, \tfrac14)$ , the entropy is $\tfrac12\log 2 + 2\cdot\tfrac14\log 4 = \tfrac32\log 2$ , and $\operatorname{erank} = 2^{3/2} = 2.83$ . That lies between 2 and 3 because the spread is uneven.

Now a small simulation. The cell runs all three measurements on three sets of 1,024 embeddings of width 128. One set is completely collapsed. One is confined to a rank-8 subspace. One is isotropic, meaning spread evenly in every direction.

Python · cell 1

The collapsed set has standard deviation 0, centered singular values that are zero up to rounding, and effective rank 1. The isotropic set sits at the reference, 0.0883 against $1/\sqrt{128} = 0.0884$ , with effective rank 125.9 of 128. The rank-8 set is the lesson. Its standard deviation, 0.0857, is within 3 percent of healthy. Its effective rank, 7.94, exposes it at once.

Warm-up

Diagnose collapse

Write collapse_report(Z, eps=1e-7). It takes a batch of embeddings Z of shape (n, d_z) with no zero rows, and it returns a dict with the three measurements above plus one more.

"std" is the per-coordinate spread. Normalize each row to unit length, take the standard deviation of each coordinate over the rows (ddof=0), and average over the coordinates.
"sv" holds all min(n, d_z) singular values of the column-centered matrix Z - Z.mean(0), in descending order.
"erank" is RankMe of the uncentered Z, with $p_k = \sigma_k/\sum_j\sigma_j + \epsilon$ and natural logarithms.
"k99" is the smallest $k$ such that the top $k$ squared centered singular values hold at least 99 percent of their total. Return 0 if the centered matrix has no spread, meaning that total is at most $10^{-12}$ times the sum of the squared entries of Z.

import numpy as np

def collapse_report(Z, eps=1e-7):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

BYOL, a network that predicts a slowly moving copy of itself

BYOL, short for Bootstrap Your Own Latent, trains a network to predict its own output for a different view of the same image. The "own output" comes from a slowly updated copy of the network. The method starts from an experiment that sounds too simple to work. Freeze a randomly initialized network and call it the target. Train a second network to predict the target's embedding of one view from the other view. The constant solution is impossible here, because the target is fixed and does not output a constant.

The trained network reached 18.8 percent linear-probe accuracy on ImageNet, while the random target itself scored 1.4 (Grill et al. 2020, Table 5). Predicting a poor representation produced a better one. Now make the improved network the new target, and repeat. BYOL runs that bootstrap continuously, at every training step.

Three pieces make it work, and each needs a name.

The online network is the one trained by gradient descent.
The target network has the same architecture but never receives gradients. After every step, its weights move a small fraction of the way toward the online weights. This update is an exponential moving average (EMA), the same one §28.2 built for MoCo's key encoder.
The predictor is a small MLP that sits only on the online side, after the projection head. It turns the online embedding into a guess of the target's embedding.

📐 BYOL

The online network, with parameters $\theta$ , has an encoder $f_\theta$ , a projector $g_\theta$ (a linear layer to 4,096, BatchNorm, ReLU, and a linear layer to $d_z = 256$ ), and a predictor $q_\theta$ of the same shape. The target network, with parameters $\xi$ , has an encoder and a projector of the same architecture and no predictor. For two views, the online network computes $\mathbf{q} = q_\theta(g_\theta(f_\theta(x^{(1)})))$ and the target computes $\mathbf{z}' = g_\xi(f_\xi(x^{(2)}))$ . The loss is

\ell = \lVert\hat{\mathbf{q}} - \hat{\mathbf{z}}'\rVert^2 = 2 - 2\cos(\mathbf{q}, \mathbf{z}'),

plus the same term with the views swapped, averaged over the batch. Gradients update $\theta$ only. After each step the target follows the online network, $\xi \leftarrow m\xi + (1-m)\theta$ . The momentum $m$ rises from $m_0 = 0.996$ to 1 over training, $m_k = 1 - (1 - m_0)\big(\cos(\pi k/K_{\text{steps}}) + 1\big)/2$ at step $k$ of $K_{\text{steps}}$ . Training uses LARS, the layer-wise learning-rate rule of §28.2, with base learning rate $0.2\times B/256$ , weight decay $1.5\times10^{-6}$ , batch 4,096, and 1,000 epochs. That takes about 8 hours for a ResNet-50 on 512 TPU v3 cores.

The equality in the loss is the identity $\lVert\hat{\mathbf{a}} - \hat{\mathbf{b}}\rVert^2 = 2 - 2\hat{\mathbf{a}}^\top\hat{\mathbf{b}}$ for unit vectors. §28.2 derived what the moving average does. It averages the online weights over a window of about $1/(1-m)$ steps, and it lags behind them by $m/(1-m)$ steps on average. At $m_0 = 0.996$ the target trails by $0.996/0.004 = 249$ steps. The cosine schedule gives $m = 0.998$ halfway through and $m = 1$ at the end. So the target changes ever more slowly as the online network settles.

Notice what is missing. Nothing pushes two different images apart. Why, then, does the online network not drift to a constant, which the target, a moving average of it, would eventually copy?

The BYOL paper gives an argument. It is worth stating carefully, because it is often overstated. Suppose the predictor is always the best possible one for the current online and target networks. Then an online network that carries any information about the image predicts the target at least as well as a collapsed one. So collapsing gains nothing. The derivation below makes this exact.

The argument is about equilibria under two assumptions, an optimal predictor and a fixed target. It is not a proof that training reaches a good place, and the BYOL authors say so. Collapsed equilibria still exist. The authors report only that training never converged to one.

∂ BYOL's argument from an optimal predictor

For squared error, the best predictor of the target embedding $\mathbf{z}'$ from the online embedding $\mathbf{z}$ is the conditional mean, $q^\star(\mathbf{z}) = \mathbb{E}[\mathbf{z}' \mid \mathbf{z}]$ (§0.4). The loss it achieves is $\mathbb{E}\sum_i\operatorname{Var}(z'_i \mid \mathbf{z})$ . By the law of total variance, $\operatorname{Var}(z'_i) = \mathbb{E}[\operatorname{Var}(z'_i \mid \mathbf{z})] + \operatorname{Var}(\mathbb{E}[z'_i \mid \mathbf{z}])$ , so

\mathbb{E}[\operatorname{Var}(z'_i \mid \mathbf{z})] \le \operatorname{Var}(z'_i).

The right side is what a constant online output achieves, since conditioning on a constant tells you nothing. More generally, replacing $\mathbf{z}$ by any function of it can only discard information, so it never lowers $\mathbb{E}[\operatorname{Var}(z'_i \mid \cdot)]$ . With an optimal predictor, the online network gains nothing by collapsing. ∎

📌 What BYOL's ablations show

From Grill et al. (2020), Table 5 and the appendix, ImageNet linear-probe accuracy of a ResNet-50 after 300 epochs, with base momentum 0.99 in these ablations.

Target momentum. A constant random target ( $m = 1$ ) gives 18.8, $m$ = 0.999 gives 69.8, 0.99 gives 72.5, and 0.9 gives 68.4. With $m = 0$ , where the target is a detached copy of the online network at every step, it collapses to 0.3. A randomly initialized network alone gives 1.4.
Removing the predictor collapses the network (0.2), even with the moving-average target.
Batch size. Accuracy is stable from 256 to 4,096 images and drops below 256. The authors attribute the drop to the BatchNorm layers in the encoder, whose batch statistics get noisy. SimCLR degrades quickly as the batch shrinks.
Augmentation. Removing color distortion costs BYOL 9.1 points and SimCLR 22.2.
The moving average can be dropped ( $m = 0$ ) if the predictor is kept near optimal. Solving for a linear predictor in closed form on each batch gives 52.5 (Section 5). Multiplying the predictor's learning rate by 10 gives 66.6, while leaving it equal to the rest of the network's gives 5.5 (Appendix I.1, Table 21).

The last line is the clue that the rest of this section develops. A slow target and a fast predictor are two ways of doing one job, which is keeping the predictor ahead of the network it predicts. In an interview, the expected answer to "what does BYOL need?" is the predictor and the asymmetry between the branches. The moving average is a stabilizer on top.

SimSiam, the stop-gradient alone

If the predictor matters more than the moving average, how much of BYOL can you delete? Chen and He (2021) deleted the target network entirely. Both branches use the same network with the same weights, an arrangement called a Siamese network. The only asymmetry left is a stop-gradient. A stop-gradient treats a value as a constant during backpropagation. The forward pass uses the value, but no gradient flows back through it. In PyTorch it is .detach().

📐 SimSiam

One network, an encoder $f_\theta$ followed by a three-layer projector $g_\theta$ of width 2,048 with BatchNorm on every layer, so $d_z = 2048$ , and a predictor $q_\theta$ (a two-layer bottleneck, 2,048 to 512 to 2,048). For two views, $\mathbf{z}^{(v)} = g_\theta(f_\theta(x^{(v)}))$ and $\mathbf{q}^{(v)} = q_\theta(\mathbf{z}^{(v)})$ for $v = 1, 2$ . With $D(\mathbf{q}, \mathbf{z}) = -\cos(\mathbf{q}, \mathbf{z})$ , the loss is

\mathcal{L} = \tfrac12 D\big(\mathbf{q}^{(1)}, \operatorname{sg}(\mathbf{z}^{(2)})\big) + \tfrac12 D\big(\mathbf{q}^{(2)}, \operatorname{sg}(\mathbf{z}^{(1)})\big),

averaged over the batch, where $\operatorname{sg}$ is the stop-gradient. There is no momentum encoder and no negative. Training uses SGD with momentum 0.9, base learning rate $0.05\times B/256$ , and weight decay $10^{-4}$ , with a default batch of 512.

In terms of architecture, SimSiam is BYOL with $m = 0$ , the setting that gave BYOL 0.3 percent. So why does one recipe fail and the other work? BYOL's own appendix gives part of the answer. With no target network and a predictor learning rate raised tenfold, BYOL reached 66.6 percent. The two recipes also differ in weight decay, $10^{-6}$ with LARS in BYOL's ablation and $10^{-4}$ with SGD in SimSiam. Weight decay and the predictor's speed are exactly the quantities the linear theory below singles out. The balance is fragile, and in practice you tune those two first.

The predictor is essential, and SimSiam's paper has an exact argument for why.

Proof. The first term sends gradient only through $\mathbf{z}^{(1)}$ , namely $\tfrac12\,\partial D(\mathbf{z}^{(1)}, \mathbf{z}^{(2)})/\partial\mathbf{z}^{(1)}$ . The second sends gradient only through $\mathbf{z}^{(2)}$ , namely $\tfrac12\,\partial D(\mathbf{z}^{(2)}, \mathbf{z}^{(1)})/\partial\mathbf{z}^{(2)}$ . Because $D$ is symmetric in its two arguments, $\partial D(\mathbf{z}^{(2)}, \mathbf{z}^{(1)})/\partial\mathbf{z}^{(2)} = \partial D(\mathbf{z}^{(1)}, \mathbf{z}^{(2)})/\partial\mathbf{z}^{(2)}$ . So the two pieces add up to exactly half the full gradient of $D(\mathbf{z}^{(1)}, \mathbf{z}^{(2)})$ , which flows through both arguments. ∎

In plain words, without a predictor the stop-gradient changes only the step size. The network then minimizes $D$ directly, and the constant solution wins. The PyTorch block near the end checks this identity with autograd.

📌 What SimSiam's experiments show

From Chen and He (2021), ImageNet linear probe with a ResNet-50.

After 100 epochs, 67.7 percent with the stop-gradient (mean of 5 trials, standard deviation 0.1) and 0.1 percent without it. Without it, the per-coordinate standard deviation of $\hat{\mathbf{z}}$ falls to 0. With it, it stays near $1/\sqrt{d_z}$ .
No predictor gives 0.1. A predictor frozen at its random initialization gives 1.5, not by collapsing but because training does not converge. A predictor whose learning rate is not decayed gives 68.1, slightly better than the baseline.
Batch sizes from 256 to 2,048 give similar results with plain SGD, 67.9 to 68.1. Batches of 128 and 64 lose 0.8 and 2.0 points. A batch of 4,096 falls to 64.0, which the authors attribute to plain SGD at a very large batch.
With batch 256, linear accuracy is 68.1, 70.0, 70.8, and 71.3 after 100, 200, 400, and 800 epochs.

🔬 SimSiam's hypothesis, an alternating optimization

Chen and He offer an explanation and call it a hypothesis. SimSiam may behave like an alternating optimization with two sets of variables. One set is the network $\theta$ . The other is a hidden target $\boldsymbol\eta_x$ for each image, which stands in for the output of the stop-gradient branch. Fix every $\boldsymbol\eta_x$ and take an SGD step on $\theta$ . Then set each $\boldsymbol\eta_x$ to the network's current average output over augmentations, and repeat.

This resembles k-means (§7.4), which alternates between moving the centers and reassigning the points. Each half-step of k-means solves a well-posed problem, and k-means does not collapse. In this reading, the predictor stands in for the average over augmentations, which one sampled view cannot provide.

A hypothesis is not a mechanism. To see a mechanism, we need a model small enough to solve.

A linear Siamese network you can solve

Tian, Chen and Ganguli (2021) found such a model. Strip the network down to one number per part, and the training dynamics can be solved by hand. The solution shows exactly what the stop-gradient, the predictor, and the weight decay each do.

Here is the model. The data are $x \sim \mathcal{N}(0, 1)$ . Each view adds independent augmentation noise of variance $\sigma^2$ , so $x^{(1)} = x + \sigma\epsilon_1$ and $x^{(2)} = x + \sigma\epsilon_2$ , with $\epsilon_1, \epsilon_2 \sim \mathcal{N}(0,1)$ . The online network is one weight $w$ . The predictor is one weight $p$ (a number, not a probability). The target is the online weight itself, as in SimSiam. The loss is the squared error between the predicted online output and the target output,

\mathcal{L}(w, p) = \tfrac12\,\mathbb{E}\big[(p\,w\,x^{(1)} - w\,x^{(2)})^2\big].

Two expectations do all the work. The first is $\mathbb{E}[(x^{(1)})^2] = \mathbb{E}[x^2] + 2\sigma\mathbb{E}[x\epsilon_1] + \sigma^2\mathbb{E}[\epsilon_1^2] = 1 + \sigma^2$ , since $x$ and $\epsilon_1$ are independent with mean 0. The second is $\mathbb{E}[x^{(1)}x^{(2)}] = \mathbb{E}[x^2] = 1$ , since every other term contains an independent zero-mean factor. Expanding the square gives $\mathcal{L} = \tfrac12 w^2\big(p^2(1+\sigma^2) - 2p + (1 + \sigma^2)\big)$ .

With the stop-gradient, the target $w\,x^{(2)}$ is held constant while we differentiate. So only the online copy of $w$ receives gradient,

\frac{\partial\mathcal{L}}{\partial w}\Big|_{\text{sg}} = \mathbb{E}\big[(p w x^{(1)} - w x^{(2)})\,p\,x^{(1)}\big] = -wp\big(1 - (1+\sigma^2)p\big), \qquad \frac{\partial\mathcal{L}}{\partial p} = -w^2\big(1 - (1+\sigma^2)p\big).

Without the stop-gradient, $w$ receives gradient through both branches, and

\frac{\partial\mathcal{L}}{\partial w}\Big|_{\text{no sg}} = w\,\mathbb{E}\big[(p x^{(1)} - x^{(2)})^2\big] = w\big(p^2(1+\sigma^2) - 2p + 1 + \sigma^2\big).

✏️ The three gradients at one point, by hand

Take $\sigma^2 = 0.25$ , $w = 1$ , $p = 0.5$ , so $1 - (1 + \sigma^2)p = 1 - 0.625 = 0.375$ . With the stop-gradient, $\partial\mathcal{L}/\partial w = -1\cdot0.5\cdot0.375 = -0.1875$ and $\partial\mathcal{L}/\partial p = -0.375$ . Both are negative, so gradient descent increases both $w$ and $p$ . Without the stop-gradient, $\partial\mathcal{L}/\partial w = 0.25\cdot1.25 - 1 + 1.25 = 0.5625$ , which is positive, so $w$ shrinks. A Monte Carlo estimate with 400,000 samples gives $-0.1885$ , $-0.3770$ , and $0.5642$ .

The same point gets two opposite verdicts. With the stop-gradient the online weight grows. Without it the weight shrinks. To see where each run ends, follow the gradients over time.

Gradient flow is the small-learning-rate limit of gradient descent (§0.6). It is written $\dot w = -\partial\mathcal{L}/\partial w$ , where the dot is the derivative in training time. Three additions bring the model closer to practice.

The predictor gets a learning rate $\alpha_p$ times that of the online weight.
Both weights get weight decay $\lambda$ .
The target becomes $\kappa w$ instead of $w$ . The value $\kappa = 1$ is SimSiam. A $\kappa$ below 1 models a moving-average target, which lags behind a growing online weight. Tian et al. call this ratio $\tau$ , a letter this book keeps for temperatures.

Finally, write $s = w^2$ for the squared online weight. It measures how much signal the online output carries, which is what a probe can use. Collapse is $s = 0$ .

∴ Without the stop-gradient, a linear Siamese network collapses, and with it, it need not

In the model above with $\sigma^2 > 0$ and target $\kappa w$ , the stop-gradient flow is

\dot p = \alpha_p\,s\,\big(\kappa - (1+\sigma^2)p\big) - \lambda p, \qquad \dot s = 2ps\big(\kappa - (1+\sigma^2)p\big) - 2\lambda s .

Conservation. $\frac{d}{dt}\big(s - p^2/\alpha_p\big) = -2\lambda\big(s - p^2/\alpha_p\big)$ . Without weight decay the quantity $c = s - p^2/\alpha_p$ is conserved. A run that starts above the parabola $s = p^2/\alpha_p$ therefore keeps $s \ge c > 0$ forever and never collapses. Its predictor moves monotonically to $p^\star = \kappa/(1+\sigma^2)$ .
Weight decay. For $\lambda > 0$ , $c(t) = c(0)e^{-2\lambda t}$ , so every run approaches the parabola $s = p^2/\alpha_p$ . On it the nonzero fixed points solve $(1+\sigma^2)p^2 - \kappa p + \lambda = 0$ , that is $p_\pm = \big(\kappa \pm \sqrt{\kappa^2 - 4\lambda(1+\sigma^2)}\big)/\big(2(1+\sigma^2)\big)$ . They exist only if $\lambda \le \kappa^2/\big(4(1+\sigma^2)\big)$ . On the parabola, $p_+$ is stable, $p_-$ is unstable and marks the edge of the collapse basin, and the collapsed point $p = 0$ is stable too. Above the threshold, $s(t) \to 0$ from every start, so collapse is the only outcome.
Without the stop-gradient ( $\kappa = 1$ ), $\dot w = -w\big(p^2(1+\sigma^2) - 2p + 1 + \sigma^2\big) - \lambda w$ . The bracket is at least $\sigma^2(2+\sigma^2)/(1+\sigma^2) > 0$ for every $p$ , so $|w(t)| \le |w(0)|\,e^{-\sigma^2(2+\sigma^2)t/(1+\sigma^2)} \to 0$ whatever the predictor does.

∂ Proof of the three items

The flow equations are the gradients above with $\kappa$ in place of the target weight 1, since $\mathbb{E}[x^{(1)}\cdot\kappa x^{(2)}] = \kappa$ . For $\dot s$ , use $\dot s = 2w\dot w$ and $\dot w = wp(\kappa - (1+\sigma^2)p) - \lambda w$ .

For item 1, $\frac{d}{dt}(s - p^2/\alpha_p) = \dot s - 2p\dot p/\alpha_p$ . The data terms are $2ps(\kappa - (1+\sigma^2)p)$ and $-2p\cdot s(\kappa - (1+\sigma^2)p)$ , which cancel. The decay terms are $-2\lambda s + 2\lambda p^2/\alpha_p = -2\lambda(s - p^2/\alpha_p)$ . With $\lambda = 0$ the derivative is 0. Since $s \ge c > 0$ , the factor $\kappa - (1+\sigma^2)p$ must vanish at a fixed point, which gives $p^\star$ . And $\dot p$ has the sign of $\kappa - (1+\sigma^2)p$ , so $p$ moves toward $p^\star$ without overshooting.

For item 2, solve the linear equation $\dot c = -2\lambda c$ . On the parabola, substitute $s = p^2/\alpha_p$ into $\dot p$ to get $\dot p = p\,\big(p(\kappa - (1+\sigma^2)p) - \lambda\big)$ . Its nonzero roots solve the stated quadratic, and they are real exactly when its discriminant $\kappa^2 - 4\lambda(1+\sigma^2)$ is nonnegative. The bracket $p(\kappa - (1+\sigma^2)p) - \lambda$ is negative for $p < p_-$ , positive between the roots, and negative above $p_+$ . So for $p > 0$ the flow pushes $p$ away from $p_-$ and toward $p_+$ or 0, which gives the stability claims. Above the threshold, use $p(\kappa - (1+\sigma^2)p) \le \kappa^2/\big(4(1+\sigma^2)\big)$ , the maximum of a downward parabola. Then $\dot s = 2s\big(p(\kappa - (1+\sigma^2)p) - \lambda\big) \le -2s\big(\lambda - \kappa^2/(4(1+\sigma^2))\big)$ , a negative multiple of $s$ , so $s$ decays exponentially from every start.

For item 3, the bracket is a quadratic in $p$ with positive leading coefficient $1+\sigma^2$ . Its minimum is at $p = 1/(1+\sigma^2)$ , where it equals $(1+\sigma^2) - 1/(1+\sigma^2) = \sigma^2(2+\sigma^2)/(1+\sigma^2)$ , positive for $\sigma^2 > 0$ . So $\frac{d}{dt}w^2 \le -2\,\frac{\sigma^2(2+\sigma^2)}{1+\sigma^2}w^2$ . A quantity whose rate of decrease is at least $k$ times itself decays at least as fast as $e^{-kt}$ (Grönwall's inequality), which gives the bound. ∎

In plain words, the loss on its own rewards shrinking. The stop-gradient changes the dynamics so that they preserve a quantity, and that quantity is a wall between the run and the collapsed axis $s = 0$ . At a non-collapsed fixed point, the predictor has learned how much an augmentation shrinks the shared signal. Its value $p^\star = \kappa/(1+\sigma^2)$ is the regression coefficient of one view on the other. Without the stop-gradient, the noise makes the two views disagree by an amount proportional to $w$ , and the cheapest fix is $w = 0$ . For $\sigma^2 = 0.25$ the weight decays at an exponential rate of at least $0.25\cdot 2.25/1.25 = 0.45$ .

Now each of the three hyperparameters has a meaning, and each is a setting you will tune in a real run.

A faster predictor (larger $\alpha_p$ ) flattens the parabola, so more starting points lie above it. BYOL and SimSiam both found fast predictors helpful. But conservation gives $s(\infty) = s(0) + (p(\infty)^2 - p(0)^2)/\alpha_p$ . So a very fast predictor does all the learning while $s$ barely grows. That fits a finding of Tian et al. on STL-10, a benchmark of 96-pixel images in 10 classes. Plugging in the optimal linear predictor at every step gave about 40 percent, against 74 to 75 for a predictor trained by gradient descent.
Weight decay makes the run forget its starting point and ties $s$ to $p$ . Then the size of the online weight reflects the augmentation strength. Too much decay collapses everything.
A lagging target ( $\kappa$ below 1) sets a modest goal $\kappa/(1+\sigma^2)$ early and raises it as $\kappa$ approaches 1. Tian et al. describe this as an automatic curriculum for the predictor.

🔬 The same theory for weight matrices

The general theory is about matrices. Tian et al. (2021) take a bias-free linear network $\mathbf{W}$ , a predictor $\mathbf{W}_p$ , and a target $\mathbf{W}_a$ , and derive the flow of all three (their Lemma 1). They prove three results.

$\mathbf{W}\mathbf{W}^\top - \alpha_p^{-1}\mathbf{W}_p^\top\mathbf{W}_p = e^{-2\lambda t}\mathbf{C}_0$ for a fixed symmetric matrix $\mathbf{C}_0$ set by the initialization, whatever the target does (their Theorem 1, the matrix form of item 1).
Without the stop-gradient, $\operatorname{vec}(\mathbf{W})$ follows $-\mathbf{H}(t)\operatorname{vec}(\mathbf{W})$ for a positive semidefinite $\mathbf{H}(t)$ , so $\mathbf{W} \to \mathbf{0}$ whenever its smallest eigenvalue stays bounded away from 0 (their Theorem 2).
A symmetric predictor's eigenvectors align with those of $\mathbf{F} = \mathbf{W}\mathbf{X}\mathbf{W}^\top$ . Here $\mathbf{X} = \mathbb{E}[\bar x\bar x^\top]$ , where $\bar x$ is the average of all augmented views of one image. Under their assumptions $\mathbf{F}$ is the correlation matrix of the online network's outputs, which are the predictor's inputs. The alignment needs a matrix built from $\mathbf{F}$ , $\mathbf{W}_p$ , and the weight decay to stay positive definite (their Theorem 3).

Once aligned, each eigen-direction obeys exactly the scalar equations above, with $s_j$ the $j$ -th eigenvalue of $\mathbf{F}$ and $p_j$ the predictor's. These results are proved for linear networks. The alignment result also assumes isotropic data and augmentations, a symmetric predictor, and a target proportional to the online weights, $\mathbf{W}_a = \kappa\mathbf{W}$ . For deep networks they are a model that predicts the ablations, not a proof.

The cell draws the flow for $\sigma^2 = 0.25$ in three regimes. The gray arrows show the direction of the flow at each point, the dashed curve is the parabola $s = p^2$ (with $\alpha_p = 1$ ), and the green line marks $p^\star = 0.8$ . The red path is one run from $(p, s) = (0.2, 1)$ .

Python · cell 2

The left run keeps $s - p^2$ at its starting value, $1 - 0.04 = 0.96$ , up to the error of the Euler steps. It ends at $p = 0.800$ , exactly $1/(1+\sigma^2)$ , with $s = 1.60$ . The middle run, with weight decay 0.05, is pulled down onto the parabola. It ends at $p = 0.747$ and $s = 0.560$ . That is essentially the larger root $p_+ = (1 + \sqrt{0.75})/2.5 = 0.746$ , with $s = p_+^2 = 0.557$ . The right run, without the stop-gradient, slides to $s = 0$ while $p$ wanders. That is collapse.

Core

SimSiam in one dimension, with and without the stop-gradient

Simulate the scalar model with $\kappa = 1$ , the one the theorem above solves. Implement simulate_siamese(w0, p0, sigma2, alpha_p=1.0, wd=0.0, stop_grad=True, dt=0.01, steps=2000). Integrate the gradient flow with explicit Euler steps of size dt. Update $w$ and $p$ together, both computed from their current values,

\dot w = -\partial\mathcal{L}/\partial w - \lambda w, \qquad \dot p = -\alpha_p\,\partial\mathcal{L}/\partial p - \lambda p,

with $\lambda$ = wd. When stop_grad is true, $\partial\mathcal{L}/\partial w$ treats the target $w\,x^{(2)}$ as a constant. Otherwise it differentiates through the target too. Return two arrays w and p of length steps + 1, with the starting values at index 0.

import numpy as np

def simulate_siamese(w0, p0, sigma2, alpha_p=1.0, wd=0.0, stop_grad=True, dt=0.01, steps=2000):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Sixteen directions at once, and which ones survive

Real data are not the same in every direction. Some directions carry a lot of variation from image to image, and some carry little. The theorem has a sharper form that says which directions a non-contrastive method keeps. Repeat the scalar algebra for one input direction $j$ along which the data have variance $\nu_j$ instead of 1. Keep the augmentation noise at $\sigma^2$ in every direction. Then $\mathbb{E}[(x^{(1)})^2] = \nu_j + \sigma^2$ and $\mathbb{E}[x^{(1)}x^{(2)}] = \nu_j$ . With the stop-gradient and $\kappa = 1$ the flow is

\dot p_j = \alpha_p s_j\big(\nu_j - (\nu_j + \sigma^2)p_j\big) - \lambda p_j, \qquad \dot s_j = 2p_js_j\big(\nu_j - (\nu_j + \sigma^2)p_j\big) - 2\lambda s_j .

The proof of item 2 of the theorem goes through word for word, with $\kappa$ replaced by $\nu_j$ and $1 + \sigma^2$ by $\nu_j + \sigma^2$ .

In plain words, with weight decay a non-contrastive network keeps exactly the directions whose signal is strong compared with the augmentation noise. It lets the rest collapse. A tiny example makes the threshold concrete. With $\sigma^2 = 0.5$ and $\lambda = 0.01$ , the bar is $4\lambda = 0.04$ . A direction with $\nu_j = 1$ scores $1/1.5 = 0.67$ and survives easily. A direction with $\nu_j = 0.1$ scores $0.01/0.6 = 0.017$ and dies.

This puts a number on an idea that runs through the whole chapter, from the augmentation graph of §28.1 onward. What the augmentations destroy, the representation learns to ignore. In practice, stronger augmentation or more weight decay means fewer surviving directions.

The cell checks the theorem in a small simulation with real matrices. It trains a 16 by 16 linear online network and a 16 by 16 linear predictor, three ways. The data have variances $\nu_j = 2\cdot0.7^j$ , the augmentation noise is $\sigma^2 = 0.5$ , and the weight decay is 0.01. The cell uses the exact expected gradients, the large-batch limit, which keeps the run deterministic and fast. The threshold $4\lambda = 0.04$ predicts that the first 8 directions survive. Along a surviving direction, the embedding's scale should be $\sqrt{s_j\nu_j} = p_{j+}\sqrt{\nu_j}$ when $\alpha_p = 1$ .

Python · cell 3

The prediction holds to within 0.002 for the first seven directions. The theory gives 1.124, 0.863, 0.645, 0.467, 0.324, 0.214, and 0.131. Training gives 1.125, 0.863, 0.646, 0.466, 0.326, 0.216, and 0.130. The network started from a random matrix that was not aligned with the data, and it was never told which directions exist. The eighth direction sits right at the threshold, $0.165^2/0.665 = 0.041$ against 0.04, and is still settling. The ninth and beyond are dying, as predicted. The moving-average target lands in the same place. Without the stop-gradient, the whole network shrinks to a scale of $10^{-71}$ . What little shape is left comes from the slowest-dying directions, so its effective rank of 1.77 describes almost nothing.

Setting the predictor directly

The theory says what a trained linear predictor converges to. So why train it at all? Tian et al. proposed to compute it directly, a method they call DirectPred. Here is the idea in plain words. Measure how the predictor's inputs are spread, find their main directions, and build a predictor that keeps those directions and scales each one by the square root of its strength.

In symbols, let $\mathbf{f}$ be the online network's output, which is the predictor's input. Its correlation matrix is $\mathbf{F} = \mathbb{E}[\mathbf{f}\mathbf{f}^\top]$ , uncentered. The matrix theory in the box above says the predictor's eigenvectors align with those of $\mathbf{F}$ . On the parabola, each eigenvalue $s_j$ of $\mathbf{F}$ is tied to the predictor's eigenvalue by $s_j = p_j^2/\alpha_p$ , so $p_j \propto \sqrt{s_j}$ . DirectPred estimates $\mathbf{F}$ with a moving average over batches, $\hat{\mathbf{F}} \leftarrow \rho\hat{\mathbf{F}} + (1-\rho)\,\mathbb{E}_{\text{batch}}[\mathbf{f}\mathbf{f}^\top]$ , where $\rho$ sets how slowly the average moves. It decomposes $\hat{\mathbf{F}} = \hat{\mathbf{U}}\operatorname{diag}(s_j)\hat{\mathbf{U}}^\top$ and sets

\mathbf{W}_p = \hat{\mathbf{U}}\operatorname{diag}\big(\sqrt{s_j} + \epsilon\max_j s_j\big)\hat{\mathbf{U}}^\top .

The small $\epsilon$ lifts weak directions so that they can still grow. The test was on STL-10 with a ResNet-18 trained for 100 epochs with a moving-average target. DirectPred reached 76.77 percent, against 74.51 for a linear predictor trained by gradient descent. The $\epsilon$ term added about another point, 77.86 at its best setting. That is small-scale evidence, but it is evidence that the linear theory captured something real. An interviewer who asks "what does the predictor learn?" will be pleased by this answer, a map aligned with the main directions of its inputs.

Research

DirectPred on a linear Siamese network

Train a linear BYOL with a moving-average target, and set its predictor by DirectPred at every step. Work in the expected-gradient (large-batch) setting of the cell above. Implement train_linear_byol(W0, nu, sigma2, predictor="directpred", steps=600, lr=0.1, m=0.99, rho=0.3, eps=0.0) and return the final online weights W. Use $\mathbf{S} = \operatorname{diag}(\nu)$ and $\mathbf{A} = \mathbf{S} + \sigma^2\mathbf{I}$ . Start with W = Wa = W0 and $\hat{\mathbf{F}} = \mathbf{W}\mathbf{A}\mathbf{W}^\top$ . Then take steps steps, each doing three things in this order.

Set the predictor. If predictor == "directpred", update $\hat{\mathbf{F}} \leftarrow \rho\hat{\mathbf{F}} + (1-\rho)\mathbf{W}\mathbf{A}\mathbf{W}^\top$ , take its eigendecomposition with np.linalg.eigh, clip negative eigenvalues to 0, and set $\mathbf{W}_p$ by the formula above. If predictor == "none", set $\mathbf{W}_p = \mathbf{I}$ .
Take a gradient step of size lr on $\tfrac12\mathbb{E}\lVert\mathbf{W}_p\mathbf{W}x^{(1)} - \mathbf{W}_a x^{(2)}\rVert^2$ with respect to $\mathbf{W}$ only, holding the target $\mathbf{W}_a$ fixed, with no weight decay.
Update the target, $\mathbf{W}_a \leftarrow m\mathbf{W}_a + (1-m)\mathbf{W}$ .

The tests compare your weights with a reference. Then they check that DirectPred keeps many directions alive while the network without a predictor keeps shrinking.

import numpy as np

def train_linear_byol(W0, nu, sigma2, predictor="directpred", steps=600, lr=0.1, m=0.99, rho=0.3, eps=0.0):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Is BatchNorm a hidden source of negatives?

BatchNorm (§0.7) normalizes each feature with the mean and variance of the whole batch. So the gradient of one image's loss flows into every other image in the batch. A widely read 2020 analysis of BYOL suggested that this coupling acts as a hidden contrastive term. In small experiments, a ResNet-18 on STL-10 trained with SGD, it reported that BYOL did no better than chance without BatchNorm in its projector and predictor.

Richemond et al. (2020) tested the idea at full scale. They replaced BatchNorm everywhere with two layers that use no batch statistics at all. Group normalization normalizes groups of channels within one image, and weight standardization normalizes each filter's weights. BYOL trained this way reached 73.9 percent, against 74.3 with BatchNorm. Removing BatchNorm with nothing in its place gave 0.1 percent. Three changes together brought that back to 65.7, with still no batch statistics. They were a better weight initialization, keeping BatchNorm's learned scale and shift as plain parameters, and a longer warmup. So batch statistics are not what prevents collapse, although they help optimization. The linear theory above has no BatchNorm either, which supports the same conclusion. This is a favorite interview question, and the answer is the experiment, not the hypothesis.

Barlow Twins, make the cross-correlation the identity

The asymmetric methods avoid collapse through their training dynamics, and it took a theorem to see why. Barlow Twins takes the other road. It writes a loss whose minimum is not a collapsed network in the first place (Zbontar et al. 2021). The name honors the neuroscientist Horace Barlow, who argued that the brain recodes its sensory input so that no two units of the code carry the same information. That is redundancy reduction. Barlow Twins adds a second demand, that each unit ignore the augmentation.

In plain words, the loss asks two things of every output unit $i$ . Unit $i$ for the first view should be perfectly correlated with unit $i$ for the second view, so the unit ignores the augmentation. And unit $i$ for one view should be uncorrelated with every other unit $j$ for the other view, so no two units duplicate each other. One matrix holds all those correlations at once, and the loss pushes it toward the identity.

📐 Barlow Twins

For the embeddings $\mathbf{Z}^{(1)}, \mathbf{Z}^{(2)} \in \R^{B\times d_z}$ of the two views (row $b$ is image $b$ ), standardize every dimension over the batch, $\tilde z_{bi} = (z_{bi} - \bar z_i)/\sqrt{v_i + \epsilon}$ , with $\bar z_i$ and $v_i$ the batch mean and population variance of dimension $i$ . The released code uses a BatchNorm layer without learned scale or shift, so $\epsilon = 10^{-5}$ . The cross-correlation matrix is

\mathbf{C}_{\text{BT}} = \frac{1}{B}\,\tilde{\mathbf{Z}}^{(1)\top}\tilde{\mathbf{Z}}^{(2)} \in \R^{d_z\times d_z},

and the loss is

\mathcal{L}_{\text{BT}} = \sum_{i=1}^{d_z}\big(1 - C_{ii}\big)^2 + \lambda_{\text{BT}}\sum_{i\ne j}C_{ij}^2,

an invariance term and a redundancy-reduction term with weight $\lambda_{\text{BT}} = 5\times10^{-3}$ (the released code uses 0.0051). The network is symmetric, with one encoder and projector, no predictor, and no stop-gradient. The projector has three linear layers, each with 8,192 outputs, so $d_z = 8{,}192$ . Training uses LARS with batch 2,048 for 1,000 epochs, about 124 hours on 32 V100 GPUs.

Each $C_{ij}$ is the correlation over the batch between dimension $i$ of the first view and dimension $j$ of the second. It lies in $[-1, 1]$ when $\epsilon = 0$ , by the Cauchy-Schwarz inequality (an inner product is at most the product of the two lengths). A positive $\epsilon$ only shrinks it. Here is a tiny batch worked through by hand.

✏️ A Barlow Twins correlation matrix, by hand

Take $B = 4$ images and $d_z = 2$ . The first-view embeddings are $(3,2), (1,2), (3,0), (1,0)$ and the second-view embeddings are $(2,4), (0,0), (2,4), (0,0)$ . The first view has mean $(2, 1)$ and population standard deviations $(1, 1)$ , so its standardized rows are $(1,1), (-1,1), (1,-1), (-1,-1)$ . The second has mean $(1, 2)$ and standard deviations $(1, 2)$ , so its standardized rows are $(1,1), (-1,-1), (1,1), (-1,-1)$ . Its two dimensions are copies of each other.

Then $C_{11} = \tfrac14(1 + 1 + 1 + 1) = 1$ and $C_{12} = \tfrac14(1\cdot1 + (-1)(-1) + 1\cdot1 + (-1)(-1)) = 1$ . Also $C_{21} = \tfrac14(1 - 1 - 1 + 1) = 0$ , and $C_{22} = 0$ the same way. So $\mathbf{C}_{\text{BT}} = \begin{pmatrix}1 & 1\\ 0 & 0\end{pmatrix}$ and $\mathcal{L}_{\text{BT}} = (1-1)^2 + (1-0)^2 + 0.005\,(1^2 + 0^2) = 1.005$ . The invariance term charges a full 1, because the second view lost what dimension 2 of the first view carries. The redundancy term charges the duplicate.

The general statement says exactly where the loss is zero, and what collapse costs.

∴ Barlow Twins is minimized exactly at the identity correlation, and collapse is not a minimum

With $\lambda_{\text{BT}} > 0$ , $\mathcal{L}_{\text{BT}} \ge 0$ , with equality if and only if $\mathbf{C}_{\text{BT}} = \mathbf{I}$ . With $\epsilon = 0$ , $\mathbf{C}_{\text{BT}} = \mathbf{I}$ holds exactly when two conditions hold. First, $\tilde{\mathbf{Z}}^{(2)} = \tilde{\mathbf{Z}}^{(1)}$ , so every dimension is invariant to the augmentation up to its own shift and scale. Second, $\tilde{\mathbf{Z}}^{(1)\top}\tilde{\mathbf{Z}}^{(1)}/B = \mathbf{I}$ , so the standardized batch is white. A completely collapsed network has $v_i = 0$ , so with $\epsilon > 0$ every standardized entry is 0, $\mathbf{C}_{\text{BT}} = \mathbf{0}$ , and the loss is $d_z$ .

∂ Proof

The loss is a sum of squares with positive weights, so it is nonnegative and vanishes exactly when every square does, $C_{ii} = 1$ and $C_{ij} = 0$ . With $\epsilon = 0$ , each standardized column has $\lVert\tilde{\mathbf{z}}_i\rVert^2 = B$ . So $C_{ii} = \tilde{\mathbf{z}}^{(1)\top}_i\tilde{\mathbf{z}}^{(2)}_i/B = 1$ is equality in Cauchy-Schwarz between two vectors of the same length, which forces $\tilde{\mathbf{z}}^{(2)}_i = \tilde{\mathbf{z}}^{(1)}_i$ . Substituting into $C_{ij} = 0$ gives $\tilde{\mathbf{z}}^{(1)\top}_i\tilde{\mathbf{z}}^{(1)}_j = 0$ for $i \ne j$ , and together with the diagonal this is $\tilde{\mathbf{Z}}^{(1)\top}\tilde{\mathbf{Z}}^{(1)}/B = \mathbf{I}$ . The converse is the same computation backwards. For a constant network, $z_{bi} - \bar z_i = 0$ , so $\tilde z_{bi} = 0/\sqrt{\epsilon} = 0$ , $\mathbf{C}_{\text{BT}} = \mathbf{0}$ , and the loss is $\sum_i 1^2 = d_z$ . ∎

In plain words, the invariance term says every output unit must report the same thing for both views. The redundancy term says no two units may report the same thing. Collapse breaks the first rule in every unit, so it costs $d_z$ instead of 0. The standardization does the work that negatives did in §28.2, because a unit with no spread cannot correlate with anything.

There is a catch that the formula hides. Centering removes one dimension, so the centered batch has rank at most $B - 1$ , and so does $\mathbf{C}_{\text{BT}}$ . It cannot equal the $d_z\times d_z$ identity unless $d_z \le B - 1$ . With the paper's $d_z = 8{,}192$ and $B = 2{,}048$ , the minimum 0 is out of reach. The loss is always a compromise between its two terms. Even so, Barlow Twins kept improving across every projector width its authors tried, while SimCLR and BYOL saturate early. In an interview, this rank argument is the expected answer to "can Barlow Twins reach zero loss?".

To train with this loss you need its gradient, which passes through the batch standardization.

∂ The gradient of the Barlow Twins loss

Write $\mathbf{A} = \tilde{\mathbf{Z}}^{(1)}$ and $\mathbf{G}_C = \partial\mathcal{L}_{\text{BT}}/\partial\mathbf{C}_{\text{BT}}$ , which has diagonal entries $-2(1 - C_{ii})$ and off-diagonal entries $2\lambda_{\text{BT}}C_{ij}$ . Since $\mathbf{C} = \mathbf{A}^\top\tilde{\mathbf{Z}}^{(2)}/B$ is linear in $\mathbf{A}$ , a small change $\mathrm{d}\mathbf{A}$ changes the loss by

\mathrm{d}\mathcal{L} = \operatorname{tr}\big(\mathbf{G}_C^\top\,\mathrm{d}\mathbf{C}\big) = \tfrac1B\operatorname{tr}\big(\mathbf{G}_C^\top\,\mathrm{d}\mathbf{A}^\top\tilde{\mathbf{Z}}^{(2)}\big) = \tfrac1B\operatorname{tr}\big(\mathrm{d}\mathbf{A}^\top\,\tilde{\mathbf{Z}}^{(2)}\mathbf{G}_C^\top\big),

using the cyclic property of the trace. So $\partial\mathcal{L}/\partial\mathbf{A} = \frac1B\tilde{\mathbf{Z}}^{(2)}\mathbf{G}_C^\top \in \R^{B\times d_z}$ , the same shape as $\mathbf{A}$ (§0.2).

Next comes the standardization. It is BatchNorm without the learned scale and shift, whose backward pass §0.7 derived. For one column with upstream gradient $\mathbf{g} = \partial\mathcal{L}/\partial\mathbf{a}_i$ and standardized values $\mathbf{a}_i$ ,

\frac{\partial\mathcal{L}}{\partial\mathbf{z}^{(1)}_i} = \frac{1}{\sqrt{v_i + \epsilon}}\Big(\mathbf{g} - \bar g\,\mathbf{1} - \mathbf{a}_i\,\overline{g\,a_i}\Big),

where $\bar g$ is the batch mean of $\mathbf{g}$ and $\overline{g\,a_i}$ the batch mean of the elementwise product. The two subtracted terms remove the parts of $\mathbf{g}$ that a shift or a rescaling of the column would absorb, so every column of the result sums to zero. Here $\bar g$ is in fact 0, because $\mathbf{g}$ is a combination of the centered columns of $\tilde{\mathbf{Z}}^{(2)}$ , but the general formula costs nothing. The second view's gradient is the same with the roles swapped, $\partial\mathcal{L}/\partial\tilde{\mathbf{Z}}^{(2)} = \frac1B\tilde{\mathbf{Z}}^{(1)}\mathbf{G}_C$ .

Interview

Barlow Twins, no loops

Implement barlow_twins(Z1, Z2, lam=5e-3, eps=1e-5) for two batches of embeddings of shape (B, d_z). Standardize each dimension of each view over the batch (axis 0), with that view's own mean and population variance and $\epsilon$ inside the square root. Form $\mathbf{C}_{\text{BT}}$ by dividing by $B$ . Return (loss, C, dZ1), where dZ1 is the gradient of the loss with respect to Z1. Use no Python loops over the batch or the dimensions.

import numpy as np

def barlow_twins(Z1, Z2, lam=5e-3, eps=1e-5):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

VICReg, three terms with three jobs

Barlow Twins ties invariance and decorrelation together in one cross-correlation matrix. It also hides the scale of the embeddings behind a standardization. VICReg separates the jobs into three terms, each computed on its own branch where possible (Bardes, Ponce and LeCun 2022). In plain words, it asks for three things. The two views of an image should be close. Every dimension should keep a minimum spread over the batch. And no two dimensions should be correlated. The name lists the terms, variance, invariance, and covariance.

📐 VICReg

For the embeddings $\mathbf{Z}, \mathbf{Z}' \in \R^{B\times d_z}$ of the two views, with rows $\mathbf{z}_b$ and columns $\mathbf{z}^j$ (dimension $j$ across the batch), define three terms.

Invariance, $s(\mathbf{Z}, \mathbf{Z}') = \frac1B\sum_{b=1}^{B}\lVert\mathbf{z}_b - \mathbf{z}'_b\rVert^2$ .
Variance, $v(\mathbf{Z}) = \frac{1}{d_z}\sum_{j=1}^{d_z}\max\big(0,\ 1 - \sqrt{\operatorname{Var}(\mathbf{z}^j) + \epsilon}\big)$ , with the unbiased batch variance (as in the paper's pseudocode and code) and $\epsilon = 10^{-4}$ . The target standard deviation is 1.
Covariance, $c(\mathbf{Z}) = \frac{1}{d_z}\sum_{i\ne j}\big[\operatorname{Cov}(\mathbf{Z})\big]_{ij}^2$ , with $\operatorname{Cov}(\mathbf{Z}) = \frac{1}{B-1}\sum_b(\mathbf{z}_b - \bar{\mathbf{z}})(\mathbf{z}_b - \bar{\mathbf{z}})^\top$ .

The loss is

\mathcal{L}_{\text{VIC}} = \lambda_{\text{inv}}\,s(\mathbf{Z}, \mathbf{Z}') + \lambda_{\text{var}}\big(v(\mathbf{Z}) + v(\mathbf{Z}')\big) + \lambda_{\text{cov}}\big(c(\mathbf{Z}) + c(\mathbf{Z}')\big),

with $\lambda_{\text{inv}} = \lambda_{\text{var}} = 25$ and $\lambda_{\text{cov}} = 1$ . (The paper uses other Greek letters, but $\lambda$ alone is weight decay in this book.) The projector, which VICReg calls the expander, has three layers of width 8,192, and training follows Barlow Twins (LARS, batch 2,048, 1,000 epochs).

The variance term is the hinge of §1.10. There it charged one minus each dimension's standard deviation whenever that standard deviation fell below 1, and it stopped a linear encoder from shrinking to zero. Here it gains a small $\epsilon$ and an average over the $d_z$ dimensions.

Each term has one job. Invariance pulls the views together. Variance forbids complete collapse. Covariance forbids dimensional collapse by decorrelating the dimensions. The variance term cannot do that alone, since every dimension could carry the same signal with unit spread.

Why hinge the standard deviation and not the variance? The paper remarks that a hinge on the variance has a vanishing gradient near collapse. Computing the gradient makes the remark exact.

∴ Why VICReg hinges the standard deviation

Fix a dimension with batch column $\mathbf{z}^j \in \R^B$ , mean $\bar z_j$ , and unbiased variance $\sigma_j^2 = \frac{1}{B-1}\lVert\mathbf{z}^j - \bar z_j\mathbf{1}\rVert^2$ . Then

\frac{\partial}{\partial\mathbf{z}^j}\sqrt{\sigma_j^2 + \epsilon} = \frac{\mathbf{z}^j - \bar z_j\mathbf{1}}{(B-1)\sqrt{\sigma_j^2 + \epsilon}}, \qquad \Big\lVert\frac{\partial}{\partial\mathbf{z}^j}\sqrt{\sigma_j^2 + \epsilon}\Big\rVert = \frac{1}{\sqrt{B-1}}\sqrt{\frac{\sigma_j^2}{\sigma_j^2 + \epsilon}},

while $\lVert\partial\sigma_j^2/\partial\mathbf{z}^j\rVert = 2\sigma_j/\sqrt{B-1}$ . The covariance term has gradient $\partial c/\partial\mathbf{Z} = \frac{4}{d_z(B-1)}\tilde{\mathbf{Z}}\operatorname{offdiag}(\operatorname{Cov}(\mathbf{Z}))$ , where $\tilde{\mathbf{Z}}$ is the centered batch and $\operatorname{offdiag}$ zeroes the diagonal.

∂ Proof

Write $\mathbf{u} = \mathbf{z}^j - \bar z_j\mathbf{1}$ . The variance is $\sigma_j^2 = \mathbf{u}^\top\mathbf{u}/(B-1)$ , and $\partial\bar z_j/\partial z_{bj} = 1/B$ . So $\partial\sigma_j^2/\partial z_{bj} = \frac{2}{B-1}\big(u_b - \frac1B\sum_c u_c\big) = \frac{2u_b}{B-1}$ , because the centered entries sum to zero. The chain rule through the square root multiplies by $1/(2\sqrt{\sigma_j^2 + \epsilon})$ , which gives the first formula. Its norm is $\lVert\mathbf{u}\rVert/((B-1)\sqrt{\sigma_j^2+\epsilon})$ with $\lVert\mathbf{u}\rVert = \sqrt{B-1}\,\sigma_j$ , which simplifies to the stated value. The norm of $2\mathbf{u}/(B-1)$ is $2\sigma_j/\sqrt{B-1}$ .

For the covariance term, $c = \frac1{d_z}\sum_{i\ne j}C_{ij}^2$ gives $\partial c/\partial\mathbf{C} = \frac{2}{d_z}\mathbf{O}$ with $\mathbf{O} = \operatorname{offdiag}(\mathbf{C})$ , which is symmetric. With $\mathbf{C} = \tilde{\mathbf{Z}}^\top\tilde{\mathbf{Z}}/(B-1)$ ,

\mathrm{d}c = \tfrac{2}{d_z(B-1)}\operatorname{tr}\big(\mathbf{O}(\mathrm{d}\tilde{\mathbf{Z}}^\top\tilde{\mathbf{Z}} + \tilde{\mathbf{Z}}^\top\mathrm{d}\tilde{\mathbf{Z}})\big) = \tfrac{4}{d_z(B-1)}\operatorname{tr}\big(\mathrm{d}\tilde{\mathbf{Z}}^\top\tilde{\mathbf{Z}}\mathbf{O}\big),

so $\partial c/\partial\tilde{\mathbf{Z}} = \frac{4}{d_z(B-1)}\tilde{\mathbf{Z}}\mathbf{O}$ . Centering multiplies by $\mathbf{I} - \frac1B\mathbf{1}\mathbf{1}^\top$ on the left, and $\mathbf{1}^\top\tilde{\mathbf{Z}} = \mathbf{0}$ , so it changes nothing. ∎

In plain words, the standard-deviation hinge pushes a shrinking dimension apart with a nearly constant force, about $1/\sqrt{B-1}$ . A variance hinge pushes with a force proportional to $\sigma_j$ . So the closer a dimension gets to collapse, the less a variance hinge resists, which is exactly backward.

Put numbers on it with $B = 257$ , so $\sqrt{B-1} = 16$ . A dimension that has shrunk to $\sigma_j = 0.1$ gets a push of $0.995/16 = 0.062$ from the standard-deviation hinge. A variance hinge would give it $0.2/16 = 0.0125$ , five times less, and the gap widens as $\sigma_j$ falls. The standard-deviation hinge keeps most of its force down to $\sigma_j \approx \sqrt{\epsilon} = 0.01$ , where it is still $1/\sqrt2$ of its full size.

The $\epsilon$ is there because the derivative of $\sqrt{\sigma_j^2}$ is undefined at $\sigma_j = 0$ . With it, the force falls off linearly once $\sigma_j$ drops well below $\sqrt{\epsilon}$ . Even there it is $1/(2\sqrt{\epsilon}) = 50$ times the variance hinge's force. An interviewer asking "why the square root?" wants this argument.

Compared with Barlow Twins, VICReg makes three different choices, and it shares one limit.

	Barlow Twins	VICReg
Scale of the embeddings	removed by standardization, so rescaling a dimension changes nothing	visible to the loss, so the variance term must hold it up
What is decorrelated	the two branches against each other, through one cross-correlation	each branch on its own, so the branches need not share weights, architectures, or even input modalities
Invariance	a correlation of 1 per dimension	a plain squared distance
Reachable minimum when $d_z$ exceeds $B - 1$	no	no

The last row holds for VICReg for a rank reason. A $d_z\times d_z$ covariance with every diagonal entry at least 1 and zeros off the diagonal is full rank. The covariance of $B$ samples has rank at most $B - 1$ . So with $d_z = 8{,}192$ and $B = 2{,}048$ , the variance and covariance terms cannot both be zero.

The second row is VICReg's practical advantage. Because each branch is regularized on its own, the paper pairs images with text and evaluates retrieval on MS-COCO, a dataset of photos with captions. Whitening methods take the hard version of the same idea. W-MSE (Ermolov et al. 2021) whitens each batch of embeddings exactly, with the matrices of §0.10, so that their covariance is the identity. Only then does it apply a mean-squared-error loss between the views, so collapse is impossible by construction.

The cell computes the Barlow Twins and VICReg terms on a healthy batch of 512 pairs in $\R^{16}$ and on four damaged versions. It shows which term answers which kind of damage.

Python · cell 4

Read the table one row at a time. Four dead dimensions cost Barlow Twins almost exactly 4 in its invariance term, because each dead dimension has $C_{ii} = 0$ . They cost VICReg through its variance term, 12.9. Eight copies of one dimension light up the redundancy and covariance terms. Noisy views raise both invariance terms. Shrinking the whole batch tenfold changes nothing for Barlow Twins, whose standardization removes scale. VICReg's variance term jumps to 44.9, because it alone watches the scale.

The widget below lets you damage a batch yourself and watch each term respond. Predict before you slide. Choose Barlow Twins and correlated dims, and push the correlation between dimensions to 0.9. Which term grows, and roughly by how much, given that there are $d_z(d_z - 1)$ off-diagonal entries weighted by 0.005? Then choose VICReg and spread, set the spread per dimension to 0.1, and predict the variance term before you look.

Each kind of damage shows up in the term built to catch it, the division of labor the cell above printed. The paper's coefficient study, which drops terms from the loss at ImageNet scale, tells the same story.

Interview

VICReg loss and its gradient

Implement vicreg(Z1, Z2, lam_inv=25.0, lam_var=25.0, lam_cov=1.0, eps=1e-4) for two batches of shape (B, d_z), exactly as the definition box writes it. Five details matter. Use the unbiased variance and covariance. Hinge $\sqrt{\operatorname{Var} + \epsilon}$ , not the variance. Keep the $1/d_z$ factors. Average the invariance over the batch. Regularize each branch separately. Return (loss, dZ1, dZ2), the loss and its gradients with respect to both inputs. Use no loops.

import numpy as np

def vicreg(Z1, Z2, lam_inv=25.0, lam_var=25.0, lam_cov=1.0, eps=1e-4):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

⚠️ The released VICReg code is not the paper's equation

The official implementation computes the invariance with F.mse_loss, a mean over all $B\,d_z$ entries, so its invariance is the paper's $s$ divided by $d_z$ . It also averages the two variance terms instead of adding them. The defaults 25, 25, 1 belong to that code. If you implement the paper's equation as printed with the same weights, the invariance term is 8,192 times heavier than in the released runs. Likewise, the Barlow Twins pseudocode in its paper standardizes with z.std(0), which in PyTorch is the unbiased estimate, while the released code uses BatchNorm, which is not. The two differ by a factor $B/(B-1)$ in every $C_{ij}$ . Before you compare numbers or reuse hyperparameters, check which normalization the code you copy actually uses.

Clustering, and codes that must be balanced

A third family prevents collapse by counting. Give the network a set of clusters, assign every image to one of them, and train the network to predict an image's assignment from another view of it. Collapse would put every image in one cluster. So add a rule that the clusters must be used equally, and the constant solution can no longer score well.

SwAV, short for swapping assignments between views (Caron et al. 2020), does the balancing inside each batch. So the method runs online, one batch at a time, with no pass over the whole dataset.

SwAV's clusters are prototypes, learned unit vectors in the embedding space, one per cluster. An image's code is a soft assignment, a probability distribution over the prototypes. The codes of a batch are chosen to agree with the embeddings as much as possible, under one rule. Every prototype must receive an equal share of the batch, a rule called equipartition. Then each view is trained to predict the other view's code.

Choosing the codes is an optimal transport problem. Think of each image as holding an equal share of mass, $1/B$ , that it must ship to the prototypes. Shipping mass from image $b$ to prototype $k$ earns the score $S_{bk}$ , the cosine similarity between them, per unit of mass. The plan must earn as much as possible, and every prototype must receive the same total, $1/K_p$ . Here is a tiny case. With 4 images and 2 prototypes, each prototype must receive half the mass, two images' worth, even if all 4 images prefer the first prototype.

📐 SwAV

Keep $K_p$ trainable prototypes, the unit columns of $\mathbf{W}_{\text{pr}} \in \R^{d_z\times K_p}$ , and let $\hat{\mathbf{Z}} \in \R^{B\times d_z}$ hold one view's unit embeddings. The codes $\boldsymbol\Gamma \in \R^{B\times K_p}$ are the soft assignment that solves

\max_{\boldsymbol\Gamma \in \mathcal{U}}\ \operatorname{tr}\big(\boldsymbol\Gamma^\top\hat{\mathbf{Z}}\mathbf{W}_{\text{pr}}\big) + \varepsilon H(\boldsymbol\Gamma), \qquad \mathcal{U} = \Big\{\boldsymbol\Gamma \ge 0 : \boldsymbol\Gamma\mathbf{1}_{K_p} = \tfrac1B\mathbf{1}_B,\ \boldsymbol\Gamma^\top\mathbf{1}_B = \tfrac{1}{K_p}\mathbf{1}_{K_p}\Big\},

with $H(\boldsymbol\Gamma) = -\sum_{bk}\Gamma_{bk}\log\Gamma_{bk}$ . Every image carries total mass $1/B$ and every prototype receives $1/K_p$ , the equipartition constraint. Rescaling each row to sum to 1 gives image $b$ 's code $\boldsymbol\gamma_b$ . Codes are computed without gradient, separately for each view, and each view predicts the other view's code (the swapped prediction),

\ell\big(\mathbf{z}^{(1)}, \boldsymbol\gamma^{(2)}\big) = -\sum_{k=1}^{K_p}\gamma^{(2)}_k\log\operatorname{softmax}\big(\mathbf{W}_{\text{pr}}^\top\hat{\mathbf{z}}^{(1)}/\tau\big)_k,

symmetrized over the two views. SwAV uses $K_p = 3{,}000$ , $\varepsilon = 0.05$ , temperature $\tau = 0.1$ , $d_z = 128$ , and three Sinkhorn iterations.

The entropy term $\varepsilon H$ keeps the codes soft, and $\varepsilon$ sets how soft. The constrained problem has a closed-form solution up to two scaling vectors, and the algorithm that finds them takes a few lines.

∴ The entropic transport plan and the Sinkhorn-Knopp iteration

With score matrix $\mathbf{S} = \hat{\mathbf{Z}}\mathbf{W}_{\text{pr}}$ , the maximizer is unique and has the form

\boldsymbol\Gamma^\star = \operatorname{diag}(\mathbf{u})\,\exp(\mathbf{S}/\varepsilon)\,\operatorname{diag}(\mathbf{v}), \qquad \mathbf{u} \in \R^{B}_{>0},\ \mathbf{v} \in \R^{K_p}_{>0},

with the exponential taken entrywise. The Sinkhorn-Knopp iteration finds $\mathbf{u}$ and $\mathbf{v}$ by alternately rescaling the columns of $\mathbf{K} = \exp(\mathbf{S}/\varepsilon)$ , so that each prototype receives $1/K_p$ , and the rows, so that each image carries $1/B$ . For a positive matrix it converges to $\boldsymbol\Gamma^\star$ (Sinkhorn and Knopp 1967).

∂ Deriving the form of the transport plan

The objective is strictly concave in $\boldsymbol\Gamma$ (a linear term plus $\varepsilon$ times the strictly concave entropy), and the constraint set is convex and compact, so the maximizer exists and is unique. Attach a multiplier $\alpha_b$ to each image's constraint and $\beta_k$ to each prototype's,

\mathcal{J} = \sum_{bk}\Gamma_{bk}S_{bk} - \varepsilon\sum_{bk}\Gamma_{bk}\log\Gamma_{bk} + \sum_b\alpha_b\Big(\tfrac1B - \sum_k\Gamma_{bk}\Big) + \sum_k\beta_k\Big(\tfrac1{K_p} - \sum_b\Gamma_{bk}\Big).

Setting $\partial\mathcal{J}/\partial\Gamma_{bk} = S_{bk} - \varepsilon(\log\Gamma_{bk} + 1) - \alpha_b - \beta_k = 0$ gives

\log\Gamma_{bk} = \frac{S_{bk}}{\varepsilon} - 1 - \frac{\alpha_b}{\varepsilon} - \frac{\beta_k}{\varepsilon}, \qquad \Gamma_{bk} = \underbrace{e^{-1/2 - \alpha_b/\varepsilon}}_{u_b}\ e^{S_{bk}/\varepsilon}\ \underbrace{e^{-1/2 - \beta_k/\varepsilon}}_{v_k}.

Every entry is positive, so the constraint $\boldsymbol\Gamma \ge 0$ is inactive, and the stationary point is the maximizer. The column constraint reads $v_k\sum_b u_bK_{bk} = 1/K_p$ , so for fixed $\mathbf{u}$ the right $\mathbf{v}$ is $v_k = (1/K_p)/(\mathbf{K}^\top\mathbf{u})_k$ . The row constraint gives $u_b = (1/B)/(\mathbf{K}\mathbf{v})_b$ . Sinkhorn-Knopp alternates these two updates. That is the same as alternately dividing each column of the current matrix by its sum and multiplying by $1/K_p$ , then each row by its sum and multiplying by $1/B$ . One global constant multiplying $\mathbf{K}$ cancels in every rescaling, which is why subtracting the largest score before exponentiating is safe.

The same entropic balanced assignment appears in §13.3, which derives it for routing tokens to experts, with tokens in place of images and experts in place of prototypes. There its job is to keep every expert busy. Here its job is to make the constant solution expensive, and the next result says exactly how expensive.

Proof. Swapping two images leaves the problem unchanged, so by uniqueness the solution has identical rows. The column constraint then forces every entry to be $1/(BK_p)$ , so every code is uniform. The loss of each image is a cross-entropy against the uniform distribution, which is at least that distribution's entropy, $\log K_p$ (§0.5). ∎

For SwAV's 3,000 prototypes, a collapsed encoder pays at least $\log 3000 = 8.01$ per image. An encoder whose predictions match confident, balanced codes can bring the loss down toward the codes' own entropy, which is near 0 when $\varepsilon$ is small. Equipartition turns the constant solution into a bad one.

The same reasoning explains SwAV's three warnings, and each is a setting you must get right in practice.

A large $\varepsilon$ weights the entropy so heavily that the codes are nearly uniform whatever the embeddings. Then the target carries no information, and the authors observed that training collapses.
A small $\varepsilon$ gives sharp codes, but $\exp(\mathbf{S}/\varepsilon)$ becomes extreme, and three iterations may not finish balancing them. The cell below shows this.
A small batch breaks the balance. The batch must be large compared with $K_p$ , or there is no meaningful balanced assignment. With batch 256 and 3,000 prototypes, SwAV keeps a queue of the last 3,840 embeddings and computes codes over it. That is a far shorter queue than MoCo's 65,536.

The cell balances 6 images over 3 prototypes, a small example you can follow entry by entry.

Python · cell 5

Four of the six images prefer prototype 1, so assigning each image to its best prototype gives totals $(4, 1, 1)$ . At $\varepsilon = 0.05$ , three iterations barely move them, to $(3.92, 1, 1.08)$ . The scores differ by up to 0.9, and $e^{0.9/0.05} = e^{18}$ dwarfs everything else. Fifty iterations reach the equipartition $(2, 2, 2)$ , and they split image 4 evenly between the first two prototypes. At $\varepsilon = 0.5$ , three iterations already balance the batch, but image 4's code is soft, $(0.45, 0.29, 0.25)$ . A small $\varepsilon$ buys sharp targets at the cost of slow balancing. With thousands of images per batch, SwAV's authors found three iterations enough in practice. The collapsed batch gets uniform codes and a loss of exactly $\log 3$ .

Interview

Sinkhorn-Knopp for balanced codes

Implement sinkhorn(scores, eps=0.05, n_iters=3) for scores of shape (B, K_p), following SwAV's released order exactly.

Form $\mathbf{Q} = \exp(\text{scores}/\varepsilon)^\top$ of shape (K_p, B), subtracting the single largest entry of scores / eps before exponentiating, and divide $\mathbf{Q}$ by its total.
Repeat n_iters times. Rescale every row (prototype) to sum to $1/K_p$ , then every column (image) to sum to $1/B$ .
Rescale every column to sum to 1, and return the transpose, codes of shape (B, K_p) whose rows sum to 1.

Use one global constant for the shift in step 1. A per-image shift would change the result after a finite number of iterations. The released code skips the shift, which changes nothing mathematically but can overflow.

import numpy as np

def sinkhorn(scores, eps=0.05, n_iters=3):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Dimensional collapse, even with negatives

Negatives prevent complete collapse. They do not prevent dimensional collapse. Jing, Vincent, LeCun and Tian (2022) trained a standard SimCLR for 100 epochs on ImageNet with 128-dimensional embeddings. A sizable part of the embedding covariance spectrum had fallen to essentially zero. They proved two mechanisms that cause it, each for linear networks.

Strong augmentation. Take a one-layer linear encoder $\mathbf{z} = \mathbf{W}x$ trained with InfoNCE, the contrastive loss of §7.2. Its gradient flow is $\dot{\mathbf{W}} = \mathbf{W}\mathbf{X}$ for a symmetric matrix $\mathbf{X}$ . That matrix is a weighted covariance of the data minus a weighted covariance of the augmentation noise. Along a direction where the augmentation varies more than the data, $\mathbf{X}$ has a negative eigenvalue. Treating $\mathbf{X}$ as fixed, as their Theorem 1 does, the weight along that direction then shrinks to zero. In plain words, pushing negatives apart along such a direction does not pay, because pulling the positive pair together along it costs more.

Implicit regularization. Now take two linear layers, $\mathbf{z} = \mathbf{W}_2\mathbf{W}_1x$ . During training, the singular vectors of adjacent layers align. Each aligned pair of singular values then grows at a rate proportional to its own size. So small directions stay small, and the product drifts toward low rank, even without strong augmentation (their Theorems 2 and 3).

The cell tests the first mechanism in a small simulation. It trains a linear encoder on 16-dimensional data of unit variance with InfoNCE over a batch of 64. The similarity is $-\lVert\mathbf{z}_i - \mathbf{z}'_j\rVert^2/2$ , as in their analysis. It runs once with mild augmentation and once with augmentation of variance 8 on dimensions 8 to 15.

Python · cell 6

With mild augmentation, all 16 singular values of $\mathbf{W}$ stay alive. The noisy half sits at 0.33 to 0.68, against 1.3 to 1.9 for the clean half. With strong augmentation, the noisy half falls to between 0.001 and 0.09, while the clean half grows past 1.6. That is a gap of at least 18 times, and up to three orders of magnitude. The encoder still has negatives, and it has still discarded half of its dimensions.

Here the discarded directions are pure nuisance, which is the point of augmentation. In a real image the line between nuisance and content is not so clean. Jing et al. argue that the lost dimensions are capacity a downstream task could have used, and the projector results below support them.

This also explains the projection head. A trainable projector between the representation and the loss can absorb the collapse, so that the representation $\mathbf{h}$ keeps more dimensions than the embedding $\mathbf{z}$ . Jing et al. confirmed it at 100 epochs with a ResNet-50.

SimCLR with a two-layer nonlinear projector reached 66.5 percent.
With a linear projector it reached 61.1, and with no projector 51.5. Without a projector, the representation's own spectrum collapsed.
Their DirectCLR applies InfoNCE directly to a fixed sub-vector of the representation, its first $d_0$ coordinates, with no trainable projector. It reached 62.7, beating the linear projector.

Garrido et al. (2023) later turned the same observation into a tool. The effective rank of the embeddings tracks downstream accuracy well enough to select hyperparameters without labels, which §28.6 takes up. An interviewer who asks why we throw the projection head away expects this answer, as well as the one from §28.2.

In PyTorch

Here are four of the losses in production form, BYOL, SimSiam, Barlow Twins, and VICReg. The stop-gradient is .detach(), and the moving-average update runs under torch.no_grad(). The block also checks two exact results against autograd. One is SimSiam's identity, that without a predictor the stop-gradient loss has half the gradient of the plain loss. The other is the VICReg covariance gradient from the theorem.

🔥 In PyTorchfor a machine with PyTorch

All the joint-embedding methods side by side

Every method in this chapter so far is a joint-embedding method. It embeds two views with encoders and compares the embeddings. The table puts the methods of §28.2 and this section together, at real scale.

✏️ One table of joint-embedding methods at real scale

ImageNet linear-probe top-1 of a ResNet-50, as reported by each paper. Epochs and views differ, so the last column is a guide, not a ranking.

Method	What prevents collapse	Large batch needed	Momentum target	Predictor	Embedding width $d_z$	Linear top-1 (epochs)
SimCLR	negatives in the batch	yes (4,096)	no	no	128	69.3 (1,000)
MoCo v2	negatives in a queue of 65,536	no (256)	yes	no	128	71.1 (800)
BYOL	predictor, stop-gradient, slow target	no (256 to 4,096)	yes	yes	256	74.3 (1,000)
SimSiam	predictor and stop-gradient	no (256)	no	yes	2,048	71.3 (800)
Barlow Twins	correlation driven to the identity	no (256 works)	no	no	8,192	73.2 (1,000)
VICReg	variance hinge and covariance penalty	no	no	no	8,192	73.2 (1,000)
SwAV	equipartitioned codes (Sinkhorn)	yes, or a queue	no	no	128	75.3 (800, multi-crop)

Each family pays a different cost for avoiding collapse. Contrastive methods pay for many negatives. The redundancy-reduction methods, Barlow Twins and VICReg, pay in memory, not compute.

The asymmetric methods, BYOL and SimSiam, pay in sensitivity to hyperparameters instead. Their success depends on the target momentum, the predictor's learning rate, and the weight decay. Those are exactly the three quantities the linear theory singles out.

Where it breaks. The mechanisms are understood for linear networks. For deep networks, the linear theorems predict the ablations, but nobody has proved that BYOL or SimSiam avoid collapse. Both also have failure modes the theory does not cover. A frozen predictor fails to converge, and BatchNorm on the predictor's output makes the loss oscillate. Barlow Twins and VICReg cannot reach their minimum when $d_z$ exceeds the batch, so the balance of their terms depends on $B$ . Clustering methods need a batch or a queue that is large compared with the number of prototypes. And every method here compares whole-image embeddings. None of them is trained to keep the spatial detail that segmentation needs. That gap is what §28.4 and §28.5 address.

The same tools reach beyond whole images. A stop-gradient and a moving-average target keep joint-embedding predictive architectures from collapsing, models that predict the embeddings of hidden image regions from the visible ones. LeCun's 2022 proposal for such architectures, A Path Towards Autonomous Machine Intelligence, suggested variance and covariance terms like VICReg's for the same purpose.

🧭 What to remember

Without negatives, a loss that only pulls views together is minimized by a constant. Collapse is complete (one point), dimensional (a low-dimensional subspace), or uniform (identical output distributions). Detect it from the embeddings, not the loss, with the coordinate standard deviation of $\hat{\mathbf{z}}$ against $1/\sqrt{d_z}$ , the centered spectrum, and the effective rank.
BYOL (predictor, stop-gradient, moving-average target) and SimSiam (predictor and stop-gradient only) avoid collapse through their training dynamics. In a linear Siamese network, the stop-gradient conserves $s - p^2/\alpha_p$ , the predictor learns $1/(1+\sigma^2)$ , and without the stop-gradient the weights decay exponentially.
With weight decay $\lambda$ , a direction along which the data have variance $\nu$ survives only if $\nu^2/(\nu + \sigma^2) \ge 4\lambda$ . What the augmentations destroy is what the representation ignores.
Barlow Twins drives the batch cross-correlation of the two views to the identity. VICReg combines an invariance term, a hinge on each dimension's standard deviation (not its variance, whose gradient vanishes near collapse), and an off-diagonal covariance penalty. Both reached 73.2 percent and need wide projectors, whose correlation matrices cost memory.
SwAV balances codes over prototypes with entropic optimal transport, solved by Sinkhorn-Knopp, so a collapsed encoder pays at least $\log K_p$ . A large $\varepsilon$ or a small batch undoes the protection.
Negatives do not prevent dimensional collapse. Strong augmentation and implicit regularization shrink directions even in SimCLR, and the projection head absorbs much of it.
Next, §28.4 pushes the moving-average target to its simplest form, a teacher with centering and sharpening, and finds that a ViT trained this way learns to segment objects.

Grill et al. (2020). Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning. BYOL, its ablations on the target and the predictor (Table 5), and the optimal-predictor argument (Section 3.2 and Appendix H).
Chen and He (2021). Exploring Simple Siamese Representation Learning. SimSiam, the stop-gradient experiment, the predictor ablation with its exact gradient argument, and the expectation-maximization hypothesis.
Tian, Chen and Ganguli (2021). Understanding Self-Supervised Learning Dynamics without Contrastive Pairs. The linear theory of this section, eigenspace alignment, and DirectPred.
Zbontar, Jing, Misra, LeCun and Deny (2021). Barlow Twins: Self-Supervised Learning via Redundancy Reduction. The cross-correlation loss and its connection to the information bottleneck.
Bardes, Ponce and LeCun (2022). VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning. The three-term loss, the coefficient study, and the comparison with every other family in its appendix.
Caron et al. (2020). Unsupervised Learning of Visual Features by Contrasting Cluster Assignments. SwAV, online codes by Sinkhorn-Knopp, and multi-crop.
Jing, Vincent, LeCun and Tian (2022). Understanding Dimensional Collapse in Contrastive Self-supervised Learning. Dimensional collapse in SimCLR, its two linear mechanisms, and DirectCLR.
Garrido, Balestriero, Najman and LeCun (2023). RankMe: Assessing the Downstream Performance of Pretrained Self-Supervised Representations by Their Rank. The effective rank as a label-free measure of representation quality.

Made for larger screens

Non-contrastive methods and collapse

What collapse looks like, and how to measure it

Measure one, the spread of each coordinate

Measure two, the spectrum

Measure three, the effective rank

Diagnose collapse

BYOL, a network that predicts a slowly moving copy of itself

SimSiam, the stop-gradient alone

A linear Siamese network you can solve

SimSiam in one dimension, with and without the stop-gradient

Sixteen directions at once, and which ones survive

Setting the predictor directly

DirectPred on a linear Siamese network

Is BatchNorm a hidden source of negatives?

Barlow Twins, make the cross-correlation the identity

Barlow Twins, no loops

VICReg, three terms with three jobs

VICReg loss and its gradient

Clustering, and codes that must be balanced

Sinkhorn-Knopp for balanced codes

Dimensional collapse, even with negatives

In PyTorch

All the joint-embedding methods side by side

Read more

Watch next