Multimodal: A Tiny Vision-Language Model.
The Concept
Take a translator and a writer. The writer only reads and produces text. They have spent years learning grammar, idiom, and how to finish sentences. They cannot look at pictures. The translator's job is different. They look at a picture, then describe it in a small handful of "words" that the writer can read. The writer trusts those words and continues the story from there.
That is the entire idea of a vision-language model. A vision encoder is the translator. A text decoder is the writer. A thin layer between them, called the projection, is the dictionary the translator uses to make sure the words come out in a vocabulary the writer recognizes.
The first version of this chapter buried this point. The blunter framing — the projection is the entire interesting part — came out of trying to explain the architecture to someone outside ML, where you cannot lean on the words "attention" or "embedding" to do the work. The chapter now opens with that framing.
A second analogy. Think of a text-only computer with no camera port. To add image input, you do not rebuild the computer. You bolt on a USB hub, and the hub converts a camera's signal into bytes the computer already knows how to read. The encoder is the camera. The projection is the USB hub. The decoder is the computer. Plug the hub in correctly, and the computer behaves as if it always had a camera.
Now the formal definitions, in plain language.
A vision encoder is a neural network whose input is an image and whose output is a list of vectors. Each vector represents one patch of the image. For a 32-by-32 image cut into 16 patches of 8-by-8 pixels, the encoder produces 16 vectors. Each vector has some dimension, say 192 numbers. In larger models the image is 224-by-224 and the patches are 14-by-14, which produces 256 vectors of dimension 768 or 1024. The shape changes; the idea does not.
When I first started thinking about image tokens on a Marunthagam pediatric-triage prototype that needed to read a photo of a malnutrition chart, I kept expecting image tokens to be doing something fundamentally different from text embeddings. They are not. They are vectors of the same dtype, the same dimensionality after projection, sitting in the same input slots. Watching the same on-device runtime dispatch a picture tensor and a word tensor through identical kernel paths drove this home. The hardware does not know one tensor is a picture and another is a word. It just shuffles bytes through buffers and dispatches kernels.
An image token is one of those output vectors. The name is deliberate. To the text decoder, the vector looks indistinguishable from a token embedding produced by looking up a word in a vocabulary table: same shape, same dtype, same place in the input sequence. The decoder does not know one is a picture and one is a word. It just sees vectors.
A projection layer is a small mapping that takes the encoder's output vectors and maps them into the decoder's embedding space. The encoder might produce 192-dimensional vectors. The decoder might expect 384-dimensional ones. The projection is a tiny multi-layer perceptron, often just two linear layers with a GELU between them, that does the dimension change and learns to put the image vectors in a region of the decoder's embedding space where they make sense as a prefix to text.
Two hundred thousand parameters. That is the whole multimodal capability. I will say that again in BREAK IT because it is the most underrated fact in modern VLM design.
The decoder's embedding space is the high-dimensional region where the decoder's normal word embeddings live. Every word in the decoder's vocabulary maps to a point in this space. Words that mean similar things end up close together. The projection's job is to put image tokens in this same neighborhood, so the decoder treats them as it would treat any other tokens at the start of a sentence.
A vision-language model, or VLM, is the full system (encoder, projection, decoder) wired together. The forward pass runs the encoder on an image to get image tokens, runs the projection to put them in the decoder's space, prepends them to a text prompt, and lets the decoder generate.
Two more terms to define before we build. Captioning is the task of producing a sentence that describes a picture: image goes in, sentence comes out. VQA, short for visual question answering, is a richer task. Image plus a text question goes in, an answer comes out. VQA is the multimodal version of instruction tuning. You already built the text-only version in Project 21: Fine-Tuning and Instruction Tuning, and the format is almost identical here.
The pattern is simple enough that it deserves a sentence in plain English. The encoder turns a picture into a small batch of vectors. The projection makes those vectors look like word embeddings. The decoder reads them as if they were words and writes a caption. The training job is to teach the projection what "looks like a word embedding" actually means for this encoder and this decoder.
Why It Matters
If you cannot accept an image as input, an enormous fraction of the world is invisible to your model. A user asks "what is this?" and points at a photo of a circuit board. A doctor asks "what do you see in this chest X-ray?" A child asks "why is the dog wearing a hat?" None of these conversations are reachable by a text-only system. Building a model that reads pixels is not a frontier capability — it is the bare minimum required to be useful in most real situations.
There is a second reason to care, and it is more subtle. The bridge from vision to language teaches you something general about how to bolt one model onto another. The same trick works for audio, for video, for sensor streams from a robot. Each modality gets its own encoder. Each encoder feeds a projection. Each projection lands in the decoder's embedding space. Once you have written one of these by hand, the others follow without conceptual surprise. This chapter is the practice case for an entire family of architectures you will see again.
A third reason is the failure mode you will see in BREAK IT. The projection layer is small. The encoder and decoder are large. A reasonable engineer might assume the projection is the easy part. Surely the hard work is the encoder and decoder, and the projection is just plumbing. That assumption is wrong, and watching it fail in this chapter is one of the cleanest demonstrations in the book that the small bridge between two big networks is where the actual multimodal capability lives.
Here is my strongest opinion in this chapter, said plainly: the projection IS the multimodal capability. Not "part of" it, not "an important component of" it. It IS it. Everything else is two pretrained networks doing what they were already doing, in their own embedding spaces, blind to each other. The projection is the only thing in the entire pipeline that has to learn the cross-modal correspondence. If it does not learn, nothing else can compensate.
The Build
build.py$python projects/29_multimodal-a-tiny-vision-language-model/build.py --tiny
"""
Project 29: Multimodal: A Tiny Vision-Language Model
Canonical runnable build for this project. Generated by tools/extract_code.py
from the book's chapter markdown — re-run that tool to regenerate.
"""
# === Step 1: Build a tiny vision encoder ===
class PatchEmbed(nn.Module):
def __init__(self, img_size=28, patch_size=4, in_chans=3, embed_dim=192):
super().__init__()
self.proj = nn.Conv2d(in_chans, embed_dim,
kernel_size=patch_size, stride=patch_size)
self.num_patches = (img_size // patch_size) ** 2
def forward(self, x):
# x: (B, 3, 28, 28)
x = self.proj(x) # (B, embed_dim, 7, 7)
x = x.flatten(2).transpose(1, 2) # (B, 49, embed_dim)
return x
class TinyViT(nn.Module):
def __init__(self, depth=4, embed_dim=192, num_heads=6,
img_size=28, patch_size=4):
super().__init__()
self.patch_embed = PatchEmbed(img_size, patch_size,
in_chans=3, embed_dim=embed_dim)
self.pos_embed = nn.Parameter(
torch.zeros(1, self.patch_embed.num_patches, embed_dim))
self.blocks = nn.ModuleList([
ViTBlock(embed_dim, num_heads) for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
def forward(self, x):
x = self.patch_embed(x)
x = x + self.pos_embed
for blk in self.blocks:
x = blk(x)
return self.norm(x)
# === Step 2: Build the projection layer ===
class Projection(nn.Module):
def __init__(self, vision_dim=192, text_dim=384):
super().__init__()
self.fc1 = nn.Linear(vision_dim, text_dim)
self.act = nn.GELU()
self.fc2 = nn.Linear(text_dim, text_dim)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
# === Step 3: Extend the GPT input path to accept vision tokens ===
class GPT(nn.Module):
def forward(self, idx=None, vision_embeds=None, targets=None):
if vision_embeds is not None:
# idx may still be present for the text portion
text_embeds = self.tok_emb(idx) if idx is not None else None
if text_embeds is not None:
x = torch.cat([vision_embeds, text_embeds], dim=1)
else:
x = vision_embeds
else:
x = self.tok_emb(idx)
T = x.shape[1]
pos = self.pos_emb[:, :T, :]
x = x + pos
for blk in self.blocks:
x = blk(x)
x = self.ln_f(x)
logits = self.lm_head(x)
if targets is not None:
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)),
targets.view(-1),
ignore_index=-100,
)
return logits, loss
return logits, None
# For a batch with 49 image tokens and a caption of length T:
# targets shape: (B, 49 + T)
# First 49 positions: -100 (ignored)
# Remaining T positions: shifted caption tokens, normal cross-entropy
targets[:, :49] = -100
# === Step 4: Pretrain on (image, caption) pairs ===
LABEL_TO_CAPTION = {
0: "a photograph of an airplane.",
1: "a photograph of an automobile.",
2: "a photograph of a bird.",
3: "a photograph of a cat.",
4: "a photograph of a deer.",
5: "a photograph of a dog.",
6: "a photograph of a frog.",
7: "a photograph of a horse.",
8: "a photograph of a ship.",
9: "a photograph of a truck.",
}
images, labels = next(loader)
captions = [LABEL_TO_CAPTION[y.item()] for y in labels]
input_ids, target_ids = encode_captions(captions, tokenizer)
vision_tokens = encoder(images) # (B, 49, 192)
vision_embeds = projection(vision_tokens) # (B, 49, 384)
# Pad targets with -100 in image positions
targets = pad_targets_with_image_mask(target_ids, num_image_tokens=49)
logits, loss = model(idx=input_ids,
vision_embeds=vision_embeds,
targets=targets)
loss.backward()
optimizer.step()
# === Step 5: Instruction-tune on visual question answering ===
VQA_TEMPLATES = {
"dog": [
("What animal is in the picture?", "A dog."),
("Is this a cat or a dog?", "A dog."),
("Describe what you see.", "I see a dog."),
],
# ... and so on
}
# === Step 6: Run inference ===
def caption(model, encoder, projection, image, tokenizer, max_new=20):
model.eval()
with torch.no_grad():
# Encode the image
vision_tokens = encoder(image.unsqueeze(0))
vision_embeds = projection(vision_tokens)
# Optional text prompt
prompt_ids = tokenizer.encode("a photograph of")
idx = torch.tensor([prompt_ids])
# Generate
for _ in range(max_new):
logits, _ = model(idx=idx, vision_embeds=vision_embeds)
next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
idx = torch.cat([idx, next_id], dim=1)
if next_id.item() == tokenizer.eos_id:
break
return tokenizer.decode(idx[0].tolist())
The Steps
6 reference filesEach step is the code for one piece, in isolation — read them as the book introduces them. They assemble into build.py above.
Step 01Build A Tiny Vision Encoderstep_01_build-a-tiny-vision-encoder.py
"""
Project 29: Step 1 — Build a tiny vision encoder
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
class PatchEmbed(nn.Module):
def __init__(self, img_size=28, patch_size=4, in_chans=3, embed_dim=192):
super().__init__()
self.proj = nn.Conv2d(in_chans, embed_dim,
kernel_size=patch_size, stride=patch_size)
self.num_patches = (img_size // patch_size) ** 2
def forward(self, x):
# x: (B, 3, 28, 28)
x = self.proj(x) # (B, embed_dim, 7, 7)
x = x.flatten(2).transpose(1, 2) # (B, 49, embed_dim)
return x
class TinyViT(nn.Module):
def __init__(self, depth=4, embed_dim=192, num_heads=6,
img_size=28, patch_size=4):
super().__init__()
self.patch_embed = PatchEmbed(img_size, patch_size,
in_chans=3, embed_dim=embed_dim)
self.pos_embed = nn.Parameter(
torch.zeros(1, self.patch_embed.num_patches, embed_dim))
self.blocks = nn.ModuleList([
ViTBlock(embed_dim, num_heads) for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
def forward(self, x):
x = self.patch_embed(x)
x = x + self.pos_embed
for blk in self.blocks:
x = blk(x)
return self.norm(x)
Step 02Build The Projection Layerstep_02_build-the-projection-layer.py
"""
Project 29: Step 2 — Build the projection layer
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
class Projection(nn.Module):
def __init__(self, vision_dim=192, text_dim=384):
super().__init__()
self.fc1 = nn.Linear(vision_dim, text_dim)
self.act = nn.GELU()
self.fc2 = nn.Linear(text_dim, text_dim)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
Step 03Extend The Gpt Input Path To Accept Vision Tokensstep_03_extend-the-gpt-input-path-to-accept-vision-tokens.py
"""
Project 29: Step 3 — Extend the GPT input path to accept vision tokens
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
class GPT(nn.Module):
def forward(self, idx=None, vision_embeds=None, targets=None):
if vision_embeds is not None:
# idx may still be present for the text portion
text_embeds = self.tok_emb(idx) if idx is not None else None
if text_embeds is not None:
x = torch.cat([vision_embeds, text_embeds], dim=1)
else:
x = vision_embeds
else:
x = self.tok_emb(idx)
T = x.shape[1]
pos = self.pos_emb[:, :T, :]
x = x + pos
for blk in self.blocks:
x = blk(x)
x = self.ln_f(x)
logits = self.lm_head(x)
if targets is not None:
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)),
targets.view(-1),
ignore_index=-100,
)
return logits, loss
return logits, None
# For a batch with 49 image tokens and a caption of length T:
# targets shape: (B, 49 + T)
# First 49 positions: -100 (ignored)
# Remaining T positions: shifted caption tokens, normal cross-entropy
targets[:, :49] = -100
Step 04Pretrain On Image Caption Pairsstep_04_pretrain-on-image-caption-pairs.py
"""
Project 29: Step 4 — Pretrain on (image, caption) pairs
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
LABEL_TO_CAPTION = {
0: "a photograph of an airplane.",
1: "a photograph of an automobile.",
2: "a photograph of a bird.",
3: "a photograph of a cat.",
4: "a photograph of a deer.",
5: "a photograph of a dog.",
6: "a photograph of a frog.",
7: "a photograph of a horse.",
8: "a photograph of a ship.",
9: "a photograph of a truck.",
}
images, labels = next(loader)
captions = [LABEL_TO_CAPTION[y.item()] for y in labels]
input_ids, target_ids = encode_captions(captions, tokenizer)
vision_tokens = encoder(images) # (B, 49, 192)
vision_embeds = projection(vision_tokens) # (B, 49, 384)
# Pad targets with -100 in image positions
targets = pad_targets_with_image_mask(target_ids, num_image_tokens=49)
logits, loss = model(idx=input_ids,
vision_embeds=vision_embeds,
targets=targets)
loss.backward()
optimizer.step()
Step 05Instruction Tune On Visual Question Answeringstep_05_instruction-tune-on-visual-question-answering.py
"""
Project 29: Step 5 — Instruction-tune on visual question answering
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
VQA_TEMPLATES = {
"dog": [
("What animal is in the picture?", "A dog."),
("Is this a cat or a dog?", "A dog."),
("Describe what you see.", "I see a dog."),
],
# ... and so on
}
Step 06Run Inferencestep_06_run-inference.py
"""
Project 29: Step 6 — Run inference
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
def caption(model, encoder, projection, image, tokenizer, max_new=20):
model.eval()
with torch.no_grad():
# Encode the image
vision_tokens = encoder(image.unsqueeze(0))
vision_embeds = projection(vision_tokens)
# Optional text prompt
prompt_ids = tokenizer.encode("a photograph of")
idx = torch.tensor([prompt_ids])
# Generate
for _ in range(max_new):
logits, _ = model(idx=idx, vision_embeds=vision_embeds)
next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
idx = torch.cat([idx, next_id], dim=1)
if next_id.item() == tokenizer.eos_id:
break
return tokenizer.decode(idx[0].tolist())
Break It
break_it.py$python projects/29_multimodal-a-tiny-vision-language-model/break_it.py --tiny
"""
Project 29: BREAK IT experiment.
Deliberately sabotages one mechanism from build.py to show what happens
when it's removed. Compare outputs to the working version.
"""
projection = Projection(vision_dim=192, text_dim=384)
for p in projection.parameters():
p.requires_grad = False
Read it in the book
This page ships the code. The full derivation — why the chain rule is mechanical not magical, what the computational graph buys you, the long-form reasoning — lives in Chapter 29 of Under the Hood.
Get the book →