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 three 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 🐱.
  • Fusion. Let a language model cross-attend to the image’s visual tokens so it can generate — describe, answer, reason — not just retrieve. A tanh gate bolts this onto a good LM without breaking it (Flamingo).

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 τ).
  • Build cross-attention (text queries, visual keys/values) and the tanh-gated fusion layer that lets a language model read an image (Flamingo), and prove the gate-zero identity — with the gates shut, the VLM is exactly the text-only LM.
  • Compress a variable number of patches to a fixed set of visual tokens with a Perceiver Resampler, and watch a tiny VLM learn to see: accuracy climbs from chance to 1.0 as the gate opens and cross-attention locks onto each image’s patch.
  • Build the early-fusion design that most modern VLMs use (LLaVA): a visual projector + splicing image tokens into the decoder’s own sequence, so plain causal self-attention reads the image — no cross-attention, no gate — and prove that spliced image tokens are just tokens.

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, and cross-attention reuses its QKᵀ/√d kernel.
  • Module 06: Transformer — the language model that the vision-language model extends with gated cross-attention.
  • 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: [0.9999999403953552, 1.0, 1.0] (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.

From Alignment to Generation: Cross-Attention VLMs

CLIP gives you a shared space, but look at what it can and cannot do. It can retrieve — given an image, find the closest caption; given text, find the closest image. It cannot generate: it will never write a novel description, answer “how many cats are in this picture?”, or reason about what it sees. For that a language model — the next-token machine you built in m06–m08 — has to somehow read the image.

The idea that makes this work is cross-attention. Self-attention (m05) lets a token look at other text tokens: the queries, keys, and values all come from one sequence. Cross-attention keeps the same softmax(QKᵀ/√d)V machinery but draws the queries from the text and the keys and values from a second sequence — the image’s visual tokens (the patch embeddings from the ViT front-end above, optionally compressed by a resampler). Now every text position can attend to the image and pull in whatever it needs.

This is exactly how Flamingo (Alayrac et al., 2022) — and the vision adapters on many chat models — turn a text-only LM into one that can see: take a frozen, already-trained LM and insert new cross-attention layers between its blocks, letting the residual stream reach over to the image.

But there is a catch, and its fix is the elegant part.

NoteKey Insight

Bolt a randomly-initialized cross-attention layer into a good LM and, at step zero, it injects noise into a stream the LM had finely tuned — training is unstable and the pretrained skill is damaged. Flamingo’s fix: start the new layer switched off, and let training turn it on gradually.

The Math: Cross-Attention & the tanh Gate

Let x \in \mathbb{R}^{T\times d} be the text stream and v \in \mathbb{R}^{M\times d} the M visual tokens. Cross-attention is the m05 formula with keys and values sourced from v:

\text{CrossAttn}(x, v) = \text{softmax}\!\left(\frac{Q K^{\top}}{\sqrt{d_k}}\right) V, \qquad Q = x W_Q,\; K = v W_K,\; V = v W_V.

Each of the T text queries scores all M visual tokens and reads a weighted mix of their values — so the output has one vector per text position, ready to add back into the residual stream.

The switch is a tanh gate. Multiply the new layer’s output by \tanh(\alpha), where \alpha is a single learnable scalar initialized to 0, before adding it to the residual — for both the cross-attention and its feed-forward:

x \;\leftarrow\; x + \tanh(\alpha_{\text{attn}})\cdot\text{CrossAttn}(\text{LN}(x), v), \qquad x \;\leftarrow\; x + \tanh(\alpha_{\text{ffn}})\cdot\text{FFN}(\text{LN}(x)).

Because \tanh(0) = 0 exactly, at initialization the gate contributes nothing: the VLM’s output is bit-for-bit identical to the original text-only LM. The image is present but ignored. As \alpha moves off 0, the image bleeds into the stream — smoothly, from zero, so the LM is never shocked. This gate-zero identity is the fusion analogue of CLIP’s “matched batch wins its diagonal,” and we test it directly below.

Step through where the gated layer sits and what the gate does:

Drag the gate below and watch how much of the image contribution actually reaches the residual. At \alpha = 0 the stream is pure LM; the image only matters once the gate opens:

The Perceiver Resampler

One loose end: a ViT emits one visual token per patch, and that count grows with resolution — and, for video, with the number of frames. Cross-attending to hundreds or thousands of visual tokens at every layer is expensive, and the count keeps changing. Flamingo fixes both with a Perceiver Resampler: a small, fixed set of learned latent queries (Flamingo uses 64) cross-attend to all the patch features and output exactly that many visual tokens — a constant-length summary of the image, no matter how big the input was.

Code: A Vision-Language Model from Scratch

Everything above is a handful of modules in fusion.py. CrossAttention is m05’s attention with the keys and values coming from a second sequence:

import torch
from fusion import CrossAttention

torch.manual_seed(0)
cross = CrossAttention(embed_dim=16, num_heads=4)
text   = torch.randn(2, 3, 16)    # 3 text positions
visual = torch.randn(2, 5, 16)    # 5 visual tokens
out, weights = cross(text, visual)
print(f"text {tuple(text.shape)}  attends to  visual {tuple(visual.shape)}")
print(f"output (one vector per text position): {tuple(out.shape)}")
print(f"attention weights (text × visual):     {tuple(weights.shape)}")
text (2, 3, 16)  attends to  visual (2, 5, 16)
output (one vector per text position): (2, 3, 16)
attention weights (text × visual):     (2, 4, 3, 5)

GatedCrossAttention wraps it with the tanh(α) gate. Freshly built, both gates are 0 — so it returns its input untouched:

from fusion import GatedCrossAttention

block = GatedCrossAttention(embed_dim=16, num_heads=4)
print(f"initial gates tanh(α): {block.gates()}")           # (0.0, 0.0)
y, _ = block(text, visual)
print(f"output == input at init? {torch.equal(y, text)}")  # True — a no-op
initial gates tanh(α): (0.0, 0.0)
output == input at init? True

VisionLanguageModel interleaves a gated cross-attention after each LM block, and exposes a text_only path (no image). Here is the headline fact, run directly: with the gates at 0, seeing the image changes nothing — the VLM is the text-only LM, bit-for-bit.

from fusion import VisionLanguageModel

torch.manual_seed(0)
vlm = VisionLanguageModel(vocab_size=10, embed_dim=16, num_heads=4, num_layers=2)
tokens = torch.randint(0, 10, (2, 5))
image  = torch.randn(2, 8, 16)      # 8 visual tokens

with_image = vlm(tokens, image)
text_only  = vlm.text_only(tokens)
print(f"logits identical with vs. without the image? {torch.equal(with_image, text_only)}")
print("→ the image is present but the shut gate ignores it — exactly.")
logits identical with vs. without the image? True
→ the image is present but the shut gate ignores it — exactly.

The PerceiverResampler compresses any number of patch features to a fixed count:

from fusion import PerceiverResampler

resampler = PerceiverResampler(embed_dim=32, num_latents=8, num_heads=4)
for n in (16, 64, 256):
    tokens_out = resampler(torch.randn(1, n, 32))
    print(f"{n:>4} patches  ->  {tuple(tokens_out.shape)}  (always 8 visual tokens)")
  16 patches  ->  (1, 8, 32)  (always 8 visual tokens)
  64 patches  ->  (1, 8, 32)  (always 8 visual tokens)
 256 patches  ->  (1, 8, 32)  (always 8 visual tokens)

Watch a VLM Learn to See

The payoff: train the tiny VLM and watch the gate open. demonstrate_vlm sets up a task that is impossible without vision — every example gets the same one-token text prompt, and the answer is determined entirely by the image. Concept c’s image marks patch c with a distinct vector (all other patches blank), and the target is answer token c. A text-only model (gates shut) gives every example identical logits, so its accuracy is pinned at chance 1/K. The only way to do better is to read the image through cross-attention.

from fusion import demonstrate_vlm

result = demonstrate_vlm(n_concepts=6, steps=500, seed=0)
============================================================
VLM: learning to read the image through cross-attention
============================================================
  K=6 concepts, 2 layers, 500 steps

    step      loss   accuracy     gate
       0   29.9692      0.000    0.000
     125    0.0002      1.000    0.310
     250    0.0001      1.000    0.311
     375    0.0001      1.000    0.312
     499    0.0000      1.000    0.313

  Accuracy: 0.000 -> 1.000  (chance = 0.167);  gate: 0.000 -> 0.313

Two things move together: accuracy climbs from chance to 1.0, and the gate tanh(α) lifts off 0 — the model literally switching the image on:

And the cross-attention earns its interpretation. Each row of the heatmap is one concept’s answer position; each column is a visual token (patch). After training, the bright cell in row c sits at column c — every concept learns to attend to exactly the patch that identifies it:

TipTry This
  1. Freeze the gate at 0: the model can only ever output one token for every concept — accuracy is stuck at 1/K. The image is unreachable without the gate.
  2. More concepts: raise n_concepts to 12. The task is harder (more patches to tell apart), but the same mechanism solves it — the diagonal still forms.
  3. Watch the gate, not just accuracy: accuracy can saturate at 1.0 while the gate keeps rising — the model sharpening how strongly it reads the image.
NoteKey Insight

This is the whole recipe for giving a language model eyes: encode the image to tokens, let the text cross-attend to them, and gate the new pathway from zero so the pretrained LM is upgraded, not overwritten. Alignment (CLIP) told the model what goes with what; fusion (cross-attention) lets it speak about what it sees.

The Simpler Fusion: Visual Tokens in the Stream

Flamingo works, but look at what it cost: a whole second attention pathway (CrossAttention), a gate to tune, and those layers threaded through the LM. What if you didn’t add any new mechanism at all?

Here is the idea that powers most current open vision-language models — LLaVA (Liu et al., 2023), and the “native multimodal” direction of GPT-4o, Gemini, and Llama 4. A patch embedding is already a vector in a d-dimensional space (PatchEmbedding). A token embedding is also a vector in a d-dimensional space (m04). So: learn a small map from one space to the other, and drop the image’s vectors into the token sequence as if they were words. The language model’s own causal self-attention (m05) reads them. No cross-attention. No gate. The only new part is a tiny projector.

This is early fusion: the image enters at the input, in the same stream as the text, instead of through side-channels bolted onto the middle of the model.

NoteKey Insight

Flamingo asks: how should text attend to a separate image stream? LLaVA asks a smaller question: how do I turn image patches into things that look like tokens? — and then lets the LM it already trusts do the rest. The image stops being a second modality the model reaches out to and becomes part of the sentence.

The Math: Project, Splice, Attend

A vision encoder turns an image into N patch features Z_v \in \mathbb{R}^{N \times d_v} (the ViT front-end from earlier, its output dimension d_v). A learned projector g maps each into the language model’s embedding dimension d:

H_v = g(Z_v), \qquad H_v \in \mathbb{R}^{N \times d}.

LLaVA’s g is a single linear layer, H_v = W Z_v; LLaVA-1.5 upgraded it to a two-layer GELU MLP — a stronger connector, still tiny. Now take the text embedding sequence, find the <image> placeholder at position p, and splice the projected image tokens in:

x = \big[\, h_1,\ \dots,\ h_{p-1},\ \underbrace{H_v}_{N\ \text{tokens}},\ h_p,\ \dots,\ h_T \,\big] \ \in \mathbb{R}^{(T+N) \times d}.

Then run the ordinary decoder over x: the same causal self-attention, the same blocks, the same LM head. Under the causal mask, every text token after the image attends back over all N image tokens — which is the fusion. There is nothing else. (A CLIP ViT-L/14 at 336px tiles into 24\times24 = 576 patches, so a real LLaVA-1.5 splices in 576 visual tokens per image.)

NoteKey Insight

Because the image tokens are just rows of the embedding sequence, “the LM reads the image” is not a trained hope — it is an identity. Running the VLM on (tokens, image) is exactly running the plain LM on the pre-spliced embeddings. The next code cell proves it with torch.allclose.

Code: Early Fusion from Scratch

Everything lives in early_fusion.py. The projector is the whole vision-specific part — one linear map (LLaVA), or a two-layer MLP (LLaVA-1.5):

import torch
from early_fusion import VisualProjector

torch.manual_seed(0)
proj = VisualProjector(d_vis=8, d_model=16, kind="mlp")   # LLaVA-1.5 connector
feats = torch.randn(2, 5, 8)                              # 5 patch features, dim 8
tokens = proj(feats)                                      # -> 5 tokens in LM space
print(f"visual features {tuple(feats.shape)}  ->  visual tokens {tuple(tokens.shape)}")
visual features (2, 5, 8)  ->  visual tokens (2, 5, 16)

The splice inserts those tokens into the text embedding sequence at the <image> slot — a plain cat, so the merged sequence grows by exactly N:

from early_fusion import splice_visual_tokens

text = torch.randn(2, 3, 16)          # 3 text embeddings, dim 16
merged = splice_visual_tokens(text, tokens, position=1)
print(f"text {tuple(text.shape)} + image {tuple(tokens.shape)}  ->  {tuple(merged.shape)}")
print(f"the image occupies merged positions [1, {1 + tokens.shape[1]})")
text (2, 3, 16) + image (2, 5, 16)  ->  (2, 8, 16)
the image occupies merged positions [1, 6)

EarlyFusionVLM ties it together: project, splice, then the unmodified causal LM. Here is the headline fact, run directly — image tokens are just tokens: feeding (tokens, image) is identical to running the plain LM over the pre-spliced embeddings, bit-for-bit.

from early_fusion import EarlyFusionVLM

torch.manual_seed(0)
vlm = EarlyFusionVLM(vocab_size=10, d_vis=8, embed_dim=16, num_heads=4, num_layers=2)
ids = torch.randint(0, 10, (2, 4))
img = torch.randn(2, 5, 8)

via_forward = vlm(ids, img, image_pos=2)                 # project + splice + LM
merged = splice_visual_tokens(vlm.token_embed(ids), vlm.projector(img), 2)
via_lm  = vlm.run_embeddings(merged)                      # plain LM over the mix
print(f"merged length: {via_forward.shape[1]}  (= 4 text + 5 image)")
print(f"forward == plain-LM-over-spliced-embeddings? {torch.allclose(via_forward, via_lm)}")
merged length: 9  (= 4 text + 5 image)
forward == plain-LM-over-spliced-embeddings? True

With no image, it is simply the language model again — nothing special-cases the text-only path:

print("text-only path == run_embeddings(text)?",
      torch.allclose(vlm.text_only(ids), vlm.run_embeddings(vlm.token_embed(ids))))
text-only path == run_embeddings(text)? True

The Unified Sequence

Step through what actually happens to one prompt. The image is encoded to patch features, the projector maps them into token space, they are spliced where the <image> marker sat, and then a single causal self-attention runs over the whole mixed sequence — the image tokens attended to like any other.

Watch Early Fusion Learn to See

The payoff, run for real. demonstrate_early_fusion sets up the same impossible task the cross-attention demo used — every example gets the same one-token text prompt, and the answer is determined entirely by the image (concept c’s image marks patch c; the target is answer token c). A model that ignores the image gives every example identical logits, so text_only accuracy is pinned at chance 1/K. The only way to do better is to read the image tokens in the stream — and there is no cross-attention and no gate doing it, just the decoder’s own self-attention.

from early_fusion import demonstrate_early_fusion

ef = demonstrate_early_fusion(n_concepts=6, steps=600, seed=0, verbose=False)
print(f"chance = 1/K = {ef['chance']:.3f}")
print(f"accuracy: start {ef['curve'][0]['accuracy']:.2f}  →  end {ef['curve'][-1]['accuracy']:.2f}")
chance = 1/K = 0.167
accuracy: start 0.00  →  end 1.00

Accuracy climbs from chance to 1.0 with no new attention mechanism — the image became part of the sentence, and the LM read it:

TipTry This
  1. Swap the projector: build the VLM with projector="linear" (the original LLaVA) instead of "mlp" and re-run — the single matrix still learns the map, a touch slower, exactly the LLaVA → LLaVA-1.5 story.
  2. Freeze the LM with vlm.freeze_lm() before training to try the true LLaVA stage-1 recipe (projector only). Our from-scratch LM has no pretraining to lean on, so it struggles — a vivid reminder that LLaVA’s projector-only magic rides on a pretrained LM that already knows how to route information.
  3. More concepts: raise n_concepts to 12 — the sequence gets longer, but the same “tokens in the stream” mechanism solves it.

Two Ways to Fuse

You have now built both fusion designs. They differ in one decision — where the image enters — and everything else follows:

Cross-attention (Flamingo) Early fusion (LLaVA)
Where the image enters New gated cross-attention layers in the middle of the LM Projected into tokens, spliced into the input sequence
New mechanism? Yes — a second attention pathway + a tanh(α) gate None — the LM’s own causal self-attention
What you train The cross-attention layers (LM often frozen) A small projector (LM frozen in stage 1, tuned in stage 2)
Sequence length Text length T (image is a side input) T + N — image tokens lengthen the sequence
Attention cost O(T^2) + O(T\!\cdot\!N) O\big((T+N)^2\big) — pays for image↔︎image too
Used by Flamingo, IDEFICS, older adapters LLaVA, Qwen-VL, InternVL, native-multimodal frontier
NoteKey Insight

Early fusion won on simplicity: reuse the transformer you already have, and the entire “adapter” is a projector you can train in hours. The price is sequence length — N visual tokens per image inflate the context and the O(L^2) attention bill, which is exactly why the PerceiverResampler (and later token-pruning tricks) matter: fewer, denser visual tokens keep early fusion affordable.

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.

WarningForgetting the tanh gate (fusion)

Bolt a fresh cross-attention layer straight into a good LM and it injects noise at step zero — training destabilizes and the pretrained skill degrades. Gate the new pathway with tanh(α), α initialized to 0, so the VLM starts identical to the LM and the image fades in. GatedCrossAttention does this; the anchor test is vlm(tokens, image) == vlm.text_only(tokens) while the gates are shut.

WarningForgetting that early-fusion image tokens cost sequence length

Cross-attention keeps the text sequence at length T; early fusion makes the image part of the sequence, so N visual tokens lengthen it to T + N and the self-attention bill grows as O\big((T+N)^2\big). A single high-res image can be 576+ tokens — dwarfing the prompt. Budget for it: resample or prune the visual tokens (PerceiverResampler) rather than splicing every patch.

WarningExpecting a projector-only fit over an untrained LM

LLaVA’s stage-1 recipe — freeze everything, train only the projector — works because the frozen LM is pretrained and already routes information well; the projector just has to speak its language. Freeze a randomly initialized LM (freeze_lm() on our from-scratch model) and the projector alone can’t drive it, which is why demonstrate_early_fusion trains the whole tiny model. The mechanism is identical; the free-lunch adapter is a property of the pretrained base.

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)

Exercise 4: The gate-zero identity

Confirm the fusion anchor for yourself: with the gates at their initial 0, a VisionLanguageModel gives the same logits with or without an image — and that it stops being true once you open a gate.

import torch
from fusion import VisionLanguageModel

torch.manual_seed(1)
vlm = VisionLanguageModel(vocab_size=8, embed_dim=16, num_heads=2, num_layers=2)
tokens = torch.randint(0, 8, (2, 4))
image = torch.randn(2, 6, 16)

# Your implementation here: show equality at init, then break it by opening a gate.
print("gates shut  → identical:", torch.equal(vlm(tokens, image), vlm.text_only(tokens)))
with torch.no_grad():
    vlm.cross_layers[0].attn_gate.fill_(1.0)   # open one gate
print("gate opened → identical:", torch.equal(vlm(tokens, image), vlm.text_only(tokens)))
gates shut  → identical: True
gate opened → identical: False

Exercise 5: Image tokens are just tokens

Prove the early-fusion anchor: running the VLM on (tokens, image) is identical to running the plain LM over the pre-spliced embedding sequence — no cross-attention to reason about, just one sequence.

import torch
from early_fusion import EarlyFusionVLM, splice_visual_tokens

torch.manual_seed(0)
vlm = EarlyFusionVLM(vocab_size=8, d_vis=6, embed_dim=16, num_heads=2, num_layers=2)
ids = torch.randint(0, 8, (2, 4))
feats = torch.randn(2, 5, 6)
pos = 2

# Your implementation here: build the merged embeddings by hand, run the LM on them,
# and compare to the VLM's own forward pass.
merged = splice_visual_tokens(vlm.token_embed(ids), vlm.projector(feats), pos)
print("forward == plain-LM-over-spliced-embeddings:",
      torch.allclose(vlm(ids, feats, image_pos=pos), vlm.run_embeddings(merged)))
forward == plain-LM-over-spliced-embeddings: True

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.
  6. Cross-attention lets the LM read the image. Keep m05’s softmax(QKᵀ/√d)V, but take the queries from the text and the keys/values from the image’s visual tokens (CrossAttention). Now every text position can attend to the picture — the step from retrieving a caption to generating one.
  7. The tanh gate upgrades an LM without breaking it. Multiply each new cross-attention layer by tanh(α) with α initialized to 0 (GatedCrossAttention): at init the VLM is bit-for-bit the text-only LM (vlm(tokens, image) == vlm.text_only(tokens)), and training opens the gate so the image fades in smoothly — the fusion analogue of CLIP’s diagonal proof.
  8. A resampler keeps the image a fixed length. A PerceiverResampler’s learned latent queries cross-attend to any number of patches and emit a constant few visual tokens, so the LM sees the same shape whatever the resolution.
  9. Early fusion needs no new mechanism. The design behind most modern VLMs (LLaVA) skips cross-attention entirely: a small projector (VisualProjector — one linear map, or a two-layer MLP in LLaVA-1.5) turns patch features into tokens, splice_visual_tokens drops them into the decoder’s sequence, and the LM’s own causal self-attention reads them. “Image tokens are just tokens” is an torch.allclose, not a slogan — and the price is N extra positions per image.

What’s Next

You now have a complete small multimodal model, and both ways to fuse: a vision front-end that turns pixels into tokens (ViT), a contrastive objective that binds vision and language into one space (CLIP), a gated cross-attention path that lets a language model read an image and speak about it (Flamingo), and the early-fusion design that just splices image tokens into the decoder’s stream (LLaVA). The frontier from here:

  • Interleaved image-text pretraining — train on documents where images and text alternate, so the model learns to reference the right image mid-sentence.
  • Native early fusion from step one — the frontier (GPT-4o, Gemini, Llama 4) trains a single transformer on mixed image-and-text tokens from scratch, rather than projecting into a frozen LM after the fact. Same “one stream” idea you built, scaled to pretraining — and increasingly, tokenized outputs too, so one model both reads and draws.
  • A real vision tower — thread the PatchEmbedding ViT and PerceiverResampler into the VLM end-to-end (this section stands in real visual tokens), and add image generation (diffusion / autoregressive pixels) for a model that both sees and draws.

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.