Cross-Attention

Cross-attention vs self-attention

In self-attention, queries, keys, and values all come from the same sequence — each position attends to every other position in the same sequence. In cross-attention, the queries come from one sequence (typically the decoder's hidden states) while the keys and values come from a different sequence (typically the encoder's output). This lets the decoder look up which parts of the source are most relevant for generating the next token.

Cross-attention is the mechanism that connects encoder and decoder in seq2seq architectures like the original Transformer (Vaswani et al. 2017), T5, and BART. It also reappears in multimodal models (LLaVA, Flamingo) to fuse visual features into a language decoder.

Mathematical formulation

Given decoder hidden states H_dec ∈ ℝ^{T_dec × d_model} and encoder output H_enc ∈ ℝ^{T_enc × d_model}:

• Q = H_dec · W_Q (from decoder)
• K = H_enc · W_K (from encoder)
• V = H_enc · W_V (from encoder)

Attention(Q, K, V) = softmax(QK^T / √d_k) · V

The result has shape T_dec × d_v — each decoder position gets a weighted combination of encoder positions. The attention weights (shape T_dec × T_enc) show which source positions the decoder is attending to when generating each target token.

Implementation in PyTorch

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CrossAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, dec_hidden, enc_output, enc_mask=None):
        # dec_hidden: (B, T_dec, d_model)
        # enc_output: (B, T_enc, d_model)
        B, T_dec, _ = dec_hidden.shape
        T_enc = enc_output.shape[1]

        Q = self.W_q(dec_hidden).view(B, T_dec, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_k(enc_output).view(B, T_enc, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_v(enc_output).view(B, T_enc, self.n_heads, self.d_head).transpose(1, 2)

        scale = math.sqrt(self.d_head)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / scale  # (B, H, T_dec, T_enc)

        if enc_mask is not None:
            scores = scores.masked_fill(enc_mask.unsqueeze(1).unsqueeze(2) == 0, -1e9)

        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)  # (B, H, T_dec, d_head)
        out = out.transpose(1, 2).contiguous().view(B, T_dec, -1)
        return self.W_o(out), attn  # return attn weights for visualisation

Role in encoder-decoder models

In the original Transformer decoder, each layer has three sublayers: (1) masked self-attention over previously generated tokens, (2) cross-attention over the encoder output, and (3) a feed-forward network. The cross-attention sublayer is what allows the model to focus on relevant source tokens when generating each target token.

Modern variants like T5 follow the same pattern. During inference, the encoder processes the input once; the decoder runs cross-attention at every step, re-attending to encoder states. This is efficient because encoder keys and values can be cached for the entire generation.

Cross-attention in multimodal models

Cross-attention is the dominant mechanism for fusing modalities. In Flamingo, visual features from a frozen vision encoder are injected into a language model via cross-attention layers inserted between transformer blocks. In LLaVA, visual tokens are projected into the language model's embedding space and prepended to the text — effectively using the language model's self-attention as a form of cross-attention. In Stable Diffusion's UNet, cross-attention fuses CLIP text embeddings with image features at multiple resolutions to condition the denoising process on the text prompt.

Attention patterns visualised

When you inspect cross-attention weights in a trained translation model, clear patterns emerge: when generating a verb in German, the model strongly attends to the corresponding verb in the English source. In summarisation models, attention heads focus on key sentences in the source document. This interpretability is one reason cross-attention is preferred over alternatives like concatenation for fusing two sequences.

import matplotlib.pyplot as plt

def plot_cross_attn(attn_weights, src_tokens, tgt_tokens, head=0):
    # attn_weights: (n_heads, T_dec, T_enc)
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.imshow(attn_weights[head].cpu().numpy(), cmap="Blues", aspect="auto")
    ax.set_xticks(range(len(src_tokens)))
    ax.set_xticklabels(src_tokens, rotation=45, ha="right")
    ax.set_yticks(range(len(tgt_tokens)))
    ax.set_yticklabels(tgt_tokens)
    ax.set_xlabel("Source (encoder)"); ax.set_ylabel("Target (decoder)")
    plt.tight_layout(); plt.show()

Gotchas

• Key-value caching: In cross-attention during generation, K and V are always derived from the encoder output (which doesn't change). Cache them after the encoder pass to avoid recomputing every decoder step.
• Padding masks: The encoder output has padding tokens for shorter sequences in a batch. Apply the encoder padding mask to cross-attention scores to prevent the decoder from attending to padding.
• Decoder-only models: GPT-style decoder-only models have no cross-attention layer — they rely entirely on self-attention over a concatenated context (prompt + response). Cross-attention is specific to encoder-decoder architectures.