Module 20: Multimodal

Introduction

Every module so far ends at a model that only ever reads token IDs: tokenize (m03), embed (m04), attend (m05), predict the next token (m06–m08). The whole book is text. This module opens the second modality — pixels — and answers how one model can relate a picture to the words that describe it.

A multimodal model processes more than one kind of input (here: images and text) in a shared representation. There are only two genuinely new ideas, and everything else is the transformer you already built:

  • Pixels become tokens. Cut an image into a grid of fixed-size patches, flatten each, and project it into the same vector space token embeddings live in. A patch is a “visual word.” This is a Vision Transformer (ViT).
  • A shared space. Encode an image and a caption to two vectors and train so matched pairs sit close and mismatched pairs far — the CLIP contrastive objective. This is what makes “a photo of a cat” land next to 🐱.

Why it matters for LLMs:

  • Modern frontier models (GPT-4o, Gemini, Claude, Llama 4) are natively multimodal — they see images by turning them into token-like patch embeddings.
  • CLIP-style alignment is the backbone of text-to-image search, zero-shot image classification, and the vision encoders bolted onto today’s chat models.

What You’ll Learn

After this module, you can:

  • Patchify an image into a sequence a transformer can read, and build a ViT patch-embedding layer (patch projection + positional embedding + CLS token).
  • Explain why images and text must be L2-normalized into a shared space, and build the CLIP contrastive loss — a symmetric cross-entropy over the batch similarity matrix — from scratch.
  • Measure alignment with image→text retrieval accuracy, and watch a tiny CLIP pull two modalities together until the similarity matrix’s diagonal lights up.
  • Reason about the temperature (logit scale) and the batch-as-negatives trick, and name the pitfalls (skipping normalization, one-sided loss, unclamped τ).

Prerequisites

This module requires familiarity with:

  • Module 04: Embeddings — patch embeddings are token embeddings for images; the positional-embedding idea carries over directly.
  • Module 05: Attention — a ViT is a transformer encoder over the patch sequence.
  • Module 07: Training — the contrastive loss is a cross-entropy, optimized exactly as you trained the GPT.
  • Module 17: Evaluation — retrieval accuracy is the metric that says whether the two modalities actually aligned.

Intuition: Patches as Visual Words

A language model never sees letters — it sees a sequence of token vectors. To feed an image to the same machinery, we need to turn it into a sequence of vectors too. The Vision Transformer’s answer is disarmingly simple: chop the image into a grid of small square patches, and treat each patch as one token.

A 32×32 image with 16×16 patches becomes a 2×2 grid — four “visual words.” Each patch (a little block of pixels) is flattened into a vector and linearly projected to the model’s embedding dimension d, exactly like a token ID is looked up into a d-dim embedding in m04. Add a positional embedding so the model knows where each patch sat (patches have no inherent order once flattened), prepend a learnable [CLS] token whose final vector summarizes the whole image, and you have a sequence the transformer reads with no other changes.

Walk the pipeline one step at a time:

The Math: Patchify & Patch Embeddings

Let an image be a tensor of shape \((C, H, W)\) — channels, height, width. With a patch size \(P\) that divides \(H\) and \(W\), patchify reshapes it into

\[ (C, H, W) \;\longrightarrow\; (N,\; P^2 C), \qquad N = \frac{H}{P}\cdot\frac{W}{P}, \]

a sequence of \(N\) patch vectors, each of length \(P^2 C\). A linear layer with weight \(W_e \in \mathbb{R}^{(P^2 C)\times d}\) projects every patch to the embedding size, we add a learned positional embedding \(E_{\text{pos}}\), and prepend a learned CLS token \(x_{\text{cls}}\):

\[ Z = \big[\, x_{\text{cls}};\; X W_e \,\big] + E_{\text{pos}} \;\in\; \mathbb{R}^{(N+1)\times d}. \]

That is the entire ViT input layer — and it is line-for-line the same shape story as m04’s token embeddings, only the “lookup” is a linear projection of pixels. From \(Z\) onward it is a standard transformer encoder (m05/m06).

Code: Patch Embeddings from Scratch

The implementation lives in multimodal.py. patchify does the reshape (it is a pure tensor permutation — no learned parameters, and lossless):

import torch
from multimodal import patchify, num_patches

image = torch.randn(3, 32, 32)          # one RGB image
patches = patchify(image, patch_size=16)
print(f"image {tuple(image.shape)}  ->  patches {tuple(patches.shape)}")
print(f"N = {num_patches(32, 32, 16)} patches, each {patches.shape[1]} numbers (16·16·3)")
image (3, 32, 32)  ->  patches (4, 768)
N = 4 patches, each 768 numbers (16·16·3)

PatchEmbedding wraps the projection, the positional embedding, and the CLS token into the layer a ViT starts with:

from multimodal import PatchEmbedding

patch_embed = PatchEmbedding(image_size=32, patch_size=16, in_channels=3, embed_dim=64)
batch = torch.randn(4, 3, 32, 32)       # 4 images
tokens = patch_embed(batch)
print(f"patch tokens: {tuple(tokens.shape)}   (4 images, 4 patches + 1 CLS, d=64)")

image_vector = patch_embed.pool(tokens) # the CLS row summarizes each image
print(f"pooled image vector: {tuple(image_vector.shape)}")
patch tokens: (4, 5, 64)   (4 images, 4 patches + 1 CLS, d=64)
pooled image vector: (4, 64)
NoteKey Insight

A patch embedding is a token embedding. Once an image is a sequence of d-dim vectors, nothing downstream knows or cares that it came from pixels — the same attention, the same blocks, the same training. Multimodality is mostly a question of getting every modality into one shared sequence of vectors.

Intuition: One Space for Two Modalities

Turning an image into tokens lets a transformer process it, but it does not yet connect the image to language. For that we need image vectors and text vectors to live in one shared space, where closeness means “these describe the same thing.”

CLIP builds that space with a beautifully simple training signal. Take a big batch of \((image, caption)\) pairs. Encode all images to vectors and all captions to vectors. Now form every possible pairing — an \(N \times N\) grid of similarities. The matched pairs are exactly the diagonal; every off-diagonal cell is a mismatch. Train to make the diagonal win: pull each image toward its own caption and push it away from the other \(N-1\) captions in the batch, which serve as free negatives.

The Math: Contrastive Alignment (CLIP)

Let \(I \in \mathbb{R}^{N\times d}\) be the batch of image embeddings and \(T \in \mathbb{R}^{N\times d}\) the paired text embeddings. First L2-normalize each row, \(\hat I_i = I_i / \lVert I_i\rVert\), so a dot product is a cosine similarity. The scaled similarity matrix is

\[ \text{logits} = s\,\hat I \hat T^{\top} \;\in\; \mathbb{R}^{N\times N}, \qquad s = \exp(t), \]

where \(t\) is a learned temperature (CLIP stores it in log space and clamps it). Row \(i\) scores image \(i\) against all \(N\) captions; the correct label is \(i\) (the diagonal). So it is just cross-entropy — done in both directions and averaged:

\[ \mathcal{L} = \tfrac12\Big( \underbrace{\text{CE}(\text{logits},\, y)}_{\text{image}\to\text{text}} + \underbrace{\text{CE}(\text{logits}^{\top},\, y)}_{\text{text}\to\text{image}} \Big), \qquad y = (0, 1, \dots, N-1). \]

The symmetry matters: the image→text term makes each image find its caption; the text→image term makes each caption find its image. Drop one and retrieval in that direction suffers.

Code: CLIP from Scratch

All of it is a handful of lines in multimodal.py. First the normalize-then-score step, and the fact that a matched batch wins its diagonal:

import torch
from multimodal import l2_normalize, clip_logits, clip_loss, contrastive_accuracy

torch.manual_seed(0)
img = torch.randn(6, 32)                 # 6 image vectors
txt = img + 0.05 * torch.randn(6, 32)    # captions that (roughly) match

print("norms before:", l2_normalize(img).norm(dim=-1)[:3].tolist(), "(all 1.0)")
logits = clip_logits(img, txt, logit_scale=10.0)
print("argmax of each row:", logits.argmax(dim=1).tolist(), "(diagonal 0..5 wins)")
norms before: [1.0, 0.9999999403953552, 0.9999999403953552] (all 1.0)
argmax of each row: [0, 1, 2, 3, 4, 5] (diagonal 0..5 wins)

The symmetric loss and the retrieval metric:

matched   = clip_loss(img, txt)                 # paired correctly
scrambled = clip_loss(img, txt.flip(0))         # every pair wrong
print(f"loss  matched={matched:.3f}   scrambled={scrambled:.3f}")
print(f"retrieval accuracy (matched):   {contrastive_accuracy(img, txt):.2f}")
print(f"retrieval accuracy (scrambled): {contrastive_accuracy(img, txt.flip(0)):.2f}")
loss  matched=0.000   scrambled=16.693
retrieval accuracy (matched):   1.00
retrieval accuracy (scrambled): 0.00

CLIPModel packages two projection encoders and the learned temperature. (Here the encoders are small linear heads over feature vectors so the demo is fast; in a real CLIP the image encoder is the ViT you built above and the text encoder is a Transformer from m06 — the contrastive machinery is identical.)

from multimodal import CLIPModel

model = CLIPModel(image_feat_dim=32, text_feat_dim=32, embed_dim=16)
print(f"learned temperature (logit scale): {model.logit_scale():.2f}  (CLIP init = 1/0.07)")
print(f"similarity matrix shape: {tuple(model(img, txt).shape)}")
learned temperature (logit scale): 14.29  (CLIP init = 1/0.07)
similarity matrix shape: (6, 6)

Watch Two Modalities Align

The payoff: train the tiny CLIP and see the shared space form. demonstrate_clip makes synthetic pairs — each concept gets an image-feature view and a text-feature view — that start unrelated, then minimizes the contrastive loss. The similarity matrix begins as noise and its diagonal lights up as matched pairs are pulled together; retrieval accuracy climbs from chance toward 1.0.

from multimodal import demonstrate_clip

result = demonstrate_clip(n_concepts=12, steps=300, seed=0)
============================================================
CLIP: aligning two modalities (synthetic)
============================================================
  N=12 pairs, embed_dim=12, 300 steps

    step      loss   top-1 acc
       0    7.9794       0.000
      75    0.0031       1.000
     150    0.0011       1.000
     225    0.0005       1.000
     299    0.0003       1.000

  Retrieval accuracy: 0.000 -> 1.000  (chance = 0.083)

Drag the slider to step through training. Watch the muddy grid resolve into a bright diagonal — that diagonal is the alignment:

TipTry This
  1. Fewer steps: rerun demonstrate_clip(steps=20) — the diagonal only half-forms, and accuracy stalls below 1.0. Alignment needs enough gradient steps.
  2. Harder batch: bump n_concepts to 32. More in-batch negatives make each classification harder — the same reason CLIP trained with enormous batch sizes.
  3. More noise: raise noise=1.5. The two views share less, so the ceiling drops.

Interactive Exploration: Temperature

The temperature \(s = \exp(t)\) scales the cosine similarities before the softmax, which controls how sharply the model must separate the matched pair from the rest. Small scale (high temperature) → a soft, forgiving distribution; large scale (low temperature) → a peaky one that demands the diagonal dominate. CLIP learns this value, and clamps it so it cannot run away and make the logits explode.

Drag the logit scale and watch the softmax over one image’s similarities to five captions (the matched caption is the first bar):

NoteKey Insight

Temperature does not change which caption is most similar — it changes how confidently the loss insists on it. Too low a scale and every pair looks equally good (no learning signal); too high and a single hard example dominates the gradient. The learned-and-clamped value is CLIP threading that needle.

Common Pitfalls

WarningForgetting to normalize

Skip the L2 step and your “cosine” similarity is a raw dot product — it grows with vector length, so a long, off-topic vector can outscore a short, on-topic one. Always normalize both towers before the similarity.

WarningUsing a one-sided loss

Only CE(logits, y) (image→text) trains images to find captions, not captions to find images — text→image retrieval will lag. The loss must be symmetric.

WarningAn unclamped learned temperature

logit_scale is learned, and gradient pressure pushes it up (sharper = lower loss). Left unclamped it can blow up and destabilize training — CLIP clamps exp(t) at 100. CLIPModel.logit_scale() does the clamp for you.

WarningPosition-embedding count must match

PatchEmbedding allocates n_patches + 1 positions (the +1 is CLS). Change the image or patch size and the count changes; a mismatch is a silent shape bug.

WarningSmall batches weaken the signal

The batch is the negative set — with N items each image sees only N−1 negatives. Tiny batches make the task too easy to learn a discriminative space, which is why real CLIP used batches in the tens of thousands.

Exercises

Exercise 1: Non-square patches / different sizes

patchify requires the patch size to divide the image. Write a check that, given an image and a patch size, either returns the patch count or explains why it doesn’t tile evenly.

from multimodal import num_patches

def can_patchify(height: int, width: int, patch_size: int):
    """Return (True, N) if it tiles evenly, else (False, reason)."""
    # Your implementation here
    if height % patch_size or width % patch_size:
        return (False, f"{patch_size} does not divide {height}x{width}")
    return (True, num_patches(height, width, patch_size))

print(can_patchify(224, 224, 16))   # (True, 196)
print(can_patchify(30, 32, 16))     # (False, ...)
(True, 196)
(False, '16 does not divide 30x32')

Exercise 2: CLS vs. mean pooling

PatchEmbedding.pool takes the CLS token. Write a mean-pool alternative (average the patch rows, excluding CLS) and compare the shapes.

import torch
from multimodal import PatchEmbedding

pe = PatchEmbedding(image_size=32, patch_size=16, embed_dim=32)
tokens = pe(torch.randn(2, 3, 32, 32))

def mean_pool(tokens):
    """Average the patch tokens (skip index 0, the CLS token)."""
    # Your implementation here
    return tokens[:, 1:].mean(dim=1)

print("CLS pool: ", tuple(pe.pool(tokens).shape))
print("mean pool:", tuple(mean_pool(tokens).shape))
CLS pool:  (2, 32)
mean pool: (2, 32)

Exercise 3: Break normalization

Reimplement clip_logits without the L2 normalization and confirm that scaling one image’s features by 10× wrongly changes its ranking.

import torch
from multimodal import clip_logits

torch.manual_seed(0)
img = torch.randn(4, 16)
txt = img.clone()
# Your implementation here: compare normalized vs un-normalized when img[0] *= 10
scaled = img.clone(); scaled[0] *= 10.0
print("normalized argmax row 0:", clip_logits(scaled, txt)[0].argmax().item(), "(still 0)")
normalized argmax row 0: 0 (still 0)

Summary

Key takeaways:

  1. A patch is a token. patchify turns a \((C, H, W)\) image into a sequence of \(N = (H/P)(W/P)\) patch vectors; a linear projection + positional embedding + CLS token (PatchEmbedding) makes it exactly the input a transformer already reads. Multimodality is mostly “get every modality into one sequence of vectors.”
  2. A shared space is learned by contrast. CLIP encodes images and text to vectors, L2-normalizes them, and trains the batch similarity matrix so the diagonal (matched pairs) wins — a symmetric cross-entropy with the in-batch mismatches as free negatives.
  3. Normalize, then score. Cosine similarity needs unit vectors; the raw dot product would rank by length. l2_normalize before clip_logits, always.
  4. Temperature is learned and clamped. The logit scale \(\exp(t)\) sets how sharply the loss separates the matched pair; CLIP learns it and clamps it at 100.
  5. The metric is retrieval accuracy. contrastive_accuracy — does each image’s nearest caption end up being its own? — climbs from chance to 1.0 as the space forms, and is the alignment made measurable.

What’s Next

You now have both halves of a multimodal model: a vision front-end that turns pixels into tokens, and a contrastive objective that binds vision and language into one space. The natural next step is fusion — letting a language model attend to image features (cross-attention, as in Flamingo) so it can describe, answer questions about, and reason over what it sees, rather than only retrieve. That turns CLIP-style alignment into a full vision-language generator.

Going Deeper

Core Papers:

Practical Resources:

  • ALIGN — Jia et al. (2021): the same contrastive recipe at billion-pair, noisy-caption scale.
  • BLIP — Li et al. (2022): bootstrapped captioning that unifies contrastive alignment with generation.