Chapter 28 · Section 28.2 · 85 min

Contrastive methods, SimCLR and MoCo

Contrastive learning trains an image encoder without labels by making it play a matching game. Take a photograph, crop it twice at random, and distort the colors of each crop. Feed both crops, along with crops of many other photographs, through the same network. Train the network so that the two crops of one photo land close together in its output space, and crops of different photos land far apart. The pairs come from how the crops were made, so nobody has to label anything.

In 2020 this recipe made pretraining without labels a serious rival to pretraining with them. Train a ResNet-50 this way on ImageNet's 1.28 million photos, without their labels. Then freeze it and fit only a linear classifier on its outputs, using the labels for that last step alone. The classifier names the right class for 69.3 percent of test images (Chen et al. 2020a). The same network trained end to end on all the labels gets about 76 percent. This section builds the two methods that defined the approach, SimCLR and MoCo, and the lessons from them that every later method in the chapter keeps.

The loss these methods use is InfoNCE, which §7.2 built for text retrieval. A batch holds $B$ matching pairs, such as questions and their answers. Each question scores all $B$ answers with a cosine similarity divided by a temperature $\tau$ , a small positive number that sharpens the scores. The loss is a softmax cross-entropy over those scores, with the question's own answer as the correct class. In plain words, each question must pick its own answer out of the batch, and every other answer is a wrong option to push away.

One experiment in the SimCLR paper makes a good puzzle. Make the color distortion in the training augmentations stronger, and the contrastive model gets better, from 59.6 to 63.2 percent. Train an ordinary supervised ResNet-50 with the same augmentations, and it gets worse, from 77.0 to 75.7 percent. Why would contrastive learning want an augmentation that supervised learning dislikes? By the end of the section you will be able to answer that exactly. On the way we adapt InfoNCE to two views of every photo, find the exact lowest value the adapted loss can reach, and derive the one new term in its gradient. The rest of the theory, the mutual-information bound and the role of the temperature, carries over from §7.2, and we recap it where the image experiments need it. We check each part of the recipe against the published experiments. Then we build MoCo, which makes the same loss work with small batches, and work out how stale its stored negatives are.

Two views of one photograph

Here is the setup in plain words. Take a photo and transform it at random twice, for example with two different crops and two different color changes. Each result is called a view. Run both views through the same network. Train the network so that the two views of one photo get similar outputs, and views of different photos do not.

§28.1 called this the joint-embedding template and showed its trap. If all you ask is that two views agree, a network that sends every image to the same vector agrees perfectly and has learned nothing. That failure is called collapse. Contrastive methods prevent it with negatives, views of other photos that the loss pushes away. The method of Chen et al. (2020a), called SimCLR, fills each part of the template with the simplest choice that works, and the field uses its names for the parts.

📐 The SimCLR framework

Take a batch of $B$ unlabeled images $x_1, \dots, x_B$ and an augmentation distribution $\mathcal{T}$ , a random recipe of crop, flip, and color distortion given below. For each image draw two independent transformations $t, t' \sim \mathcal{T}$ and form the views $\tilde x_i = t(x_i)$ and $\tilde x_{i+B} = t'(x_i)$ . The batch then holds $2B$ views, the first views in rows $1$ to $B$ and the second views in rows $B+1$ to $2B$ . Each view goes through three steps,

\mathbf{h}_i = f_\theta(\tilde x_i) \in \R^{d}, \qquad \mathbf{z}_i = g_\theta(\mathbf{h}_i) \in \R^{d_z}, \qquad \hat{\mathbf{z}}_i = \mathbf{z}_i / \lVert \mathbf{z}_i \rVert .

Here $f_\theta$ is the encoder, $\theta$ the trainable parameters, and $\mathbf{h}_i$ the representation. For SimCLR the encoder is a ResNet-50 without its classification layer, so $\mathbf{h}$ is its $d = 2048$ dimensional pooled feature. The map $g_\theta$ is a small network called the projection head, with output width $d_z = 128$ . The loss sees only the unit vectors $\hat{\mathbf{z}}_i$ . After training, $g_\theta$ is thrown away, and $\mathbf{h} = f_\theta(x)$ is the representation every downstream task reads.

In plain words, the network gets two distorted copies of each photo and must embed them so that each copy can pick out its twin among all the other views. The encoder is a ResNet-50, a convolutional network of 50 layers built from residual blocks (§0.8). It ends by averaging its feature maps over all positions, and it has about 24 million parameters below its classifier.

Two notes on notation. The SimCLR paper numbers the views of image $k$ as $2k - 1$ and $2k$ . We use $i$ and $i + B$ , as most implementations do, because then the index of a view's twin is one line of code. And Chapter 27 wrote a ViT's pooled image vector as $\mathbf{z}$ . The self-supervised literature calls that vector $\mathbf{h}$ and keeps $\mathbf{z}$ for the projection head's output, and so do we.

A representation is graded with the protocols of §28.1, and the main one is the linear probe. Freeze the encoder and compute $\mathbf{h}$ for every labeled ImageNet training image. Train only a linear classifier on those features. Report its top-1 accuracy, the fraction of test images whose highest-scoring class is the right one. ImageNet here means the standard 1.28 million training photos in 1,000 classes. Every percentage in this section is that number unless it says otherwise.

The loss over 2B views

SimCLR's loss is InfoNCE with one change. In §7.2 the questions and the answers were two separate sets of vectors. Here there is one set, the $2B$ views, and each view plays both parts. The view that is asking is the anchor, its twin is the positive, and every other view is a negative. That one change forces four details.

Every view is an anchor. Its candidates are the other $2B - 1$ views of the batch, one positive and $2B - 2$ negatives (both views of each of the other $B - 1$ images).
A view's cosine with itself is always 1. It must be removed from the softmax, or it would win every row.
Each positive pair is scored in both directions, once with each view as the anchor.
The negatives cost nothing extra, because they are the other anchors. A batch of 4,096 images gives each view 8,190 negatives.

📐 NT-Xent, the normalized temperature-scaled cross-entropy

Let $\hat{\mathbf{Z}} \in \R^{2B\times d_z}$ hold the unit embeddings of the $2B$ views as rows, with rows $i$ and $i+B$ the two views of image $i$ . The positive of view $i$ is

i^+ = \begin{cases} i + B, & i \le B,\\ i - B, & i > B,\end{cases} \qquad\text{and the positive-pair matrix is}\qquad \Pi_{ij} = \mathbb{1}[\,j = i^+\,],\quad \boldsymbol\Pi \in \lbrace 0,1\rbrace^{2B\times 2B}.

With temperature $\tau > 0$ and logits $\mathbf{S} = \hat{\mathbf{Z}}\hat{\mathbf{Z}}^\top/\tau \in \R^{2B\times 2B}$ , the loss of anchor $i$ and the batch loss are

\ell_i = -\log\frac{\exp(S_{i i^+})}{\sum_{j \ne i}\exp(S_{ij})}, \qquad \mathcal{L}_{\text{NT-Xent}} = \frac{1}{2B}\sum_{i=1}^{2B}\ell_i .

This is equation 1 of the SimCLR paper, with its batch size $N$ written $B$ . In code, set the diagonal of $\mathbf{S}$ to $-\infty$ and take a cross-entropy over each full row with target column $i^+$ . Write $\mathbf{A}$ for the row softmax of that masked $\mathbf{S}$ , so $A_{ii} = 0$ and each row of $\mathbf{A}$ sums to 1.

In plain words, each of the $2B$ views plays a guessing game with $2B - 1$ options ("which of the other views is my twin?"), and the loss is the average surprise at the right answer. The name lists the three ingredients, normalized embeddings, a temperature, and a cross-entropy. The loss is older than the name. The same softmax over one positive and many negatives appears in Sohn (2016), Wu et al. (2018), and van den Oord et al. (2018). SimCLR's contribution was the recipe around it.

Interviewers often ask you to write this loss from memory. People slip in two places, the mask on the diagonal and the index of the positive. Here is the loss on four views, small enough to check by hand.

✏️ NT-Xent by hand

Two images, two views each, embedded on the unit circle. Image 1's views sit at 0 and 30 degrees, image 2's at 120 and 150 degrees. In the order $(a_1, a_2, b_1, b_2)$ , first views then second views, the angles are $(0°, 120°, 30°, 150°)$ and the cosine matrix is

\hat{\mathbf{Z}}\hat{\mathbf{Z}}^\top = \begin{pmatrix} 1 & -0.5 & 0.866 & -0.866\\ -0.5 & 1 & 0 & 0.866\\ 0.866 & 0 & 1 & -0.5\\ -0.866 & 0.866 & -0.5 & 1\end{pmatrix}.

Take $\tau = 0.5$ . Anchor $a_1$ has its positive $b_1$ at cosine $0.866$ , and its negatives $a_2$ and $b_2$ at $-0.5$ and $-0.866$ . Divide the top and bottom of the softmax by the positive's term,

\ell_{a_1} = \log\!\big(1 + e^{(-0.5 - 0.866)/0.5} + e^{(-0.866 - 0.866)/0.5}\big) = \log(1 + 0.0651 + 0.0313) = 0.092 .

Anchor $a_2$ has its positive $b_2$ at $0.866$ , but one negative, $b_1$ , sits at cosine 0. So $\ell_{a_2} = \log(1 + e^{-2.732} + e^{-1.732}) = \log(1.242) = 0.217$ . By symmetry $\ell_{b_1} = 0.217$ and $\ell_{b_2} = 0.092$ , and the batch loss is the mean, $0.154$ . The two inner views, at 30 and 120 degrees, are the ones that could be confused. At $\tau = 1$ the same embeddings cost $0.438$ , and at $\tau = 0.1$ only $0.0001$ .

Here is the loss in NumPy, line by line as in the definition, checked on the hand example.

Python · cell 1

Where the loss starts and how low it can go

Two numbers frame every NT-Xent training curve. The first is where training starts, and §7.2 already found it for InfoNCE. (Losses here use the natural logarithm, so they are measured in nats.) A random encoder gives nearly equal logits, so each row's softmax is close to uniform and the loss is about the log of the number of candidates. Here each anchor has $2B - 1$ candidates, so $\ell_i \approx \log(2B - 1)$ . That is $\log 511 = 6.24$ at $B = 256$ and $\log 8191 = 9.01$ at $B = 4096$ .

The second number is the lowest the loss can ever go. §7.2 showed that at $\tau = 1$ it cannot reach zero, using an ideal encoder with every negative at cosine $-1$ . No encoder can do that, because many vectors cannot all point away from each other. With all $2B$ views in one set we can find the exact floor. Cosines lie between $-1$ and $1$ , so every logit lies between $-1/\tau$ and $1/\tau$ , and the floor depends on $\tau$ .

∂ Proof of the NT-Xent lower bound, by Jensen's inequality twice

Write $c_{ij} = \hat{\mathbf{z}}_i^\top\hat{\mathbf{z}}_j$ , let $p_i = c_{i i^+}$ be anchor $i$ 's positive cosine, and let $\mathcal{N}(i)$ be its $n = 2B - 2$ negatives. Dividing the softmax by the positive's term, as in the hand example,

\ell_i = \log\Big(1 + \sum_{j \in \mathcal{N}(i)} e^{(c_{ij} - p_i)/\tau}\Big).

Step 1, average the negatives. The exponential is convex, so by Jensen's inequality $\frac{1}{n}\sum_{j} e^{x_j} \ge e^{\bar x}$ for the mean $\bar x$ . With $\bar c_i$ the mean negative cosine of anchor $i$ , this gives $\ell_i \ge \varphi\big((\bar c_i - p_i)/\tau\big)$ with $\varphi(u) = \log(1 + n e^{u})$ .

Step 2, average the anchors. The function $\varphi$ is increasing and convex (its second derivative is $n e^{u}/(1 + n e^{u})^2 > 0$ ), so by Jensen again $\frac{1}{2B}\sum_i \varphi(u_i) \ge \varphi(\bar u)$ , and

\mathcal{L} \;\ge\; \varphi\Big(\frac{1}{\tau}\Big[\frac{1}{2B}\sum_i \bar c_i - P\Big]\Big), \qquad P = \frac{1}{2B}\sum_i p_i .

Step 3, the negatives cannot all be far. Expand $\big\lVert\sum_k \hat{\mathbf{z}}_k\big\rVert^2 = \sum_{i}\sum_{j} c_{ij} \ge 0$ . The diagonal contributes $2B$ and the positive pairs contribute $\sum_i p_i = 2BP$ , so the negative cosines sum to at least $-2B - 2BP$ . There are $2B \cdot n$ of them, so

\frac{1}{2B}\sum_i \bar c_i \;\ge\; -\frac{1 + P}{n} .

Step 4, put it together. The argument of $\varphi$ is at least $\frac{1}{\tau}\big[-\frac{1+P}{n} - P\big]$ , which decreases as $P$ grows, and $P \le 1$ . At $P = 1$ it equals $-\frac{1}{\tau}\big(1 + \frac{2}{n}\big) = -\frac{B}{(B-1)\tau}$ , and since $\varphi$ is increasing this gives the bound. Equality needs every inequality to be tight. Step 4 needs $P = 1$ (twins coincide), Step 3 needs $\sum_k \hat{\mathbf{z}}_k = \mathbf{0}$ , and the two Jensen steps need every negative cosine to be equal. $B$ unit vectors with equal pairwise cosines summing to zero form a regular simplex, with cosine $-1/(B-1)$ , and such a simplex spans $B - 1$ dimensions. $\square$

In plain words, the temperature sets a floor under the loss. The loss can approach zero only when $e^{1/\tau}$ is much larger than the number of negatives, roughly when $\tau$ is well below $1/\log(2B)$ . At $B = 4096$ that threshold is $0.11$ .

✏️ The floor at SimCLR's scale

Take $B = 4096$ and $\tau = 1$ . The loss starts at $\log 8191 = 9.01$ nats, and the floor is $\log(1 + 8190\,e^{-4096/4095}) = 8.01$ nats, only one nat lower. A loss of at least 8.01 means that the softmax gives the positive a probability of at most $e^{-8.01}$ , about 0.03 percent, in the sense of a geometric mean over anchors. The positive can still be the top candidate, but the softmax can never be confident about it.

The SimCLR paper measured exactly this (its Table 5). It scores each anchor against the 4,096 views in the other half of the batch, its positive and 4,095 negatives. At $\tau = 1$ the contrastive accuracy, how often the positive is the top candidate, was 59.1 percent. Yet the entropy of that softmax stayed at 8.3 nats, essentially the $\log 4096 = 8.32$ of a uniform distribution over the 4,096 candidates. Real ImageNet embeddings in $d_z = 128$ dimensions cannot form a simplex of 4,096 points, so their true minimum is higher still. At $\tau = 0.1$ the floor drops to 0.32.

Python · cell 2

At $\tau = 1$ the random encoder sits at 6.25, right at $\log(2B-1)$ . At $\tau = 0.1$ it starts higher, at 6.74. Random directions in many dimensions are nearly orthogonal, and their cosines spread by about $1/\sqrt{128} = 0.088$ around zero (§1.3). A small temperature magnifies that spread tenfold. A softmax that is confidently wrong costs more than a uniform one, as §7.2 showed. The simplex meets the bound to six decimals, so the bound cannot be improved.

This tells you what a healthy training curve looks like. A loss that starts far from $\log(2B - 1)$ points to a bug in the mask or the indexing. A loss that stalls near the floor at a large $\tau$ is not a bug at all.

The widget shows the full $2B \times 2B$ matrix of cosines for a batch of four images. Click a row to see its softmax, the probability the anchor gives each candidate. The lower bars split the negatives' probability among them. Each negative's share is also its share of the gradient that pushes the negatives away, as the next part shows.

The gradient pushes every view twice

The gradient tells you what the loss does to the embeddings, and "implement NT-Xent with its backward pass" is a standard interview question. §7.2 derived the InfoNCE gradient as expected minus observed, a pull toward the positive and a push away from the softmax-weighted average of the candidates. NT-Xent differs from it in one way that trips people up. There the questions and the answers were different matrices, so each received one gradient term. Here the same matrix $\hat{\mathbf{Z}}$ sits on both sides of $\mathbf{S} = \hat{\mathbf{Z}}\hat{\mathbf{Z}}^\top/\tau$ . So every view receives a gradient twice. It gets one as an anchor, from its own row, and another as a candidate, from every other anchor's row.

Recall the rule for a softmax cross-entropy from §0.2. The gradient with respect to the logits is the predicted probabilities minus the one-hot target. Applied to every row at once, that gives the theorem.

∴ The gradient of NT-Xent

With $\mathbf{A}$ the row softmax of the masked logits (zero diagonal) and $\boldsymbol\Pi$ the positive-pair matrix,

\frac{\partial\mathcal{L}}{\partial\mathbf{S}} = \frac{\mathbf{A} - \boldsymbol\Pi}{2B}, \qquad \frac{\partial\mathcal{L}}{\partial\hat{\mathbf{Z}}} = \frac{1}{2B\tau}\Big[(\mathbf{A} - \boldsymbol\Pi) + (\mathbf{A} - \boldsymbol\Pi)^\top\Big]\hat{\mathbf{Z}} = \frac{1}{2B\tau}\big(\mathbf{A} + \mathbf{A}^\top - 2\boldsymbol\Pi\big)\hat{\mathbf{Z}} .

If $\hat{\mathbf{z}}_i = \mathbf{z}_i/\lVert\mathbf{z}_i\rVert$ and $\mathbf{g}_i$ is row $i$ of the matrix above, then $\partial\mathcal{L}/\partial\mathbf{z}_i = \big(\mathbf{g}_i - (\hat{\mathbf{z}}_i^\top\mathbf{g}_i)\hat{\mathbf{z}}_i\big)/\lVert\mathbf{z}_i\rVert$ . This last step, the gradient through normalization, is the one derived in §7.2. It removes the part of $\mathbf{g}_i$ that points along $\hat{\mathbf{z}}_i$ , since that part would only change the vector's length.

∂ Proof of the NT-Xent gradient

Row $i$ of the loss is a softmax cross-entropy over the entries $S_{ij}$ , $j \ne i$ , with target $i^+$ . By the "predicted minus true" rule its gradient is $(A_{ij} - \Pi_{ij})/(2B)$ , the $1/(2B)$ coming from the mean. The masked diagonal entry is the constant $-\infty$ , so it has no gradient, and $A_{ii} = \Pi_{ii} = 0$ makes the formula hold there too.

Now $S_{ij} = \hat{\mathbf{z}}_i^\top\hat{\mathbf{z}}_j/\tau$ depends on $\hat{\mathbf{z}}_k$ when $k = i$ and when $k = j$ . Write $\mathbf{G} = \partial\mathcal{L}/\partial\mathbf{S}$ . The chain rule gives

\frac{\partial\mathcal{L}}{\partial\hat{\mathbf{z}}_k} = \frac{1}{\tau}\sum_j G_{kj}\,\hat{\mathbf{z}}_j + \frac{1}{\tau}\sum_i G_{ik}\,\hat{\mathbf{z}}_i ,

which is row $k$ of $(\mathbf{G} + \mathbf{G}^\top)\hat{\mathbf{Z}}/\tau$ . The positive-pair relation is symmetric ( $i^{++} = i$ ), so $\boldsymbol\Pi^\top = \boldsymbol\Pi$ , which gives the last form. $\square$

In plain words, a gradient step pulls view $k$ toward its twin twice, once through its own row and once through its twin's row. It pushes view $k$ away from the softmax-weighted average of the candidates in its own row. And it pushes view $k$ away from every other anchor that mistook it for that anchor's twin, in proportion to the probability that anchor gave it.

Forgetting the transpose term is the most common bug. The code still trains, slowly and wrongly, so always check a hand-written gradient against finite differences. One check is free. Every row of $\mathbf{A} - \boldsymbol\Pi$ sums to zero, because each softmax row sums to one.

✏️ The gradient by hand

Continue the hand example at $\tau = 0.5$ . Row $a_1$ of the softmax is $(0, 0.059, 0.912, 0.029)$ over $(a_1, a_2, b_1, b_2)$ , so row $a_1$ of $\mathbf{A} - \boldsymbol\Pi$ is $(0, 0.059, -0.088, 0.029)$ . Add column $a_1$ , which records how much the other anchors chose $a_1$ . That gives row $a_1$ of $\mathbf{A} + \mathbf{A}^\top - 2\boldsymbol\Pi$ as $(0, 0.112, -0.283, 0.057)$ . Divide by $2B\tau = 2$ and multiply by the unit vectors,

\mathbf{g}_{a_1} = \tfrac12\big(0.112\,\hat{\mathbf{z}}_{a_2} - 0.283\,\hat{\mathbf{z}}_{b_1} + 0.057\,\hat{\mathbf{z}}_{b_2}\big) = (-0.175,\ -0.008).

Almost all of it points along $\hat{\mathbf{z}}_{a_1} = (1, 0)$ itself, and the gradient through normalization removes that part. What survives is $(0, -0.008)$ . A gradient step moves against the gradient, so it turns $a_1$ slightly counterclockwise, toward its twin at 30 degrees.

The inner views get a sideways gradient 33 times larger, of norm 0.26. A step turns $a_2$ toward its twin at 150 degrees and $b_1$ toward its twin at 0 degrees, away from each other. That is exactly where the loss was.

Interview

NT-Xent loss and gradient over 2B views, no loops

Implement nt_xent(Z, tau) and return (loss, dZ).

The input Z has shape (2B, d) and holds unnormalized embeddings. Rows i and i + B are the two views of image i.
Normalize the rows, then compute $\mathcal{L}_{\text{NT-Xent}}$ as defined above. Leave each view's similarity with itself out of its softmax, take the positive of row i to be row (i + B) % (2B), and average over all $2B$ anchors.
dZ is the gradient of the loss with respect to the unnormalized Z.

Use no Python loops over rows, and keep everything finite at tau = 0.001.

import numpy as np

def nt_xent(Z, tau):
    """Z: (2B, d), rows i and i + B are two views of image i. Returns (loss, dZ)."""
    raise NotImplementedError

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

What the augmentations must destroy

Now the puzzle from the start. The loss asks the network to tell which views came from the same photograph. It does not care how the network tells. The network will find and use any cue that two crops of one photo share and crops of different photos do not. The content of the photo is such a cue, which is what we want. But so is its color distribution.

A beach photo is mostly blue and sand, and a forest photo is mostly green. Two random crops of one photo have nearly the same color histogram, the fraction of pixels at each intensity in each color channel. Crops of different photos do not. SimCLR's Figure 6 shows this. Histograms of different crops of one image look alike and differ across images, until color distortion is added. So a network can solve the contrastive task by matching histograms, which takes no understanding of shapes at all.

Supervised learning has no such shortcut. The labels define its task, and a color histogram does not predict "golden retriever" well enough to get far. There, color distortion only removes information that is sometimes useful, so it hurts a little. In contrastive learning the augmentations define the task. They say which differences between two inputs should be ignored. If color is not among them, color becomes the answer. That resolves the opening puzzle.

📐 The SimCLR augmentation policy for ImageNet

Each view is made by these steps, in order.

Random resized crop. A crop of area fraction $a \sim U(a_-, a_+) = U(0.08, 1)$ of the image, with aspect ratio between $3/4$ and $4/3$ (torchvision draws it log-uniformly), resized to $224\times224$ pixels.
Horizontal flip with probability 0.5.
Color jitter with probability 0.8 and strength $s$ (default $s = 1$ ). The brightness, contrast, and saturation factors are drawn from $[1 - 0.8s, 1 + 0.8s]$ , and the hue shifts by up to $0.2s$ . In torchvision this is ColorJitter(0.8s, 0.8s, 0.8s, 0.2s). The TensorFlow code of the original SimCLR shifts brightness additively by up to $0.8s$ instead.
Grayscale with probability 0.2.
Gaussian blur with probability 0.5, standard deviation $\sigma \sim U(0.1, 2.0)$ , and a kernel 10 percent of the image side.

The supervised augmentations of §27.4 are not part of this policy. RandAugment applies random image operations tuned for classifiers, and Mixup and CutMix blend two images and their labels, which here do not exist. The evidence for this policy comes from SimCLR's own experiments. No single transformation is enough on its own. Cropping plus color distortion is the strongest pair. Blur adds a little (63.2 to 64.5 percent at 100 epochs), and so does flipping (removing it drops 64.5 to 63.4).

The widget makes two views of one synthetic photo with the crop, flip, color jitter, and grayscale steps. It then plays the matching game 128 times with color alone. Each time a view's twin hides among one fresh view of each of 8 photos, and the widget counts how often the twin has the closest color histogram.

Warm-up

The color-histogram shortcut

Measure the shortcut yourself. The setup provides 60 synthetic photographs and a function that crops them.

IMAGES is a list of 60 arrays, each $24\times24\times3$ with values in $[0, 1]$ .
make_views(jitter, seed) returns two lists of 60 views, A and B, where A[i] and B[i] are two crops of photo i. Crops have different sizes and are not resized.

Implement two functions.

color_histogram(view, bins) returns a vector of length 3 * bins. It holds the histogram of each color channel over bins equal bins on $[0, 1]$ , concatenated in the order red, green, blue. A value of exactly 1 goes in the last bin. Divide each channel's counts by the number of pixels, so each channel's histogram sums to 1.
histogram_match_accuracy(views_a, views_b, bins) returns the fraction of indices i for which views_b[i] is the view in views_b closest to views_a[i], by L1 distance between histograms.

Then compare the accuracy with and without jitter.

import numpy as np

def color_histogram(view, bins):
    """view: (h, w, 3) in [0, 1]. Returns (3 * bins,), each channel's histogram summing to 1."""
    raise NotImplementedError

def histogram_match_accuracy(views_a, views_b, bins):
    """Fraction of i whose nearest neighbor (L1) of views_a[i] among views_b is views_b[i]."""

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

The histogram shortcut is not hypothetical, and we can watch a contrastive learner take it. The next cell builds a tiny world of "photos", each described by 16 numbers. Eight are content features, the photo's class (one of 4) plus details of that particular photo. The other eight are color features, the photo's own palette, unrelated to its class. A crop changes the content features a lot, since different crops show different parts. Without jitter, the color features barely change between views.

The encoder is a linear map down to only $d_z = 4$ dimensions, so it must choose what to keep. We train it with NT-Xent and the gradient derived above. Then we freeze it and fit two linear probes, one for the class and one for the color features.

Python · cell 3

✏️ A contrastive learner takes the color shortcut

Without jitter the encoder spends its four dimensions on color. The norm of its weights on the color features is 3.77, against 0.12 on content. The class probe scores 23.5 percent, which is chance. A probe for the color features recovers half their variance ( $R^2 = 0.50$ ), the most that four dimensions can hold of eight equally varied features.

With strong jitter the encoder switches to content, with weight norms 5.97 on content and 0.58 on color. The class probe jumps to 91.8 percent and color is gone ( $R^2 = 0.00$ ).

Now look at the final losses, 2.64 without jitter and 4.71 with it. The shortcut solves the contrastive task better and gives useless features. We will meet the same pattern with the temperature, the projection head, and BatchNorm. A low contrastive loss is not evidence of a good representation.

The second half of that result is a warning. Jitter made the features blind to color, and for some tasks color is the signal, as with flowers, birds, ripe fruit, or lesions in medical images. Ericsson et al. (2021) compared 13 self-supervised models on 40 downstream tasks (§28.1). They found that the best ones keep less color information than supervised models do.

Tian et al. (2020) turned the lesson into a principle they call InfoMin. Good views should share as little information as possible beyond what the downstream task needs, since the encoder may use anything the views share. They showed that stronger augmentation lowers the information shared between views and raises downstream accuracy, up to a point, and they reached 73.0 percent with a ResNet-50.

The projection head, and why it is thrown away

SimCLR computes the loss on $\mathbf{z} = g_\theta(\mathbf{h})$ but evaluates $\mathbf{h}$ , the layer below. Why have a head at all, and why throw it away after training?

Start with a simple fact. A head that is only a linear map cannot help a linear probe directly.

So if a head helps, it must help by changing what the encoder learns. The experiments say it does, a lot. SimCLR proposed an explanation and tested it. The loss asks $\mathbf{z}$ to be invariant to the augmentations, so the head learns to throw away whatever they change, such as color, orientation, and the exact crop. The layer below, $\mathbf{h}$ , does not have to, and keeps more.

To test this, SimCLR trained small MLPs to predict which transformation had been applied, once from $\mathbf{h}$ and once from $\mathbf{z}$ . For this test the pretraining augmentations also included rotations, noise, and edge filtering.

The head is a buffer that absorbs the invariance, so the representation underneath stays rich. In practice, always probe the encoder's output and never the head's. An interviewer who asks "why throw the head away?" wants this answer, with the rotation numbers as evidence.

🔬 Where to cut, across methods

Bordes et al. (2022) showed that the effect is general. Across self-supervised methods, reading features from below the last projector layers rather than from the top is worth more than 30 points of ImageNet accuracy in some settings. They call this Guillotine Regularization. The best layer to cut at depends on the training setup, the data, and the downstream task. SimCLR v2 (Chen et al. 2020b) moved the cut the other way. With a deeper three-layer head, fine-tuning from the head's first layer rather than from $\mathbf{h}$ improved learning from 1 percent of the labels by as much as 14 percent relative.

The head has a second role, which §28.3 takes up. Even with negatives, contrastive embeddings can collapse onto a low-dimensional subspace. Jing et al. (2022) measured this in SimCLR. Its 128-dimensional $\mathbf{z}$ uses only part of its dimensions. When SimCLR is trained without a projector, the collapse moves into the 2,048-dimensional $\mathbf{h}$ itself, so the head also shields the representation.

Temperature and normalization

§7.2 showed that the temperature decides which negatives get the push. A small $\tau$ concentrates the gradient on the few negatives closest to the anchor, the hard negatives. A large $\tau$ spreads it evenly over all of them. SimCLR measured what that does to image features, and the floor above explains one end of the curve.

Read the table from both ends. At $\tau = 1$ the softmax cannot become confident, because of the floor. So every negative keeps pushing about equally, and the hard ones get no special attention. At $\tau = 0.05$ almost all the push goes to the few nearest negatives. Many of those are false negatives, other photos of the same class, and the loss pushes them away hardest (the pitfall of §7.2).

The unnormalized runs are the shortcut lesson again. Without normalization the network can sharpen its softmax just by making its vectors longer, since a longer vector acts like a smaller temperature. That raises the contrastive accuracy without improving the features.

The best temperature is not universal. MoCo, built later in this section, used 0.07. MoCo v2 found 0.2 better once it added an MLP head. MoCo v3's ResNet-50 recipe uses $\tau = 1$ with an extra prediction head and reaches 73.8 percent. Tune $\tau$ together with the rest of the recipe, and expect an interviewer to ask what goes wrong at each extreme.

Big batches and the LARS optimizer

More negatives make a harder game. They also raise a ceiling. §7.2 showed that InfoNCE with $N$ candidates can certify at most $\log N$ nats of information shared by the two views, here $\log(2B - 1)$ . SimCLR deliberately keeps no store of old embeddings, so its negatives come only from the batch. It measured what batch size buys.

📌 Batch size and epochs (SimCLR, Table B.1)

ResNet-50 linear probe, with the learning rate scaled linearly with the batch.

At 100 epochs, batches of 256, 1,024, 4,096, and 8,192 give 57.5, 62.8, 64.6, and 64.8.
At 800 epochs they give 66.6, 68.5, 68.9, and 69.0.
With square-root scaling, $\eta = 0.075\sqrt{B}$ (equal to the linear rule $0.3 \times B/256$ at $B = 4096$ ), the batch-256 run at 100 epochs rises from 57.5 to 62.8.

In longer runs with square-root scaling (the paper's Figure B.2), gains saturate near a batch of 8,192, while training longer keeps helping up to 3,200 epochs. The defaults are batch 4,096, 100 epochs for ablations, the LARS optimizer with learning rate 4.8, weight decay $10^{-6}$ , 10 warmup epochs, and cosine decay.

Big batches help most when training is short. The gap mostly closes with more epochs, because each image then meets many more negatives over the course of training.

A batch of 4,096 at learning rate 4.8 is unstable with plain momentum SGD (§0.6), so SimCLR uses LARS (You et al. 2017) at every batch size. The idea is that different layers need different step sizes. The ratio of a layer's gradient norm to its weight norm varies wildly between layers and over training. LARS rescales each layer's step to a fixed fraction of that layer's own weight norm, with a per-layer factor called the trust ratio.

📐 LARS, layer-wise adaptive rate scaling

For each layer $\ell$ with weights $\mathbf{w}_\ell$ , stochastic gradient $\mathbf{g}_\ell$ , weight decay $\lambda$ , trust coefficient $\eta_{\text{trust}}$ (0.001 in SimCLR), global learning rate $\eta$ , and momentum $\mu$ , LARS sets the local learning rate, or trust ratio,

\eta_\ell = \eta_{\text{trust}}\,\frac{\lVert\mathbf{w}_\ell\rVert}{\lVert\mathbf{g}_\ell\rVert + \lambda\lVert\mathbf{w}_\ell\rVert}

(taken to be 1 when $\lVert\mathbf{w}_\ell\rVert = 0$ or the denominator is 0) and updates

\mathbf{v}_\ell \leftarrow \mu\,\mathbf{v}_\ell + \eta\,\eta_\ell\,(\mathbf{g}_\ell + \lambda\mathbf{w}_\ell), \qquad \mathbf{w}_\ell \leftarrow \mathbf{w}_\ell - \mathbf{v}_\ell .

Biases and BatchNorm gains and shifts are excluded from both the trust ratio and the weight decay. They get plain momentum SGD, $\mathbf{v} \leftarrow \mu\mathbf{v} + \eta\,\mathbf{g}$ . This is You et al.'s Algorithm 1 with the exclusions used by SimCLR and BYOL.

∴ LARS bounds every layer's relative step

For an adapted layer with nonzero weights and a zero momentum buffer, $\lVert\Delta\mathbf{w}_\ell\rVert \le \eta\,\eta_{\text{trust}}\,\lVert\mathbf{w}_\ell\rVert$ , whatever the scale of the gradient. Equality holds when $\lambda = 0$ or when $\mathbf{g}_\ell$ points the same way as $\mathbf{w}_\ell$ .

Proof. If the denominator is 0, then $\mathbf{g}_\ell = \mathbf{0}$ and $\lambda = 0$ , so the step is zero. Otherwise $\lVert\Delta\mathbf{w}_\ell\rVert = \eta\,\eta_\ell\,\lVert\mathbf{g}_\ell + \lambda\mathbf{w}_\ell\rVert \le \eta\,\eta_\ell\,(\lVert\mathbf{g}_\ell\rVert + \lambda\lVert\mathbf{w}_\ell\rVert) = \eta\,\eta_{\text{trust}}\lVert\mathbf{w}_\ell\rVert$ , by the triangle inequality and the definition of $\eta_\ell$ . The triangle inequality is tight exactly when one vector is a nonnegative multiple of the other. $\square$

In plain words, at SimCLR's settings each layer moves by at most $4.8 \times 0.001 = 0.48$ percent of its own norm per step, before momentum, whether its gradient is tiny or huge. With momentum $\mu = 0.9$ the buffer adds up as many as $1/(1 - \mu) = 10$ such steps, so about 5 percent. That is what makes a learning rate of 4.8 safe. A layer with small weights and large gradients, which plain SGD at that rate would blow up, gets a proportionally small step.

Core

LARS with exclusions

Implement lars_step(params, grads, bufs, lr, wd, trust, momentum, exclude), one LARS step exactly as in the definition above.

params, grads, and bufs are dicts from parameter names to NumPy arrays, where bufs holds the momentum buffers.
exclude is a set of names, the biases and normalization parameters.
For an array not in exclude, use the trust ratio $\eta_{\text{trust}}\lVert\mathbf{w}\rVert/(\lVert\mathbf{g}\rVert + \lambda\lVert\mathbf{w}\rVert)$ . Use a ratio of 1 if $\lVert\mathbf{w}\rVert = 0$ or the denominator is 0.
For an excluded array, use plain momentum SGD without weight decay.

Norms are taken over the whole array. Return (new_params, new_bufs) as new dicts, and leave the inputs unchanged.

import numpy as np

def lars_step(params, grads, bufs, lr, wd, trust, momentum, exclude):
    """params, grads, bufs: dicts name -> array. exclude: set of names. Returns (new_params, new_bufs)."""
    raise NotImplementedError

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

BatchNorm can leak the pairing

The last ingredient of SimCLR is the one that silently breaks contrastive learning when you spread training over many devices. Recall from §0.7 that BatchNorm, in training mode, normalizes each feature with the mean and variance of the current batch. Training on many devices usually uses data parallelism. Each device holds a copy of the model and processes its own slice of the batch, so it normally computes those statistics over its own slice only.

There are three fixes.

Global BatchNorm, used by SimCLR, computes the mean and variance over all devices together (torch.nn.SyncBatchNorm in PyTorch).
Shuffled BatchNorm, used by MoCo, applies to the second encoder that MoCo uses to make its keys (built below). It shuffles the batch before splitting it across devices for that encoder, and puts the outputs back in order afterward. So the two views of an image are normalized with statistics from different slices.
No batch statistics. Hénaff et al. (2019) replaced BatchNorm with LayerNorm, which normalizes each example on its own.

In an interview, "my contrastive accuracy is near 100 percent and my probe is bad" should make you think of this pitfall first.

The negatives must also be shared across devices. Each device collects the embeddings of the whole global batch with an all-gather, an operation that gives every device the tensors of every other device. Then it computes the loss for its own anchors. The gather must pass gradients back to the device that produced each embedding, or every device's negatives become constants. The collectives themselves are in §10.2. §7.2 shows how to compute a giant contrastive batch in chunks with exactly the full-batch gradient, a trick called gradient caching.

Here is how the pieces look in PyTorch. The block is shown, not run, because it needs torchvision and several processes.

✏️ SimCLR at real scale

A ResNet-50 costs about 4.1 billion multiply-accumulates per forward pass on a 224-pixel image, so about 8.2 GFLOPs, and a backward pass costs about twice that. SimCLR runs two views of each of the 1.28 million ImageNet images forward and backward every epoch. That is about $1.28\times10^{6} \times 2 \times 3 \times 8.2\times10^{9} = 6.3\times10^{16}$ FLOPs per epoch.

The paper reports about 1.5 hours per 100 epochs on 128 TPU v3 cores at batch 4,096. So the hardware sustained roughly $6.3\times10^{18}/5400\ \text{s} \approx 1.2$ PFLOP/s, our estimate from those two numbers. The logit matrix is $8192\times8192$ , 268 MB in float32, small next to the activations of 8,192 ResNet passes.

On ordinary GPUs the batch is the problem. The MoCo v2 authors estimated that the same loss at batch 4,096 needs 93 GB per GPU, on a machine with 8 V100 GPUs of 16 GB each. That is why they kept a batch of 256 and found the negatives elsewhere.

MoCo, a dictionary larger than the batch

SimCLR's negatives cost a huge batch. He et al. (2020) found a way to get many negatives with a small one. They looked at contrastive learning as a dictionary lookup. The query is one view of an image. The keys are encoded views of many images. The loss asks the query to find its own key, the one made from the other view of the same image.

Their hypothesis was that good features need a dictionary with two properties. It should be large, so there are many negatives. And it should be consistent, with all keys made by nearly the same encoder, so that comparing the query with them is fair. The two other mechanisms of the time each gave up one of the two.

In the end-to-end mechanism, which SimCLR later scaled up, the dictionary is the batch. It is consistent, but its size is the batch size, which memory limits.
The memory bank of Wu et al. (2018), met in §28.1, stores one feature for every image in the dataset. It is large. But each stored feature was computed when its image was last seen, up to a whole epoch earlier, by an encoder that has changed since.

MoCo keeps both properties with two pieces. A queue holds the keys of recent batches, and a slowly moving copy of the encoder makes them.

📐 MoCo, momentum contrast

MoCo has a query encoder $f_\theta$ , trained by gradient descent, and a key encoder of the same architecture with parameters $\xi$ , never trained by gradients. In MoCo v1 both are ResNet-50s whose final layer outputs 128 dimensions, normalized to unit length. For each image of a batch of $B$ , two views give a query $\mathbf{q} = f_\theta(x^{(1)})$ and its positive key $\mathbf{k}^+ = f_\xi(x^{(2)})$ . A queue $\mathcal{Q} \in \R^{K_q\times d_z}$ holds the keys $\mathbf{k}^-_1, \dots, \mathbf{k}^-_{K_q}$ of the most recent batches, and the loss of one query is the InfoNCE

\mathcal{L}_{\mathbf{q}} = -\log\frac{\exp(\mathbf{q}^\top\mathbf{k}^+/\tau)}{\exp(\mathbf{q}^\top\mathbf{k}^+/\tau) + \sum_{j=1}^{K_q}\exp(\mathbf{q}^\top\mathbf{k}^-_j/\tau)},

averaged over the batch. After the gradient step on $\theta$ , the key encoder moves toward it by an exponential moving average, the momentum update

\xi \leftarrow m\,\xi + (1 - m)\,\theta ,

with momentum $m \in [0, 1)$ . Then the batch's keys are added to the queue and the oldest batch's keys are dropped. MoCo's defaults are $K_q = 65{,}536$ , $m = 0.999$ , $\tau = 0.07$ , a batch of 256 on 8 GPUs, SGD with learning rate 0.03, and 200 epochs, which took about 53 hours for a ResNet-50.

In plain words, the queue makes the number of negatives independent of the batch size, and the slow key encoder keeps the queue's keys comparable with each other. Each query sees one fresh positive and 65,536 negatives, while backpropagation runs through only the 256 queries.

Three details matter in practice. The positive comes from the current batch and is not in the queue. MoCo v1 uses only the queue as negatives, not the other keys of the batch. And the keys get no gradient, so the key encoder's forward pass runs under no_grad. We rename two of the paper's symbols to avoid clashes. Its queue length $K$ becomes $K_q$ , since $K$ is a number of classes in this book, and its key-encoder parameters $\theta_k$ become $\xi$ .

The exponential moving average is the same one §0.6 used to average weights for evaluation, with $m$ in the role of that section's $\beta$ . Here it produces the encoder that makes the targets. The non-contrastive methods of §28.3 and the self-distillation of §28.4 use it the same way. Three numbers describe it.

∴ The window, lag, and half-life of a momentum encoder

With the update order of MoCo (step $\theta_t$ , then $\xi_t = m\,\xi_{t-1} + (1 - m)\,\theta_t$ ), after the start-up transient

\xi_t = (1 - m)\sum_{j \ge 0} m^{j}\,\theta_{t-j},

an average whose weights sum to 1. Its window is $1/(1 - m)$ steps. The mean age of the parameters in it, its lag behind $\theta_t$ , is $m/(1 - m)$ steps. Half of its weight lies on the last $\ln 2/(-\ln m) \approx 0.693/(1 - m)$ steps, its half-life.

Proof. Unroll the recursion. The initial value is multiplied by $m^t$ , which vanishes. The weights sum to $(1 - m)\sum_j m^j = 1$ . The mean age $\sum_j j\,(1 - m)m^j = m/(1 - m)$ was derived in §0.6 by differentiating the geometric series. The total weight on ages $j \ge s$ is $(1 - m)\sum_{j \ge s} m^j = m^s$ , which equals $1/2$ at $s = \ln 2/(-\ln m)$ , and $-\ln m \approx 1 - m$ for $m$ near 1. $\square$

For $m = 0.999$ the window is 1,000 steps, the lag 999 steps, and the half-life 693 steps. In practice this means a key encoder that averages roughly the last thousand versions of the query encoder, so it changes very little from one step to the next.

How old are the keys in the queue

Now we can put a number on MoCo's consistency argument. At step $t$ the query encoder has parameters $\theta_t$ . The positive key was just computed with $\xi_t$ . The keys added $j$ steps ago were computed with $\xi_{t-j}$ , for $j = 1, \dots, K_q/B$ . How far behind $\theta_t$ is each of them?

∴ How stale are the keys in the queue

Measured by mean parameter age behind the current query encoder $\theta_t$ , the positive key lags by $m/(1 - m)$ steps, and the keys added $j$ steps ago lag by

j + \frac{m}{1 - m}, \qquad j = 1, \dots, \frac{K_q}{B}.

If the parameters drift at a constant velocity, $\theta_t = \theta_0 + t\,\mathbf{v}$ , then after the transient $\theta_t - \xi_{t-j} = \big(j + \frac{m}{1-m}\big)\mathbf{v}$ exactly.

Proof. By the previous theorem $\xi_{t-j} = (1 - m)\sum_{i \ge 0} m^i\,\theta_{t-j-i}$ , an average of parameters whose ages relative to step $t$ are $j + i$ , with weights $(1 - m)m^i$ . Its mean age is $j + m/(1 - m)$ . Under constant drift, $\theta_t - \theta_{t-j-i} = (j + i)\mathbf{v}$ , so averaging with those weights gives $(j + m/(1-m))\mathbf{v}$ . $\square$

In plain words, every key, the positive included, comes from an encoder that trails the query encoder by the lag $m/(1-m)$ . Keys that sat in the queue for $j$ steps trail by $j$ more. When the lag is large compared with the queue's length in batches, all keys are about equally old, and the dictionary is consistent.

✏️ Key ages with MoCo's defaults

The queue holds $65{,}536/256 = 256$ batches. With $m = 0.999$ the lag is 999 steps, so every key used in a step comes from an encoder between 999 and 1,255 steps old. The oldest is only 26 percent staler than the freshest, so the dictionary is nearly consistent.

Now set $m = 0$ , which copies the query encoder at every step. The positive key comes from the current encoder, with lag 0, while the negatives are 1 to 256 steps old. The one fresh key in every softmax is the positive, and the ratio of ages is unbounded. MoCo tried $m = 0$ with a smaller queue of 4,096 keys (16 batches), and training failed, with a loss that oscillates. That the age gap causes the failure is a hypothesis supported by the experiment, not a proof. The theorem says how inconsistent the keys are, and the experiment says inconsistency hurts.

Python · cell 4

The widget draws the age of every key used in one step, its lag behind the query encoder, as a bar. The queue holds MoCo's 65,536 keys, so you can see how momentum and batch size change the spread. A dashed line carries the positive's age across the queue, so each negative's extra age shows above it.

With a batch of 1,024 the queue holds only 64 batches, so the oldest key lags $64 + 999 = 1{,}063$ steps and the ratio falls to 1.06. With MoCo's defaults the ratio $(256 + \text{lag})/\text{lag}$ stays below 1.5 only when the lag exceeds 512 steps, which needs $m > 512/513 \approx 0.998$ .

📌 MoCo's ablations and results (MoCo, Section 4)

ResNet-50, ImageNet linear probe.

With a queue of 4,096 keys, momentum $m = 0, 0.9, 0.99, 0.999, 0.9999$ gives failure (the training loss oscillates and does not converge), 55.2, 57.8, 59.0, and 58.9.
A memory bank with the same InfoNCE loss and 65,536 negatives reaches 58.0, 2.6 points below MoCo.
Negatives from the batch alone, as in SimCLR, match MoCo at small dictionary sizes, but on 8 V100 GPUs of 32 GB the batch cannot go past 1,024.
MoCo reaches 60.6 with a ResNet-50 and 68.6 with a ResNet-50 four times as wide.

MoCo was also trained on Instagram-1B, 940 million public Instagram photos (batch 1,024 on 64 GPUs, 1.25 million iterations, about 1.4 epochs, about 6 days). Pretrained on ImageNet or on Instagram-1B, MoCo outperformed ImageNet-supervised pretraining on 7 detection and segmentation tasks on PASCAL VOC, COCO, and other datasets.

The transfer result mattered more than the linear probe. MoCo's linear accuracy, 60.6, was below SimCLR's 69.3, which came months later and used 1,000 epochs. But MoCo showed self-supervised pretraining beating supervised pretraining at the tasks most practitioners cared about, finding and outlining objects in photos, in seven detection and segmentation tasks.

Interview

One MoCo step

Implement moco_step(q, k, queue, ptr, theta, xi, m, tau), one training step of MoCo without the encoders. The inputs are these.

q and k are (B, d) arrays of unit-length queries and their positive keys. Use them as given, without renormalizing.
queue is (K_q, d), with K_q a multiple of B, and ptr is the index of the oldest batch in the queue.
theta and xi are flat arrays holding the parameters of the query encoder (already updated by this step's optimizer) and of the key encoder.

Return (loss, dq, new_queue, new_ptr, new_xi).

loss is the mean over the batch of $\mathcal{L}_{\mathbf{q}}$ . The logits are [q.k+, q.queue^T] / tau, with the positive in column 0 and only the queue as negatives.
dq is the gradient of loss with respect to q. The keys are constants.
After computing the loss, write the batch's keys into rows ptr to ptr + B of a copy of the queue, and advance the pointer by B modulo K_q.
new_xi is the momentum update of the key encoder.

Compute everything stably.

import numpy as np

def moco_step(q, k, queue, ptr, theta, xi, m, tau):
    """q, k: (B, d) unit. queue: (K_q, d). Returns (loss, dq, new_queue, new_ptr, new_xi)."""
    raise NotImplementedError

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

Research

How consistent is a queue?

Make the staleness theorem of this section concrete, and check it by simulation.

key_lag(m, K_q, B) returns (pos, neg, ratio). Here pos is the positive key's mean lag $m/(1-m)$ . The array neg holds the lags of the $K_q/B$ batches in the queue, with entry j - 1 for the batch added j steps ago. And ratio is neg[-1] / pos, or np.inf when pos is 0.
simulate_lag(m, n_slots, steps, noise, rng, D=8) simulates MoCo's two encoders. Start with theta = xi = zeros(D) and a unit drift direction v = ones(D) / sqrt(D). At each of steps steps, first move the query encoder, theta = theta + v + noise * rng.normal(size=D). Then apply the momentum update to xi, and record xi as the encoder of that step's keys. After the last step, return an array of n_slots + 1 measured lags. Entry j is the projection $(\theta_T - \xi_{T-j})^\top\mathbf{v}$ for the keys encoded j steps before the end, so entry 0 is the positive key.

With noise = 0 and enough steps, the simulation should reproduce the closed form.

import numpy as np

def key_lag(m, K_q, B):
    """Returns (positive lag, array of the K_q // B negative-slot lags, oldest / positive ratio)."""
    raise NotImplementedError

def simulate_lag(m, n_slots, steps, noise, rng, D=8):
    """Returns n_slots + 1 measured lags, entry 0 for the positive key."""

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

MoCo v2, MoCo v3, and unstable contrastive ViTs

SimCLR and MoCo appeared three months apart, and each had something the other lacked. SimCLR had the MLP head and strong augmentation. MoCo had the queue that made small batches work. Within weeks the MoCo authors combined them.

📌 MoCo v2 (Chen, Fan, Girshick and He, 2020)

ResNet-50, batch 256, 8 GPUs (Chen et al. 2020c).

Replacing MoCo's linear output layer with SimCLR's two-layer MLP head lifts the linear probe from 60.6 to 62.9 at $\tau = 0.07$ , and to 66.2 at the better $\tau = 0.2$ .
Adding blur alone gives 63.4. Yet that model transfers better to object detection on PASCAL VOC than the MLP model (56.8 against 56.4 AP, the standard detection score, average precision). So linear accuracy and detection transfer do not move together.
Both changes give 67.3, and a cosine learning-rate schedule gives 67.5 at 200 epochs. SimCLR at the same 200 epochs and batch 256 gets 61.9, and at batch 8,192 gets 66.6.
With 800 epochs MoCo v2 reaches 71.1, above SimCLR's 69.3 at 1,000 epochs and batch 4,096.

At batch 256, MoCo needs 5.0 GB per GPU and 53 hours for 200 epochs, against 7.4 GB and 65 hours for the SimCLR-style loss, because MoCo backpropagates through the query encoder only.

MoCo v3 (Chen et al. 2021) asked what happens with Vision Transformers, and the answer reshaped the recipe. A Vision Transformer, or ViT, cuts an image into square patches, turns each patch into a vector with one linear layer, and runs a Transformer over those vectors (§27.2). At a batch of 4,096 the queue's extra negatives no longer helped, so MoCo v3 dropped the queue and uses the other keys of the batch as negatives. It kept the momentum encoder.

Its momentum ablation tells the staleness story from the other side. For ViT-B trained for 300 epochs, momentum $m = 0, 0.9, 0.99, 0.999$ gives 74.3, 75.6, 76.5, and 75.0. Without a queue, $m = 0$ no longer fails, because every key comes from the current batch and one encoder. The momentum encoder still adds 2.2 points.

🔬 The rest of the MoCo v3 recipe

MoCo v3 symmetrizes the loss, $\mathcal{L} = \text{ctr}(\mathbf{q}_1, \mathbf{k}_2) + \text{ctr}(\mathbf{q}_2, \mathbf{k}_1)$ . Here $\text{ctr}$ is InfoNCE with the $B$ keys of the other view as candidates, scaled by the constant $2\tau$ . The authors note that the scale only makes the loss less sensitive to $\tau$ at a fixed learning rate and weight decay.

It adds a prediction head on the query side only, borrowed from BYOL, a two-layer MLP after a three-layer projector. Both have hidden width 4,096, output width 256, and BatchNorm in every layer. In MoCo v3 the symbol $\mathbf{q}$ is the predictor's output, the online vector compared with a target, the role it keeps in §28.3. Its ResNet-50 recipe uses LARS, $\tau = 1$ , and $m = 0.996$ raised toward 1 on a cosine schedule.

The real finding of MoCo v3 was instability. Self-supervised ViTs trained with large batches and learning rates showed sudden dips in accuracy during training. The dips recovered partly and left the final number a point or two lower. Run-to-run variation was only 0.1 to 0.3 points, so the damage was easy to miss.

The authors measured k-NN accuracy every few tens of iterations to see the dips. Each dip followed a spike in the gradient norm. The spike appeared first in the patch projection, the first layer of the ViT that turns image patches into vectors (§27.1), and tens of iterations later in the last layers.

📌 MoCo v3 and the frozen patch projection

Freezing the patch projection at its random initialization, with a stop-gradient right after it so that no gradient reaches it, stabilized training. For ViT-B/16 (the base size, with 16-pixel patches) at 100 epochs and learning rate $1.5\times10^{-4}$ , accuracy rose by 1.7 points, from 71.7 to 73.4. The gain grew with the learning rate, 0.4, 0.6, and 1.7 points at $0.5$ , $1.0$ , and $1.5\times10^{-4}$ . The same trick helped SimCLR (69.3 to 70.1) and BYOL (69.7 to 71.0) with ViTs.

Final MoCo v3 results, with AdamW, batch 4,096, and a 40-epoch warmup, are 73.4 (ViT-S, the small size), 76.7 (ViT-B), 77.6 (ViT-L, large), and 78.1 (ViT-H, huge) on the ImageNet linear probe, and 73.8 for a ResNet-50 at 800 epochs.

The authors were careful about what this shows. Freezing the first layer narrows the set of solutions and changes no architecture, so the problem is in optimization, not in capacity. The first layer is only the easiest place to intervene, since it is the one layer that is not a Transformer block, and they called for a more basic fix. Later work supplied several. These include a lower learning rate and longer warmup, normalizing the queries and keys inside attention (QK-norm, §5.7), and the stability measures of DINOv2 in §28.4. The practical lesson is to monitor training at high resolution, because instability can hide inside a final number that looks fine.

In PyTorch

Production NT-Xent is four lines. Normalize, mask the diagonal with masked_fill, and call F.cross_entropy with targets $(i + B) \bmod 2B$ , which handles the stable log-softmax and the mean. The block checks autograd against the closed-form gradient of the theorem.

Then it runs a few steps of MoCo's two moving parts, the momentum update under torch.no_grad() and a ring-buffer queue. The loss rises over the first steps. The queue starts with random keys, which are easy negatives, and real keys replace them.

🔥 In PyTorchfor a machine with PyTorch

Where contrastive learning falls short

Contrastive learning set the template for self-supervised vision, and its weaknesses set the agenda for the rest of the chapter. There are four.

Negatives cost a big batch or a queue. SimCLR needs thousands of images per step, with all-gathers and LARS. MoCo needs a second encoder and a queue whose keys must be kept consistent. Either way, the method's quality depends on how many negatives you can afford.

Some negatives are not negative. Every other image in the batch is pushed away, including images of the same class. ImageNet has 1,000 roughly balanced classes, so each negative shares the anchor's class with probability about 1/1000. That is about 8 of SimCLR's 8,190 negatives per anchor at batch 4,096, and about 66 of MoCo's 65,536. At small $\tau$ these false negatives get the largest pushes.

The embeddings can still partly collapse. Negatives stop every image from mapping to one point. They do not stop the embeddings from filling only a low-dimensional subspace. Strong augmentation is one cause. Along a direction where the augmentation varies more than the data does, the encoder's weights shrink, as they did on the color features in our simulation under strong jitter. When that happens along many directions, the embeddings fill fewer dimensions. §28.3 names this dimensional collapse, shows how to measure it, and explains its causes.

The features are global. NT-Xent compares one vector per image. So the encoder is rewarded for what identifies the whole image, not for where things are in it. That shows up as weaker dense features, the per-patch features used for segmentation and depth, than the masked methods of §28.5 give.

The table puts the mechanisms side by side.

Method	Negatives from	Negatives per query	Batch	Key encoder	BatchNorm fix	ResNet-50 linear probe
Memory bank (Wu et al. 2018)	features of every image, stored when last seen	4,096 sampled	256	none, stored features	not needed	54.0 (58.0 reimplemented with InfoNCE and 65,536 negatives)
SimCLR	the other views in the batch	$2B - 2 = 8{,}190$	4,096	the same network, with gradients	global BatchNorm	69.3 (1,000 epochs)
MoCo v1	a queue of past keys	65,536	256	momentum, $m = 0.999$	shuffled BatchNorm	60.6 (200 epochs)
MoCo v2	a queue, plus the MLP head and blur	65,536	256	momentum, $m = 0.999$	shuffled BatchNorm	71.1 (800 epochs)
MoCo v3	the other view's keys in the batch	$B - 1 = 4{,}095$ per direction	4,096	momentum, $m = 0.99$ for ViT	global BatchNorm (SyncBN), only in the heads for ViT	73.8 (800 epochs); ViT-B 76.7

These comparisons mix training lengths and architectures, so read them as a history of recipes rather than a controlled experiment.

The same loss reappears across the book. CLIP trains it across two kinds of data, an image encoder against a text encoder, with a learned temperature (§30.2). And wav2vec 2.0 uses a contrastive loss to pick the true quantized speech frame among distractors (§35.4).

🧭 What to remember

SimCLR makes a positive pair from one photo with two random augmentations, encodes both with $f_\theta$ and a projection head $g_\theta$ , and trains with NT-Xent. That is InfoNCE over $2B$ views with the self-similarity masked, $2B - 2$ negatives per view, and every pair counted in both directions.
The loss starts near $\log(2B - 1)$ and can never fall below $\log(1 + (2B - 2)e^{-B/((B-1)\tau)})$ , so $\tau$ must be well below $1/\log(2B)$ for confident matches. The gradient on the unit embeddings is $(\mathbf{A} + \mathbf{A}^\top - 2\boldsymbol\Pi)\hat{\mathbf{Z}}/(2B\tau)$ , followed by the gradient through normalization.
The augmentations define the task. Crops alone let the network match color histograms, so contrastive learning needs strong color distortion that supervised learning dislikes, and the features become blind to whatever the augmentations vary.
The projection head absorbs the invariance, and the layer below it, $\mathbf{h}$ , is the representation, better than $\mathbf{z}$ by more than 10 points in SimCLR. A low contrastive loss, from a shortcut, missing normalization, or BatchNorm leakage, does not mean a good representation.
Large batches need LARS, which bounds every layer's step to about $\eta\,\eta_{\text{trust}}$ of its weight norm, and BatchNorm that cannot leak the pairing (global, shuffled, or replaced).
MoCo replaces the big batch with a queue of 65,536 keys made by a momentum encoder, $\xi \leftarrow m\xi + (1-m)\theta$ , whose keys lag by $j + m/(1-m)$ steps, nearly consistent at $m = 0.999$ . MoCo v2 added SimCLR's head and blur, and MoCo v3 dropped the queue and exposed the instability of contrastive ViTs.
Negatives cost a big batch or a queue, and they push apart images that share a class. §28.3 asks whether we can drop them, and what then stops every image from collapsing to one point.

Chen, Kornblith, Norouzi & Hinton (2020a). A Simple Framework for Contrastive Learning of Visual Representations. The SimCLR paper. Its ablations of augmentation (Table 1, Figures 5 and 6), the projection head (Table 3), the loss and temperature (Tables 4 and 5), and batch size (Table B.1) are the evidence behind most of this section.
He, Fan, Wu, Xie & Girshick (2020). Momentum Contrast for Unsupervised Visual Representation Learning. MoCo. Section 3 explains the dictionary view, the queue, the momentum update, and shuffled BatchNorm, and Section 4 has the momentum and mechanism ablations and the transfer results.
Chen, Fan, Girshick & He (2020c). Improved Baselines with Momentum Contrastive Learning. MoCo v2, a three-page note that isolates the MLP head, the temperature, and blur, with the memory and time cost of each mechanism.
Chen, Xie & He (2021). An Empirical Study of Training Self-Supervised Vision Transformers. MoCo v3, the instability of contrastive ViTs, and the frozen random patch projection.
Chen, Kornblith, Swersky, Norouzi & Hinton (2020b). Big Self-Supervised Models are Strong Semi-Supervised Learners. SimCLR v2, with bigger encoders, a deeper head, fine-tuning from the head's middle, and distillation for semi-supervised learning.
You, Gitman & Ginsburg (2017). Large Batch Training of Convolutional Networks. The LARS optimizer and the layer-wise trust ratio.
Tian, Sun, Poole, Krishnan, Schmid & Isola (2020). What Makes for Good Views for Contrastive Learning?. InfoMin, the argument that views should share only what the task needs.
Bordes, Balestriero, Garrido, Bardes & Vincent (2022). Guillotine Regularization: Why removing layers is needed to improve generalization in Self-Supervised Learning. Why the projector is cut off, across methods, and how the best cut depends on the task.

Made for larger screens

Contrastive methods, SimCLR and MoCo

Two views of one photograph

The loss over 2B views

Where the loss starts and how low it can go

The gradient pushes every view twice

NT-Xent loss and gradient over 2B views, no loops

What the augmentations must destroy

The color-histogram shortcut

The projection head, and why it is thrown away

Temperature and normalization

Big batches and the LARS optimizer

LARS with exclusions

BatchNorm can leak the pairing

MoCo, a dictionary larger than the batch

How old are the keys in the queue

One MoCo step

How consistent is a queue?

MoCo v2, MoCo v3, and unstable contrastive ViTs

In PyTorch

Where contrastive learning falls short

Read more

Watch next