Attention mechanism

The attention mechanism lets a model dynamically focus on the most relevant parts of the input when producing each output, rather than compressing everything into one fixed vector. It is like how you glance back at specific words in a sentence when answering a question about it. Each output gets a weighted combination of all input representations, where higher weights mean greater relevance.

Queries represent what the current token is 'looking for'; Keys represent what each input token 'advertises' about itself; Values hold the actual content to be retrieved. The dot product of a Query with each Key gives a relevance score, which is softmaxed into weights, and the output is a weighted sum of the Values. It works like a search engine: your query matches document keys to retrieve document values.

Early encoder-decoder models compressed the entire input into one fixed-size context vector, which became a bottleneck for long sentences since a single vector cannot store everything. Attention was introduced so the decoder can look at all encoder hidden states at every step and focus on the most relevant ones dynamically. This dramatically improved performance on tasks like machine translation, especially for longer inputs.

Softmax converts the raw attention scores (dot products of queries and keys) into a probability distribution that sums to 1, ensuring each output is a proper weighted average of the value vectors. Positions with higher scores receive more weight and therefore more influence on the output. Without softmax, the weights would not sum to 1 and could be arbitrarily large or negative, making the combination uninterpretable and numerically unstable.

Self-attention computes attention between tokens within the same sequence — every token attends to every other token in the same input to build a richer, context-aware representation. In encoder-decoder attention, the decoder's query attends to the encoder's keys and values (two separate sequences). Self-attention allows a model to capture relationships between all words simultaneously, regardless of their distance, within a single sentence or document.

The formula is Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V, where d_k is the dimension of the key vectors. First, Q·Kᵀ computes pairwise similarity scores between all queries and keys. Dividing by √d_k prevents these scores from becoming too large; softmax turns them into weights; and the final multiplication with V produces context-weighted output vectors.

import numpy as np

def scaled_dot_product_attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)          # shape: (n_q, n_k)
    weights = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)  # softmax
    return weights @ V                         # shape: (n_q, d_v)

When the key dimension d_k is large, dot products between random query and key vectors grow in magnitude (their variance scales with d_k), pushing the softmax into near-zero-gradient saturation regions. Dividing by √d_k keeps the scores in a moderate range so the softmax output remains spread out and gradients flow well during training. This scaling is theoretically motivated: if q and k are vectors of i.i.d. values with mean 0 and variance 1, then q·k has mean 0 and variance d_k, so dividing by √d_k gives variance 1.

Multi-head attention runs h attention operations in parallel, each with its own learned projection matrices for Q, K, and V. Each 'head' can specialize in capturing a different type of relationship — for example, one head for syntactic dependencies, another for semantic similarity. The outputs of all heads are concatenated and linearly projected, giving representations that integrate multiple perspectives without significantly increasing computation.

Unlike RNNs that process tokens in order, self-attention treats all tokens symmetrically regardless of position — so 'cat chases dog' and 'dog chases cat' would produce identical representations without positional information. Positional encoding adds a unique signal to each token's embedding based on its position, usually via sine and cosine functions of different frequencies, added element-wise. This gives the model the ability to distinguish token order without breaking parallelism.

In RNNs, information from an early token must pass through every intermediate hidden state to reach a late token, causing gradients to vanish or explode over many steps. Attention creates a direct, one-step connection from any token to any other token, regardless of distance, via a single dot-product computation. This means a word at position 1 and a word at position 100 are equally 'close' in terms of information flow and gradient propagation.

The attention score measures how much token i should 'pay attention to' token j when computing its output representation — a higher score means token j is more relevant to understanding token i. It is computed as the dot product of token i's query vector with token j's key vector, optionally scaled. After softmax, these scores become attention weights that are used to take a weighted average of the value vectors.

The context vector is a weighted sum of all encoder hidden states, where the weights (attention scores) indicate how relevant each encoder state is at the current decoding step. Unlike the single fixed context vector in basic seq2seq models, the attention-based context vector is computed fresh at every decoding step and focuses on different parts of the input. This allows the decoder to 'look back' at different source words when generating each target word.

Soft attention computes a weighted average over all input positions using differentiable softmax weights, making the entire model trainable end-to-end with gradient descent. Hard attention selects exactly one (or a discrete subset of) positions to attend to, making it non-differentiable and requiring reinforcement learning or other tricks to train. Soft attention is used in nearly all modern architectures like Transformers because its gradients flow naturally.

In self-attention, the queries, keys, and values all come from the same sequence, enabling tokens to build context from one another within that sequence. In cross-attention, queries come from one sequence (e.g., the decoder's current state) while keys and values come from a different sequence (e.g., encoder outputs), allowing information to flow between the two. Cross-attention is the mechanism by which the Transformer decoder 'reads' the encoded source input.

Self-attention has O(n²·d) time complexity and O(n²) space complexity because it computes pairwise dot products between all n tokens, producing an n×n attention matrix. This quadratic scaling is the primary bottleneck for processing very long sequences such as entire books or genomic data. Research into efficient attention (Longformer, Performer, Flash Attention) aims to reduce this toward near-linear complexity.

Masked self-attention prevents each decoder position from attending to any future positions by setting those attention scores to −∞ before the softmax, making those weights effectively zero. This is necessary because during training the entire target sequence is fed in at once for efficiency, but the model must not 'cheat' by reading tokens it is supposed to predict. The mask enforces the autoregressive property: each predicted token can only depend on previously generated tokens.

When batching sequences of different lengths, shorter sequences are padded with a special token. The padding mask adds a very large negative value (e.g., −1×10⁹) to attention scores for padding positions before the softmax, which causes their weights to become virtually zero. This ensures the model learns to ignore padding tokens and prevents meaningless pad positions from influencing the representations of real tokens.

Bahdanau (additive) attention, introduced in 2014 (published at ICLR 2015) for neural machine translation, computes a score between decoder state s and encoder hidden state h using a small neural network: score(s, h) = vᵀ · tanh(W₁·s + W₂·h). It is called 'additive' because it sums the linear projections of both inputs before applying a nonlinearity. This was one of the first practical attention mechanisms and showed that decoders could selectively focus on source words.

Luong (multiplicative) attention, also from 2015, uses simpler score functions: a direct dot product (sᵀ·h), a scaled dot product, or a bilinear form (sᵀ·W·h) — all without a tanh nonlinearity. Unlike Bahdanau attention, which uses the previous decoder state as the query, Luong attention uses the current decoder hidden state. Its simpler computation makes it generally faster while achieving comparable translation quality.

Attention weight heatmaps (rows = queries, columns = keys) show which input tokens each output token attends to most, revealing learned linguistic patterns. In machine translation, these often show word alignment (e.g., the French output word attending strongly to its English source counterpart). In BERT, different heads have been found to capture syntactic roles, coreference (pronoun → antecedent), and positional patterns, all without explicit supervision.

The total parameter count and FLOPs are roughly the same since each head uses d_model/h dimensions, but multi-head attention gains expressive power because each head projects Q, K, V into a different learned subspace. One head might track subject-verb agreement, another pronoun coreference, another local context — simultaneously. This 'divide and specialize' strategy consistently outperforms a single head of the same full dimension on benchmarks.

The variable `weights` is the attention weight matrix (shape n×n), where entry weights[i][j] represents how much query token i attends to key token j after softmax normalization. Each row sums to 1, forming a probability distribution over all input tokens for that query. Multiplying `weights @ V` then produces the output — a weighted mixture of value vectors where more relevant values contribute more.

import numpy as np
def self_attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    weights = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)
    return weights @ V

Using a range of frequencies — from very high (small i) to very low (large i) — encodes position at multiple scales simultaneously, like how a clock's second hand, minute hand, and hour hand each represent position at a different timescale. High-frequency dimensions distinguish nearby positions precisely; low-frequency dimensions give coarse global structure. Additionally, using both sin and cos at each frequency allows the model to represent relative positions as linear transformations, which is useful for tasks requiring positional reasoning.

The encoder uses fully bidirectional self-attention — every token can attend to every other token in both directions, producing rich, globally-informed representations of the input. The decoder uses two types of attention per layer: causal (masked) self-attention, where each output token attends only to previous output tokens to preserve autoregressive generation; and cross-attention, where decoder queries attend to encoder keys and values to incorporate source information. This design lets the decoder generate output one token at a time while using the complete encoded input.

During training, the full target sequence is fed in simultaneously (for GPU parallelism), but the model must learn to predict each token using only the tokens that precede it — the same constraint it faces at inference. Without causal masking, attention scores at position t would include token t itself and all future tokens, meaning the model could trivially copy the answer, learning nothing. Applying −∞ masks to future positions before softmax enforces this constraint and ensures training and inference behaviors match.

Absolute positional encoding assigns a fixed embedding to position 1, 2, 3, … in the sequence, which can struggle when a model is tested on sequences longer than those seen in training. Relative positional encoding instead encodes the offset between two tokens (e.g., 'j is 3 positions before i'), which captures the same relational structure regardless of where in the sequence the pair appears, enabling better length generalization. Models like Transformer-XL, T5, and RoPE use forms of relative encoding.

In an RNN, computing hidden state h_t requires h_{t-1}, creating a sequential dependency chain that forces step-by-step processing — an n-token sequence needs n sequential steps. Self-attention computes all query-key dot products simultaneously as a single matrix multiplication QKᵀ, and all weighted sums as another (weights·V), both of which GPUs can execute in one parallel forward pass across the entire sequence. This makes Transformers dramatically faster to train on long sequences because all positions are processed in parallel, not one after another.

Sparse attention restricts each token to attend only to a chosen subset of positions — such as a local window of nearby tokens, strided global tokens, or a combination — instead of all n positions. This reduces the O(n²) cost of full attention to roughly O(n·k) where k is the number of attended positions, making processing of very long sequences (documents, audio, genomic sequences) feasible. Models like Longformer and BigBird use sparse patterns combined with a few global tokens that attend everywhere to preserve long-range capacity.

During autoregressive generation, each new token must attend to all previous tokens. Without caching, the keys and values for all previous tokens would be recomputed from scratch at every step, costing O(n·d) per new token. KV caching stores the computed K and V tensors for all past tokens in memory so only the new token's K and V need to be computed and appended, reducing per-step attention cost to O(d). This is why generating 1,000 tokens is much faster with caching than without.

Studies (e.g., Voita et al. 2019, Clark et al. 2019) found that some heads track syntactic dependency relations, such as attending from a verb to its subject or from a word to its syntactic head in a dependency parse. Other heads focus on coreference, linking pronouns back to their antecedents (e.g., 'it' strongly attending to 'the cat' several sentences earlier). Yet other heads specialize in positional patterns — always attending to the immediately next or previous token — without being explicitly trained to do so.

Standard attention must write the full n×n attention matrix to GPU high-bandwidth memory (HBM), which for long sequences becomes the dominant memory and speed bottleneck — not the FLOPs. Flash Attention uses a tiling strategy: it splits Q, K, V into small blocks that fit in the fast on-chip SRAM, computes attention incrementally in tiles (using the online softmax trick to avoid materializing the full matrix), and writes only the final output back to HBM. This achieves mathematically identical results to standard attention while being 2–4× faster and using O(n) instead of O(n²) memory.

BERT uses fully bidirectional self-attention — at every layer, each token can attend to all other tokens in both directions — making it ideal for tasks that require understanding the full context, like text classification, named entity recognition, and question answering. GPT uses unidirectional (causal/left-to-right) self-attention — each token can only attend to itself and previous tokens — making it suited for autoregressive text generation. This architectural choice is the root reason BERT excels at encoding/understanding while GPT excels at generation.

In the decoder's cross-attention layer, the Queries (Q) are derived from the decoder's own previous sub-layer output — representing the decoder's current 'need' based on what it has generated so far. The Keys (K) and Values (V) both come from the encoder's final output — the encoded representations of the source sentence. This setup lets each decoder position selectively retrieve relevant source information; the encoder output is computed once and reused at every decoder layer and every generation step.

The first limitation is quadratic memory: the n×n attention matrix requires O(n²) space, so for n=10,000 and even a single attention head, this is 100 million entries — quickly exhausting GPU VRAM with multi-head, multi-layer models. The second is quadratic time: computing all pairwise dot products costs O(n²·d) FLOPs, making training and inference prohibitively slow for long documents or sequences. These bottlenecks have driven research into efficient variants: Flash Attention addresses memory via tiling, Longformer and BigBird address both via sparse attention patterns, and Performer approximates full attention in O(n·d) via kernel methods.