"Dog bites man" is not news. "Man bites dog" is. The two sentences contain exactly the same words, and every reader knows they mean opposite things. Here is a surprising fact about the attention layer we have built. Without some extra ingredient, it cannot tell them apart. It would assign "dog" the same output vector in both sentences, and the same for "bites" and "man". We will prove this, because it is the cleanest example of a symmetry in deep learning and the reason every Transformer adds positional information to its inputs.

First, though, we generalize. Everything so far has been self-attention, where queries, keys, and values all come from the same sequence. Nothing forces that. The query side and the key-value side can be different sequences entirely, which is cross-attention, the original form of attention in translation and today the standard way to connect one modality to another. The section ends with two ideas that change how you see attention, as message passing on a graph, and as an operation on sets that the Perceiver turned into a general-purpose architecture.

Where the queries, keys, and values come from

Self-attention. One sequence $\mathbf{X} \in \R^{T\times d}$ supplies everything, $\mathbf{Q} = \mathbf{X}\mathbf{W}_Q$, $\mathbf{K} = \mathbf{X}\mathbf{W}_K$, $\mathbf{V} = \mathbf{X}\mathbf{W}_V$. Each token gathers information from the other tokens of its own sequence. The score matrix is square, $T\times T$.

Cross-attention. The queries come from one sequence $\mathbf{X} \in \R^{T_q\times d}$, and the keys and values from another, $\mathbf{Z} \in \R^{T_k\times d_z}$, often called the context or memory,

In plain words, each position of $\mathbf{X}$ asks a question, and the answers are looked up in $\mathbf{Z}$.

Three facts follow from the shapes. The score matrix is $T_q\times T_k$, generally not square. The output has one row per query, so cross-attention returns a sequence as long as $\mathbf{X}$, enriched with information from $\mathbf{Z}$. And the two sequences can have different widths, $d$ and $d_z$, because $\mathbf{W}_K$ and $\mathbf{W}_V$ translate the context into the query's key space. An image encoder with 1,024-dimensional features can feed a language model of width 4,096 this way.

This is Bahdanau attention (§3.6) in modern notation, with a scaled dot product in place of the additive score. The decoder state is the query, and the encoder states are the keys and values. In the original Transformer, each decoder layer has three sublayers in order, a causal self-attention over the target tokens generated so far, then a cross-attention whose queries come from the decoder and whose keys and values come from the final encoder output, then a feed-forward network (Vaswani et al. 2017). The cross-attention needs no causal mask, because the whole source sentence is known before decoding starts. It does need the source's padding mask. And since the encoder output never changes during generation, each layer computes its cross-attention keys and values once per source and reuses them for every generated token.

The same block appears all over modern AI. Speech recognizers like Whisper cross-attend from text tokens to audio features. Text-to-image diffusion models cross-attend from image patches to the prompt's text embeddings, which is how the prompt steers the picture. Llama 3's vision adapter feeds image-encoder features into the language model through a series of cross-attention layers (Llama Team 2024), and vision-language models in the Flamingo family do the same with gated cross-attention layers inserted into a frozen language model.

import numpy as np

def cross_attention(X, Z, W_Q, W_K, W_V, W_O, h, z_mask=None):
    raise NotImplementedError

Attention does not know the order of its inputs

Now the symmetry. Reorder the tokens of a sequence, and ask what happens to the outputs of self-attention. Let's first check numerically on "dog bites man" with random embeddings and random weights, and no position information.

Each word gets exactly the same output in both sentences, just in a different row. The outputs are permuted along with the inputs. Any summary that ignores row order, such as the mean, is therefore identical for the two sentences, so a classifier built on it would give "dog bites man" and "man bites dog" the same label. This property has a name.

Definition. A function $f$ from $T\times d$ matrices to $T\times d'$ matrices is permutation equivariant if $f(\mathbf{P}\mathbf{X}) = \mathbf{P}f(\mathbf{X})$ for every $T\times T$ permutation matrix $\mathbf{P}$. It is permutation invariant if $f(\mathbf{P}\mathbf{X}) = f(\mathbf{X})$.

Here a permutation matrix has exactly one 1 in every row and every column. If $\pi$ is a permutation of $\{1, \dots, T\}$ and $\mathbf{P}$ has $P_{i,\pi(i)} = 1$, then row $i$ of $\mathbf{P}\mathbf{X}$ is row $\pi(i)$ of $\mathbf{X}$, so $\mathbf{P}\mathbf{X}$ is the same tokens in a new order. Permutation matrices are orthogonal, $\mathbf{P}^\top\mathbf{P} = \mathbf{I}$, because undoing a reordering is another reordering.

The proof

Theorem. Self-attention without a mask, $f(\mathbf{X}) = \operatorname{softmax}\!\big(\mathbf{X}\mathbf{W}_Q(\mathbf{X}\mathbf{W}_K)^\top/\sqrt{d_k}\big)\,\mathbf{X}\mathbf{W}_V$ with a row-wise softmax, is permutation equivariant.

We need one lemma about the softmax first.

Lemma. For any matrix $\mathbf{S} \in \R^{T\times T}$ and permutation matrix $\mathbf{P}$, the row-wise softmax satisfies $\operatorname{softmax}(\mathbf{P}\mathbf{S}\mathbf{P}^\top) = \mathbf{P}\operatorname{softmax}(\mathbf{S})\mathbf{P}^\top$.

Proof of the lemma. Entry $(i, j)$ of $\mathbf{P}\mathbf{S}\mathbf{P}^\top$ is $\sum_{a,b}P_{ia}S_{ab}P_{jb} = S_{\pi(i)\pi(j)}$, so permuting this way reorders both the rows and the columns of $\mathbf{S}$ by $\pi$. Apply the softmax to row $i$,

The middle step holds because as $l$ runs over all positions, so does $\pi(l)$, and a sum does not depend on the order of its terms. ∎

Proof of the theorem. Replace $\mathbf{X}$ by $\mathbf{P}\mathbf{X}$. The projections act row by row, so $\mathbf{Q}' = \mathbf{P}\mathbf{X}\mathbf{W}_Q = \mathbf{P}\mathbf{Q}$, and likewise $\mathbf{K}' = \mathbf{P}\mathbf{K}$ and $\mathbf{V}' = \mathbf{P}\mathbf{V}$. The scores become

By the lemma, $\mathbf{A}' = \mathbf{P}\mathbf{A}\mathbf{P}^\top$. Finally,

using $\mathbf{P}^\top\mathbf{P} = \mathbf{I}$. ∎

In plain words, reordering the tokens reorders the queries, keys, and values in the same way, the attention matrix gets its rows and columns shuffled consistently, and every token ends up with exactly the output it had before, sitting in its new slot.

The result extends to a whole Transformer. Multi-head attention is a sum of equivariant heads followed by a row-wise $\mathbf{W}_O$, so it is equivariant. The MLP, LayerNorm, and residual connection each act on every row separately, and any row-wise map $g$ satisfies $g(\mathbf{P}\mathbf{X}) = \mathbf{P}g(\mathbf{X})$. A composition of equivariant maps is equivariant, since $f_2(f_1(\mathbf{P}\mathbf{X})) = f_2(\mathbf{P}f_1(\mathbf{X})) = \mathbf{P}f_2(f_1(\mathbf{X}))$. So an encoder with no positional information is permutation equivariant from end to end, and anything it outputs after mean pooling is permutation invariant. It treats its input as a set of tokens (strictly, a multiset, since repeated tokens are allowed).

Cross-attention has a slightly different symmetry. Permuting the context with $\mathbf{P}_z$ gives scores $\mathbf{S}\mathbf{P}_z^\top$ (the columns are reordered), the row-wise softmax commutes with that reordering, and the output is $\operatorname{softmax}(\mathbf{S})\mathbf{P}_z^\top\mathbf{P}_z\mathbf{V} = \mathbf{O}$. Cross-attention is invariant to the order of its context and equivariant in its queries. Exercise 4.6a tested exactly this.

Why position must be added, and what breaks the symmetry

For sets, equivariance is exactly right. A point cloud, the atoms of a molecule, the objects in a scene, or a bag of retrieved documents has no natural order, and a model that depended on the order you happened to list them in would be wrong. For language it is a disaster, since order carries meaning. The fix is to break the symmetry on purpose by making each token's vector depend on its position, for instance by adding a position vector $\mathbf{p}_t$ to the embedding at position $t$. Then permuting the tokens does not permute the position vectors, which stay attached to slots, and the proof above no longer applies. How to choose those vectors (sinusoidal, learned, or relative, as in rotary embeddings) is the subject of §5.4 and §5.5.

There is a subtlety that catches people out in interviews. A causal mask also breaks the symmetry. The proof used $\mathbf{S}' = \mathbf{P}\mathbf{S}\mathbf{P}^\top$. With a fixed mask $\mathbf{M}$ the masked scores become $\mathbf{P}\mathbf{S}\mathbf{P}^\top + \mathbf{M}$, which equals $\mathbf{P}(\mathbf{S}+\mathbf{M})\mathbf{P}^\top$ only if $\mathbf{P}\mathbf{M}\mathbf{P}^\top = \mathbf{M}$, and for the causal mask only the identity permutation does that. A decoder therefore knows something about position even with no position encoding at all. The cleanest way to see it is that the first token can attend only to itself, the second to two tokens, and so on. A head whose scores are all equal attends uniformly to the $t$ tokens it can see, and if the first token's value is distinctive, the output at position $t$ is exactly $1/t$ times that value. Position is written into the output.

This is not just a curiosity. Haviv et al. (2022) trained causal language models with no positional encoding at all and found them competitive with models that had one, and probes showed the models had learned absolute position internally. Kazemnejad et al. (2023) found that, on a suite of small-scale reasoning and arithmetic tasks, such "NoPE" decoders generalized to sequences longer than those seen in training better than decoders with absolute, rotary, or ALiBi encodings. Whether that holds for large language models is not settled. Encoders like BERT have no such escape, since without a causal mask they are fully equivariant and need explicit positions.

Attention as message passing on a graph

Here is another way to read $\mathbf{o}_i = \sum_j a_{ij}\mathbf{v}_j$. Draw each token as a node. Draw an edge from $j$ to $i$ whenever $i$ may attend to $j$. A layer of attention then performs one round of message passing. Every node $j$ prepares a message $\mathbf{v}_j$, every node $i$ assigns its incoming edges the weights $a_{ij}$ (a softmax over its neighbors only), and every node replaces its state with the weighted sum of its incoming messages. Self-attention without a mask runs on the complete graph. Each mask from §4.4 is just a different edge set.

This view connects Transformers to graph neural networks. A Graph Attention Network (Veličković et al. 2018) is attention restricted to the edges of a given graph, with its own (additive) score function and multiple heads. Joshi (2020) put the connection plainly, that a Transformer is a graph neural network on the complete graph of its tokens. The complete graph has no structure of its own, so whatever the model knows about order has to arrive as node features, which is the job of positional encodings. It also explains the cost. Message passing costs one message per edge, and the complete graph on $T$ nodes has $T^2$ edges, which is where the quadratic cost of §4.7 comes from. Sparse attention patterns (§12.6) are cheaper graphs, chosen so that information still reaches everywhere within a few layers.

import numpy as np

def graph_attention(X, adj, W_Q, W_K, W_V):
    raise NotImplementedError

Sets, learned queries, and the Perceiver

If attention is naturally an operation on sets, we can design for that instead of fighting it. Three ideas build on each other.

Pooling a set with a learned query. To summarize a set into one vector, cross-attend to it with a single learned query vector $\mathbf{q}$ that is a parameter of the model, $\mathbf{o} = \sum_j \operatorname{softmax}_j(\mathbf{q}^\top\mathbf{k}_j/\sqrt{d_k})\,\mathbf{v}_j$. By the cross-attention result above, the output is invariant to the order of the set. Zaheer et al. (2017) proved that for sets drawn from a countable universe every permutation-invariant function can be written as $\rho\big(\sum_j\phi(\mathbf{x}_j)\big)$ (the Deep Sets form), and attention pooling is a learned, input-dependent weighting of that sum. Lee et al. (2019) built the Set Transformer from such blocks, using a few learned "seed" queries to pool, and a few learned "inducing points" to cut the cost of self-attention over large sets.

Learned queries as slots. Use several learned queries, and cross-attention returns several summaries, one per query, each free to specialize. DETR (Carion et al. 2020) detects objects with 100 learned "object queries" that cross-attend to image features, and each query learns to claim one object. The Q-Former and similar resamplers in vision-language models compress hundreds of image patches into a fixed, small number of tokens the same way.

The Perceiver. Self-attention over $M$ inputs costs $O(M^2)$ per layer, hopeless for raw inputs like the 50,176 pixels of a $224\times224$ image. Jaegle et al. (2021) proposed a small array of $N \ll M$ learned latent vectors that cross-attend to the input, at cost $O(MN)$, followed by a deep stack of self-attention among the latents, at $O(N^2)$ per layer, independent of $M$. The input is touched only through cross-attention, so the depth of the network is decoupled from the size of its input, and one architecture handles images, audio, point clouds, and video given suitable position features. Perceiver IO (Jaegle et al. 2021b) adds a cross-attention at the output, with one query per desired output element, so the outputs can be structured however a task needs.

import numpy as np

def attention_pool(X, q, W_K, W_V, mask=None):
    return (X @ W_V).mean(axis=1)      # placeholder: plain mean pooling

In PyTorch

PyTorch's attention functions take the query and the key-value sequences separately, so cross-attention needs nothing new. Here is an encoder-decoder layer's cross-attention with a source padding mask, followed by a numerical check of permutation equivariance for a whole encoder layer.

Read more

• Vaswani et al. (2017). Attention Is All You Need. Section 3.2.3 describes the three uses of attention in the encoder-decoder model.
• Phuong & Hutter (2022). Formal Algorithms for Transformers. Pseudocode for encoder-decoder, encoder-only, and decoder-only models, with every cross-attention spelled out.
• Lee et al. (2019). Set Transformer and Zaheer et al. (2017), Deep Sets. Permutation-invariant architectures, and pooling a set with learned queries.
• Jaegle et al. (2021). Perceiver and Perceiver IO. Latent cross-attention for arbitrary, very large inputs and outputs.
• Veličković et al. (2018). Graph Attention Networks and Joshi (2020), Transformers are Graph Neural Networks. Attention as message passing.
• Haviv et al. (2022). Transformer Language Models without Positional Encodings Still Learn Positional Information and Kazemnejad et al. (2023), The Impact of Positional Encoding on Length Generalization in Transformers. How a causal mask leaks position, and how far models without positional encodings can go.
• Carion et al. (2020). End-to-End Object Detection with Transformers. DETR's learned object queries.

Watch next

▶

▶