Self-Attention vs Cross-Attention

This is the hands-on post, and it is intentionally code-heavy. The previous posts built understanding from concept through mechanics; here you assemble a single, runnable multi-head attention module that serves BOTH flavors, then call it two ways and watch the score shapes differ.

The design choice that makes this work is taking the query input and the key/value input as separate arguments. Pass the same tensor for both and you get self-attention; pass the decoder state and the encoder memory and you get cross-attention. The most valuable thing you can do is run it and inspect the returned attention weights — square for self, rectangular for cross — to see the distinction with your own eyes.

The module's constructor is deliberately built to support both flavors. It stores the head count and the per-head dimension d_k, then creates three linear layers: a query projection, a combined key-value projection of width two-times-d_model, and an output projection. Separating the query projection from the key-value projection is exactly what lets the module accept a different input for the query than for the keys and values.

This is a small but important departure from a self-attention-only module, which often fuses all three projections into one because they share an input. Here, because cross-attention feeds Q and K/V from different sequences, the query projection must be applied to a different tensor than the key-value projection — so they are kept separate from the start.

The first half of the forward pass shows the key design decision: it takes two inputs, x_q for the query and x_kv for the keys and values, and crucially reads their lengths independently as Tq and Tk. It projects the query from x_q and the key and value from x_kv, then reshapes each into multiple heads. Because Tq and Tk are read separately, the query and key sequences are free to differ in length.

That independence is the whole point. In a self-attention call the two inputs are the same tensor so Tq equals Tk, but the code never assumes that. The score matrix it computes, q times k-transpose, is therefore Tq by Tk in general — square only when the caller happens to pass the same sequence for both. This is how one module transparently handles both flavors.

The second half of the forward pass finishes the computation in a flavor-agnostic way. It optionally applies the mask by filling masked positions with negative infinity, softmaxes the scores into weights, blends the values, transposes and reshapes to merge the heads back into a single vector per query position, and projects out. It returns both the output and the attention weights, whose shape is batch by heads by Tq by Tk.

That returned weight shape is the diagnostic that tells you which flavor ran. If Tq equals Tk the map is square — self-attention; if they differ it is rectangular — cross-attention. Returning the weights rather than discarding them is deliberate: inspecting them is how you confirm, at a glance, that your call did what you intended and aligned the sequences you meant to align.

This is the self-attention call. You instantiate the layer once, then pass the same tensor x as both the query input and the key/value input, with no mask. Because both inputs are the five-token sequence, the output has shape batch by five by sixty-four — the same layout, re-mixed — and the attention weights come back as batch by eight by five by five: square, one map per head relating every token to every other token in the sequence.

The square shape is the signature of self-attention. Printing it confirms that passing one tensor twice produced a within-sequence relation. This is the simplest possible demonstration that self-attention is just the general attention module called with matching query and key/value inputs — no special module, no special flag, just the same tensor in both slots.

This is the cross-attention call, using the very same layer object. Now the query input is the decoder state with five positions and the key/value input is the encoder output with seven positions. The output still has five positions — one per query — but the attention weights come back as batch by eight by five by seven: rectangular, a five-by-seven alignment map between target and source for each head.

The rectangular shape is the signature of cross-attention, and that it emerged from the identical layer used for the self-attention call is the entire lesson of this post made tangible. Nothing about the module changed; only the inputs did. The five-by-seven map is the alignment you would visualize in translation, showing which of the seven source positions each of the five target positions attended to.

This comparison summarizes the two calls and their differing outputs. Calling the layer with x, x gives a five-by-five square map where queries, keys, and values all come from x and tokens relate to tokens. Calling it with dec, enc gives a five-by-seven rectangular map where the query comes from the decoder and the keys and values from the encoder, aligning target to source.

The table reinforces that the difference is entirely at the call site. One layer, two call patterns, two shapes of attention map. Internalizing this means you will never again wonder whether you need a separate 'cross-attention class' — you do not. You need one well-designed attention module and the discipline to feed it the right inputs.

This snippet wires both flavors into a tiny but complete decoder block, which is where the two come together in a real architecture. The block holds two attention instances — one for self-attention, one for cross-attention — plus a feed-forward network. In the forward pass, the target y first attends to itself with a causal mask, then cross-attends to the encoder memory, then passes through the feed-forward layer, each step added back as a residual.

The ordering mirrors the real Transformer decoder: understand the generated target first, then pull in the source, then process. Note that the self-attention call passes y for both inputs plus the causal mask, while the cross-attention call passes y as the query and memory as the key/value with no causal mask. This block is the smallest piece of code that uses both flavors in their correct, distinct roles.

These shapes-and-gotchas are where practitioners lose time when handling both flavors. The headline point is that the query and key/value inputs can have different lengths — Tq versus Tk — which is precisely what enables cross-attention; code that hard-codes a single sequence length will break on the rectangular case. The attention weights are shaped batch by heads by Tq by Tk, square only when the call is self-attention.

The remaining traps carry over from attention in general: pass a causal mask for decoder self-attention but never for cross-attention, ensure d_model divides evenly by the head count, and remember that a reshape after a transpose may require calling contiguous first. None of these raise conceptual errors in your understanding of the two flavors, but each is a concrete trap that turns a correct mental model into broken code.

This snippet shows that, in practice, you rarely hand-roll either flavor. PyTorch's fused scaled_dot_product_attention takes query, key, and value as separate arguments, so you call it with the same tensor three times (plus is_causal) for self-attention, or with the decoder query and encoder key/value for cross-attention. The full nn.MultiheadAttention module likewise accepts query, key, and value separately, so the same module handles both by what you pass.

The reason to learn the from-scratch version first is that these built-ins are black boxes until you know what they compute. Once you have seen that self versus cross is purely a matter of which tensors fill the query, key, and value slots, the built-ins' separate-argument signatures make perfect sense — they are designed for exactly this, and in real projects you should prefer them for speed and correctness.

This recap consolidates the build: one module takes separate x_q and x_kv inputs; self-attention is calling it with the same tensor twice; cross-attention is calling it with the target and the source; the self map is square while the cross map is rectangular; and a decoder block runs self-attention, then cross-attention, then the feed-forward sublayer.

With a working module that serves both flavors — and whose score shapes you have verified change with the inputs — the only thing left is to learn the traps. The final post catalogs the self-versus-cross mistakes that produce wrong alignments without ever raising an error.

The teaser points to the common-mistakes post. You now have a single attention module that runs both flavors and a decoder block that uses them in their correct roles; the next step is hardening your instincts against the silent failures — feeding K and V from the wrong sequence, masking cross-attention causally, forgetting the source padding mask, and over-trusting the alignment maps.