← All tutorials

Non-Contrastive Self-Supervised Learning

On this page

As in Contrastive Learning, training pulls the embeddings of two random edits of a photo, its views, together, but here with no negatives (other images pushed away). Every training step, one network learns while its target, the other view’s embedding it must match, comes from a momentum encoder (a copy whose weights follow a moving average of the network’s) or from the same network with no gradient flowing through it (a stop-gradient). Each method adds one part that stops every image from getting the same output.

1 Pulling pairs together with no negatives should collapse

If the only signal is make the two views match, one vector for every image matches every pair and says nothing; this is collapse. Negatives forbid it; without them, nothing in the loss does.

In a subtler dimensional collapse, the embeddings spread along only a few directions of the space and stay flat along the rest; Failure Modes and Checks plots it.

2 BYOL predicts the output of a slowly updated copy of itself

BYOL (Bootstrap Your Own Latent, Grill et al. 2020; latent means the embedding space) trains one network to predict a momentum encoder’s embedding of the other view, with no negatives, and still does not collapse. Its trained online network turns view 1 into a prediction p₁ through an encoder, a projector (a small head, as in SimCLR) and a predictor (a second small head that guesses the target). The target network, an encoder and projector with no predictor, turns view 2 into the target z₂. It gets no gradients and is an exponential moving average (EMA) of the online weights.

The target follows MoCo’s momentum rule. Each step it keeps a share m of its old weights, and m rises from 0.996 to 1.

The snippet computes eq line 2 for both view orders on toy vectors; cosine_similarity scales both to unit length itself.

import torch
import torch.nn.functional as F

def byol_loss(p, z):
    cos = F.cosine_similarity(p, z, dim=-1)
    return (2 - 2 * cos).mean()               # batch mean

B, D = 8, 256
p1, p2 = torch.randn(B, D), torch.randn(B, D)   # online predictions
z1, z2 = torch.randn(B, D), torch.randn(B, D)   # targets
L = byol_loss(p1, z2) + byol_loss(p2, z1)

1. 1The loss in Section 2 eq line 2, 2 − 2 cos.
2. 2Toy predictions p and targets z, the two ends of Figure 3.
3. 3Section 2 eq line 1 plus the same with the views swapped.

With no negatives to collect, BYOL needs no big batch.

3 Three parts hold collapse off

Removing any one of three parts collapses BYOL. BYOL’s authors removed the predictor and changed how the target updates; a later method, SimSiam (Section 4), removed the stop-gradient.

1. The stop-gradient treats the target’s output as fixed, so the two sides cannot walk together to a constant.
2. The predictor, on the online side only, learns to guess the target from the online output. A constant output gives it nothing to guess from, so the authors argue training moves away from the constant output (BYOL sec. 3.2).
3. The EMA target lags the online network: the slow target changes little per step, so the authors argue the predictor can stay near the best prediction of it (BYOL sec. 3.2).

4 SimSiam shows a stop-gradient and a predictor are enough

SimSiam (Simple Siamese networks, Chen and He 2021; Siamese means two branches with the same weights) uses one network for both views, with no momentum encoder and no negatives. As in BYOL, each view gives an output z and a prediction p, and each prediction is compared with the other view’s output behind a stop-gradient, written sg.

Removing the predictor collapses it too (0.1%). The SimSiam paper compares training to k-means clustering, which holds cluster centers fixed while fitting points, then updates the centers; the stop-gradient does the job of the fixed-centers step.

5 DINO matches a student to a centered, sharpened teacher

DINO (self-distillation with no labels, Caron et al. 2021; the student learns from a teacher built from itself) trains a student network to copy a teacher, a momentum encoder of the student (MoCo’s momentum rule). Two edits to the teacher’s output, centering and sharpening, stop collapse. Each network has K outputs, 65,536 in DINO. A softmax turns them into shares that sum to 1, and dividing by a temperature first sets how peaked the shares are, as in Knowledge Distillation but set low for the teacher. The student matches the teacher’s shares by cross entropy.

6 BYOL, SimSiam and DINO in practice

This snippet runs one BYOL step, then moves the target toward the online network by EMA.

import copy, math
import torch
import torch.nn as nn
import torch.nn.functional as F

def byol_loss(p, z):
    return (2 - 2 * F.cosine_similarity(p, z, dim=-1)).mean()

def mlp(d_in, d_out=256):
    return nn.Sequential(nn.Linear(d_in, 4096), nn.BatchNorm1d(4096), nn.ReLU(), nn.Linear(4096, d_out))

encoder = nn.Sequential(nn.Flatten(), nn.Linear(3 * 32 * 32, 512))
online = nn.Sequential(encoder, mlp(512))       # encoder + projector
predictor = mlp(256)
target = copy.deepcopy(online).requires_grad_(False)
params = [*online.parameters(), *predictor.parameters()]
opt = torch.optim.SGD(params, lr=0.2, momentum=0.9)
k, K_STEPS, M_BASE = 0, 1000, 0.996             # step k

x1, x2 = torch.rand(64, 3, 32, 32), torch.rand(64, 3, 32, 32)
p1, p2 = predictor(online(x1)), predictor(online(x2))
with torch.no_grad():
    z1, z2 = target(x1), target(x2)
L = byol_loss(p1, z2) + byol_loss(p2, z1)
opt.zero_grad()
L.backward()
opt.step()

m = 1 - (1 - M_BASE) * (math.cos(math.pi * k / K_STEPS) + 1) / 2
with torch.no_grad():
    for p_t, p_o in zip(target.parameters(), online.parameters()):
        p_t.mul_(m).add_(p_o, alpha=1 - m)

1. 1The loss of Section 2 eq line 2, as in the snippet there.
2. 2The two-layer MLP used for the projector and the predictor.
3. 3Online network, predictor, target copy and optimizer, as in Figure 3; LARS (an SGD variant that scales each layer’s step to its weight size) is not in PyTorch, so SGD stands in.
4. 4Toy views, predictions, no-gradient targets and the loss, Section 2 eq lines 1 and 2.
5. 5Updates the online side only, the side trained by gradients in Figure 3.
6. 6Cosine schedule for m and the target EMA, the EMA arrow in Figure 3.

A SimSiam step has no target network; z.detach() is the stop-gradient sg of Section 4 eq lines 2 and 3.

import torch
import torch.nn as nn
import torch.nn.functional as F

def D(p, z):
    return -F.cosine_similarity(p, z.detach(), dim=-1).mean()

online = nn.Sequential(nn.Flatten(), nn.Linear(3 * 32 * 32, 256))   # stand-in
predictor = nn.Linear(256, 256)
x1, x2 = torch.rand(64, 3, 32, 32), torch.rand(64, 3, 32, 32)
z1, z2 = online(x1), online(x2)
p1, p2 = predictor(z1), predictor(z2)
L = D(p1, z2) / 2 + D(p2, z1) / 2

1. 1D of Section 4 eq line 1, with detach() as sg.
2. 2One network gives outputs z and predictions p for both views.
3. 3Eq lines 2 and 3, each view once predictor and once fixed target.

A DINO step trains the student, then moves the teacher toward it by EMA and the center c toward the batch mean.

import copy
import torch
import torch.nn as nn
import torch.nn.functional as F

def dino_loss(s_out, t_out, c, tau_s=0.1, tau_t=0.04):
    P_t = F.softmax((t_out - c) / tau_t, dim=-1)     # center, sharpen
    log_P_s = F.log_softmax(s_out / tau_s, dim=-1)
    return -(P_t * log_P_s).sum(dim=-1).mean()

K_OUT, LAM, M_C, WD = 1024, 0.996, 0.9, 0.04   # toy K; DINO uses 65,536
student = nn.Sequential(nn.Flatten(), nn.Linear(3 * 32 * 32, K_OUT))
teacher = copy.deepcopy(student).requires_grad_(False)
c = torch.zeros(K_OUT)
opt = torch.optim.AdamW(student.parameters(), lr=5e-4, weight_decay=WD)

x1, x2 = torch.rand(64, 3, 32, 32), torch.rand(64, 3, 32, 32)
s1, s2 = student(x1), student(x2)
with torch.no_grad():
    t1, t2 = teacher(x1), teacher(x2)
L = (dino_loss(s1, t2, c) + dino_loss(s2, t1, c)) / 2
opt.zero_grad()
L.backward()
opt.step()
with torch.no_grad():
    for p_t, p_s in zip(teacher.parameters(), student.parameters()):
        p_t.mul_(LAM).add_(p_s, alpha=1 - LAM)
    c = M_C * c + (1 - M_C) * torch.cat([t1, t2]).mean(dim=0)

1. 1Loss with centered, sharpened teacher shares, Section 5 eq line 1.
2. 2Student, teacher copy, center c and optimizer, the DINO panel of Figure 1; DINO raises the weight decay WD from 0.04 to 0.4 on a cosine schedule.
3. 3Toy views, both networks on both views, then the loss on the shares of Section 5 eq line 1.
4. 4Updates the student only; the teacher sits behind the stop-gradient in Figure 1.
5. 5Teacher EMA with rate λ, the EMA copy in Figure 1.
6. 6Moves c, the center in Section 5 eq line 1, 0.1 of the way to the batch mean of teacher outputs.

From the BYOL, SimSiam and DINO papers and main_dino.py.

Masked Image Modeling comes next, which hides most of an image and rebuilds it.

7 References

• Grill, J.-B., et al. (2020). Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning. NeurIPS 2020. arXiv:2006.07733.
• Chen, X., and He, K. (2021). Exploring Simple Siamese Representation Learning. CVPR 2021. arXiv:2011.10566.
• Caron, M., et al. (2021). Emerging Properties in Self-Supervised Vision Transformers. ICCV 2021. arXiv:2104.14294.
• He, K., et al. (2020). Momentum Contrast for Unsupervised Visual Representation Learning. CVPR 2020. arXiv:1911.05722.