Projection Layer for Modality Alignment

A projection layer is a small network that maps vectors from one embedding space into another. Here it maps vision-encoder tokens into a text decoder's embedding space. The encoder (lessons 58-59) produces tokens of dimension vision_hidden = 768. The decoder you want to bolt on top has embedding dimension text_hidden = 512 (any other number is just as plausible) and expects text-shaped tokens. Image tokens aren't text-shaped: they live in a basis the encoder learned during vision-only pretraining, with no relationship to the decoder's word vectors.

A two-layer MLP projection (linear, GELU, linear) bridges the gap. It has about 768 * 1024 + 1024 * 512 = 1.3M parameters, small enough to train in minutes on a single GPU, and it is the piece that has to learn during the alignment phase. The vision encoder stays frozen. The text embedding table stays frozen. Only the projection moves. LLaVA shipped this recipe in 2023, BLIP-2 reframed it as a Q-Former, and open-weight VLMs since have adopted it in some form.

Pooling before projection

Pooling collapses a sequence of token vectors into one vector. The vision encoder emits 197 tokens and the text side has a single caption-level embedding, so alignment needs one image-level vector per sample. CLS pooling is the simplest: take the first token from the encoder and project it. Mean pooling over all 197 tokens is another option. Either one reduces 197 vectors to one.

Why two layers and not one

A single linear projection can rotate and rescale but cannot fix the basis if the two spaces have curvature mismatches. A GELU between two linear layers adds one non-linear bend, which is empirically enough to align CLIP-style features to language model embeddings. Deeper projections are extensions (LLaVA-NeXT used GLU, Qwen-VL used a stack of attention layers). The two-layer MLP is the canonical baseline, and it is what BLIP-2's Q-Former projection head ships with under the hood.

About 1.3M parameters for a 768 -> 1024 -> 512 head.

Cosine alignment loss

Cosine alignment loss is 1 - cos_sim(image, text). Align does not mean image_emb == text_emb. It means image_emb points in the same direction as text_emb in the joint space. The loss ranges from 0 (perfectly aligned) to 2 (opposite), and training drives it toward zero per pair. Lesson 62 generalizes to a contrastive batch (InfoNCE) where each image must be closer to its own caption than to any other caption in the batch. I use the per-pair version here so the dynamics are visible.

Frozen encoder is the trick

Freezing a model means setting requires_grad=False on its parameters so training leaves them alone. The vision encoder here is a small ViT (hidden=768, depth=4, heads=8, mlp_ratio=2.0) with roughly 20M parameters, and the text table has another couple of million (4096 * 512). A production encoder like CLIP ViT-L is far larger. Training all of them from a mock corpus is a non-starter. With both frozen, the projection's 1.3M parameters are what change, and a few hundred steps on synthetic pairs is enough to drive the loss down. Adapter-based VLMs share this operational shape: the heavy parts stay frozen and the light bridge trains.

code/main.py implements:

• MLPProjector(in_dim, hidden_dim, out_dim), two-layer linear MLP with GELU activation.
• MockTextEmbedding(vocab_size, dim), a frozen embedding table with deterministic init from a seed.
• make_pair(seed, vocab_size, max_len), which synthesizes one paired (image, caption) sample. Captions are short id sequences; the caption embedding is mean-pooled over token embeddings.
• cosine_alignment_loss(image_emb, text_emb), the per-pair 1 - cos_sim objective.
• A training loop that runs the projection for 200 steps over 32 synthetic pairs (cycled), with the vision encoder and text table frozen, and prints the loss every 25 steps.

Run it:

python3 code/main.py

Expected output: the loss drops from around 1.07 to about 0.80 within 200 steps, which shows the projection alone can pull image tokens toward the text space. The final cosine similarity is printed too, as a single number: the mean over the last batch (one pair per step).

The same pattern shows up across open-weight VLMs:

• LLaVA 1.5. Two-layer GELU MLP projection from CLIP-ViT-L hidden to LLaMA embedding dim. Frozen vision encoder, frozen LLM, train only the projection (then unfreeze the LLM in stage two).
• BLIP-2. Q-Former takes 32 learned query tokens through cross-attention against image tokens, then projects to the LLM embedding dim. The projection head at the very end of Q-Former is the analog of the MLP here.
• MiniGPT-4. Single linear projection from BLIP-2 Q-Former output to Vicuna embedding dim.
• Qwen-VL. Cross-attention adapter with several layers, but the final piece is again a projection to the LM embedding dim.

The shape varies but the role stays the same: pool image tokens, project to text embedding dim, train alone.

code/test_main.py holds 8 tests covering:

• projector output shape matches the configured out_dim
• projector gradient flows after one backward pass
• frozen text embedding table has zero requires_grad parameters
• text pooling ignores padding: a caption padded with id 0 pools to the same vector as the unpadded one
• cosine loss is zero on identical vectors and is 2 on anti-parallel vectors (two tests)
• the training loop reduces loss between its first and last step (40 steps on 4 pairs)
• make_pair returns the expected image and token-id shapes, with ids inside the vocabulary

Run them:

python3 -m unittest code/test_main.py

1. Replace CLS pooling with mean pooling over the 196 patch tokens and compare final loss after 200 steps. Mean pooling usually trains faster on synthetic data; CLS is more sample-efficient on natural images.

2. Add a learned scalar temperature to the cosine loss (cos / tau) and observe what happens when tau is too small (gradient noise) or too large (loss plateaus high).

3. Swap the two-layer MLP for a single linear layer and quantify the loss gap. The non-linearity matters more on natural image features and less on synthetic ones.

4. Add a small L2 penalty on the projector weights and watch how it interacts with cosine alignment (cosine is scale-invariant, so the penalty mostly shrinks unused directions).

5. Persist projector weights, then reload and run inference without the vision encoder backward pass to verify that only the projector is needed at deploy time.

• LLaVA paper for the two-stage train (project, then unfreeze LM).
• BLIP-2 paper for Q-Former as a learnable projection alternative.
• Qwen-VL technical report for cross-attention adapters as deeper projection heads.