Cross Attention — Module 10.1

How This Lesson Fits the Module & Engineering Practice

Self-attention contextualizes one sequence. Many tasks need a second stream: a decoder reading an encoder, a captioner reading image tokens, or a multimodal model aligning text to vision features. That pattern is cross-attention.

In the original transformer, decoder layers stack masked self-attention, then cross-attention into encoder memory, then a feed-forward network. Engineers who confuse self and cross wiring will mis-implement seq2seq and retrieval-style blocks.

Learning Objectives

By the end of this lesson, students should be able to:

Define cross-attention in terms of separate query and key/value sources.
Compare self-attention and cross-attention side by side.
Describe where cross-attention sits in an encoder–decoder stack.
Implement a cross-attention forward pass in PyTorch.
Reason about sequence-length shapes when source and target lengths differ.
Avoid common wiring mistakes (swapping Q vs K/V sources).

Definition

Cross-attention is attention where queries come from one sequence (the target) and keys and values come from another sequence (the source / memory). The target positions retrieve and mix information from the source.

The Cross-Attention Formula

Let \(X_{\mathrm{tgt}} \in \mathbb{R}^{m \times d}\) be the target and \(X_{\mathrm{src}} \in \mathbb{R}^{n \times d}\) the source:

\[Q = X_{\mathrm{tgt}} W_Q,\quad K = X_{\mathrm{src}} W_K,\quad V = X_{\mathrm{src}} W_V\]

\[\mathrm{CrossAttn} = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\]

The score matrix is \(m \times n\): each of \(m\) target positions distributes mass over \(n\) source positions. Output length matches the target (\(m\)), not the source.

Self-Attention vs Cross-Attention (Clear Comparison)

Aspect	Self-attention	Cross-attention
Inputs	One tensor \(X\)	Two tensors: target + source
Q from	\(X\)	Target
K, V from	\(X\)	Source / encoder memory
Score shape	\(n \times n\)	\(m \times n\) (lengths may differ)
Intuition	“Talk to yourself”	“Look up the other sequence”
Classic home	Encoder; decoder self-attn	Decoder–encoder attention

When to use self

Build contextual token reps in one modality/stream.
Language modeling within a single context window.
Encoder-only classifiers (BERT-style).

When to use cross

Machine translation: decoder reads encoder.
Image captioning: text queries vision tokens.
Any retrieve-and-condition design.

Encoder–Decoder Placement

1. Encode source

Run the encoder with self-attention to get memory.

2. Decoder self-attn

Target tokens attend among themselves (often masked).

3. Cross-attn

Target queries attend to encoder keys/values.

4. FFN + residual

Finish the block with FFN, residuals, and norms.

Related module pages: Query, Key, Value, Attention Score, Scaled Dot-Product Attention.

PyTorch: Cross-Attention Forward

import math
import torch
from torch import nn

class CrossAttention(nn.Module):
    def __init__(self, d_model: int, d_k: int | None = None):
        super().__init__()
        d_k = d_k or d_model
        self.d_k = d_k
        self.W_q = nn.Linear(d_model, d_k, bias=False)
        self.W_k = nn.Linear(d_model, d_k, bias=False)
        self.W_v = nn.Linear(d_model, d_k, bias=False)

    def forward(self, x_tgt, x_src):
        # x_tgt: (B, m, d)  x_src: (B, n, d)
        Q = self.W_q(x_tgt)
        K = self.W_k(x_src)
        V = self.W_v(x_src)
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)  # (B, m, n)
        weights = torch.softmax(scores, dim=-1)
        return weights @ V  # (B, m, d_k)

tgt = torch.randn(2, 5, 32)   # shorter target
src = torch.randn(2, 12, 32)  # longer source
out = CrossAttention(32)(tgt, src)
print(out.shape)  # torch.Size([2, 5, 32])

Practical Notes

Benefits

Aligns sequences of different lengths cleanly.
Reusable memory: encoder runs once, decoder queries often.
Natural fit for translation and multimodal conditioning.

Caveats

Still \(O(m n)\) attention cost per layer.
Wrong Q/K/V wiring silently breaks learning.
Needs good encoder representations to retrieve from.

Common Mistake

Feeding the source into the query projection and the target into keys/values. Queries must come from the sequence that needs information; keys/values must come from the sequence that provides it.

Misconception

“Cross-attention replaces self-attention.” In a full decoder block you usually need both: self-attention for target coherence and cross-attention for source grounding. See also multi-head attention, which can wrap either pattern.

Knowledge Check

Short Answer: In cross-attention, which sequence produces the queries? Answer: The target sequence (the one requesting information).
True/False: Keys and values in cross-attention come from the source / memory sequence. Answer: True.
Multiple Choice: If target length is 5 and source length is 12, score shape per batch item is: (a) \(5\times5\), (b) \(5\times12\), (c) \(12\times12\). Answer: (b).
Short Answer: How does self-attention differ from cross-attention? Answer: Self uses one sequence for Q/K/V; cross uses target for Q and source for K/V.
True/False: Cross-attention output length matches the source length. Answer: False; it matches the target (query) length.
Multiple Choice: Classic home of cross-attention: (a) decoder looking at encoder, (b) only BatchNorm, (c) only max-pooling. Answer: (a).
Short Answer: Name one application of cross-attention. Answer: Machine translation, image captioning, or any encoder–decoder / multimodal conditioning setup.
True/False: Cross-attention and self-attention use the same scaled-dot-product formula once Q, K, V exist. Answer: True.
Multiple Choice: Which module page is the natural sibling contrast? (a) Self Attention, (b) Dropout only, (c) Stemming. Answer: (a).
Short Answer: What comes next after understanding cross-attention in this module track? Answer: Multi-head attention (parallel heads + output projection).

Key Takeaways

Cross-attention: Q from target, K/V from source; score matrix is \(m \times n\).
It is the main mechanism for decoder–encoder and many multimodal links.
Do not confuse it with self-attention; both often appear in one decoder block.
Shape checks (output length = target length) catch wiring bugs early.
Next: Multi-Head Attention.

Trainer’s Guide

Hands-on idea: Run the PyTorch snippet with unequal lengths and have students predict out.shape before printing.

Discussion prompt: In a vision–language model, which modality should provide Q vs K/V for caption generation—and why?

Recap: Cross-attention lets a target sequence retrieve from a source sequence via separate Q vs K/V streams. Continue with Multi-Head Attention.