Pretext Tasks: Learning Without Labels

"They handed me a photo cut into nine squares and shuffled like a deck of cards, and said: put it back. I had never been told what a dog was, but to win this game I had to learn that ears go above legs. The label was free. The understanding was not."

A Jigsaw-Solving Convolutional Network

In the previous chapters you trained networks against human labels: a class index from ImageNet, a bounding box, a per-pixel mask. This section steps off that path entirely. We will train a network with no human labels at all, by inventing a supervision signal that can be computed automatically from any image. The plan is to introduce the general framework, build the simplest pretext task (rotation prediction) end to end in PyTorch, survey two more (jigsaw and colorization) that illustrate different design choices, and then pin down the protocol that measures whether the learned features are any good. By the end you will understand why self-supervision works at all, and why some pretext tasks teach far more than others. This is the foundation that Section 25.2 sharpens into contrastive learning. The illustration below captures the central bargain: free data, self-made labels.

1. The Self-Supervised Recipe Beginner

The entire field rests on one move. Take an unlabeled image $x$, apply a transformation $t$ whose parameters you choose, and obtain a transformed input $\tilde{x} = t(x)$. Now ask the network to predict something about $t$ from $\tilde{x}$. Because you chose $t$, you know the correct answer, so you have a label for free. The network is trained with ordinary supervised loss (usually cross-entropy) against that automatically generated label. The crucial point is that this is genuine supervised learning; the only difference from Chapter 20 is the source of the labels.

Why would this teach anything useful? Because for many well-chosen transformations, the only way to predict $t$ reliably is to recognize what is in the image. To tell that a photo of a dog has been rotated ninety degrees, you must already have an internal sense of which way a dog normally faces; sky is up, legs point down, faces are upright. The network cannot get this from low-level cues alone, so it is pushed to build a representation that encodes object orientation, parts, and layout. That representation, captured in the network's intermediate layers, is the prize. The pretext task is scaffolding we throw away.

2. Rotation Prediction, End to End Beginner

The cleanest pretext task to implement is RotNet (Gidaris et al., 2018). Rotate each image by one of four angles, $0$, $90$, $180$, or $270$ degrees, and train the network to classify which of the four rotations was applied. That is a four-way classification problem with labels generated on the fly. The transformation $t$ is a discrete rotation; the label is the rotation index. Figure 25.1.1 shows the setup: one image becomes four training examples, each tagged with the angle it carries.

The code below builds the full pipeline against a torchvision backbone. The key piece is the data transform that produces the four rotations and their labels; everything else is a standard supervised training step you have seen since Chapter 18.

import torch
import torch.nn as nn
import torchvision

def make_rotations(images):
    """Given a batch (B, C, H, W), return 4B rotated images and their rotation labels."""
    rots, labels = [], []
    for k in range(4):                       # k = number of 90-degree turns
        rots.append(torch.rot90(images, k, dims=(2, 3)))  # rotate H,W plane
        labels.append(torch.full((images.size(0),), k, dtype=torch.long))
    return torch.cat(rots, dim=0), torch.cat(labels, dim=0)

# A ResNet-18 backbone with a fresh 4-way head replacing the 1000-way ImageNet head.
backbone = torchvision.models.resnet18(weights=None)
backbone.fc = nn.Linear(backbone.fc.in_features, 4)   # 4 rotation classes
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(backbone.parameters(), lr=0.1, momentum=0.9)

images = torch.randn(8, 3, 64, 64)          # 8 unlabeled images, no class labels needed
rot_imgs, rot_labels = make_rotations(images)         # -> (32, 3, 64, 64), (32,)
logits = backbone(rot_imgs)
loss = criterion(logits, rot_labels)        # supervised loss against free labels
loss.backward(); optimizer.step()
print("rotated batch:", rot_imgs.shape, "labels:", rot_labels[:8].tolist())
# rotated batch: torch.Size([32, 3, 64, 64]) labels: [0, 0, 0, 0, 0, 0, 0, 0]

Notice that no .targets or class annotations appear anywhere. The supervision comes entirely from make_rotations. Once this network is trained on a large unlabeled corpus, we throw away backbone.fc (the rotation head) and keep the convolutional trunk as a feature extractor. Subsection 4 below explains how we measure whether those features are good.

3. Jigsaw Puzzles and Colorization Intermediate

Rotation is the simplest pretext task, but it is one of many, and the variety illustrates the design space. The jigsaw task (Noroozi and Favaro, 2016) cuts the image into a grid of tiles, shuffles them by one of a fixed set of permutations, and asks the network to predict which permutation was used. To reassemble the puzzle the network must understand spatial relationships between parts: a face tile belongs above a body tile, a wheel sits below a fender. Because the number of permutations of nine tiles is enormous, the task uses a curated subset of a few hundred permutations chosen to be maximally distinct, turning it into a manageable classification problem.

Colorization (Zhang et al., 2016) takes the opposite tack: it is a generative pretext task. The network receives a grayscale image and must predict the missing color channels. To color a banana yellow and grass green it must recognize the objects, since color is a semantic property, not a local one. Colorization predicts a dense output (a color per pixel) rather than a single class, so it teaches features useful for dense downstream tasks. The contrast with rotation is instructive: rotation is a global, four-way classification; colorization is a dense, per-pixel regression. Both invent their labels from the data, but they push the representation in different directions.

A recurring danger unites all three. A pretext task is only as good as the shortcut it forbids. If the network can predict the rotation by reading a camera watermark in the corner, or solve the jigsaw by matching the chromatic aberration at tile edges, or colorize by memorizing a texture statistic, it will do exactly that and learn nothing about content. The illustration below shows this loophole in cartoon form, and the practical example below is a real instance of the same failure.

4. Measuring Success: The Linear Probe Intermediate

Since the pretext accuracy is not the goal, we need a separate yardstick for representation quality. The standard tool is the linear probe, an evaluation protocol inherited directly from the transfer-learning discussion of Chapter 21. Freeze the pretrained backbone so its weights cannot change, extract features for a labeled dataset, and train a single linear classifier (one fully connected layer) on top of those frozen features. The accuracy of that linear classifier measures how linearly separable the learned representation is: if a simple linear boundary can carve the frozen features into the right classes, the features have already done the hard work of organizing the data semantically.

Formally, if the frozen encoder maps an image to a feature vector $z = f_\theta(x)$ with $\theta$ fixed, the linear probe fits only a weight matrix $W$ and bias $b$ by minimizing

where $\mathcal{L}_{\text{CE}}$ is cross-entropy and $y_i$ are the downstream labels. Only $W$ and $b$ are learned; $\theta$ never moves. The two other common protocols are fine-tuning (unfreeze the whole backbone and train end to end on the downstream labels, which measures the representation as a starting point rather than as a fixed feature) and k-nearest-neighbors (classify each test feature by a majority vote of its nearest training features, which needs no training at all and is a pure measure of feature-space structure). The code below implements the linear probe directly.

import torch
import torch.nn as nn

@torch.no_grad()
def extract_features(encoder, loader, device):
    """Run the frozen encoder over a labeled loader and stack features + labels."""
    encoder.eval()
    feats, labels = [], []
    for x, y in loader:
        feats.append(encoder(x.to(device)).flatten(1).cpu())  # frozen forward pass
        labels.append(y)
    return torch.cat(feats), torch.cat(labels)

def linear_probe(train_feats, train_labels, num_classes, epochs=100):
    """Train ONLY a linear layer on frozen features; the encoder is untouched."""
    clf = nn.Linear(train_feats.size(1), num_classes)
    opt = torch.optim.Adam(clf.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(epochs):
        opt.zero_grad()
        loss = loss_fn(clf(train_feats), train_labels)        # only clf has gradients
        loss.backward(); opt.step()
    return clf

# Usage sketch: freeze the RotNet trunk, drop its rotation head, probe on real labels.
encoder = backbone                       # from Section 2
encoder.fc = nn.Identity()               # expose 512-d features instead of the 4-way head
for p in encoder.parameters():
    p.requires_grad_(False)              # freeze: this is what makes it a linear PROBE
# feats, labels = extract_features(encoder, labeled_loader, "cuda")
# clf = linear_probe(feats, labels, num_classes=10)
print("encoder frozen:", not any(p.requires_grad for p in encoder.parameters()))
# encoder frozen: True

The linear probe is what lets us compare RotNet, jigsaw, colorization, and the contrastive methods of Section 25.2 on a level field. It is also why the field converged on the methods it did: contrastive and masked-modeling approaches produce far more linearly separable features than the early pretext tasks, which is exactly the gap the three numbers above expose and the next two sections close.

# Load a pretrained self-supervised ViT as a frozen feature extractor.
# num_classes=0 strips the classification head, exposing the raw embedding.
import timm
# A self-supervised ViT backbone, features only (no classifier), weights downloaded:
encoder = timm.create_model("vit_small_patch16_224.dino", pretrained=True, num_classes=0)
features = encoder(torch.randn(1, 3, 224, 224))   # (1, 384) frozen embedding