Cross-Attention

Cross-attention lets one sequence attend to another, the bridge that connects an encoder's representation to a decoder's generation.

Attending across two sequences

In self-attention, queries, keys, and values come from the same sequence. In cross-attention, the queries come from one sequence and the keys and values come from another. This lets the first sequence pull information from the second. It is the mechanism that connects the encoder and decoder in an encoder-decoder transformer, and more broadly any architecture where one modality or stream must condition on another.

In the encoder-decoder transformer

Kronos motion — cross section

Each decoder block contains three sublayers: masked self-attention over the tokens generated so far, cross-attention where decoder queries attend to encoder keys and values, and a feedforward layer. The cross-attention step is how the decoder consults the source: while generating each output token, it reads the encoded input, focusing on the source positions most relevant to what it is producing.

python

import numpy as np
def cross_attention(dec, enc, Wq, Wk, Wv):
    Q = dec @ Wq            # queries from decoder
    K = enc @ Wk; V = enc @ Wv   # keys/values from encoder
    s = Q @ K.T / np.sqrt(K.shape[-1])
    s -= s.max(-1, keepdims=True)
    w = np.exp(s); w /= w.sum(-1, keepdims=True)
    return w @ V

Beyond translation

Cross-attention is a general conditioning tool. In text-to-image diffusion models, image features attend to text embeddings so the generated image follows the prompt. In multimodal models, a language decoder attends to visual features from an image encoder. In retrieval-augmented systems, a generator attends to retrieved documents. Wherever an output must be conditioned on a separate source, cross-attention is a natural fit.

Shapes and asymmetry

Cross-attention is asymmetric: the query sequence and the key/value sequence can have different lengths and even different meanings. The output has the same length as the query sequence, since there is one attended result per query. This asymmetry is what allows a decoder of one length to condition on an encoder output of another length.

Relation to self-attention

Mechanically, cross-attention is identical to self-attention except for where queries versus keys and values originate. Both use the same scaled dot-product with softmax, and both are usually multi-headed. Understanding self-attention gives cross-attention almost for free; the only new idea is that the two roles are filled by two different sequences.

Queries from one sequence, keys and values from another.
Connects encoder and decoder in transformers.
Conditions outputs on text, images, or retrieved data.
Same math as self-attention, different sources.