Attention Mechanism Explained

Why Attention Was Needed

Before attention, sequence-to-sequence models (used for translation, summarization, and other tasks) processed the entire input through an encoder and compressed it into a single fixed-size vector. The decoder then generated the output from this single vector. This created a bottleneck: all the information in a 500-word paragraph had to fit into a vector of perhaps 512 dimensions. Important details were inevitably lost, and performance degraded severely on long inputs.

Attention, introduced by Bahdanau et al. in 2014 for machine translation, solved this by letting the decoder look back at all encoder positions when generating each output word. When translating the French word "chat" to English, the decoder can focus on the encoder positions corresponding to "chat" rather than trying to recover that information from a compressed summary of the entire sentence. Performance improved dramatically on long sentences.

Queries, Keys, and Values

Modern attention (as used in transformers) formalizes the mechanism using three vectors for each position: a query (Q), a key (K), and a value (V).

The analogy is a search engine. The query represents what you are looking for. The keys represent what each stored item contains. The values represent the actual information stored at each position. To find relevant information, you compare your query against all keys, determine which keys match best, and retrieve the corresponding values.

Each input position is projected through three separate learned weight matrices to produce its Q, K, and V vectors. The attention score between position i and position j is the dot product of position i's query with position j's key. High dot products mean the query and key are aligned, indicating relevance. The scores are divided by the square root of the key dimension (to prevent numerical instability) and passed through softmax to produce weights that sum to 1. The output for position i is the weighted sum of all positions' value vectors.

The weight matrices W_Q, W_K, and W_V are learned during training. The network discovers what constitutes a useful query, what constitutes a useful key, and what information to store in values, all through gradient descent on the prediction task. Different layers learn different types of attention patterns because they learn different projection matrices.

Self-Attention vs. Cross-Attention

Self-attention (also called intra-attention) computes attention within a single sequence. Each position attends to all other positions in the same sequence. This is what transformer encoder and decoder blocks use internally. In the sentence "The cat sat on the mat," self-attention lets "sat" attend to "cat" (to learn the subject-verb relationship) and "mat" attend to "on" (to learn the prepositional relationship).

Cross-attention computes attention between two different sequences. The queries come from one sequence (typically the decoder) and the keys and values come from another (typically the encoder). In machine translation, cross-attention lets each output word attend to the relevant input words. Cross-attention is also used in multimodal models where text attends to image features or vice versa.

The computation is identical in both cases; the only difference is whether Q, K, and V come from the same sequence or from different sequences.

Multi-Head Attention

A single attention computation captures one type of relationship. But language involves many simultaneous relationship types: syntactic structure, semantic similarity, coreference, positional proximity, and more. Multi-head attention runs multiple attention computations in parallel, each with its own learned projections.

If the model dimension is 768 and there are 12 attention heads, each head operates on 768/12 = 64 dimensions. The 12 heads compute attention independently, and their outputs are concatenated (producing a 768-dimensional vector) and projected through a final weight matrix. The total parameter count is similar to a single attention computation at full dimension, but the model can capture 12 different types of relationships simultaneously.

Analysis of trained transformers reveals that heads specialize. In BERT, specific heads consistently track syntactic dependencies (connecting verbs to subjects across intervening clauses). Others track positional patterns (attending primarily to the previous or next token). Others capture semantic relationships (connecting semantically related words). This division of labor is not engineered; it emerges as the most efficient strategy the model discovers during training.

Causal Masking

In decoder models (GPT, Claude), each position should only attend to previous positions, not future ones. During training, the model processes the entire sequence at once for efficiency, but it must not "cheat" by looking at future tokens when predicting the current one.

Causal masking enforces this constraint by setting the attention scores for all future positions to negative infinity before the softmax. After softmax, these positions have zero weight, effectively making them invisible. Position 5 can attend to positions 1 through 5, but not positions 6 onward. This creates the autoregressive property: each prediction depends only on past context.

Encoder models (BERT) do not use causal masking because they process the full sequence bidirectionally. Each position can attend to every other position, which gives encoders richer representations for understanding tasks but makes them unsuitable for generation (where future tokens do not exist yet).

Computational Cost and Optimizations

Self-attention's main limitation is its quadratic complexity. Computing attention scores between all pairs of n positions requires n^2 operations and n^2 memory. For a sequence of 100,000 tokens, this means 10 billion pairwise comparisons per attention layer. This cost limits context window sizes and makes long-document processing expensive.

FlashAttention reduces the memory overhead by computing attention in tiles that fit in GPU SRAM, avoiding the need to materialize the full n x n attention matrix in slower GPU memory. It does not change the algorithmic complexity but reduces wall-clock time by 2-4x.

Sparse attention patterns restrict each position to attending to a subset of positions. Local attention (attend to nearby positions only), strided attention (attend to every kth position), and global attention (a few designated positions attend to everything) can be combined to approximate full attention at O(n * sqrt(n)) or O(n * log(n)) cost.

Linear attention approximates the softmax attention with a kernel function, reducing complexity to O(n). The approximation introduces some accuracy loss but enables processing very long sequences that would be impossible with standard attention.

Despite these optimizations, the quadratic cost of attention remains a fundamental constraint of the transformer architecture, driving ongoing research into alternative sequence processing mechanisms like state-space models that achieve O(n) complexity natively.

Attention Beyond Language

Attention has proven valuable far beyond its original NLP application. Vision transformers (ViT) apply self-attention to image patches, letting each patch attend to every other patch and capturing global image relationships that CNNs (with their local receptive fields) can miss. Graph attention networks (GATs) use attention to weight messages between graph nodes. Protein structure prediction (AlphaFold) uses attention to model interactions between amino acid residues. Audio processing uses attention to capture long-range dependencies in spectrograms.

The versatility of attention comes from its content-based, dynamic nature. Unlike convolutional filters (which are the same for every input) or recurrent states (which compress information into a fixed-size vector), attention computes different connection patterns for every input, adapting its behavior to the specific content it is processing.