Lab 05 — Multimodal Bridge (CLIP → LLM Projection)

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

LLaVA's core insight: a 2-layer MLP is all it takes to connect a frozen vision encoder to a frozen LLM. This lab builds that bridge at toy scale, where the architecture, the freezing discipline, and the loss masking are identical to the 7B-parameter version.

What you build

  • Projector — the trainable 2-layer MLP mapping vision-feature space → LLM embedding space (the only component that trains; <0.25% of total parameters)
  • MultimodalModel — frozen vision tower + frozen causal LM with visual tokens prepended to the text sequence
  • loss — next-token cross-entropy on text positions only, with the visual-token causal shift and prompt masking (ignore_index=-100)
  • generate — greedy image-conditioned decoding

Frozen TinyVisionEncoder / TinyCausalLM stand-ins are provided with the same interfaces as CLIP ViT-L/14 and TinyLlama, so the tests run in seconds with no downloads; swapping in the real models is the extension.

Key concepts

ConceptWhat to understand
Patch tokens, not CLSLLaVA feeds all patch tokens to the LM — VQA needs spatial detail the pooled embedding discards
Freezing disciplinerequires_grad_(False) + vision under no_grad — optimizer never sees frozen params, no grad buffers allocated
The causal shiftLogits at the LAST visual position predict the FIRST text token; off-by-one here is the classic multimodal bug (a test targets it)
Loss maskingPrompt/question positions set to -100 — grade the answer, not the echo
Stage-1 vs stage-2LLaVA trains projector-only on captions first, then unfreezes the LM for instruction tuning

Files

FilePurpose
lab.pySkeleton with # TODO markers
solution.pyComplete reference; python solution.py shows projector-only training working
test_lab.py12 tests: freezing, grad flow, logit alignment, masking, image→text information flow
requirements.txttorch, pytest

Run

pip install -r requirements.txt
pytest test_lab.py -v
LAB_MODULE=solution pytest test_lab.py -v

Suggested TODO order

  1. Projector — two lines of architecture, but know why MLP > single linear
  2. MultimodalModel.__init__ — the freezing; check test_only_projector_is_trainable
  3. visual_tokens → forward — concatenation order matters
  4. loss — draw the position layout on paper before coding the slice
  5. generate

Success criteria

  • All 12 tests pass — especially test_loss_alignment_uses_image_conditioned_prediction (the off-by-one catcher) and test_trained_model_uses_the_image (proves information actually flows image → projector → LM)
  • You can state the trainable-parameter fraction and why that makes stage-1 training cheap
  • You can explain what breaks if you detach() the projector output instead of the vision output (no gradient reaches the projector — silent non-training)

Full-scale extension (the original spec)

Swap stand-ins for openai/clip-vit-large-patch14 (frozen) and TinyLlama-1.1B (frozen); train the projector on the LLaVA-Instruct-150K caption subset; evaluate VQAv2 accuracy (target >50%); compare single-linear vs MLP projector. Then: Q-Former connector (BLIP-2), and measure the NPU-relevant cost — visual tokens inflate the KV cache by N_patches at every layer.

Interview Q&A

Q: Why freeze both towers in stage 1? The projector starts random — its gradient noise would destroy pretrained representations (the LM would "catastrophically adapt" to garbage visual tokens). Train the bridge until visual tokens live in the LM's embedding distribution, then optionally unfreeze for instruction tuning.

Q: Why does LLaVA prepend visual tokens rather than cross-attend (Flamingo-style)? Prepending reuses the LM unchanged (no new cross-attention layers to train), is trivially compatible with KV caching and existing serving stacks, and scales with patch count. Cross-attention isolates modalities better and keeps text-only throughput intact, at the cost of new trained layers — a classic deployment-vs-architecture tradeoff.

Q: What is the inference cost of visual tokens on an edge NPU? Each of the N visual tokens occupies KV cache in every layer (N × layers × 2 × d_kv) and participates in every attention computation — for 576 CLIP tokens this often dominates a short-question VQA prompt. Mitigations: token pruning/merging, pooled visual tokens, Q-Former-style compression to a fixed small budget.

References