Finally Teaching a Network to See Without Labels

Abstract

A small CIFAR-10 study of how SimCLR, DeepCluster, and BYOL try to learn useful image representations without class labels, including the limits of one failed BYOL implementation.

Self-supervised learning interested me during my PhD because it offered a way to learn from structure in the data rather than hand-written labels. The framing I retained from a Yann LeCun talk on self-supervision and energy-based models was simple: construct a prediction target from the input itself.

This post asks a narrower question for image models: what prevents a joint-embedding objective from mapping every input to the same representation? I implemented the training logic for SimCLR, DeepCluster, and BYOL in plain PyTorch around the same ResNet-18 backbone, then trained them on CIFAR-10. SimCLR receives the full walkthrough; the other two experiments are documented in collapsed appendices.

The saved runs produced strong class-discriminative features for SimCLR, a weaker but useful DeepCluster representation, and a poor BYOL representation. Those are observations from particular configurations, not a ranking of the algorithms. The experiment also connects to later work in which colleagues and I applied self-supervised clustering to wireless-spectrum activity (Milosheski et al. 2023).

Families of Self-Supervised Image Methods

Joint-embedding methods must avoid a trivial constant representation. Different families constrain that solution in different ways:

• Contrastive methods such as SimCLR (Chen et al. 2020) and MoCo (He et al. 2020) compare positive pairs with other images. SimCLR uses the rest of the batch as negatives; MoCo maintains a queue populated by a momentum encoder.
• Self-distillation methods such as BYOL (Grill et al. 2020) and DINO (Caron et al. 2021) train an online or student network against a stop-gradient target. Predictors, momentum updates, centering, and sharpening provide method-specific asymmetries and regularization.
• Clustering methods such as DeepCluster (Caron et al. 2018) alternate between clustering features and predicting the resulting pseudo-labels. Balanced sampling limits domination by large clusters, although empty or degenerate clusters require separate monitoring.
• Joint-embedding predictive methods such as I-JEPA (Assran et al. 2023) predict representations of masked regions from visible context and use a momentum-updated target encoder.

Method	Family	Main training signal	Implemented here
SimCLR	Contrastive	In-batch negatives	Yes
MoCo	Contrastive	Queued negatives	No
BYOL	Self-distillation	Predictor and EMA target	Yes, appendix
SimSiam (Chen and He 2021)	Self-distillation	Predictor and stop-gradient	No
DINO	Self-distillation	EMA teacher, centering, and sharpening	No
DeepCluster	Clustering	K-means pseudo-labels and balanced sampling	Yes, appendix
I-JEPA	Joint-embedding predictive	Masked prediction and EMA target	No

Experimental Scope

All three runs use a ResNet-18 with a CIFAR stem: a 3x3 stride-1 convolution and no initial max-pool. Self-supervised pretraining uses the 50,000 CIFAR-10 training images with batches of 256. Labels enter only during the downstream linear-probe and k-NN evaluations.

The comparison is illustrative rather than controlled. SimCLR and BYOL use two strongly augmented views, while DeepCluster uses Sobel-filtered inputs and lighter augmentation during pseudo-label training. Each reported result comes from one saved run in optimized, nondeterministic CUDA mode. The notebook records an RTX 3090 for the SimCLR run, but it does not preserve exact software versions or repeated-seed uncertainty.

import gzip
import math
import os
import random
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Final

import numpy as np
import torch
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from torch import Tensor, nn, optim
from torch._inductor import config as inductor_config
from torch.nn import functional as F
from torch.utils import data
from torchvision import datasets, models
from torchvision import transforms as VT

IMAGENET_NORMALIZE = {"mean": (0.485, 0.456, 0.406), "std": (0.229, 0.224, 0.225)}
DATA_ROOT = Path.home() / "datasets" / "cifar10"
SEED: Final[int] = 42
WIDTH: Final[int] = 8  # inches; at dpi=200, renders 2x this blog's ~800px column, no InlineBackend indirection
GOLDEN_RATIO: Final[float] = (1 + math.sqrt(5)) / 2


DPI: Final[int] = 200  # 2x the ~800px content column, so figures stay crisp on HiDPI screens
FIGURES = Path("./figures")


def save_fig(fig: Figure, name: str, lossless: bool = True) -> None:
    """Write a figure into figures/ as WebP. Charts keep lossless; photographic panels do not."""
    FIGURES.mkdir(parents=True, exist_ok=True)
    opts = {"lossless": True} if lossless else {"quality": 90}
    fig.savefig(FIGURES / name, dpi=DPI, pil_kwargs={"method": 6, **opts})
    plt.close(fig)


@dataclass
class Arguments:
    num_workers: int = len(os.sched_getaffinity(0))


args = Arguments()
args

Arguments(num_workers=16)

def set_random_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def set_torch_mode(use_optimized: bool, debug: bool) -> None:
    if use_optimized:
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = False
        torch.backends.cudnn.allow_tf32 = True
        torch.set_float32_matmul_precision("medium")
    else:
        torch.use_deterministic_algorithms(True)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        torch.set_float32_matmul_precision("high")
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

    # torch.compile is otherwise chatty during autotuning: it prints a full
    # kernel-comparison table for every op it benchmarks.
    inductor_config.trace.log_autotuning_results = False

    if not debug:
        torch.autograd.set_detect_anomaly(False)
        torch.autograd.profiler.profile(False)
        torch.autograd.profiler.emit_nvtx(False)

    print(f"Using {'optimized' if use_optimized else 'deterministic'} mode")
    print(f"Debug functionalities are {'enabled' if debug else 'disabled'}")


set_torch_mode(use_optimized=True, debug=False)
set_random_seed(SEED)

Using optimized mode
Debug functionalities are disabled

def make_cifar_resnet(in_channels: int = 3) -> nn.Module:
    """ResNet-18 with a CIFAR-style stem: 3x3 stride-1 conv, no initial maxpool.

    The stock torchvision resnet18() stem (7x7 stride-2 conv + maxpool) is tuned for
    224x224 ImageNet inputs; on CIFAR-10's 32x32 images it throws away most of the
    spatial resolution before the first residual block even runs.
    """
    resnet = models.resnet18()
    resnet.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1, bias=False)
    resnet.maxpool = nn.Identity()
    return resnet


class SobelFilter(nn.Module):
    """Fixed grayscale-then-Sobel edge filter (no learned parameters).

    DeepCluster uses this ahead of its backbone so that k-means clusters on shape and
    texture instead of finding the trivial shortcut of clustering by color.
    """

    def __init__(self) -> None:
        super().__init__()
        gray_weight = torch.tensor([0.2989, 0.5870, 0.1140]).view(1, 3, 1, 1)
        kernel_x = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]])
        sobel_weight = torch.stack([kernel_x, kernel_x.t()]).unsqueeze(1)
        self.register_buffer("gray_weight", gray_weight)
        self.register_buffer("sobel_weight", sobel_weight)

    def forward(self, x: Tensor) -> Tensor:
        x = F.conv2d(x, self.gray_weight)
        return F.conv2d(x, self.sobel_weight, padding=1)


class ProjectionHead(nn.Module):
    """Two-layer MLP with BatchNorm after each linear layer, used by the SimCLR and BYOL projectors
    and by the BYOL predictor. This shared design is a notebook choice, not an exact reproduction of every paper."""

    def __init__(self, in_dim: int = 512, hidden_dim: int = 2048, out_dim: int = 2048) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, out_dim),
            nn.BatchNorm1d(out_dim),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.net(x)


def nt_xent_loss(z0: Tensor, z1: Tensor, temperature: float = 0.5) -> Tensor:
    """SimCLR's NT-Xent / InfoNCE loss between two batches of augmented-view projections."""
    batch_size = z0.shape[0]
    z = F.normalize(torch.cat([z0, z1], dim=0), dim=1)

    similarity = z @ z.t() / temperature
    similarity.fill_diagonal_(-torch.inf)

    positive_idx = torch.arange(batch_size, device=z.device)
    targets = torch.cat([positive_idx + batch_size, positive_idx])
    return F.cross_entropy(similarity, targets)


def byol_loss(p: Tensor, z: Tensor) -> Tensor:
    """BYOL's negative cosine similarity between an online prediction and a stop-gradient target projection."""
    p = F.normalize(p, dim=1)
    z = F.normalize(z, dim=1)
    return 2 - 2 * (p * z).sum(dim=1).mean()


@torch.no_grad()
def update_moving_average(target: nn.Module, online: nn.Module, tau: float) -> None:
    """In-place EMA update of `target`'s parameters and buffers towards `online`, used by BYOL's target network."""
    for target_param, online_param in zip(target.parameters(), online.parameters(), strict=True):
        target_param.mul_(tau).add_(online_param, alpha=1 - tau)
    for target_buffer, online_buffer in zip(target.buffers(), online.buffers(), strict=True):
        target_buffer.copy_(online_buffer)


def byol_tau_schedule(step: int, total_steps: int, tau_base: float) -> float:
    """BYOL paper's cosine ramp: tau starts at `tau_base` and rises to 1.0 by the final step, so
    the target network mixes in less of the online network's per-step change as training
    progresses - a near-frozen target late in training gives the predictor stable regression
    targets, unlike the fixed tau used in earlier attempts (see train_byol's docstring)."""
    progress = min(step / total_steps, 1.0)
    return 1.0 - (1.0 - tau_base) * (math.cos(math.pi * progress) + 1) / 2

def build_augment(strength: str = "full") -> VT.Compose:
    """Augmentation pipeline shared by all three methods, at three preset strengths.

    "full": crop + color jitter + grayscale + blur + flip - SimCLR/BYOL's paired views,
    where augmentation *is* the training signal, not just regularization.
    "light": crop + flip only - DeepCluster's pseudo-label training pass.
    "none": no augmentation - DeepCluster's feature-extraction/clustering pass (clustering
    on noisy augmented features would make cluster assignments unstable run to run) and eval.
    """
    normalize = VT.Normalize(mean=IMAGENET_NORMALIZE["mean"], std=IMAGENET_NORMALIZE["std"])

    if strength == "none":
        return VT.Compose([VT.ToTensor(), normalize])

    transform_steps = [VT.RandomResizedCrop(32, scale=(0.2, 1.0)), VT.RandomHorizontalFlip()]

    if strength == "full":
        transform_steps += [
            VT.RandomApply([VT.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
            VT.RandomGrayscale(p=0.2),
            VT.RandomApply([VT.GaussianBlur(kernel_size=3)], p=0.5),
        ]
    elif strength != "light":
        raise ValueError(f"Unknown augmentation strength: {strength!r}")

    transform_steps += [VT.ToTensor(), normalize]
    return VT.Compose(transform_steps)


class TwoViews:
    """Applies a transform twice, producing the (x0, x1) positive pair used by SimCLR/BYOL."""

    def __init__(self, transform: VT.Compose) -> None:
        self.transform = transform

    def __call__(self, x: Any) -> tuple[Tensor, Tensor]:
        return self.transform(x), self.transform(x)


@torch.no_grad()
def extract_embeddings(
    forward_fn: Any, loader: data.DataLoader, device: torch.device
) -> tuple[np.ndarray, np.ndarray]:
    """Runs `forward_fn` (already in eval mode) over a loader, returning (embeddings, labels) as numpy arrays."""
    all_embeddings, all_labels = [], []
    for samples, labels in loader:
        embeddings = forward_fn(samples.to(device)).flatten(start_dim=1)
        all_embeddings.append(embeddings.cpu())
        all_labels.append(labels)
    return torch.cat(all_embeddings).float().numpy(), torch.cat(all_labels).numpy()


def save_checkpoint(
    path: Path,
    model: nn.Module,
    optimizer: optim.Optimizer,
    epoch: int,
    scheduler: Any = None,
    loss_history: list[float] | None = None,
    best_loss: float | None = None,
) -> None:
    """Gzip-compressed checkpoint. Model weights don't compress dramatically (fairly high-entropy
    floats), but it's free - no extra dependency, no precision loss - so there's no reason not to.

    Also saves RNG state (Python/NumPy/torch CPU/torch CUDA) plus loss_history/best_loss, so a
    resumed run continues the same random stream instead of replaying identical batches and
    augmentation choices from a fixed seed - see load_checkpoint's `restore_rng`.
    """
    checkpoint = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "epoch": epoch,
        "rng_state": {
            "random": random.getstate(),
            "numpy": np.random.get_state(),
            "torch": torch.get_rng_state(),
            "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
        },
    }
    if scheduler is not None:
        checkpoint["scheduler"] = scheduler.state_dict()
    if loss_history is not None:
        checkpoint["loss_history"] = loss_history
    if best_loss is not None:
        checkpoint["best_loss"] = best_loss
    with gzip.open(path, "wb") as f:
        torch.save(checkpoint, f)


def load_checkpoint(
    path: Path,
    model: nn.Module,
    optimizer: optim.Optimizer = None,
    scheduler: Any = None,
    restore_rng: bool = False,
) -> dict[str, Any]:
    """Auto-detects gzip (via its magic bytes) so this loads both the new compressed checkpoints
    and the plain ones already on disk from before compression was added.

    `restore_rng=True` resets Python/NumPy/torch RNG state to exactly where training left off -
    only meaningful when actually resuming a training loop (not for eval-only loading, where the
    notebook's own fixed-seed `set_random_seed()` call should stay in charge instead). Returns
    the epoch/loss_history/best_loss a resuming train_* function needs; old checkpoints predating
    this - like the first SimCLR run's `epoch_200.pt` - simply lack those keys and default sanely.

    `weights_only=False`: these checkpoints carry RNG state (Python/NumPy tuples, not just
    tensors), which PyTorch's default `weights_only=True` refuses to unpickle. Safe here since
    every checkpoint this loads is one this notebook produced itself, not a third-party file.
    """
    with path.open("rb") as f:
        is_gzip = f.read(2) == b"\x1f\x8b"
    opener = gzip.open if is_gzip else open
    with opener(path, "rb") as f:
        checkpoint = torch.load(f, map_location="cpu", weights_only=False)
    model.load_state_dict(checkpoint["model"])
    if optimizer is not None and "optimizer" in checkpoint:
        optimizer.load_state_dict(checkpoint["optimizer"])
    if scheduler is not None and "scheduler" in checkpoint:
        scheduler.load_state_dict(checkpoint["scheduler"])
    if restore_rng and "rng_state" in checkpoint:
        rng = checkpoint["rng_state"]
        random.setstate(rng["random"])
        np.random.set_state(rng["numpy"])
        torch.set_rng_state(rng["torch"])
        if rng["cuda"] is not None and torch.cuda.is_available():
            torch.cuda.set_rng_state_all(rng["cuda"])
    return {
        "epoch": checkpoint["epoch"],
        "loss_history": checkpoint.get("loss_history", []),
        "best_loss": checkpoint.get("best_loss", float("inf")),
    }


def load_resume_state(
    resume_from: Path | None,
    model: nn.Module,
    optimizer: optim.Optimizer = None,
    scheduler: Any = None,
) -> tuple[int, list[float], float]:
    """Shared resume-state loading for all three train_* functions - they only differ in which
    of optimizer/scheduler they pass through (DeepCluster passes neither, since it rebuilds both
    fresh every reclustering round regardless of resume). Returns (start_epoch, loss_history,
    best_loss), at fresh-start defaults (0, [], inf) if `resume_from` is None.
    """
    if resume_from is None:
        return 0, [], float("inf")
    resumed = load_checkpoint(resume_from, model, optimizer, scheduler, restore_rng=True)
    return resumed["epoch"], resumed["loss_history"], resumed["best_loss"]

import warnings


class UnNormalize:
    def __init__(self, mean: Iterable[float], std: Iterable[float]):
        self.mean = mean
        self.std = std

    def __call__(self, tensor: Tensor) -> Tensor:
        for t, m, s in zip(tensor, self.mean, self.std, strict=True):
            t.mul_(s).add_(m)
        return tensor


reverse_transform = VT.Compose(
    [
        UnNormalize(mean=IMAGENET_NORMALIZE["mean"], std=IMAGENET_NORMALIZE["std"]),
        VT.ToPILImage(),
    ]
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# CIFAR10's unpickling triggers a noisy VisibleDeprecationWarning on newer NumPy (align=0);
# it's torchvision's own internal call, nothing this notebook controls, so just suppress it.
with warnings.catch_warnings():
    warnings.simplefilter("ignore", category=UserWarning)

    ### Self-supervised pretraining dataset: SimCLR/BYOL need two full-strength augmented views per image
    pretrain_transform = TwoViews(build_augment("full"))
    pretrainset = datasets.CIFAR10(root=DATA_ROOT, train=True, download=False, transform=pretrain_transform)
    pretrain_loader = data.DataLoader(
        pretrainset, batch_size=256, shuffle=True, drop_last=True, pin_memory=True, num_workers=args.num_workers
    )

    ### Classification training (linear probe) and evaluation: no augmentation
    eval_transform = build_augment("none")
    trainset = datasets.CIFAR10(root=DATA_ROOT, train=True, download=False, transform=eval_transform)
    trainset_loader = data.DataLoader(
        trainset, batch_size=256, shuffle=True, drop_last=True, pin_memory=True, num_workers=args.num_workers
    )

    testset = datasets.CIFAR10(root=DATA_ROOT, train=False, download=False, transform=eval_transform)
    testset_loader = data.DataLoader(
        testset, batch_size=256, shuffle=False, drop_last=False, num_workers=args.num_workers
    )

    ### Visualization testset (raw PIL images, for the image-mapped t-SNE plot)
    vis_testset = datasets.CIFAR10(root=DATA_ROOT, train=False, download=False, transform=None)

SimCLR

For a batch of \(N\) images, SimCLR draws two augmented views per image and encodes all \(2N\) views. Its normalized temperature-scaled cross-entropy loss, NT-Xent, treats the other view of the same image as the positive and the remaining \(2N-2\) views as negatives (Chen et al. 2020):

\[ \mathcal{L}_i = -\log \frac{\exp(\operatorname{sim}(z_i,z_i')/\tau)}{\exp(\operatorname{sim}(z_i,z_i')/\tau) + \sum_{j\neq i}\left[\exp(\operatorname{sim}(z_i,z_j)/\tau) + \exp(\operatorname{sim}(z_i,z_j')/\tau)\right]}. \]

Here \(\operatorname{sim}\) is cosine similarity and \(\tau=0.5\). A constant representation cannot distinguish the positive from the negatives. For \(N=2\), if every similarity equals \(1\), the anchor loss is \(-\log(1/3)\approx1.10\). If the positive similarity is \(0.8\) and the two negative similarities are \(0.1\) and \(0.05\), the same loss is about \(0.39\).

The augmentations define which information should remain invariant. This run combines random resized crops, horizontal flips, color jitter, grayscale conversion, and Gaussian blur. The saved checkpoint contains 199 logged epochs and ends at an NT-Xent loss of 4.4750, which is also its minimum stored epoch loss. A previous scheduler configuration plateaued near the same value, but two runs are not enough to establish a general loss floor.

from tqdm import tqdm


class SimCLRModel(nn.Module):
    """Backbone + projection head. The projected output feeds the NT-Xent loss; the backbone
    output (before the projector) is what every evaluation and visualization cell uses."""

    def __init__(self) -> None:
        super().__init__()
        resnet = make_cifar_resnet()
        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
        self.projector = ProjectionHead(512, 2048, 2048)

    def forward(self, x: Tensor) -> Tensor:
        return self.projector(self.backbone(x).flatten(start_dim=1))


def train_simclr(
    model: SimCLRModel,
    loader: data.DataLoader,
    epochs: int,
    device: torch.device,
    checkpoint_path: Path = Path("./checkpoints/simclr/best.pt.gz"),
    warmup_epochs: int = 10,
    resume_from: Path | None = None,
) -> list[float]:
    """Saves whenever the epoch loss improves and also stores the final epoch. A lower
    NT-Xent training loss is the checkpoint criterion here, not proof of better downstream features.
    Also unconditionally saves the true final-epoch state to "final.pt.gz" once training completes,
    for comparison against "best.pt.gz" (see train_byol's docstring for why that comparison matters).

    `resume_from`, if given, restores model/optimizer/scheduler/RNG state and continues from the
    saved epoch, rather than restarting the shuffle order, augmentation randomness, and LR
    schedule from scratch.
    """
    model.to(device)
    model = torch.compile(model)
    optimizer = optim.Adam(model.parameters(), lr=1e-3, fused=device.type == "cuda")
    # ReduceLROnPlateau never triggered in practice (its eps=0.5 threshold was looser than the
    # loss's actual epoch-to-epoch noise) - a real warmup + cosine decay replaces it here, same
    # as train_byol.
    warmup = optim.lr_scheduler.LinearLR(optimizer, start_factor=0.01, total_iters=warmup_epochs)
    cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs - warmup_epochs)
    scheduler = optim.lr_scheduler.SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warmup_epochs])
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)

    start_epoch, loss_history, best_loss = load_resume_state(resume_from, model, optimizer, scheduler)

    for epoch in tqdm(range(start_epoch, epochs), desc="SimCLR", initial=start_epoch, total=epochs):
        model.train()
        epoch_loss = 0.0
        for (x0, x1), _ in loader:
            x0, x1 = x0.to(device), x1.to(device)
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
                loss = nt_xent_loss(model(x0), model(x1))

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        epoch_loss /= len(loader)
        loss_history.append(epoch_loss)
        scheduler.step()

        if epoch_loss < best_loss:
            best_loss = epoch_loss
            save_checkpoint(checkpoint_path, model, optimizer, epoch + 1, scheduler, loss_history, best_loss)

    save_checkpoint(
        checkpoint_path.parent / "final.pt.gz", model, optimizer, epochs, scheduler, loss_history, best_loss
    )
    return loss_history


simclr_model = SimCLRModel()

# A previous scheduler configuration plateaued near the retained run; see the prose for scope.
# The best-loss checkpoint is loaded below, while downstream evaluation provides the useful check.

# simclr_loss_history = train_simclr(simclr_model, pretrain_loader, epochs=200, device=device)

simclr_checkpoint = load_checkpoint(Path("./checkpoints/simclr/best.pt.gz"), simclr_model)
print(
    f"SimCLR pretraining: {len(simclr_checkpoint['loss_history'])} epochs, "
    f"final NT-Xent loss {simclr_checkpoint['loss_history'][-1]:.4f}, "
    f"best {min(simclr_checkpoint['loss_history']):.4f}"
)

fig, ax = plt.subplots(figsize=(WIDTH, WIDTH / GOLDEN_RATIO), dpi=200, constrained_layout=True)
ax.plot(
    range(len(simclr_checkpoint["loss_history"])),
    simclr_checkpoint["loss_history"],
    linewidth=1.5,
    alpha=0.80,
    clip_on=False,
)
ax.set_xlabel("epoch", fontsize=9)
ax.set_ylabel("NT-Xent loss", fontsize=9)
ax.set_title("SimCLR training loss", fontsize=10)
ax.tick_params(labelsize=8)
ax.autoscale(axis="x", tight=True)  # epoch 0 starts flush against the y-axis, no left margin
ax.grid(color="0.8", linestyle=":", linewidth=1)
save_fig(fig, "simclr-loss.webp")

SimCLR pretraining: 199 epochs, final NT-Xent loss 4.4750, best 4.4750

Figure 1: Stored epoch-average NT-Xent loss for the SimCLR run.

The DeepCluster and BYOL implementations use the same backbone but different objectives and input pipelines. Their results appear in the appendices.

Evaluating the Representations

The linear probe freezes the backbone and trains one labeled linear classifier on all 50,000 CIFAR-10 training images. The k-NN evaluation stores labeled training embeddings and classifies test embeddings by cosine distance. Neither evaluation is label-free; the labels measure how linearly or locally class-discriminative the frozen representation is.

The protocol has an important limitation. It evaluates the CIFAR-10 test set after every probe epoch, and the code passes test accuracy to a learning-rate scheduler. The scheduler’s unusually large eps=0.5 makes every proposed learning-rate change too small to apply, so it did not alter these stored runs, but the test set is still repeatedly inspected. A future run should use a validation split during training and reserve the test set for one final evaluation.

The notebook also lacks supervised and random-backbone baselines. The reported accuracies compare these three saved representations with one another, not with training from scratch or with a formal sample-efficiency baseline.

from sklearn.metrics import accuracy_score, f1_score


class LinearProbeModel(nn.Module):
    """Combines a frozen backbone with a linear_probe's trained classifier head into one
    callable module - used by the attribution-map cells, which need a single forward pass
    from image to class logits."""

    def __init__(self, backbone: nn.Module, classifier: nn.Module) -> None:
        super().__init__()
        self.backbone = backbone
        self.classifier = classifier

    def forward(self, x: Tensor) -> Tensor:
        return self.classifier(self.backbone(x).flatten(start_dim=1))


def linear_probe(
    backbone: nn.Module,
    train_loader: data.DataLoader,
    test_loader: data.DataLoader,
    epochs: int,
    device: torch.device,
    label: str = "",
) -> dict[str, Any]:
    """Freezes `backbone` and trains a single linear layer on top of it, plotting ACC/F1 per epoch."""
    backbone = backbone.to(device).eval()
    backbone = torch.compile(backbone)
    for param in backbone.parameters():
        param.requires_grad = False

    classifier = nn.Linear(512, 10).to(device)
    optimizer = optim.Adam(classifier.parameters(), lr=1e-3, fused=device.type == "cuda")
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, eps=0.5)

    acc = f1 = 0.0
    acc_history: list[float] = []
    f1_history: list[float] = []
    for _epoch in tqdm(range(epochs), desc="linear probe"):
        classifier.train()
        for inputs, targets in train_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            with torch.no_grad():
                features = backbone(inputs).flatten(start_dim=1)
            loss = F.cross_entropy(classifier(features), targets)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        classifier.eval()
        labels_true, labels_pred = [], []
        with torch.no_grad():
            for inputs, targets in test_loader:
                features = backbone(inputs.to(device)).flatten(start_dim=1)
                labels_true.append(targets)
                labels_pred.append(torch.argmax(classifier(features), dim=1).cpu())

        y_true = torch.cat(labels_true).numpy()
        y_pred = torch.cat(labels_pred).numpy()
        acc = accuracy_score(y_true, y_pred)
        f1 = f1_score(y_true, y_pred, average="weighted")
        acc_history.append(acc)
        f1_history.append(f1)
        scheduler.step(acc)

    epochs_range = range(epochs)
    fig, ax = plt.subplots(figsize=(WIDTH, WIDTH / GOLDEN_RATIO), dpi=200, constrained_layout=True)
    ax.plot(epochs_range, [a * 100 for a in acc_history], label="ACC", linewidth=1.5, alpha=0.80, clip_on=False)
    ax.plot(epochs_range, [f * 100 for f in f1_history], label="F1", linewidth=1.5, alpha=0.80, clip_on=False)
    ax.set_xlabel("epoch", fontsize=9)
    ax.set_ylabel("%", fontsize=9)
    ax.set_title(f"{label} linear probe" if label else "linear probe", fontsize=10)
    ax.tick_params(labelsize=8)
    ax.autoscale(axis="x", tight=True)  # epoch 0 starts flush against the y-axis, no left margin
    ax.legend(prop={"size": 8}, framealpha=0.6)
    ax.grid(color="0.8", linestyle=":", linewidth=1)
    save_fig(fig, f"probe-{label.lower()}.webp")

    return {"accuracy": acc, "f1": f1, "classifier": classifier, "model": LinearProbeModel(backbone, classifier)}

from sklearn.neighbors import KNeighborsClassifier


def knn_evaluate(
    train_embeddings: np.ndarray,
    train_labels: np.ndarray,
    test_embeddings: np.ndarray,
    test_labels: np.ndarray,
    k: int = 20,
) -> float:
    """k-NN accuracy on frozen embeddings - DINO's convention for a training-free evaluation."""
    knn = KNeighborsClassifier(n_neighbors=k, metric="cosine", n_jobs=-1)
    knn.fit(train_embeddings, train_labels)
    return accuracy_score(test_labels, knn.predict(test_embeddings))

from sklearn.model_selection import train_test_split


def data_efficiency_probe(
    backbone: nn.Module,
    train_dataset: datasets.CIFAR10,
    test_loader: data.DataLoader,
    device: torch.device,
    fractions: tuple[float, ...] = (0.01, 0.1, 1.0),
    epochs: int = 20,
) -> dict[float, float]:
    """Linear probe accuracy trained on a stratified fraction of the labeled train set, per fraction.

    This is the actual "fewer labels needed" argument for SSL: if a good representation was
    learned without labels, a linear probe on top of it should need far fewer labeled examples
    to reach a given accuracy than training from scratch would.
    """
    targets = np.array(train_dataset.targets)
    results = {}

    for fraction in fractions:
        if fraction >= 1.0:
            subset = train_dataset
        else:
            indices, _ = train_test_split(
                np.arange(len(train_dataset)), train_size=fraction, stratify=targets, random_state=0
            )
            subset = data.Subset(train_dataset, indices)

        subset_loader = data.DataLoader(subset, batch_size=256, shuffle=True, num_workers=args.num_workers)
        metrics = linear_probe(backbone, subset_loader, test_loader, epochs=epochs, device=device)
        results[fraction] = metrics["accuracy"]
        print(f"fraction {fraction:.0%}: ACC {metrics['accuracy'] * 100:.1f}%")

    return results

The stored SimCLR evaluation reports:

### SimCLR - linear probe and k-NN evaluation. The printed line below is the record:
### hard-coding the figures in a comment only guarantees they drift from the run.
simclr_backbone = simclr_model.backbone
simclr_linear = linear_probe(
    simclr_backbone, trainset_loader, testset_loader, epochs=64, device=device, label="SimCLR"
)
simclr_train_emb, simclr_train_labels = extract_embeddings(simclr_backbone.eval(), trainset_loader, device)
simclr_test_emb, simclr_test_labels = extract_embeddings(simclr_backbone.eval(), testset_loader, device)
simclr_knn = knn_evaluate(simclr_train_emb, simclr_train_labels, simclr_test_emb, simclr_test_labels)
print(
    f"SimCLR: linear probe ACC {simclr_linear['accuracy']:.1%} / F1 {simclr_linear['f1']:.1%}, "
    f"k-NN ACC {simclr_knn:.1%}"
)
# simclr_data_efficiency = data_efficiency_probe(simclr_backbone, trainset, testset_loader, device=device)

SimCLR: linear probe ACC 87.6% / F1 87.6%, k-NN ACC 86.0%

Figure 2: Test accuracy and weighted F1 recorded after each SimCLR linear-probe epoch.

The final SimCLR probe reaches 87.6% accuracy and 87.6% weighted F1, while 20-nearest-neighbor classification reaches 86.0% accuracy. These results show that this checkpoint retains substantial CIFAR-10 class information. They do not isolate the effect of negatives or show that SimCLR is more label-efficient than a supervised baseline.

Inspecting the Embeddings

t-SNE preserves selected local neighborhoods rather than global distances, and it can make clusters appear more distinct than they are in the original feature space (Maaten and Hinton 2008). In the SimCLR projection below, many vehicle thumbnails occupy a different region from animal thumbnails, and several classes form locally coherent groups. I treat this as a qualitative complement to the probe metrics, not independent proof of representation quality.

### Image-mapped t-SNE embeddings for SimCLR - actual CIFAR-10 thumbnails placed at their
### frozen test-set t-SNE coordinates, instead of plain dots. Self-contained (duplicates the
### DeepCluster appendix's tsne_pipeline/plot_embeddings setup below) rather than reordering
### cells, so this cell only depends on cell 10's already-loaded simclr_model checkpoint, not
### the appendix.
from matplotlib import offsetbox
from PIL.Image import Image as PILImage
from sklearn.manifold import TSNE
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

tsne_pipeline = Pipeline([("scaler", StandardScaler()), ("tsne", TSNE(n_jobs=-1, random_state=42))])

vis_samples = np.asarray([x for x, _ in vis_testset])


def plot_embeddings(
    x1: np.ndarray,
    x2: np.ndarray,
    samples: Iterable[np.ndarray | PILImage | Tensor],
    min_distance: float = 0.1,
    zoom: float = 0.5,
    ax: Any | None = None,
) -> Any:
    if ax is None:
        ax = plt.gca()

    X = np.stack([x1, x2], axis=-1)
    shown_images = np.array([[np.inf, np.inf]])
    for i in range(X.shape[0]):
        dist = np.sum((X[i] - shown_images) ** 2, 1)
        if np.min(dist) < min_distance:
            continue

        sample = samples[i]
        if isinstance(sample, Tensor):
            sample = sample.permute(1, 2, 0).numpy()
        if isinstance(sample, PILImage):
            sample = np.asarray(sample)

        shown_images = np.r_[shown_images, [X[i]]]
        im = offsetbox.OffsetImage(sample, zoom=zoom)
        ab = offsetbox.AnnotationBbox(offsetbox=im, xy=X[i], frameon=False)
        ax.add_artist(ab)

    return ax


simclr_backbone_gpu = simclr_model.backbone.to(device).eval()
simclr_embeddings, _ = extract_embeddings(simclr_backbone_gpu, testset_loader, device)
simclr_projection = tsne_pipeline.fit_transform(simclr_embeddings)

fig, ax = plt.subplots(figsize=(WIDTH, WIDTH), dpi=200)
ax.axis("off")
ax.scatter(simclr_projection[:, 0], simclr_projection[:, 1], marker=".", edgecolors="none", alpha=0.0)
plot_embeddings(simclr_projection[:, 0], simclr_projection[:, 1], samples=vis_samples, ax=ax)
save_fig(fig, "tsne-images-simclr.webp", lossless=False)

Figure 3: CIFAR-10 thumbnails placed at their SimCLR t-SNE coordinates.

Coloring the same projection by the held-out labels makes local class mixing easier to inspect:

### Class-colored t-SNE embeddings for SimCLR - same projection as the image-mapped version
### above, colored by true label instead of shown as thumbnails. Each class also gets its own
### marker shape, not just a color, since color alone isn't accessible to colorblind readers.
import matplotlib as mpl

classes = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
markers = ["o", "s", "^", "v", "D", "P", "X", "*", "p", "h"]
cmap = mpl.colormaps.get_cmap("tab10")
vis_labels = np.asarray([label for _, label in vis_testset])

fig, ax = plt.subplots(figsize=(WIDTH, WIDTH), dpi=200, constrained_layout=True)
for class_idx in sorted(set(vis_labels)):
    ax.scatter(
        simclr_projection[vis_labels == class_idx, 0],
        simclr_projection[vis_labels == class_idx, 1],
        color=cmap(class_idx),
        marker=markers[class_idx],
        s=8,
        edgecolors="none",
        alpha=0.8,
        label=classes[class_idx],
    )
ax.axis("off")
ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
save_fig(fig, "tsne-classes-simclr.webp")

Figure 4: The SimCLR t-SNE projection colored and marked by CIFAR-10 class.

The Guided Backpropagation panels below visualize gradients of the true-class logit for four examples. Such saliency maps are not reliable explanations of which features caused a prediction; some methods can remain visually similar after model parameters are randomized (Adebayo et al. 2018). I include them only as qualitative gradient visualizations.

### GuidedBackprop attribution for SimCLR, a few examples spanning different classes.
import warnings
from copy import deepcopy

from captum.attr import GuidedBackprop

example_classes = [0, 3, 5, 8]  # airplane, cat, dog, ship - two from each supercluster above
targets_array = np.asarray(testset.targets)
example_indices = [int(np.flatnonzero(targets_array == label)[0]) for label in example_classes]

simclr_probe_model = deepcopy(simclr_linear["model"]).eval()
gbp = GuidedBackprop(simclr_probe_model)

fig, axes = plt.subplots(
    nrows=2, ncols=len(example_indices), figsize=(WIDTH, WIDTH / 2), dpi=200, constrained_layout=True
)
for col, idx in enumerate(example_indices):
    sample_norm, sample_label = testset[idx]
    sample_image = reverse_transform(sample_norm.clone())

    # requires_grad_(True) avoids captum's "did not already require gradients" warning; the
    # other one it raises (temporary backward hooks on ReLU) is unavoidable and just suppressed.
    sample_input = sample_norm.unsqueeze(0).to(device).requires_grad_(True)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=UserWarning)
        attr = gbp.attribute(sample_input, target=sample_label)
    attr = attr.squeeze().permute(1, 2, 0).cpu().numpy()
    attr = np.sum(np.abs(attr), axis=-1)

    axes[0, col].imshow(sample_image, interpolation="none")
    axes[0, col].set_title(classes[sample_label])
    axes[0, col].axis("off")

    axes[1, col].imshow(attr, cmap="bwr", interpolation="none")
    axes[1, col].axis("off")

save_fig(fig, "attribution-guidedbackprop-simclr.webp")

Figure 5: Guided Backpropagation gradients for four SimCLR probe examples.

Later Directions

SimCLR obtains many negatives through large batches. MoCo reduces that requirement with a queue of embeddings from a momentum encoder (He et al. 2020). Methods without explicit negatives take a different route: BYOL uses an online predictor and EMA target (Grill et al. 2020), while DINO adds teacher-output centering and sharpening (Caron et al. 2021).

DeepCluster also led to online clustering approaches. SwAV replaces repeated offline k-means with swapped predictions of online cluster assignments (Caron et al. 2020). DINOv2 combines image-level self-distillation with patch-level objectives and large-scale data curation (Oquab et al. 2023). I-JEPA instead predicts target-region representations from visible context (Assran et al. 2023). These methods are referenced to locate the three notebook experiments in the broader design space; they are not evaluated here.

Conclusion

The saved checkpoints support a narrow result. SimCLR produced the strongest class-discriminative representation in this setup, DeepCluster retained a moderate signal, and my BYOL implementation performed poorly. The mechanisms explain what each objective is trying to do, but one run per configuration cannot establish why the accuracies differ.

The BYOL failure is therefore an implementation and configuration result, not evidence that BYOL needs ImageNet-scale data or that predictor-plus-EMA training fails on CIFAR-10. The most useful lesson was methodological: a low self-distillation loss is not enough to diagnose representation quality, and neither t-SNE nor an attribution map substitutes for downstream metrics and appropriate baselines.

Appendix

DeepCluster

NoteShow the DeepCluster write-up

DeepCluster alternates between two steps (Caron et al. 2018): cluster backbone features with k-means, then train the backbone and a new linear head to predict those cluster assignments. The cluster identifiers are arbitrary pseudo-labels rather than semantic classes.

A degenerate solution can concentrate samples in a few clusters. The implementation uses inverse-frequency sampling so large populated clusters do not dominate the gradient, and it records cluster sizes for inspection. Its Sobel preprocessing removes color before feature extraction, following the original method’s emphasis on shape and texture.

I tried reclustering every epoch and every five epochs. I retained the five-epoch checkpoint because its downstream metrics were better in my notes and it avoided four expensive clustering passes out of five. The notebook does not preserve the first run’s metrics or either run’s NMI output, so it cannot support a quantitative cadence comparison from stored evidence alone.

import warnings

from sklearn.cluster import MiniBatchKMeans
from sklearn.decomposition import PCA


class DeepClusterDataset(data.Dataset):
    """Wraps CIFAR-10 (base transform=None) so the applied transform and pseudo-label targets
    can be swapped out each reclustering round without touching the underlying data."""

    def __init__(self, base: data.Dataset) -> None:
        self.base = base
        self.transform: VT.Compose | None = None
        self.pseudo_labels = np.zeros(len(base), dtype=np.int64)

    def __len__(self) -> int:
        return len(self.base)

    def __getitem__(self, index: int) -> tuple[Tensor, int]:
        image, _ = self.base[index]
        return self.transform(image), int(self.pseudo_labels[index])


class DeepClusterModel(nn.Module):
    """Backbone (fed Sobel-filtered input) + a classification head over pseudo-labels.
    The head is reinitialized every reclustering round, since cluster identities are
    arbitrary and don't carry meaning across rounds."""

    def __init__(self, num_clusters: int) -> None:
        super().__init__()
        self.sobel = SobelFilter()
        resnet = make_cifar_resnet(in_channels=2)
        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
        self.cluster_head = nn.Linear(512, num_clusters)

    def reset_head(self, num_clusters: int) -> None:
        device = next(self.backbone.parameters()).device
        self.cluster_head = nn.Linear(512, num_clusters).to(device)

    def features(self, x: Tensor) -> Tensor:
        return self.backbone(self.sobel(x))

    def forward(self, x: Tensor) -> Tensor:
        return self.cluster_head(self.features(x).flatten(start_dim=1))


def train_deepcluster(
    model: DeepClusterModel,
    dataset: DeepClusterDataset,
    epochs: int,
    num_clusters: int,
    device: torch.device,
    train_transform: VT.Compose,
    cluster_transform: VT.Compose,
    checkpoint_path: Path = Path("./checkpoints/deepcluster/best.pt.gz"),
    resume_from: Path | None = None,
    recluster_every: int = 1,
) -> tuple[list[float], list[np.ndarray]]:
    """Saves (overwriting) whenever the epoch loss improves on the best seen so far. Cross-entropy
    against pseudo-labels doesn't have BYOL's collapse-via-low-loss failure mode, but it's still
    only a proxy - the cluster-size histogram (below) is the real diagnostic for degenerate
    solutions, not this loss. Also unconditionally saves the true final-epoch state to
    "final.pt.gz" once training completes, for comparison against "best.pt.gz" (see train_byol's
    docstring for why that comparison matters) - and this comparison mattered here too: see the
    run log below, best.pt.gz stuck on an early-epoch fluke just like BYOL's attempt 2.

    `recluster_every` controls how many epochs pass between reclustering rounds (feature
    extraction + PCA + k-means + a fresh classification head). The first attempt used 1 (matches
    the original DeepCluster paper's cadence). `recluster_every=5` matches the actual cadence
    used in the author's own DeepCluster-based paper (Milosheski, Cerar et al., "Self-supervised
    learning for clustering of wireless spectrum activity," Computer Communications 2023) - their
    published code reclusters every 5 epochs rather than every 1, giving the backbone several
    epochs of gradient signal per cluster assignment before labels shift again. On non-reclustering
    epochs, the previous round's pseudo-labels, classification head, and optimizer are reused and
    continue training rather than being reset - resetting a head that's still training on the same
    labels would throw away progress for no reason.

    `resume_from`, if given, restores the backbone/RNG state and continues from the saved epoch.
    The classification head and its optimizer are rebuilt fresh either way, since a fresh head
    every reclustering round is already this method's normal behavior (cluster identities aren't
    meaningful across rounds), resumed or not - only the backbone needs to carry over.
    """
    model.to(device)
    # model.features(x) during reclustering (below) is called directly, bypassing forward(),
    # so it needs its own compile - torch.compile(model) alone only covers model(images).
    model = torch.compile(model)
    model.features = torch.compile(model.features)
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)

    start_epoch, loss_history, best_loss = load_resume_state(resume_from, model)
    cluster_size_history = []
    optimizer = None
    train_loader = None

    for epoch in tqdm(range(start_epoch, epochs), desc="DeepCluster", initial=start_epoch, total=epochs):
        if epoch % recluster_every == 0:
            # 1. Extract features on an unaugmented view - clustering on noisy features would
            #    make cluster assignments unstable from one round to the next.
            model.eval()
            dataset.transform = cluster_transform
            cluster_loader = data.DataLoader(
                dataset, batch_size=256, shuffle=False, pin_memory=True, num_workers=args.num_workers
            )
            features, _ = extract_embeddings(model.features, cluster_loader, device)

            # 2. PCA-whiten, then k-means for pseudo-labels.
            features = PCA(n_components=256, whiten=True, random_state=0).fit_transform(features)
            pseudo_labels = MiniBatchKMeans(n_clusters=num_clusters, n_init="auto", random_state=0).fit_predict(
                features
            )
            dataset.pseudo_labels = pseudo_labels

            cluster_sizes = np.bincount(pseudo_labels, minlength=num_clusters)
            cluster_size_history.append(cluster_sizes)

            # 3. Fresh head each reclustering round.
            model.reset_head(num_clusters)
            optimizer = optim.Adam(model.parameters(), lr=1e-3, fused=device.type == "cuda")

            # 4. Inverse-frequency sampling, so a few large clusters don't dominate the gradient.
            sample_weights = 1.0 / cluster_sizes[pseudo_labels]
            sampler = data.WeightedRandomSampler(sample_weights, num_samples=len(dataset), replacement=True)

            dataset.transform = train_transform
            train_loader = data.DataLoader(
                dataset, batch_size=256, sampler=sampler, pin_memory=True, num_workers=args.num_workers
            )

        model.train()
        epoch_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
                loss = F.cross_entropy(model(images), labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        epoch_loss /= len(train_loader)
        loss_history.append(epoch_loss)

        if epoch_loss < best_loss:
            best_loss = epoch_loss
            save_checkpoint(
                checkpoint_path, model, optimizer, epoch + 1, loss_history=loss_history, best_loss=best_loss
            )

    save_checkpoint(
        checkpoint_path.parent / "final.pt.gz",
        model,
        optimizer,
        epochs,
        loss_history=loss_history,
        best_loss=best_loss,
    )
    return loss_history, cluster_size_history


# Same noisy torchvision/NumPy VisibleDeprecationWarning as the main dataset setup above.
with warnings.catch_warnings():
    warnings.simplefilter("ignore", category=UserWarning)
    deepcluster_dataset = DeepClusterDataset(
        datasets.CIFAR10(root=DATA_ROOT, train=True, download=False, transform=None)
    )
deepcluster_model = DeepClusterModel(num_clusters=100)

# Attempt 1 (recluster_every=1, the original paper's cadence) - see the appendix prose for the
# full story. best.pt.gz got stuck on an early-epoch fluke (same pattern as BYOL's below);
# final.pt.gz is the checkpoint that matters. Superseded by attempt 2; checkpoints not kept on disk.
# deepcluster_loss_history, deepcluster_cluster_sizes = train_deepcluster(
#     deepcluster_model,
#     deepcluster_dataset,
#     epochs=200,
#     num_clusters=100,
#     device=device,
#     train_transform=build_augment("light"),
#     cluster_transform=build_augment("none"),
#     recluster_every=1,
# )

# Attempt 2 (recluster_every=5, adopted as the default) - see the appendix prose for why. Same
# best.pt.gz-vs-final.pt.gz fluke pattern as attempt 1; final.pt.gz is again the checkpoint that
# matters.
# deepcluster_loss_history, deepcluster_cluster_sizes = train_deepcluster(
#     deepcluster_model,
#     deepcluster_dataset,
#     epochs=200,
#     num_clusters=100,
#     device=device,
#     train_transform=build_augment("light"),
#     cluster_transform=build_augment("none"),
#     recluster_every=5,
# )

deepcluster_checkpoint = load_checkpoint(Path("./checkpoints/deepcluster/final.pt.gz"), deepcluster_model)

fig, ax = plt.subplots(figsize=(WIDTH, WIDTH / GOLDEN_RATIO), dpi=200, constrained_layout=True)
ax.plot(
    range(len(deepcluster_checkpoint["loss_history"])),
    deepcluster_checkpoint["loss_history"],
    linewidth=1.5,
    alpha=0.80,
    clip_on=False,
)
ax.set_xlabel("epoch", fontsize=9)
ax.set_ylabel("cross-entropy loss", fontsize=9)
ax.set_title("DeepCluster training loss", fontsize=10)
ax.tick_params(labelsize=8)
ax.autoscale(axis="x", tight=True)  # epoch 0 starts flush against the y-axis, no left margin
ax.grid(color="0.8", linestyle=":", linewidth=1)
save_fig(fig, "deepcluster-loss.webp")

Figure 6: Stored cross-entropy loss for the retained DeepCluster run.

### DeepCluster evaluation for the retained recluster_every=5 checkpoint.
### Historical comparison metrics are not preserved in a stored output.

### DeepClusterModel's backbone expects Sobel-filtered input, so wrap sobel+backbone together -
### nn.Sequential accepts the same 3-channel RGB loaders as SimCLR/BYOL and applies Sobel internally.
deepcluster_backbone = nn.Sequential(deepcluster_model.sobel, deepcluster_model.backbone)
deepcluster_linear = linear_probe(
    deepcluster_backbone, trainset_loader, testset_loader, epochs=64, device=device, label="DeepCluster"
)
deepcluster_train_emb, deepcluster_train_labels = extract_embeddings(
    deepcluster_backbone.eval(), trainset_loader, device
)
deepcluster_test_emb, deepcluster_test_labels = extract_embeddings(deepcluster_backbone.eval(), testset_loader, device)
deepcluster_knn = knn_evaluate(
    deepcluster_train_emb, deepcluster_train_labels, deepcluster_test_emb, deepcluster_test_labels
)
print(
    f"DeepCluster: linear probe ACC {deepcluster_linear['accuracy']:.1%} / "
    f"F1 {deepcluster_linear['f1']:.1%}, k-NN ACC {deepcluster_knn:.1%}"
)
# deepcluster_data_efficiency = data_efficiency_probe(deepcluster_backbone, trainset, testset_loader, device=device)

DeepCluster: linear probe ACC 64.8% / F1 64.7%, k-NN ACC 58.6%

Figure 7: Test accuracy and weighted F1 recorded after each DeepCluster linear-probe epoch.

The retained DeepCluster checkpoint reaches 64.8% linear-probe accuracy, 64.7% weighted F1, and 58.6% k-NN accuracy. Those values show class-discriminative information well above random guessing, but the missing random-backbone baseline prevents a stronger claim about how much the pretraining added.

Comparing the Visualizations

The following t-SNE projections are fitted independently for SimCLR and DeepCluster. Their axes, orientation, cluster sizes, and distances are therefore not comparable across panels. Within each panel, the label colors suggest more local class separation for SimCLR, consistent with its higher probe accuracy.

The attribution figures use one airplane and target its true-class logit. Guided Backpropagation and a one-sample GradientShap estimate are both sensitive visual summaries, not causal explanations or quantitative comparisons (Adebayo et al. 2018).

from sklearn.manifold import TSNE
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

tsne_pipeline = Pipeline([("scaler", StandardScaler()), ("tsne", TSNE(n_jobs=-1, random_state=42))])

### Raw images and true labels for the visualization test set - method-independent, no trained model needed
vis_samples = np.asarray([x for x, _ in vis_testset])
vis_labels = np.asarray([label for _, label in vis_testset])

### t-SNE projection of each method's frozen test-set embeddings.
method_backbones = {
    "SimCLR": simclr_model.backbone,
    "DeepCluster": nn.Sequential(deepcluster_model.sobel, deepcluster_model.backbone),
}

method_projections = {}
for name, backbone in method_backbones.items():
    embeddings, _ = extract_embeddings(backbone.eval(), testset_loader, device)
    method_projections[name] = tsne_pipeline.fit_transform(embeddings)

from matplotlib import offsetbox  # noqa: F811 - reimported for cell self-containment
from PIL.Image import Image


def plot_embeddings(
    x1: np.ndarray,
    x2: np.ndarray,
    samples: Iterable[np.ndarray | Image | Tensor],
    labels: Iterable[int | str] | None = None,
    min_distance: float = 0.1,
    zoom: float = 0.5,
    ax: Any | None = None,
) -> Any:
    if ax is None:
        ax = plt.gca()

    # Requires matplotlib >= 1.0
    X = np.stack([x1, x2], axis=-1)

    shown_images = np.array([[np.inf, np.inf]])
    for i in range(X.shape[0]):
        dist = np.sum((X[i] - shown_images) ** 2, 1)
        if np.min(dist) < min_distance:
            # don't show points that are too close
            continue

        sample = samples[i]

        if isinstance(sample, Tensor):
            sample = sample.permute(1, 2, 0).numpy()

        if isinstance(sample, Image):
            sample = np.asarray(sample)

        shown_images = np.r_[shown_images, [X[i]]]
        im = offsetbox.OffsetImage(sample, zoom=zoom)
        ab = offsetbox.AnnotationBbox(offsetbox=im, xy=X[i], frameon=False)
        ax.add_artist(ab)

    return ax

### Image-mapped t-SNE embeddings, side by side across SimCLR and DeepCluster (the favorite
### visualization from the original notebook).
fig, axes = plt.subplots(ncols=2, figsize=(WIDTH, WIDTH / 2), dpi=200, constrained_layout=True)
for ax, (name, X) in zip(axes, method_projections.items(), strict=True):
    ax.axis("off")
    ax.set_title(name)
    ax.scatter(X[:, 0], X[:, 1], marker=".", edgecolors="none", alpha=0.0)
    plot_embeddings(X[:, 0], X[:, 1], samples=vis_samples, ax=ax)
save_fig(fig, "tsne-images-comparison.webp", lossless=False)

Figure 8: CIFAR-10 thumbnails in independently fitted SimCLR and DeepCluster t-SNE projections.

import matplotlib as mpl

classes = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
cmap = mpl.colormaps.get_cmap("tab10")

### Class-colored t-SNE embeddings, side by side across SimCLR and DeepCluster.
fig, axes = plt.subplots(ncols=2, figsize=(WIDTH, WIDTH / 2), dpi=200, constrained_layout=True)
for ax, (name, X) in zip(axes, method_projections.items(), strict=True):
    for class_idx in sorted(set(vis_labels)):
        ax.scatter(
            X[vis_labels == class_idx, 0],
            X[vis_labels == class_idx, 1],
            color=cmap(class_idx),
            marker=".",
            edgecolors="none",
            alpha=0.8,
            label=classes[class_idx],
        )
    ax.axis("off")
    ax.set_title(name)
axes[-1].legend(bbox_to_anchor=(1.05, 1), loc="upper left")
save_fig(fig, "tsne-classes-comparison.webp")

Figure 9: Independent SimCLR and DeepCluster t-SNE projections colored by class.

from copy import deepcopy

from captum.attr import GradientShap, GuidedBackprop
from captum.attr import visualization as viz

### GuidedBackprop attribution, one sample, compared across SimCLR and DeepCluster.
import warnings

method_probe_models = {
    "SimCLR": simclr_linear["model"],
    "DeepCluster": deepcluster_linear["model"],
}

sample_idx = 1555
sample_norm, sample_label = testset[sample_idx]
sample_image = reverse_transform(sample_norm.clone())

fig, axes = plt.subplots(ncols=3, figsize=(WIDTH, WIDTH / 3), dpi=200, constrained_layout=True)
axes[0].imshow(sample_image, interpolation="none")
axes[0].set_title("original")
axes[0].axis("off")

for ax, (name, probe_model) in zip(axes[1:], method_probe_models.items(), strict=True):
    probe_model.eval()
    gbp = GuidedBackprop(deepcopy(probe_model))

    # requires_grad_(True) avoids captum's "did not already require gradients" warning; the
    # other one it raises (temporary backward hooks on ReLU) is unavoidable and just suppressed.
    sample_input = sample_norm.unsqueeze(0).to(device).requires_grad_(True)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=UserWarning)
        attr = gbp.attribute(sample_input, target=sample_label)
    attr = attr.squeeze().permute(1, 2, 0).cpu().numpy()
    attr = np.sum(np.abs(attr), axis=-1)

    ax.imshow(attr, cmap="bwr", interpolation="none")
    ax.set_title(name)
    ax.axis("off")

save_fig(fig, "attribution-guidedbackprop-comparison.webp")

Figure 10: Guided Backpropagation gradients for one true-class logit in the SimCLR and DeepCluster probes.

### GradientShap attribution, same sample, compared across SimCLR and DeepCluster.
torch.manual_seed(0)
np.random.seed(0)

rand_img_dist = torch.cat([sample_norm.unsqueeze(0) * 0, sample_norm.unsqueeze(0) * 1]).to(device)

for name, probe_model in method_probe_models.items():
    gradient_shap = GradientShap(probe_model)
    attributions_gs = gradient_shap.attribute(
        sample_norm.unsqueeze(0).to(device), n_samples=1, stdevs=0.0001, baselines=rand_img_dist, target=sample_label
    )
    print(name)
    fig, _ = viz.visualize_image_attr_multiple(
        np.transpose(attributions_gs.squeeze().cpu().detach().numpy(), (1, 2, 0)),
        np.transpose(VT.PILToTensor()(sample_image).squeeze().cpu().detach().numpy(), (1, 2, 0)),
        ["original_image", "heat_map"],
        ["all", "absolute_value"],
        cmap="viridis",
        show_colorbar=True,
        use_pyplot=False,
    )
    save_fig(fig, f"attribution-gradientshap-{name.lower()}.webp")

SimCLR
DeepCluster

Figure 11: One-sample GradientShap visualization for the SimCLR probe.

Figure 12: One-sample GradientShap visualization for the DeepCluster probe.

The t-SNE panels offer a visual check of local neighborhoods, while the attribution panels show how two gradient methods render one selected input. I do not use either set to infer why SimCLR outperformed DeepCluster; the probe and k-NN results are the more direct evidence.

BYOL Failure Investigation

NoteShow the BYOL investigation

BYOL trains an online backbone, projector, and predictor against a stop-gradient target network updated by exponential moving average (Grill et al. 2020). The predictor and slowly changing target create an asymmetric learning problem without explicit negatives.

My implementation produced highly aligned backbone features in four attempts. I monitored average pairwise cosine similarity and per-dimension standard deviation because the BYOL training loss can decrease even when inputs receive nearly identical representations. Those diagnostic values survive only in comments and notes, not stored cell outputs, so the reproducible evidence in this notebook is the final checkpoint’s downstream evaluation.

The attempts changed several plausible settings:

1. A fixed EMA momentum, flat learning rate, and plateau scheduler.
2. Warmup followed by cosine learning-rate decay and a lower fixed momentum.
3. AdamW with weight decay in addition to the second configuration.
4. A cosine EMA-momentum ramp from \(0.996\) toward \(1.0\), matching the schedule described by BYOL.

All four remained poor by my diagnostics. This narrows the local debugging history, but it does not rule out implementation differences or other hyperparameters. The notebook uses AdamW rather than the paper’s large-batch LARS setup, a different projector and predictor design, CIFAR-10, and one backbone size. Any of those choices could matter.

A second lesson concerns checkpoint selection. The lowest BYOL training loss is not necessarily the most useful representation, so the evaluation below deliberately loads the final checkpoint. A future rerun should store collapse diagnostics every epoch and compare with a randomly initialized backbone.

class BYOLModel(nn.Module):
    """Online backbone, projector, and predictor with an EMA-updated target backbone and projector.
    The predictor and momentum target introduce asymmetry; their presence does not by itself prove
    that a particular implementation will avoid a degenerate representation."""

    def __init__(self, tau_base: float = 0.996) -> None:
        super().__init__()
        # Base tau for the cosine ramp (see byol_tau_schedule) - the value at step 0, rising to
        # 1.0 by the final training step. Not used directly as a fixed EMA rate; train_byol
        # computes the actual per-step tau from this and passes it to update_target.
        self.tau_base = tau_base

        online_resnet = make_cifar_resnet()
        self.online_backbone = nn.Sequential(*list(online_resnet.children())[:-1])
        self.online_projector = ProjectionHead(512, 2048, 2048)
        self.predictor = ProjectionHead(2048, 2048, 2048)

        target_resnet = make_cifar_resnet()
        self.target_backbone = nn.Sequential(*list(target_resnet.children())[:-1])
        self.target_projector = ProjectionHead(512, 2048, 2048)
        self.target_backbone.load_state_dict(self.online_backbone.state_dict())
        self.target_projector.load_state_dict(self.online_projector.state_dict())
        for param in list(self.target_backbone.parameters()) + list(self.target_projector.parameters()):
            param.requires_grad = False

    def train(self, mode: bool = True) -> "BYOLModel":
        # The target network is never trained directly (only EMA-updated), so it always
        # stays in eval mode - otherwise its BatchNorm running stats would drift from forward
        # passes that never contribute a gradient.
        super().train(mode)
        self.target_backbone.eval()
        self.target_projector.eval()
        return self

    def forward_online(self, x: Tensor) -> Tensor:
        return self.predictor(self.online_projector(self.online_backbone(x).flatten(start_dim=1)))

    @torch.no_grad()
    def forward_target(self, x: Tensor) -> Tensor:
        return self.target_projector(self.target_backbone(x).flatten(start_dim=1))

    @torch.no_grad()
    def update_target(self, tau: float) -> None:
        update_moving_average(self.target_backbone, self.online_backbone, tau)
        update_moving_average(self.target_projector, self.online_projector, tau)


def train_byol(
    model: BYOLModel,
    loader: data.DataLoader,
    epochs: int,
    device: torch.device,
    checkpoint_path: Path = Path("./checkpoints/byol/best.pt.gz"),
    warmup_epochs: int = 10,
    weight_decay: float = 1e-6,
    resume_from: Path | None = None,
) -> list[float]:
    """Saves (overwriting) whenever the epoch loss improves on the best seen so far, plus an
    unconditional save of the true final-epoch state to "final.pt.gz" once training completes.

    Caution: the self-distillation loss is not a sufficient checkpoint-quality metric. Similar
    outputs can also reduce it, so compare final and best-loss checkpoints with representation
    diagnostics and downstream evaluation before choosing one.

    `model.tau_base` sets the EMA momentum's starting point; the actual per-step tau follows
    the original BYOL paper's cosine ramp (byol_tau_schedule) up to 1.0 by the final step,
    rather than staying fixed for the whole run as in attempts 1-3.

    `resume_from`, if given, restores model/optimizer/scheduler/RNG state and continues from the
    saved epoch, rather than restarting the shuffle order, augmentation randomness, and LR
    schedule from scratch.
    """
    model.to(device)
    # BYOL has no single forward() the training loop calls (forward_online/forward_target are
    # separate entry points), so torch.compile(model) wouldn't touch either - compile the two
    # methods directly instead.
    model.forward_online = torch.compile(model.forward_online)
    model.forward_target = torch.compile(model.forward_target)
    online_params = (
        list(model.online_backbone.parameters())
        + list(model.online_projector.parameters())
        + list(model.predictor.parameters())
    )
    # AdamW (decoupled weight decay), not Adam(weight_decay=...) - plain Adam's weight decay is
    # really L2-via-gradient, which interacts oddly with Adam's per-parameter adaptive learning
    # rates and gives a weaker regularization effect than AdamW's decoupled version.
    optimizer = optim.AdamW(online_params, lr=1e-3, weight_decay=weight_decay, fused=device.type == "cuda")
    # ReduceLROnPlateau never triggered in practice (its eps=0.5 threshold was looser than the
    # loss's actual epoch-to-epoch noise) - a real warmup + cosine decay replaces it here.
    warmup = optim.lr_scheduler.LinearLR(optimizer, start_factor=0.01, total_iters=warmup_epochs)
    cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs - warmup_epochs)
    scheduler = optim.lr_scheduler.SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warmup_epochs])
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)

    start_epoch, loss_history, best_loss = load_resume_state(resume_from, model, optimizer, scheduler)

    steps_per_epoch = len(loader)
    total_steps = epochs * steps_per_epoch
    step = start_epoch * steps_per_epoch

    for epoch in tqdm(range(start_epoch, epochs), desc="BYOL", initial=start_epoch, total=epochs):
        model.train()
        epoch_loss = 0.0
        for (x0, x1), _ in loader:
            x0, x1 = x0.to(device), x1.to(device)
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
                p0, p1 = model.forward_online(x0), model.forward_online(x1)
                z0, z1 = model.forward_target(x0), model.forward_target(x1)
                loss = byol_loss(p0, z1) / 2 + byol_loss(p1, z0) / 2

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            model.update_target(byol_tau_schedule(step, total_steps, model.tau_base))
            step += 1
            epoch_loss += loss.item()

        epoch_loss /= len(loader)
        loss_history.append(epoch_loss)
        scheduler.step()

        if epoch_loss < best_loss:
            best_loss = epoch_loss
            save_checkpoint(checkpoint_path, model, optimizer, epoch + 1, scheduler, loss_history, best_loss)

    save_checkpoint(
        checkpoint_path.parent / "final.pt.gz", model, optimizer, epochs, scheduler, loss_history, best_loss
    )
    return loss_history


byol_model = BYOLModel(tau_base=0.996)

# Four attempts (fixed vs. ramped tau, three tau values, three LR schedules, Adam vs. AdamW
# with/without weight decay) all collapsed to a similar degree - see the BYOL appendix for the
# full story.
# byol_loss_history = train_byol(byol_model, pretrain_loader, epochs=200, device=device)

# Load the final checkpoint so the downstream cells evaluate the retained run rather than select
# a checkpoint by self-distillation loss alone.

byol_checkpoint = load_checkpoint(Path("./checkpoints/byol/final.pt.gz"), byol_model)

fig, ax = plt.subplots(figsize=(WIDTH, WIDTH / GOLDEN_RATIO), dpi=200, constrained_layout=True)
ax.plot(
    range(len(byol_checkpoint["loss_history"])),
    byol_checkpoint["loss_history"],
    linewidth=1.5,
    alpha=0.80,
    clip_on=False,
)
ax.set_xlabel("epoch", fontsize=9)
ax.set_ylabel("BYOL loss", fontsize=9)
ax.set_title("BYOL training loss", fontsize=10)
ax.tick_params(labelsize=8)
ax.autoscale(axis="x", tight=True)  # epoch 0 starts flush against the y-axis, no left margin
ax.grid(color="0.8", linestyle=":", linewidth=1)
save_fig(fig, "byol-loss.webp")

Figure 13: Stored self-distillation loss for the retained BYOL run.

The final checkpoint is evaluated with the same linear-probe and k-NN code as the other two methods.

### BYOL linear-probe and k-NN evaluation for the retained final checkpoint.
### The result is weak relative to the other methods, but needs a random-backbone baseline.

byol_backbone = byol_model.online_backbone
byol_linear = linear_probe(byol_backbone, trainset_loader, testset_loader, epochs=64, device=device, label="BYOL")
byol_train_emb, byol_train_labels = extract_embeddings(byol_backbone.eval(), trainset_loader, device)
byol_test_emb, byol_test_labels = extract_embeddings(byol_backbone.eval(), testset_loader, device)
byol_knn = knn_evaluate(byol_train_emb, byol_train_labels, byol_test_emb, byol_test_labels)
print(f"BYOL: linear probe ACC {byol_linear['accuracy']:.1%} / F1 {byol_linear['f1']:.1%}, k-NN ACC {byol_knn:.1%}")
# byol_data_efficiency = data_efficiency_probe(byol_backbone, trainset, testset_loader, device=device)

BYOL: linear probe ACC 18.4% / F1 15.5%, k-NN ACC 25.1%

Figure 14: Test accuracy and weighted F1 recorded after each BYOL linear-probe epoch.

### t-SNE of the retained BYOL embeddings; this visualization is not a collapse diagnostic.
byol_embeddings, _ = extract_embeddings(byol_backbone.eval(), testset_loader, device)
byol_projection = tsne_pipeline.fit_transform(byol_embeddings)

fig, ax = plt.subplots(figsize=(WIDTH, WIDTH), dpi=200, constrained_layout=True)
ax.axis("off")
ax.set_title("BYOL (collapsed)")
ax.scatter(byol_projection[:, 0], byol_projection[:, 1], marker=".", edgecolors="none", alpha=0.0)
plot_embeddings(byol_projection[:, 0], byol_projection[:, 1], samples=vis_samples, ax=ax)
save_fig(fig, "tsne-byol-collapse.webp", lossless=False)

Figure 15: CIFAR-10 thumbnails in the t-SNE projection of the retained BYOL representation.

The BYOL checkpoint reaches 18.4% linear-probe accuracy, 15.5% weighted F1, and 25.1% k-NN accuracy. These are far below the other two checkpoints but above a 10% uniform-random guess. Without a random-backbone baseline, they do not establish complete collapse or show how much class information pretraining removed or retained. The t-SNE image is also not a collapse diagnostic: near-identical high-dimensional points can be spread apart by numerical noise and the projection procedure.

References

Adebayo, Julius, Justin Gilmer, Michael Muelly, Ian Goodfellow, Moritz Hardt, and Been Kim. 2018. “Sanity Checks for Saliency Maps.” Advances in Neural Information Processing Systems 31: 9505–15. https://proceedings.neurips.cc/paper/2018/hash/294a8ed24b1ad22ec2e7efea049b8737-Abstract.html.

Assran, Mahmoud, Quentin Duval, Ishan Misra, et al. 2023. “Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

Caron, Mathilde, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. 2018. “Deep Clustering for Unsupervised Learning of Visual Features.” Proceedings of the European Conference on Computer Vision (ECCV).

Caron, Mathilde, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. 2020. “Unsupervised Learning of Visual Features by Contrasting Cluster Assignments.” Advances in Neural Information Processing Systems (NeurIPS).

Caron, Mathilde, Hugo Touvron, Ishan Misra, et al. 2021. “Emerging Properties in Self-Supervised Vision Transformers.” Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV).

Chen, Ting, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. 2020. “A Simple Framework for Contrastive Learning of Visual Representations.” Proceedings of the 37th International Conference on Machine Learning (ICML).

Chen, Xinlei, and Kaiming He. 2021. “Exploring Simple Siamese Representation Learning.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

Grill, Jean-Bastien, Florian Strub, Florent Altché, et al. 2020. “Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning.” Advances in Neural Information Processing Systems (NeurIPS).

He, Kaiming, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. 2020. “Momentum Contrast for Unsupervised Visual Representation Learning.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

Maaten, Laurens van der, and Geoffrey Hinton. 2008. “Visualizing Data Using t-SNE.” Journal of Machine Learning Research 9 (86): 2579–605. https://www.jmlr.org/papers/v9/vandermaaten08a.html.

Milosheski, Ljupcho, Gregor Cerar, Blaž Bertalanič, Carolina Fortuna, and Mihael Mohorčič. 2023. “Self-Supervised Learning for Clustering of Wireless Spectrum Activity.” Computer Communications 212: 353–65.

Oquab, Maxime, Timothée Darcet, Théo Moutakanni, et al. 2023. “DINOv2: Learning Robust Visual Features Without Supervision.” arXiv Preprint arXiv:2304.07193.