The Attention Mechanism

In a vanilla encoder-decoder, an encoder (often a recurrent network) reads the input sequence $x_1, \dots, x_n$ and produces a single context vector $c$ — typically the encoder's final hidden state. The decoder then generates the output sequence conditioned only on $c$ and its own previous outputs. The problem is structural: $c$ has fixed dimensionality regardless of $n$, so as $n$ grows, $c$ must represent more information with the same capacity. Empirically, translation BLEU scores for such models fall off a cliff past roughly 20–30 tokens. Attention removes the bottleneck by giving the decoder access to the entire sequence of encoder hidden states $h_1, \dots, h_n$, and letting it compute, at each decoding step $t$, a context vector specific to that step:

where $\alpha_{t,i}$ are non-negative weights summing to 1 — the attention distribution — computed from how relevant encoder position $i$ is to the decoder's current need at step $t$. The modern, general formulation abstracts this into three roles, borrowed from information retrieval:

• Query $q$: a representation of "what am I looking for right now" (e.g., the decoder's current hidden state).
• Key $k_i$: a representation of "what does item $i$ offer," used to compute relevance against the query.
• Value $v_i$: the actual content of item $i$ that gets retrieved if it's deemed relevant.

Relevance is scored by a compatibility function $f(q, k_i)$, turned into a probability distribution via softmax, and used to weight the values:

The most widely used compatibility function is the scaled dot product, $f(q, k_i) = \dfrac{q \cdot k_i}{\sqrt{d_k}}$, giving the equation that underlies Transformers:

Here $Q, K, V$ are matrices stacking queries, keys, and values as rows, so the whole computation is expressed with two matrix multiplications and a softmax — no loops, fully batchable on a GPU.

Step-by-step mechanism (single query, additive or dot-product attention): 1. Obtain encoder representations $h_1, \dots, h_n$ (from an RNN, CNN, or embedding layer). 2. At decoder step $t$, form a query $q_t$ (e.g., the decoder's hidden state $s_{t-1}$). 3. Project or reuse $h_i$ as keys $k_i$ and values $v_i$ (in vanilla seq2seq attention, $k_i = v_i = h_i$; in Transformers, $k_i, v_i$ are learned linear projections of $h_i$). 4. Compute alignment scores $e_{t,i} = f(q_t, k_i)$ for every $i$. 5. Normalize with softmax to get weights $\alpha_{t,i}$. 6. Compute the context vector $c_t = \sum_i \alpha_{t,i} v_i$. 7. Combine $c_t$ with the decoder state (e.g., concatenate then project) to produce the output distribution at step $t$. 8. Repeat for every decoder step; each step gets its own $\alpha_{t,\cdot}$, so different output words can attend to different input words.

Multi-head attention runs $H$ independent copies of this process in parallel, each on a different learned linear projection of $Q, K, V$ into a lower-dimensional subspace, then concatenates the results and projects back: ``` function SCALED_DOT_PRODUCT_ATTENTION(Q, K, V, mask=None):

# Q: (n_q, d_k), K: (n_kv, d_k), V: (n_kv, d_v) scores = MATMUL(Q, TRANSPOSE(K)) / sqrt(d_k) # (n_q, n_kv) if mask is not None: scores = scores + mask # -inf where disallowed weights = SOFTMAX(scores, axis=-1) # rows sum to 1 output = MATMUL(weights, V) # (n_q, d_v) return output, weights function MULTI_HEAD_ATTENTION(X_q, X_kv, H, d_model, mask=None):

# X_q: (n_q, d_model), X_kv: (n_kv, d_model) d_k = d_model / H heads = [] for i in 1 .. H: Q_i = MATMUL(X_q, W_Q[i]) # (n_q, d_k) K_i = MATMUL(X_kv, W_K[i]) # (n_kv, d_k) V_i = MATMUL(X_kv, W_V[i]) # (n_kv, d_k) head_i, _ = SCALED_DOT_PRODUCT_ATTENTION(Q_i, K_i, V_i, mask) heads.append(head_i) concat = CONCAT(heads, axis=-1) # (n_q, d_model) return MATMUL(concat, W_O) # (n_q, d_model) ``` Two important variants controlled by the `mask` argument:

• Cross-attention (decoder attends to encoder): queries come from the decoder, keys/values from the encoder; no masking needed.
• Causal self-attention (used in decoders/autoregressive models): queries, keys, and values all come from the same sequence, but a mask sets $\text{scores}[t, i] = -\infty$ for $i > t$, so position $t$ cannot see future tokens.

Time complexity. For a single attention operation with query length $n_q$, key/value length $n_{kv}$, and feature dimension $d$:

• Computing $QK^\top$ costs $O(n_q \cdot n_{kv} \cdot d)$.
• Softmax over the score matrix costs $O(n_q \cdot n_{kv})$.
• Multiplying weights by $V$ costs $O(n_q \cdot n_{kv} \cdot d)$.

Total: $O(n_q n_{kv} d)$ time. For self-attention where $n_q = n_{kv} = n$, this is $O(n^2 d)$ — quadratic in sequence length. This is the well-known efficiency tradeoff of attention-based architectures: unlike an RNN's $O(n \cdot d^2)$ per-layer cost (linear in sequence length but with recurrence preventing parallelism across time), attention is parallelizable across all positions simultaneously but scales quadratically with $n$. For multi-head attention with $H$ heads of dimension $d/H$ each, the total cost across heads is still $O(n^2 d)$ (the per-head cost $O(n^2 \cdot d/H)$ summed over $H$ heads), so multi-head attention has essentially the same asymptotic cost as single-head attention with full dimension $d$ — the benefit is representational (multiple relevance patterns), not asymptotic.

Space complexity. Storing the full attention weight matrix requires $O(n_q \cdot n_{kv})$ memory, in addition to $O(nd)$ for the representations themselves. For long sequences this $O(n^2)$ memory term dominates and is the practical bottleneck that motivates sparse/linear attention variants (not covered here, but worth knowing exist).

Correctness argument. Three properties must hold for attention to behave as a well-defined weighted retrieval: 1. The weights form a valid probability distribution. Softmax guarantees $\alpha_i \geq 0$ and $\sum_i \alpha_i = 1$ for any real-valued scores, since $\text{softmax}(e)_i = \exp(e_i)/\sum_j \exp(e_j)$ is a ratio of positive numbers normalized to sum to 1. This means $c_t$ is a genuine convex combination of the values — it always lies inside their convex hull, so it can never "invent" values outside the range of what's observed, only reweight them. 2. Scaling by $\sqrt{d_k}$ prevents saturation. If query and key components are independent random variables with mean 0 and variance 1, then $q \cdot k = \sum_{j=1}^{d_k} q_j k_j$ has mean 0 and variance $d_k$ (variance of a sum of $d_k$ independent terms each with variance 1). Without correction, as $d_k$ grows, dot products grow in magnitude, pushing softmax inputs into regions where gradients vanish (softmax saturates towards one-hot). Dividing by $\sqrt{d_k}$ restores unit variance, keeping the softmax in a well-conditioned regime — this is a variance-normalization argument, not a heuristic. 3. Attention is differentiable end-to-end. Every operation (matrix multiply, softmax, weighted sum) has a well-defined gradient, so the entire mechanism can be trained with backpropagation jointly with the rest of the network — this is precisely why it superseded earlier hard/discrete alignment models that required separate training signals (e.g., HMM-style alignment in classical statistical machine translation).

Honest tradeoffs. Attention removes the fixed-vector bottleneck and enables long-range dependency modeling in $O(1)$ path length between any two positions (versus $O(n)$ for RNNs), which is why it dramatically improves long-sequence performance. But this comes at $O(n^2)$ time and memory, versus $O(n)$ for a plain RNN pass — for very long sequences (documents, genomes, high-resolution images) this quadratic cost becomes prohibitive, motivating an entire research subfield of efficient/sparse attention.

Consider a toy machine-translation decoder generating the French word for "bank" in the sentence "The river bank was steep," where the encoder has produced 4 hidden states $h_1, h_2, h_3, h_4$ corresponding to "river," "bank," "was," "steep" (dimension $d=4$, values shown for illustration). Suppose the decoder's query at this step, $q_t$, is designed to look for "the noun modified by a geography-related word."

(Scores computed with $d_k = 4$; softmax applied across the column.) Step through the computation: 1. Scores. The query strongly resembles $k_1$ (river) and $k_2$ (bank) since it's seeking geography-related context, giving them the highest raw scores 1.8 and 2.6. 2. Softmax normalization. Exponentiating: $e^{1.8}=6.05,\, e^{2.6}=13.46,\, e^{-1.0}=0.37,\, e^{0.4}=1.49$. Sum $= 21.37$. Dividing each by the sum gives $\alpha \approx [0.283, 0.630, 0.017, 0.070]$ (matches the table up to rounding). 3. Weighted context. With values $v_i = h_i$ (say $v_2$ encodes "riverbank" semantics strongly in its first coordinate, e.g. $v_2 = [2.0, 0.5, 0.1, 0.3]$ and the others are smaller in that coordinate), the context vector $c_t = \sum_i \alpha_i v_i$ is dominated by $v_2$ (weight 0.63) with a meaningful contribution from $v_1$ (weight 0.28) — correctly disambiguating "bank" as riverbank rather than financial institution, because the model attended heavily to "river" and "bank" and almost ignored "was"/"steep." 4. Decoding. This $c_t$ is concatenated with the decoder's hidden state and passed through the output projection + softmax over the French vocabulary, favoring "rive" (riverbank) over "banque" (financial bank). Contrast this with a fixed-vector encoder-decoder: it would have compressed all four words into one $c$ used for every decoding step, unable to sharpen focus onto "river" and "bank" specifically when disambiguating this particular output word — the attention mechanism's per-step, per-query flexibility is exactly what fixes this.

The attention mechanism as used in deep sequence models was introduced to solve the fixed-vector bottleneck of encoder-decoder machine translation, with the query-key-value abstraction later crystallizing the idea and enabling the fully attention-based Transformer.

Attention is the computational core of the Transformer, which underlies nearly every large-scale language model in production, including GPT-family, BERT-family, and T5-family systems used for chat assistants, code generation, search ranking, and document summarization. Machine translation systems (e.g., production NMT pipelines at major search and cloud providers) moved from RNN-with-attention to fully self-attentional Transformers for both quality and training-throughput reasons, since attention parallelizes across sequence positions on GPUs/TPUs far better than the sequential dependency of recurrence. Speech recognition systems use encoder-decoder attention to align variable-length audio frames to variable-length transcripts. Vision Transformers apply the same QKV mechanism to image patches, treating an image as a sequence of tokens, unifying vision and language architectures under one attention-based paradigm. Multimodal systems (image captioning, visual question answering) use cross-attention to let a text decoder query visual features directly, exactly as the original attention mechanism let a translation decoder query encoder states. More broadly, viewing attention as a learned, weighted aggregation over a set of neighbors — rather than over a fixed grid or sequence order — is the same conceptual leap that motivates message passing in Graph Neural Networks: An Introduction, where nodes aggregate information from their neighbors with learned, data-dependent weights instead of a fixed convolutional kernel.

Ask the tutor: How does multi-head attention actually differ in what it learns compared to using one large single head, and can you show a case where multiple heads capture qualitatively different relationships?