Self-Supervised Learning: SimCLR, BYOL, MAE, and CLIP

A collection of images can contain far more examples than anyone has labeled. Self-supervised learning constructs training targets from the data: another view of an image, its missing patches, or a representation computed by a second network. The resulting encoder can then be evaluated on a labeled downstream task. CLIP uses an additional source of supervision, human-written text paired with images; it belongs in this comparison because it also learns transferable representations without requiring a fixed set of training classes.

What should the training task teach?

Early work used automatically generated targets: predict an applied rotation, solve a jigsaw of shuffled patches, or colorize a grayscale photograph. These pretext tasks produced useful representations, but success on the training task need not require the information a downstream task needs. A model might exploit image borders or low-level statistics instead of learning object structure.

The methods below choose different targets and different constraints on how the model can reach them. SimCLR compares views against other images; BYOL predicts another network’s output; MAE reconstructs hidden pixels; CLIP matches images to text. Their training losses measure success at those tasks. Transfer evaluation determines whether the learned features are useful elsewhere. The code illustrates the objectives and update rules with small tensors, not full image-pretraining systems. Run the Python blocks in order; they require PyTorch.

The reconstruction objective used by MAE is introduced in Autoencoders and Representation Learning.

Contrastive learning: SimCLR

SimCLR samples two augmentations of each image. A shared backbone maps each view to a feature vector \(h\), and a small projection network maps \(h\) to \(z\), where the contrastive loss is applied. For an anchor view in a batch of \(N\) images, its other view is the positive; the remaining \(2N-2\) views are treated as negatives. The denominator below contains \(2N-1\) candidates, including the positive and excluding the anchor itself.

\[\mathcal{L}_{i,j}=-\log\frac{\exp(\text{sim}(z_i,z_j)/\tau)}{\sum_{k\ne i}\exp(\text{sim}(z_i,z_k)/\tau)}\]

This NT-Xent objective is a form of InfoNCE: cross-entropy for identifying the matching view. Here \(\mathrm{sim}\) is cosine similarity and \(\tau>0\) is a temperature. Smaller temperatures make a given similarity difference produce a larger logit difference; the choice must be tuned with the rest of the training setup. For two images ordered as [a1, b1, a2, b2], the target indices are [2, 3, 0, 1]. Each anchor has one positive and two negatives.

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

def nt_xent(z1, z2, temperature=0.1, reduction="mean"):
    """Corresponding rows are paired views; inputs have shape (N, D)."""
    if z1.ndim != 2 or z1.shape != z2.shape or min(z1.shape) < 1:
        raise ValueError("Expected matching, nonempty (N, D) embeddings")
    if z1.size(0) < 2 or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Use at least two images and a finite positive temperature")
    n = z1.size(0)
    z = F.normalize(torch.cat([z1, z2]), dim=1)
    logits = z @ z.T / temperature
    logits.fill_diagonal_(float("-inf"))
    targets = (torch.arange(2 * n, device=z.device) + n) % (2 * n)
    return F.cross_entropy(logits, targets, reduction=reduction)

Both views pass through the same backbone and projection head. Backpropagating this loss updates both modules. The inputs here should be finite floating-point embeddings on the same device and with the same dtype. With only one image, the positive would be the only candidate and the loss would be zero regardless of the representation, so this example rejects that case.

Augmentation defines the desired invariances. SimCLR’s experiments found the combination of cropping and color distortion effective. Matching color statistics alone can make the pretext task too easy. However, an augmentation can also remove useful information: strong color changes would be unsuitable when flower color determines the target label, and a crop can remove the object entirely. Choose transformations that usually preserve what the downstream task needs.

Transfer uses the backbone features. The original SimCLR experiments obtained better linear evaluation from features before the projection head. The head lets the contrastive objective act on a transformed version of those features. This is a useful design choice, not a guarantee that every downstream task prefers the same layer; specify which representation you evaluate.

SimCLR reported benefits from larger batches, but 4,096 images is not a minimum requirement. The implementation above builds a square matrix over all views, so its similarity storage grows quadratically with batch size. MoCo instead maintains a queue of previously computed embeddings and uses a slowly updated encoder to limit their inconsistency over time. This allows the number of comparison examples to exceed the current batch size.

BYOL: predicting a target representation

Matching positive pairs alone admits a collapsed representation: every input can map to the same nonzero vector. BYOL learns useful representations without negatives in its reported experiments, using an online predictor and a slowly updated target. Those training dynamics need to be distinguished from the set of solutions admitted by its loss.

The online branch contains a backbone, projector, and predictor. The target contains a backbone and projector, initialized as copies of their online counterparts; it has no predictor. For one view, the online prediction is compared with the target projection of the other view, and the directions are swapped and added. Only the online branch receives gradients. After its optimizer step, each corresponding target parameter is updated as \(\theta_t\leftarrow\tau\theta_t+(1-\tau)\theta_o\), with \(0\leq\tau\leq1\). For example, 0.99 retains 99% of the old target parameter. This EMA coefficient is separate from SimCLR’s temperature, despite the shared symbol. For transfer, retain the online backbone and set aside the projector and predictor.

The following helper handles parameter averaging for modules without buffers. Batch-normalized networks also have running statistics and require a specified policy for those states; parameter averaging alone is not a complete implementation for them.

def byol_loss(p1, z2, p2, z1):
    """Sum the two directional losses, then average over examples."""
    def distance(p, z):
        return (F.normalize(p, dim=1) - F.normalize(z, dim=1)).square().sum(1)
    return (distance(p1, z2.detach()) + distance(p2, z1.detach())).mean()

@torch.no_grad()
def update_target(online, target, tau=0.996):
    """EMA for matching backbone/projector modules without state buffers."""
    if not 0 <= tau <= 1:
        raise ValueError("tau must lie in [0, 1]")
    source, dest = dict(online.named_parameters()), dict(target.named_parameters())
    if source.keys() != dest.keys() or any(source[k].shape != dest[k].shape for k in source):
        raise ValueError("Online and target structures must match; exclude the predictor")
    if list(online.buffers()) or list(target.buffers()):
        raise ValueError("This example needs an explicit policy for state buffers")
    for name, parameter in dest.items():
        parameter.mul_(tau).add_(source[name], alpha=1 - tau)

An aligned nonzero constant output still gives zero loss and zero gradient, even with the predictor and stop-gradient in place. If corresponding online and target parameters already agree, EMA leaves them equal. These components therefore do not mathematically exclude collapse. BYOL’s successful training and its ablations concern the behavior reached from particular initializations and training setups. A low training loss alone cannot establish a useful representation; inspect feature variation across inputs and evaluate transfer.

Here is one complete update using small MLPs without normalization buffers. online stands in for the backbone plus projector. The separately defined predictor is optimized but never copied into the target. Gaussian perturbations stand in for image augmentations; this step checks the wiring and does not establish that the learned features avoid collapse.

from copy import deepcopy
from torch import nn

torch.manual_seed(7)
online = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
predictor = nn.Sequential(nn.Linear(3, 8), nn.ReLU(), nn.Linear(8, 3))
target = deepcopy(online).requires_grad_(False)
optimizer = torch.optim.SGD(list(online.parameters()) + list(predictor.parameters()), lr=0.01)
x = torch.randn(16, 4)
v1, v2 = x + 0.1 * torch.randn_like(x), x + 0.1 * torch.randn_like(x)
with torch.no_grad():
    z1, z2 = target(v1), target(v2)
loss = byol_loss(predictor(online(v1)), z2, predictor(online(v2)), z1)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
update_target(online, target)
print("target has gradients", any(p.grad is not None for p in target.parameters()))
# target has gradients False

SimSiam demonstrated successful training without a momentum target in its own architecture and optimization setup, retaining a predictor and stop-gradient. It does not show that adding stop-gradient alone makes an arbitrary representation-learning system avoid collapse.

Masked autoencoding: MAE

A masked autoencoder divides an image into patches, hides a randomly selected subset, and predicts their pixels. The encoder sees only visible patches with their positional information. A lightweight decoder receives their encoded features plus mask tokens at the missing positions, adds positional information, and predicts a patch at every position. Reconstruction error is averaged over the hidden patches. The encoder is retained for downstream tasks; the reconstruction decoder can be discarded.

The original MAE used 75% masking, leaving 49 visible patches out of a 196-patch image. This makes reconstruction rely on less local evidence while reducing encoder work. The paper reported substantial speedups for its setup; the visible-token fraction alone does not determine end-to-end speed because decoder and other costs remain. The ratio is a design choice, not a general constant for images. The original method also used image augmentations, so it does not remove preprocessing choices.

The small example below uses four one-number patches to isolate masking and loss selection. A real ViT encodes visible patches and supplies features to the decoder; its predictions are substituted here by fixed numbers. The placeholder sequence restores spatial order before decoding. Only errors at positions 0, 1, and 3 count: the mean is (1 + 4 + 9) / 3, approximately 4.667. The visible patch at position 2 contributes no direct reconstruction loss. MAE can also normalize the pixel targets within each patch, which changes the reconstruction target; this example uses unnormalized values.

patches = torch.tensor([[[10.], [20.], [30.], [40.]]])
order = torch.tensor([2, 0, 3, 1])
visible = patches[:, order[:1]]
mask_tokens = torch.zeros(1, 3, 1)
restored = torch.cat([visible, mask_tokens], dim=1)[:, order.argsort()]
mask = torch.ones(1, 4, dtype=torch.bool)
mask[:, order[:1]] = False
prediction = torch.tensor([[[11.], [22.], [999.], [43.]]], requires_grad=True)
patch_error = (prediction - patches).square().mean(dim=-1)
masked_loss = patch_error[mask].mean()
masked_loss.backward()
print("restored positions", restored.flatten().tolist())
print("masked loss", round(masked_loss.item(), 3))
print("visible prediction gradient", prediction.grad[0, 2, 0].item())
# restored positions [0.0, 0.0, 30.0, 0.0]
# masked loss 4.667
# visible prediction gradient 0.0

CLIP: supervision from the internet

CLIP jointly trains image and text encoders on paired images and captions. Their projections have a shared embedding dimension. Each image is trained to identify its paired caption among the batch’s captions, and each caption to identify its paired image. This is natural-language supervision: existing captions avoid collecting a fixed class label for every image, but pairing, filtering, and data quality still matter. Different rows may also describe the same concept; the one-positive-per-row loss below nevertheless treats them as negatives.

The diagonal of the score matrix contains the recorded pairs. Row-wise cross-entropy evaluates image-to-text matching; transposing it evaluates text-to-image matching. logit_scale is a trainable scalar holding the logarithm of the inverse temperature, so its exponential multiplies cosine similarities. It and the embeddings must remain finite, on a compatible device and dtype; a full training implementation also controls this scale to avoid numerical overflow.

def clip_loss(image_emb, text_emb, logit_scale):
    """Matched (N, D) embeddings and a scalar log inverse-temperature."""
    if image_emb.ndim != 2 or image_emb.shape != text_emb.shape or image_emb.size(0) < 2:
        raise ValueError("Use matching (N, D) embeddings with at least two pairs")
    image_emb = F.normalize(image_emb, dim=1)
    text_emb = F.normalize(text_emb, dim=1)
    logits = logit_scale.exp() * (image_emb @ text_emb.T)
    targets = torch.arange(len(logits), device=logits.device)
    return 0.5 * (F.cross_entropy(logits, targets)
                  + F.cross_entropy(logits.T, targets))

For zero-shot classification, encode candidate descriptions such as “a photo of a cat” and “a photo of a dog.” Normalize each text embedding and each image embedding, compute their dot products, and choose the largest score. No labeled examples of this downstream task are used to fit a classifier. This ability does not guarantee accuracy on every candidate label set, and it does not mean the pretraining data lacked human supervision or examples of those concepts. Softmax scores depend on the candidate set and temperature; they are not automatically calibrated probabilities.

# These vectors illustrate scoring; they are not outputs of a trained CLIP model.
image_features = torch.tensor([[1., 0.2], [0.1, 1.]])
text_features = torch.eye(2)
scores = F.normalize(image_features, dim=1) @ F.normalize(text_features, dim=1).T
print("selected candidate indices", scores.argmax(dim=1).tolist())
# selected candidate indices [0, 1]

Prompt wording and training-data coverage can affect zero-shot performance. If prompts are selected using downstream labels, report that validation step. CLIP-style text encoders are also used to condition some image generators, but a shared image–text space alone is not a generator: a separate generative model learns how to use the text representation.

Evaluating a self-supervised model

A linear probe fixes the representation and learns a linear readout. Fine-tuning allows the representation itself to change, so the rankings can differ. Neither contrastive learning nor masked reconstruction guarantees a win under either protocol. Probe optimization, available labels, architecture, and pretraining budget all affect the comparison; repeated runs help distinguish a ranking from sampling and optimization noise.

For a frozen probe, disable encoder parameter gradients and use evaluation mode so dropout and running normalization statistics do not change the features. Fit the classifier on training labels, choose its settings on validation data, and score an untouched test set. A k-NN evaluation also needs labeled reference examples, and its neighbor count can be tuned. Low-shot evaluation describes the label budget and can use either probing or fine-tuning. For a strictly inductive test, split the data before pretraining or augmentation and exclude test examples even when unlabeled. If test inputs are deliberately used during pretraining, report that different evaluation protocol.

The patch encoder and positional information are explained in Vision Transformers and Modern Vision Models.

When to use it

Domain-specific unlabeled data provides a reason to test self-supervised pretraining when existing features transfer poorly. Compare it with an available pretrained checkpoint under the same downstream label budget. Continued self-supervised training of that checkpoint is another option; starting from scratch and reusing a checkpoint are not the only two choices.

Corpus size alone cannot settle the choice. Domain relevance, duplication, augmentation quality, model capacity, and compute affect the result. A small pilot can compare validation performance and training cost before a larger run. Use the protocol you will actually deploy: frozen features when the encoder must stay fixed, or fine-tuning when adaptation is affordable.

How dense word vectors are learned in the first place is covered in Word Embeddings: Word2Vec, Negative Sampling, GloVe, and Bias.

Exercises

1. More comparison candidates. Generate paired embeddings once, then use nested batches of 8, 64, and 512 images. Report the average NT-Xent loss and the loss for the same first eight images’ views in each batch. Which comparison must increase when candidates are added, and why?

torch.manual_seed(0)
base = torch.randn(512, 128)
other = base + 0.5 * torch.randn_like(base)
for n in (8, 64, 512):
    losses = nt_xent(base[:n], other[:n], temperature=0.5, reduction="none")
    anchors = torch.cat([torch.arange(8), n + torch.arange(8)])
    print(f"N={n:3d} all={losses.mean().item():.4f} fixed={losses[anchors].mean().item():.4f}"
          f" equal_logits={math.log(2 * n - 1):.4f}")
# N=  8 all=1.1982 fixed=1.1982 equal_logits=2.7081
# N= 64 all=3.1135 fixed=3.0926 equal_logits=4.8442
# N=512 all=5.1646 fixed=5.1463 equal_logits=6.9305

2. Linear probe versus fine-tuning. Explain why a model can rank higher under a frozen linear probe and lower after fine-tuning. Which experiment would guide a deployment that permits only a linear classifier?

3. A collapsed BYOL solution. Show that identical nonzero constant predictions and targets have zero loss. Does detaching the target eliminate that solution?

p = torch.tensor([[1., 0.]]).repeat(4, 1).requires_grad_()
z = p.detach().clone().requires_grad_()
collapsed_loss = byol_loss(p, z, p, z)
collapsed_loss.backward()
print("collapsed loss", collapsed_loss.item())
print("prediction gradient norm", p.grad.norm().item())
print("target gradient", z.grad)
# collapsed loss 0.0
# prediction gradient norm 0.0
# target gradient None

References

• Noroozi and Favaro (2016). Unsupervised Learning of Visual Representations by Solving Jigsaw Puzzles. ECCV.
• Gidaris, Singh, and Komodakis (2018). Unsupervised Representation Learning by Predicting Image Rotations. ICLR.
• van den Oord, Li, and Vinyals (2018). Representation Learning with Contrastive Predictive Coding. arXiv preprint.
• Chen, Kornblith, Norouzi, and Hinton (2020). A Simple Framework for Contrastive Learning of Visual Representations. ICML.
• He, Fan, Wu, Xie, and Girshick (2020). Momentum Contrast for Unsupervised Visual Representation Learning. CVPR.
• Grill et al. (2020). Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning. NeurIPS.
• Chen and He (2021). Exploring Simple Siamese Representation Learning. CVPR.
• Caron et al. (2021). Emerging Properties in Self-Supervised Vision Transformers. ICCV.
• Radford et al. (2021). Learning Transferable Visual Models from Natural Language Supervision. ICML.
• He et al. (2022). Masked Autoencoders Are Scalable Vision Learners. CVPR.