Lab 01 — ViT Patch Embedder + CLIP Contrastive Encoder + LLaVA Projector

Phase: 04 — Vision-Language Models Difficulty: ⭐⭐⭐☆☆ (the math is small; the fusion plumbing is what trips people) Time: 3–4 hours

A vision-language model is three machines bolted together: a Vision Transformer that turns pixels into tokens, a CLIP objective that teaches an image encoder and a text encoder to live in the same embedding space, and a connector that splices image tokens into an LLM's token stream so the language model can "read" the picture. This lab builds all three from scratch on nested lists — patchify + patch_embed (ViT), cosine_similarity + clip_similarity_matrix + info_nce_loss + clip_zero_shot_classify (CLIP), and Projector + assemble_multimodal_sequence (LLaVA) — so the mechanism that lets GPT-4V, Claude, and LLaVA see is muscle memory, not magic.

What you build

• patchify(image, patch_size) — ViT's tokenizer: chop an H×W grayscale image into non-overlapping P×P patches in raster order, each flattened to a vector; reject non-divisible sizes. An image becomes (H/P)·(W/P) tokens — "an image is worth 16×16 words."
• patch_embed(patches, W) — the linear projection that lifts each flat patch into a d-dim token, with an optional prepended learned [CLS] vector.
• cosine_similarity / l2_normalize — the similarity metric CLIP lives on, with the zero-norm guard interviewers love to probe.
• clip_similarity_matrix(image_embs, text_embs, temperature) — the temperature-scaled cosine logits: entry [i][j] is how much image i matches text j, divided by temperature.
• info_nce_loss(sim_matrix) — the symmetric InfoNCE loss (image→text + text→image cross-entropy, diagonal = positive), with the log-sum-exp trick so large logits never overflow. The soul of the lab: low when matched pairs sit on the diagonal, high when you shuffle.
• clip_zero_shot_classify(image_emb, class_text_embs) — classify by argmax cosine similarity to text prompts: no training, no fixed label set.
• Projector(in_dim, out_dim, hidden, seed) — LLaVA's connector: a seeded 2-layer GELU MLP that maps vision features into the LLM's token-embedding width.
• assemble_multimodal_sequence(text_token_embs, image_token_embs, idx) — splice the projected image tokens into the text stream at the <image> placeholder, yielding one fused sequence of length len(text) - 1 + len(image).

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why patchify of a 4×4 image with patch_size=2 gives 4 patches of length 4, and why a non-divisible size must raise.
• You can state why cosine_similarity is 1 for identical, -1 for opposite, 0 for orthogonal vectors, and why a zero-norm input returns 0 instead of dividing by zero.
• You can explain why test_info_nce_low_when_pairs_matched and test_info_nce_high_when_pairs_shuffled are the whole point — the loss is low only when the diagonal (the matched pair) dominates its row and column.
• You can explain why info_nce_loss([[1000,0],[0,1000]]) is ~0 and does not overflow (max-subtraction / log-sum-exp).
• You can explain why assemble_multimodal_sequence returns length len(text) - 1 + len(image) and where in the sequence the image tokens land.
• Given any "how does an LLM see an image?" prompt you can sketch ViT → CLIP-style encoder → projector → LLM, and name the token-budget cost.

How this maps to the real stack

Limits of the miniature (be honest in the interview): real ViTs add learned position embeddings and run a full transformer over the patch tokens — we stop at the linear patch embed; CLIP trains the encoders for hours on hundreds of millions of pairs — we just score fixed vectors; the projector here is randomly initialized, not trained (training it, with the encoder frozen and LoRA on the LLM, is Phase 05); and we model grayscale images as nested lists, so there is no 3-channel Conv2d, no interpolation of position embeddings for new resolutions, and no O(T²) attention cost — the parts that make high-resolution VLMs expensive.

Extensions (build these on real hardware)

• Add learned position embeddings and a [CLS] token to patch_embed, then run a one-layer self-attention block over the patch tokens (reuse your Phase 02 attention) to make it a real (tiny) ViT.
• Add a real temperature parameter that is learned (CLIP's logit_scale) and show the gradient of info_nce_loss w.r.t. it (reuse your Phase 03 autograd).
• Swap InfoNCE for SigLIP's pairwise sigmoid loss and compare batch-size sensitivity.
• Implement Flamingo-style cross-attention fusion (image tokens as keys/values the LLM attends to) and contrast its token-budget cost with the LLaVA projection approach.
• Compute the KV-cache inflation: an image at 336×336 with 14×14 patches is 576 tokens — plug that into your Phase 00 kv_cache_bytes and show how many images blow the context budget.
• Load a real openai/clip-vit-base-patch32 and check your zero-shot classification picks the same class on a handful of images.

Interview / resume

• Talking points: "How does an LLM actually see an image?" "What is contrastive pretraining and why does the temperature matter?" "Projection vs cross-attention fusion — when would you pick each?" "How does adding an image change the KV-cache and the cost?"
• Resume bullet: Built a from-scratch vision-language pipeline — ViT patch embedding, CLIP-style contrastive encoder with a numerically stable symmetric InfoNCE loss and zero-shot text-prompt classification, and a LLaVA-style MLP connector that fuses projected image tokens into an LLM's token stream — clarifying the exact mechanism behind multimodal models.