Cross-Attention Fusion

Early fusion concatenates image tokens and text tokens into one sequence. Chameleon and Emu3 take this path. Late fusion keeps the two streams separate: the text decoder runs on text-only tokens and reaches into the image stream through cross-attention. This lesson's code adds it to every block; Flamingo inserts it every K blocks of a frozen language model, and later Flamingo-shaped decoders copied the design.

Late fusion has two advantages. The text stream stays clean, so the model preserves text-only capabilities. And the image stream is computed once per image and reused for every decode step, which keeps generation cheap even for long captions. The cost is one extra attention sub-layer per block.

Mask shapes

The two attentions inside a decoder block need different masks:

The self-attention path validates its mask shape, so passing the wrong mask there surfaces as a ValueError instead of a silently broken loss curve. Cross-attention takes no mask argument at all.

Why no mask on cross-attention

The image is fully observed before any text is generated, and image patches have no temporal order. Token t of the caption may attend to any patch. Some Flamingo variants add a per-sample masking pattern when interleaving multiple images and text segments, but for a single image plus a caption, cross-attention sees everything.

Key/value caching

A KV cache holds the image keys and values, computed once at the start of the decode. Each new text token reads from the cache without recomputing them. This is what makes captioning fast at inference: the heavy ViT runs once, and the cross-attention reuses its keys and values at every step. The code exposes the cache, and the tests cover the cache-hit path. One limit to know: VisionLanguageDecoder.forward builds the cache inside a single call, so the demo proves the cached and uncached logits match but does not carry a cache across separate decode steps. A real generation loop would call build_kv_cache once and reuse it.

Block composition

A decoder block runs pre-LN -> self-attention -> residual -> pre-LN -> cross-attention -> residual -> pre-LN -> feed-forward -> residual. That is three sub-layers, each with its own LayerNorm. The Flamingo paper added a learned gate on cross-attention so the model could opt out of the image path and start training from text-only behavior, which helps stability. The canonical baseline, which this code uses, has no gate.

class DecoderBlock(nn.Module):
    def forward(self, x, memory, text_mask, kv_cache=None):
        x = x + self.self_attn(self.ln1(x), mask=text_mask)
        x = x + self.cross_attn(self.ln2(x), memory, kv_cache=kv_cache)
        x = x + self.ffn(self.ln3(x))
        return x

code/main.py implements:

• CrossAttention(cfg), taking a DecoderConfig (hidden size, heads, depth, vocab sizes), a multi-head cross-attention with separate q and kv projections.
• CausalSelfAttention(cfg), the masked self-attention from a standard decoder.
• DecoderBlock, composing the three sub-layers with pre-LN residuals.
• VisionLanguageDecoder, four-layer decoder fed by a mock vision encoder output and a small text embedding table.
• causal_mask(length) returning a (length, length) lower-triangular boolean tensor.
• A demo that feeds a batch of two text sequences of length 10 with image memory of length 197 and prints output shape, the self-attention mask shape, and the cross-attention output norm per position.

Run it:

python3 code/main.py

The decoder produces a (2, 10, text_vocab) logits tensor, and the mask shape is (10, 10). The KV-cache check confirms the cached and uncached logits agree to within 1e-4.

Cross-attention shows up in two production families:

• Flamingo and IDEFICS. Insert a cross-attention sub-layer every K language model blocks, with a frozen LM. The vision-language adapter is the cross-attention block plus its gate.
• BLIP-2. The Q-Former uses cross-attention from a fixed set of 32 query tokens into the image features, then projects the queries into the LM embedding space.

The block built here maps directly onto both, and so does the mask discipline: causal on self-attention, none on cross-attention.

code/test_main.py covers:

• causal mask is lower-triangular and matches expected boolean shape
• cross-attention output shape is (B, Nt, hidden) regardless of key length
• KV-cache path matches uncached path to float tolerance
• a batch-size mismatch between text and image streams, or a wrong-shaped causal mask, raises a clear ValueError
• a full decoder forward pass produces the right batch and sequence shape

Run them:

python3 -m unittest code/test_main.py

1. Add a learned tanh gate to the cross-attention residual (the Flamingo trick) and verify training converges from a near-zero initial gate. The gate starts at 0; the model recovers text-only behavior before mixing the image stream in.

2. Implement interleaved attention where the same decoder consumes multiple images plus multiple text segments. Build the per-sample cross-attention mask that prevents text segment 2 from attending to image 1.

3. Profile the cross-attention vs the self-attention layer at Nt=64, Nv=576 (raise max_text_len in DecoderConfig above its default of 32 first) (a 24x24 grid at higher resolution). The cross-attention cost is Nt * Nv and dominates at high image resolution.

4. Add a query-side dropout on the cross-attention map and measure caption diversity on the demo (caption sample variance increases with dropout in the cross map).

5. Swap the cross-attention layer for a Q-Former-style attention block where a fixed 32-token query pool attends to image features once per layer.

• Flamingo (2022) for the canonical late-fusion design with gated cross-attention.
• BLIP-2 (2023) for the Q-Former, which is a cross-attention block dressed as a learned query pool.
• IDEFICS (2023) for an open-weight reproduction of the Flamingo recipe.