zero2robot · Phase 5 · Practitionerch5.2-align · align.py

Chapter 5.2

Why AlignedContrastive Vision-Language Pretraining

By the end you can

  1. Map an image and the WORDS that describe it into ONE shared embedding space with two from-scratch towers — a tiny ViT image tower (re-contained: patch-embed a 64x64 frame into 64 tokens, a CLS token, learned positions, a few pre-norm attention blocks) and a trivial text tower (word-level tokenizer -> embed -> one block -> masked mean-pool) — each projecting to a small L2-normalized vector.
  2. Train the alignment with SYMMETRIC InfoNCE over a batch's image-text cosine matrix (cross-entropy in BOTH directions, image->text AND text->image) plus a LEARNED temperature — using NO labels, only which caption came with which frame. Flag the free-tier reality: contrastive wants a big batch (128-256) for in-batch negatives, which fits here only because the towers are tiny, and these near-blank 64x64 sim frames need background subtraction so the ViT sees the block, not the table.
  3. Measure alignment as RETRIEVAL: type "the block is near the top left corner" and rank held-out frames by cosine. Show the contrastive encoder beats BOTH a ch5.1-style supervised-probe encoder AND a random-init encoder — the DIRECTION (aligned >> random; aligned > supervised on the fine metric), seed-robust, never an exact number.
  4. Diagnose the classic InfoNCE bug the learner writes themselves — dropping the NEGATIVES (pull each matched pair's cosine to 1 but never push non-matches apart) — predict its retrieval, run it (--break noneg), and self-explain why, with nothing repelling non-matches, the embedding space collapses so retrieval stays above chance but is measurably worse. And refute the misconception "contrastive learning needs labels": the negatives come free from PAIRING (every other caption in the batch), and the supervised baseline that DOES use quadrant labels still loses to the label-free contrastive tower.

See it work

live · P2

Why aligned — one instruction, two encoders

Pick an instruction and watch which held-out scenes light up by cosine similarity. The aligned tower (symmetric InfoNCE) returns the scenes that really sit in that corner; the random-init tower returns noise. Same rankings, side by side — that gap is what contrastive pretraining bought, from pairing alone, no labels.

…or choose one below

instruction“the block is near the top left corner”

Aligned encodercontrastive · symmetric InfoNCE
bottom leftbottom righttop lefttop right12345
5 / 5 top matches in the top left corner
Random-init encoderuntrained · no shared space
bottom leftbottom righttop lefttop right12345
1 / 5 top matches in the top left corner
in the named corner (a hit)wrong corner (a miss)the other 244 held-out scenes
Instruction: "the block is near the top left corner". Aligned contrastive encoder: 5 of its top 5 scenes sit in the top left corner — the retrieved scenes cluster there. Random-init encoder: 1 of 5 — the retrieved scenes are scattered across corners. Contrastive pretraining built a shared image-text space from pairing alone; the untrained encoder did not.

Across all four instructions the aligned tower puts 20 / 20 of its top matches in the named corner; the random tower only 6 / 20 — and it returns the same few high-norm scenes for different instructions (frame 189 is in all four of its top-5 lists). The rock is aligned ≫ random: contrastive learned a real shared image-text space from pairing alone, no labels. These are tiny towers over near-blank 64×64 sim frames — a mechanism demo, not a strong perception model; a web-scale CLIP/SigLIP is the reference. Real precomputed rankings from align.py (seed 0, cpu, 249 held-out scenes); poster reads with JS off.

Open in Colabsoon

Free-tier notebook — the button goes live when the course repository is published.

From "what do I see?" to "does this match what I said?"

Chapter 5.1 trained a tiny ViT with a supervised probe: you handed it a label — which quadrant the block sits in — and a linear head learned to read that label off the pixels. It works, and it is the honest way to ask does my encoder see anything? But a label is a thin description. "Top-left" is one of four buckets; the caption "the block is far in the top-left corner" says more, and there is no fixed list of captions to enumerate. This chapter asks the harder question a real robot needs: can an image and the words that describe it land in the same space, so that a sentence you have never turned into a class can still pull up the right picture?

That is contrastive vision-language pretraining — the idea behind CLIP, and behind the vision backbones real VLAs run on (SigLIP inside SmolVLA and OpenVLA). The surprise, and the thing to hold onto, is the supervision: there is none, in the usual sense. We never tell the model what a frame contains. We only tell it which caption came with which frame — the pairing — and let it sort out the rest. Open align.py. Seven regions: setup, data, language (the text tower), vision (the image ViT), contrastive (the loss), eval (retrieval), and report.

The data: cheap captions from sim state

align.py#datasha256:e54913d646…
# ~2-4k (frame, caption) pairs from PushT. We replay the scripted expert in-process (ch1.7# pattern) and render 64x64 frames; the block sweeps the table as the expert pushes it to# center, so the cache spans every quadrant and radius. The CAPTION is generated from the# block's position in SIM STATE — a cheap label, no human annotation (the shared-arc trick):# a quadrant word (2x2 from the sign of tee_xy) and a near/far word (radius, split at the# median). Contrastive pretraining never SEES the quadrant/near-far labels — it only knows# which caption came with which frame; the labels only train the SUPERVISED baseline + SCORE.QUAD_WORDS = ["bottom left", "bottom right", "top left", "top right"]  # index by quadrant id 0..3NF_WORDS = ["near", "far"]                                             # index by the near/far bitTEMPLATES = [                                                          # paraphrases, one picked per frame    "the block is {nf} the {q} corner",    "the tee sits {nf} the {q}",    "push toward the block {nf} the {q}",    "the t piece is {nf} the {q} region",]  def collect_frames(episodes: int, seed: int, stride: int):    """Replay the PushT expert; return (frames uint8 (N,64,64,3), tee_xy (N,2))."""    env = PushTEnv()    frames, tee = [], []    for e in range(episodes):        obs = env.reset(seed + e)        expert = ScriptedExpert(noise=0.0, seed=seed + e)        step, done = 0, False        while not done:            action = expert.action(env)            if step % stride == 0:                       # subsample near-duplicate frames                frames.append(env.render_frame(IMG_HW, IMG_HW))                tee.append(obs[2:4].copy())              # obs idx 2,3 = tee_x, tee_y (PushT obs layout)            obs, _, done, _ = env.step(action)            step += 1    return np.asarray(frames, np.uint8), np.asarray(tee, np.float32)  frames, tee = collect_frames(args.episodes, args.seed, args.frame_stride)num_pairs = len(frames)# Background subtraction: a raw frame is ~98% constant table, so center on the MEAN frame so the# ViT sees the block (a linear probe on raw pixels is at chance; on centered pixels near-perfect).frame_mean = (frames.astype(np.float32) / 127.5 - 1.0).mean(0)   # (64, 64, 3) — subtracted in ImageTowerradius = np.hypot(tee[:, 0], tee[:, 1])radius_split = float(np.median(radius))                  # data-driven near/far boundary (deterministic given frames)quadrant = ((tee[:, 1] >= 0).astype(np.int64) * 2 + (tee[:, 0] >= 0).astype(np.int64))  # 0..3 (top=+y, right=+x)near_far = (radius >= radius_split).astype(np.int64)     # 0=near, 1=farcls = quadrant * 2 + near_far                            # 0..7 fine class (quadrant x near/far)variant = (args.seed + np.arange(num_pairs)) % len(TEMPLATES)   # one paraphrase per frame, seed-deterministiccaptions = [TEMPLATES[variant[i]].format(nf=NF_WORDS[near_far[i]], q=QUAD_WORDS[quadrant[i]])            for i in range(num_pairs)]

We reuse PushT. Replaying the scripted expert sweeps the block across the whole table, so a few thousand cached 64×64 frames span every quadrant and every distance from the center. The caption for each frame is generated from the block's position in the simulator state — a cheap label, no human annotation: a quadrant word (top left, …) from the sign of the block's coordinates, and a near/far word from its radius. A handful of paraphrase templates keeps the wording from being a single fixed string.

Read the boundary carefully, because it is the whole point of the chapter: the contrastive model never sees these quadrant or near/far labels. It only ever sees this frame goes with this sentence. The labels exist for two jobs that stand outside the contrastive learner — to train the supervised baseline we will race against, and to score retrieval at the end. Pairing is not labeling, and the difference is the lesson.

Two towers into one space

An image and a sentence are different kinds of object; to compare them we need to map both into a single vector space where "close" means "these belong together." That takes two encoders — towers — that meet in the middle.

The text tower is deliberately trivial (the interesting one is the image ViT):

align.py#languagesha256:4260f4c5c8…
# The whole tokenizer, from scratch (re-contained from ch1.7 — chapters do not import each# other): a fixed WORD-LEVEL vocab (no BPE, no HF tokenizers), 4 special ids, pad/truncate.# The corpus is CLOSED, so the vocab is fixed (no OOV). A real CLIP/SigLIP uses a 30k+ subword# tokenizer; here the point is the mechanism (text -> ids -> a learned embedding), laid bare.class Tokenizer:    def __init__(self, corpus: list[str]) -> None:        words = sorted({w for text in corpus for w in text.split()})        self.itos = ["<pad>", "<unk>", "<bos>", "<eos>"] + words   # ids 0..3 are the specials        self.stoi = {w: i for i, w in enumerate(self.itos)}     @property    def vocab_size(self) -> int:        return len(self.itos)     def encode(self, text: str) -> np.ndarray:        ids = [self.stoi["<bos>"]] + [self.stoi.get(w, 1) for w in text.split()] + [self.stoi["<eos>"]]        ids = ids[:MAX_TOKENS] + [PAD_ID] * (MAX_TOKENS - len(ids))        return np.asarray(ids[:MAX_TOKENS], dtype=np.int64)  # The vocab is the closure of the templates over every (quadrant, near/far) fill.corpus = [t.format(nf=nf, q=q) for t in TEMPLATES for nf in NF_WORDS for q in QUAD_WORDS]tokenizer = Tokenizer(corpus)tokens = np.stack([tokenizer.encode(c) for c in captions])       # (N, MAX_TOKENS) int64  class TextTower(nn.Module):    """The TRIVIAL tower: embed the token ids, run ONE attention block so words can talk,    masked-mean-pool to one vector, and project into the shared space. Kept deliberately    small — the interesting tower is the image ViT; text just needs to name the scene."""     def __init__(self, vocab_size: int, dim: int, heads: int, embed_dim: int) -> None:        super().__init__()        self.embed = nn.Embedding(vocab_size, dim, padding_idx=PAD_ID)        self.pos = nn.Parameter(0.02 * torch.randn(1, MAX_TOKENS, dim))        self.block = Block(dim, heads)        self.norm = nn.LayerNorm(dim)        self.proj = nn.Linear(dim, embed_dim)     def forward(self, tok: torch.Tensor) -> torch.Tensor:        pad = tok == PAD_ID                                      # (B, L) True where padded        x = self.block(self.embed(tok) + self.pos, pad)        keep = (~pad).float()[:, :, None]                        # mask pads OUT of the mean        pooled = (x * keep).sum(1) / keep.sum(1).clamp(min=1.0)        return self.proj(self.norm(pooled))                      # (B, embed_dim), pre-L2-norm

We re-contain ch1.7's word-level tokenizer — string in, fixed-length id array out, no transformers, no BPE — then embed the ids, run one attention block so the words can talk, masked-mean-pool to a single vector, and project. That is all a caption this simple needs. A real CLIP text encoder is a full Transformer over a 30,000-token subword vocab; ours is a couple dozen words. The mechanism is the same; the scale is not.

The image tower is the tiny ViT from 5.1, re-derived here (chapters re-contain their backbones — we do not import 5.1):

align.py#visionsha256:45d26d533e…
# The image tower is a tiny ViT — the SAME shape ch5.1 built and the SAME attention block# ch1.8 uses, re-derived here (chapters re-contain their backbone; do NOT import ch5.1).class Block(nn.Module):    """One pre-norm transformer block: multi-head self-attention (Q/K/V from nn.Linear +    a scaled-dot-product softmax) then a per-token MLP. key_pad masks padded text keys;    the image tower has no padding and passes None."""     def __init__(self, dim: int, heads: int) -> None:        super().__init__()        self.heads = heads        self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim)        self.qkv = nn.Linear(dim, 3 * dim)        self.proj = nn.Linear(dim, dim)        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))     def forward(self, x: torch.Tensor, key_pad: torch.Tensor | None) -> torch.Tensor:        B, L, dim = x.shape        h, hd = self.heads, dim // self.heads        qkv = self.qkv(self.ln1(x)).reshape(B, L, 3, h, hd).permute(2, 0, 3, 1, 4)        q, k, v = qkv[0], qkv[1], qkv[2]                          # each (B, h, L, hd)        scores = (q @ k.transpose(-2, -1)) / math.sqrt(hd)       # (B, h, L, L)        if key_pad is not None:            scores = scores.masked_fill(key_pad[:, None, None, :], float("-inf"))        x = x + self.proj((scores.softmax(-1) @ v).transpose(1, 2).reshape(B, L, dim))        return x + self.mlp(self.ln2(x))  class ImageTower(nn.Module):    """Patch-embed a 64x64 frame into 8x8 = 64 tokens, prepend a CLS token, add learned    positions, run the attention blocks, read the fused scene off CLS, project to the    shared space. This IS ch5.1's ViT; here its CLS output must MATCH language, not a label."""     def __init__(self, dim: int, depth: int, heads: int, embed_dim: int, frame_mean: np.ndarray) -> None:        super().__init__()        n_tok = (IMG_HW // PATCH) ** 2                            # 64 patches        self.patch = nn.Conv2d(3, dim, PATCH, stride=PATCH)      # (B,3,64,64) -> (B,dim,8,8)        self.cls = nn.Parameter(torch.zeros(1, 1, dim))        self.pos = nn.Parameter(0.02 * torch.randn(1, n_tok + 1, dim))        self.blocks = nn.ModuleList([Block(dim, heads) for _ in range(depth)])        self.norm = nn.LayerNorm(dim)        self.proj = nn.Linear(dim, embed_dim)        self.register_buffer("frame_mean", torch.from_numpy(frame_mean))  # background, subtracted below     def forward(self, images_uint8: torch.Tensor) -> torch.Tensor:        x = images_uint8.to(torch.float32) / 127.5 - 1.0         # (B,64,64,3) in [-1,1]        x = (x - self.frame_mean).permute(0, 3, 1, 2)            # center on background, then (B,3,64,64)        x = self.patch(x).flatten(2).transpose(1, 2)             # (B, 64, dim)        x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], dim=1) + self.pos        for blk in self.blocks:            x = blk(x, None)        return self.proj(self.norm(x[:, 0]))                     # (B, embed_dim), pre-L2-norm

Patch-embed the frame into 8×8 = 64 tokens, prepend a CLS token, add learned positions, run a few pre-norm attention blocks, read the fused scene off CLS, and project. One detail earns its line and is easy to skip: background subtraction. A 64×64 top-down PushT frame is ~98% constant table — the block is a handful of pixels. Feed that raw to a tiny ViT and the CLS token is dominated by the blank background; every frame maps to nearly the same point and nothing learns. (Measure it: a linear probe on raw pixels sits at chance; on background-subtracted pixels it is near-perfect — the signal was always there, drowned by the constant.) So we subtract the mean frame first. It is not a trick; it is what makes a tiny encoder on near-blank sim frames learnable at all, and worth remembering the next time a "working" architecture refuses to train on flat inputs.

The loss: pull the pair, push everything else

Cross-entropy down the rows (each image picks its caption) plus cross-entropy across the columns (each caption picks its frame), over the $B\times B$ matrix of cosines scaled by a learned temperature $\tau$ — and here it is in code:

L=12Bi=1B[logexp(sii/τ)jexp(sij/τ)+logexp(sii/τ)jexp(sji/τ)],sij=ziimgzjtxt\mathcal{L} = -\frac{1}{2B}\sum_{i=1}^{B}\left[ \log\frac{\exp(s_{ii}/\tau)}{\sum_{j}\exp(s_{ij}/\tau)} + \log\frac{\exp(s_{ii}/\tau)}{\sum_{j}\exp(s_{ji}/\tau)} \right], \qquad s_{ij} = z^{\mathrm{img}}_i \cdot z^{\mathrm{txt}}_j
align.py#contrastivesha256:00cd5df1a3…
class Aligner(nn.Module):    """Two towers + a LEARNED temperature. encode_* returns L2-normalized embeddings so a    dot product IS a cosine. logit_scale is stored in log space and init to CLIP's 1/0.07. The    SUPERVISED baseline builds one with embed_dim=NUM_QUAD, so its image tower IS the classifier."""     def __init__(self, vocab_size: int, dim: int, depth: int, heads: int, embed_dim: int,                 frame_mean: np.ndarray) -> None:        super().__init__()        self.image = ImageTower(dim, depth, heads, embed_dim, frame_mean)        self.text = TextTower(vocab_size, dim, heads, embed_dim)        self.logit_scale = nn.Parameter(torch.tensor(math.log(1.0 / 0.07)))     def encode_image(self, imgs: torch.Tensor) -> torch.Tensor:        return F.normalize(self.image(imgs), dim=-1)      # onto the unit sphere: a dot product IS a cosine     def encode_text(self, tok: torch.Tensor) -> torch.Tensor:        return F.normalize(self.text(tok), dim=-1)  def info_nce(img_e: torch.Tensor, txt_e: torch.Tensor, logit_scale: torch.Tensor,             negatives: bool) -> torch.Tensor:    """Symmetric InfoNCE over a batch's B x B cosine matrix. The B matched pairs are the    positives (the diagonal); every OTHER caption in the batch is a NEGATIVE — that is why a    bigger batch is a harder, better lesson, and why contrastive needs no labels (the negatives    come free from PAIRING). Cross-entropy is applied in BOTH directions — image->text AND    text->image — the symmetric pull CLIP uses, scaled by the learned temperature.    THE BUG (--break noneg): keep the positive pull but DROP the negatives (just push each pair's    cosine to 1). With nothing pushing non-matches apart the space collapses and retrieval,    though still above chance, is measurably worse — negatives are the whole point."""    if not negatives:        return (1.0 - (img_e * txt_e).sum(-1)).mean()        # positives only -> collapse (the deliberate failure)    scale = logit_scale.exp().clamp(max=100.0)               # clamp like CLIP so temperature can't blow up    logits = scale * img_e @ txt_e.t()                       # (B, B): logits[i,j] = cos(image_i, text_j) / temp    labels = torch.arange(len(logits), device=logits.device)  # the positive for row i is column i    return 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.t(), labels))  def train_contrastive(model: Aligner, imgs: torch.Tensor, tok: torch.Tensor, epochs: int,                      negatives: bool, freeze_image: bool, tag: str) -> float:    if freeze_image:                                         # the supervised baseline aligns text to a FROZEN image tower        for p in model.image.parameters():            p.requires_grad_(False)    opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=args.lr)    loss_val = float("nan")    for epoch in range(epochs):        total, nb = 0.0, 0        for batch in torch.randperm(len(imgs), generator=gen).split(args.batch_size):            loss = info_nce(model.encode_image(imgs[batch]), model.encode_text(tok[batch]),                            model.logit_scale, negatives)            opt.zero_grad()            loss.backward()            opt.step()            total, nb = total + loss.item(), nb + 1        loss_val = total / nb        if args.rerun:            rr.set_time("epoch", sequence=epoch)            rr.log(f"contrastive/{tag}/loss", rr.Scalars([loss_val]))            rr.log(f"contrastive/{tag}/temperature", rr.Scalars([1.0 / model.logit_scale.exp().clamp(max=100.0).item()]))    return loss_val  def supervised_pretrain(model: Aligner, imgs: torch.Tensor, quad: torch.Tensor, epochs: int) -> None:    """ch5.1's recipe, re-contained: train the image tower to CLASSIFY the block quadrant.    With embed_dim=NUM_QUAD the tower's output IS the 4-way logits — no extra head. We freeze    it and align text to it; it recovers the quadrant but LOSES the near/far half of the    caption the contrastive tower keeps. Supervised != aligned."""    opt = torch.optim.Adam(model.image.parameters(), lr=args.lr)    for _ in range(epochs):        for batch in torch.randperm(len(imgs), generator=gen).split(args.batch_size):            loss = F.cross_entropy(model.image(imgs[batch]), quad[batch])            opt.zero_grad()            loss.backward()            opt.step()

Here is the engine. Take a batch of B (frame, caption) pairs, L2-normalize every embedding onto the unit sphere so a dot product is a cosine, and form the B×B matrix of image-to-caption cosines. If the space were perfect, that matrix would be bright on its diagonal — each frame matches its caption — and dark everywhere else. Symmetric InfoNCE makes it so: a cross-entropy down the rows (each image should pick its caption out of all B) plus a cross-entropy across the columns (each caption should pick its frame), scaled by a learned temperature. Two forces, every step: pull each matched pair together, push every mismatched pair apart.

Where do the "mismatched pairs" come from? For free — every other caption in the batch is a negative for your frame. That is why contrastive learning wants a big batch (128–256): more captions in the batch means more negatives, a harder and more informative lesson. It fits the free-tier floor here only because the towers are tiny — flag that as a real constraint, not a detail. And it is why the "no labels" claim holds: the negatives are handed to you by the pairing, not by any annotation.

Alignment is a number: retrieval

A shared space is a claim; retrieval is the measurement that makes it honest.

align.py#evalsha256:c8a0888a00…
# ALIGNMENT = RETRIEVAL. Split the cache into disjoint held-out GALLERY (images) and QUERY# (captions) sets so no caption can trivially retrieve its own frame. For each query caption,# rank gallery frames by cosine, take the top-1: is it the RIGHT class? retrieval@1 (fine,# 8-way) demands quadrant AND near/far; (quad, 4-way) is coarser. Chance ~1/8 and ~1/4.@torch.no_grad()def retrieval(model: Aligner, gal_imgs: torch.Tensor, qry_tok: torch.Tensor,              gal_cls: torch.Tensor, qry_cls: torch.Tensor) -> tuple[float, float]:    img_e = model.encode_image(gal_imgs)                     # (G, D)    top1 = (model.encode_text(qry_tok) @ img_e.t()).argmax(dim=1)   # best-matching frame per query    fine = (gal_cls[top1] == qry_cls).float().mean().item()    quad = ((gal_cls[top1] // 2) == (qry_cls // 2)).float().mean().item()    return fine, quad  # Held-out split: shuffle (seeded), last 30% held-out, halved into disjoint gallery / query;# the rest (70%) trains the towers. No held-out frame is ever trained on.perm = rng.permutation(num_pairs)n_hold = max(4, int(0.30 * num_pairs))train_idx, hold = perm[:-n_hold], perm[-n_hold:]gal_idx, qry_idx = hold[: len(hold) // 2], hold[len(hold) // 2:]imgs_t = torch.from_numpy(frames).to(device)tok_t = torch.from_numpy(tokens).to(device)quad_t = torch.from_numpy(quadrant).to(device)cls_t = torch.from_numpy(cls).to(device)tr_imgs, tr_tok, tr_quad = imgs_t[train_idx], tok_t[train_idx], quad_t[train_idx]gal_imgs, gal_cls = imgs_t[gal_idx], cls_t[gal_idx]qry_tok, qry_cls = tok_t[qry_idx], cls_t[qry_idx]print(f"pairs: {num_pairs} ({len(train_idx)} train / {len(gal_idx)} gallery / {len(qry_idx)} query), "      f"vocab {tokenizer.vocab_size}, radius_split {radius_split:.3f}, break={args.break_mode or 'none'}") # THREE encoders, one retrieval task. (1) ALIGNED: both towers, symmetric InfoNCE.negatives = args.break_mode != "noneg"torch.manual_seed(args.seed)aligned = Aligner(tokenizer.vocab_size, args.dim, args.depth, args.heads, args.embed_dim, frame_mean).to(device)loss_aligned = train_contrastive(aligned, tr_imgs, tr_tok, args.epochs, negatives, False, "aligned")at1_aligned, quad_aligned = retrieval(aligned, gal_imgs, qry_tok, gal_cls, qry_cls) # (2) RANDOM: both towers random-init, NEVER trained — the floor. Retrieval ~ chance.torch.manual_seed(args.seed + 1)random_model = Aligner(tokenizer.vocab_size, args.dim, args.depth, args.heads, args.embed_dim, frame_mean).to(device)at1_random, quad_random = retrieval(random_model, gal_imgs, qry_tok, gal_cls, qry_cls) # (3) SUPERVISED: ch5.1's probe. embed_dim=NUM_QUAD makes the image tower the 4-way quadrant# classifier; pretrain it on LABELS, freeze it, align text to it. Uses labels the contrastive# tower never saw — and still loses, because a label-shaped space can't hold the whole caption.torch.manual_seed(args.seed + 2)supervised = Aligner(tokenizer.vocab_size, args.dim, args.depth, args.heads, NUM_QUAD, frame_mean).to(device)supervised_pretrain(supervised, tr_imgs, tr_quad, args.sup_epochs)train_contrastive(supervised, tr_imgs, tr_tok, args.epochs, True, True, "supervised")at1_supervised, quad_supervised = retrieval(supervised, gal_imgs, qry_tok, gal_cls, qry_cls)temp_final = 1.0 / aligned.logit_scale.exp().clamp(max=100.0).item()  # the learned temperatureprint(f"retrieval@1 a/s/r  fine {at1_aligned:.3f}/{at1_supervised:.3f}/{at1_random:.3f}  "      f"quad {quad_aligned:.3f}/{quad_supervised:.3f}/{quad_random:.3f}")

Split the cache into a disjoint held-out gallery (images) and query (captions). For each query — "the block is near the top left corner" — rank every gallery frame by cosine and take the top one: is it actually a top-left, near frame? Average over queries and you have retrieval@1. We score it two ways: fine (the retrieved frame must match on quadrant and near/far — the whole caption) and quad (quadrant only — the coarser bar). Chance is ~1/8 and ~1/4.

Then we run the same measurement over three encoders:

  • aligned — the contrastive towers above.
  • random — the same two towers, never trained. The floor.
  • supervised — 5.1's probe: a ViT trained on the 4-way quadrant label, then frozen and aligned to text. This is the "contrastive needs labels?" foil — it uses them.

The result (measured, seed-robust in direction — never trust the exact number on rendered frames across machines): aligned ≫ random by a wide margin — contrastive built a real shared space from pairing alone, while random init retrieves noise. That is the rock. On the finer question, aligned > supervised: both nail the quadrant (on the quad score they roughly tie — the probe was trained on quadrant, so say so), but the supervised encoder compressed the scene to a 4-way label and lost the near/far half of the caption. Contrastive kept the whole sentence. The encoder that saw no labels wins on the richer question — which is the misconception, dead on the page.

Misconception: "contrastive learning needs labels." It does not. The only signal is which caption rode along with which frame. Our supervised baseline literally does use quadrant labels and still loses. Negatives — the thing that makes contrastive work — come from the batch, not from an annotator.

Break it: forget the negatives

You will write this loss yourself, and there is one bug almost everyone writes on the first try. Contrastive learning is pull and push. It is tempting to write only the pull: maximize the cosine of each matched pair and stop. --break noneg is exactly that — drop the negatives, keep the positive term:

# correct: cross-entropy over the B x B matrix  (pull the diagonal, push the rest)
# noneg:   (1.0 - (img_e * txt_e).sum(-1)).mean()   # pull only — no push

Predict before you run (exercise 1). Then measure: retrieval@1 falls sharply — in the reference run from ~0.83 to ~0.38 — yet stays above random. (Those figures are a provisional draft, measured at a reduced config — episodes 70, epochs 25, not the default 80/30 — so read them for their direction, not the exact percentages.) That "plausible but worse" signature is the tell. Ask yourself what stops the buggy model from mapping every frame and every caption to the same vector: the loss would be zero and retrieval useless. Nothing stops it — that is representational collapse, and it is what the negatives were preventing. (Retrieval clings above chance only because the tiny surviving structure from initialization and the shared background subtraction is not fully washed out.) Two forces, not one; the push is not optional.

Read the real thing

Open openai/CLIP, model.py. You will recognize it immediately: two towers projecting to a shared, L2-normalized space; logits_per_image and logits_per_text; two cross-entropies against arange labels; and a learned logit_scale — the same clamped log-temperature we use, exp()'d. Everything you built here is there at scale: a big ViT and a real Transformer text encoder instead of our two toys, hundreds of millions of image-text pairs instead of a few thousand PushT frames, but the objective is line-for-line ours.

Then look at what modern VLAs actually ship: google-research/big_vision, models/proj/image_text/two_towers.py. SigLIP swaps the softmax InfoNCE for a per-pair sigmoid loss. The reason is exactly the batch constraint we hit: a softmax over in-batch negatives couples the whole batch, while the sigmoid scores each pair independently, which decouples the loss and scales to enormous batches far better. That is why SmolVLA's vision backbone is a SigLIP model — and why the "big batch" note in this chapter is not a toy detail but the thing the next design decision turns on.

What's next

You now have an encoder whose space language can point into. That is the missing piece under a real VLA: 5.1 asked whether an encoder sees, 5.2 aligned what it sees to what you say, and the perception arc closes (5.3) by putting these features to work — a pretrained, language-aligned backbone is precisely what lets a policy follow an instruction it was never explicitly trained on. The from-scratch towers here are a mechanism demo on near-blank frames; the payoff is understanding, cheaply and completely, the objective that a web-scale CLIP/SigLIP spends millions of pairs to earn.

Practice

practice · candidate exercisesdrafted by the exercise generator, pending human promotion. Answers reveal only after you predict — honor system.

  1. predict-then-run

    Exercise 1

    SUGGESTED exercise candidate (humans promote) — predict-then-run, ch5.2.

    THE LEARNER-GENERATED FAILURE. Contrastive learning has two moves in every batch: PULL each matched (frame, caption) pair together, and PUSH every non-matching pair apart. The push comes free — every OTHER caption in the batch is a negative for your frame. A very common bug when you write InfoNCE yourself is to keep the pull and forget the push: you maximize the cosine of each matched pair and never repel the negatives. align.py's --break noneg is exactly that buggy loss:

    # correct: cross-entropy over the B x B cosine matrix (pull the diagonal, push the rest)
    # buggy:   (1.0 - (img_e * txt_e).sum(-1)).mean()   <- pull only, NO negatives
    

    PREDICT before you run: you train the aligned encoder twice at a reduced config — once with the correct symmetric InfoNCE, once with --break noneg — and compare retrieval@1 (fine). Which row is right?

    • A) noneg >= correct — pulling the pairs together is all retrieval needs.

    • B) correct much higher; noneg lower but still ABOVE random — without negatives the space partially COLLAPSES (nothing keeps different scenes apart), so retrieval still beats chance but is measurably worse.

    • C) noneg near 0 (below random) — forgetting negatives makes it worse than an untrained encoder.

    Record your answer in PREDICTION, then run (this TRAINS twice at a reduced scale — a few minutes on CPU). Estimated learner time: 15 minutes.

    Predict, then commit

    Pick the outcome you expect from the options above. The answer and the local run command reveal only after you commit — predicting after you know teaches nothing.

  2. predict-then-run

    Exercise 2

    SUGGESTED exercise candidate (humans promote) — predict-then-run, ch5.2.

    Objective tested: WHY contrastive, not just supervised. align.py trains three encoders and scores each by retrieval: type an instruction, rank held-out frames by cosine.

    • ALIGNED: contrastive InfoNCE — no labels, only which caption rode with which frame.
    • SUPERVISED: ch5.1's probe — a ViT trained on the 4-way QUADRANT label, then frozen and aligned to text (this IS the "contrastive needs labels?" foil — it uses them).
    • RANDOM: untrained init — the floor. Captions name BOTH the quadrant AND near/far ("the block is FAR the top left corner"). The retrieval@1 (FINE) score demands the retrieved frame match on quadrant AND near/far.

    PREDICT before you run: how do the three FINE retrieval@1 scores order?

    • A) aligned > supervised > random — contrastive keeps the WHOLE caption; the supervised probe only kept its 4-way quadrant label, so it loses the near/far half; random is chance.

    • B) supervised > aligned > random — labels beat no-labels; supervision always wins.

    • C) aligned ~ supervised >> random — both learn the full caption equally well.

    (Hint: also print the QUAD-only score. On quadrant alone the supervised probe ~ ties aligned — it was TRAINED on the quadrant. The gap opens on the FINE score, where the caption says more than the label ever did.) Record PREDICTION, then run (TRAINS three encoders at a reduced config — a few minutes on CPU). Estimated learner time: 15 minutes.

    Predict, then commit

    Pick the outcome you expect from the options above. The answer and the local run command reveal only after you commit — predicting after you know teaches nothing.

  3. code-completion

    Exercise 3

    SUGGESTED exercise candidate (humans promote) — code-completion, ch5.2.

    Implement the heart of the chapter: SYMMETRIC InfoNCE over a batch's image-text cosine matrix. You are given a batch of L2-normalized image embeddings and text embeddings and a learned log-temperature. Build the B x B matrix of scaled cosines, then average the cross-entropy in BOTH directions — image->text (each row's correct caption) AND text->image (each column's correct frame) — against the arange labels.

    This is the CLIP training objective (openai/CLIP model.py: logits_per_image / logits_per_text, two cross-entropies). The self-contained check compares your loss to a reference on a fixed fixture, within meta.yaml's abs_tol.

    Fill in symmetric_info_nce. Do NOT change the signature.

    Run it locally:

    pytest curriculum/phase5_practitioner/ch5.2_align/exercises/suggested/checks.py -k ex3
  4. code-completion

    Exercise 4

    SUGGESTED exercise candidate (humans promote) — code-completion, ch5.2.

    Implement the MEASUREMENT: retrieval@1. Given text-query embeddings and a gallery of image embeddings (both L2-normalized, so a dot product is a cosine), rank the gallery for each query, take the top-1, and report the fraction whose CLASS matches the query's class. This is exactly how align.py turns "a shared embedding space" into a number you can trust.

    The self-contained check compares your accuracy to a reference on a fixed fixture, within meta.yaml's abs_tol. Fill in retrieval_at1. Do NOT change the signature.

    Run it locally:

    pytest curriculum/phase5_practitioner/ch5.2_align/exercises/suggested/checks.py -k ex4
wall-clock · rendered from wallclock.csvone source · every tier
cpu-laptopexpected wall-clock on cpu-laptop: ~0.69 min (measured)measured
mpswall-clock on mps: not yet measuredpending
t4wall-clock on t4: not yet measuredpending
4090wall-clock on 4090: not yet measuredpending

Colophon · provenance

The code on this page is not pasted — each panel is included by region straight fromalign.py, and these fingerprints are its sha256, the same ones check_prose_code_drift re-checks on every PR. Edit a shown region without re-rendering and CI turns red.

#setup
570f9e328c6b11374d0ea6cfeb32ab86aa022c254114ae1a33ac7f4e6ef8e91c
#data
e54913d6461f1f2ac71c7ffed802b8254448284c7c17ea72a97233ba936df243
#language
4260f4c5c8eea9a55e06455c02ee61a5c062d9f33dd1d2ff0f84cc9f55df6157
#vision
45d26d533eee1211be32f9412cfd2fc86bd87603f93f293e8646bf5e0d508342
#contrastive
00cd5df1a3492e9dc3dbfed8f609a7e580e996d6e69eb44e04f50e9d87560c36
#eval
c8a0888a0020d22de913c4fe34fd03b6106f8c553d0fad13680eebd2d3f7fdca
#report
f555c970662e8da15b497dbcd9670ccf7b13cf3c7df34b269f26cf3bdb736e4c