# Attention mechanism

Attention lets a neural network weigh which parts of the input matter for each output. Learn queries, keys and values, self-attention and multi-head attention.

## Overview

Attention is a way for a neural network to decide, for each piece of output it produces, which parts of its input matter most. Instead of squeezing a whole sentence, image or table into one fixed summary, the network keeps every element available and builds a weighted mix of them. The weights are computed on the fly from the data, so they change from one input to the next.

Attention is the core operation of the [transformer](/tools/transformers) architecture, which underlies modern language models, [vision transformers](/tools/vit), speech models such as Whisper and many models for tables. When people say a model "attends" to a word, they mean that word received a large weight in this mix.

An everyday comparison is reading with a highlighter. When you answer a question about a long article, you do not reread every word with equal care. You scan for the passages that relate to the question and lean on those. Attention gives a network the same ability, and it learns during training what kind of passages are worth highlighting for which questions.

## What it is

Modern attention is built from three projections of the input and a few simple operations.

| Part | Role |
|---|---|
| Query (Q) | What the current position is looking for |
| Key (K) | What each position offers, used to judge relevance |
| Value (V) | The information each position passes on if chosen |
| Scores | Dot products between a query and every key |
| Scaling | Division by the square root of the key dimension to keep scores stable |
| Softmax | Turns scores into positive weights that sum to one |
| Weighted sum | Combines the values using those weights |
| Mask | Blocks certain positions, such as future tokens or padding |
| Heads | Several attention operations in parallel, each with its own projections |
| Output projection | Mixes the heads' results back into one vector |

Two broad forms are common. In **self-attention**, queries, keys and values all come from the same sequence, so each element looks at the others in its own input. In **cross-attention**, queries come from one sequence and keys and values from another, for example a decoder looking at an encoded source sentence, or an image model looking at a text prompt.

## Where it came from

Attention in its modern neural form appeared in machine translation. Early encoder-decoder models built from recurrent networks such as the [LSTM](/tools/lstm) compressed the whole source sentence into one vector, and quality dropped for long sentences. In 2014, Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio proposed letting the decoder look back at all encoder states and learn a soft alignment between source and target words. Their paper, "Neural Machine Translation by Jointly Learning to Align and Translate", appeared at ICLR 2015. Later in 2015 Minh-Thang Luong and colleagues compared simpler, multiplicative scoring functions.

In 2017, Ashish Vaswani and colleagues at Google published "Attention Is All You Need". It removed recurrence entirely and built the transformer from self-attention and feed-forward layers. It introduced scaled dot-product attention and multi-head attention in the form used today. The transformer trained in parallel far better than recurrent networks and became the base of models such as [BERT](/tools/bert) and the [GPT family](/tools/gpt-architecture).

Since then, much work has targeted the cost of attention. Multi-query attention (2019) and grouped-query attention (2023) share keys and values across heads to shrink memory during generation. [FlashAttention](/tools/flashattention) (2022) computes exact attention while moving far less data through GPU memory.

## How it works

### Step 1: project the input

Each input position, such as a token embedding, is multiplied by three learned matrices to produce a query vector, a key vector and a value vector. For a sequence of n positions, this gives matrices Q, K and V with n rows each.

### Step 2: score every pair

For each query, the model takes the dot product with every key. A high dot product means the two vectors point in similar directions, which the model has learned to treat as relevance. This produces an n by n score matrix.

### Step 3: scale and mask

Scores are divided by the square root of the key dimension, d. Without this, large dimensions give large dot products, and the softmax becomes so sharp that gradients nearly vanish. A mask then sets forbidden scores to minus infinity. In a decoder, a causal mask stops each token from seeing tokens that come after it.

### Step 4: softmax and combine

A softmax across each row turns scores into weights between zero and one that add up to one. Each output is the weighted sum of the value vectors. Written compactly, attention equals softmax(QKᵀ / √d) V.

```python
import torch
import torch.nn.functional as F

q = k = v = torch.randn(1, 8, 128, 64)   # batch, heads, tokens, head dim
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
```

### Step 5: many heads

Multi-head attention runs several of these operations in parallel with separate, smaller projections. One head might track which noun a pronoun refers to, another nearby words, another punctuation, although heads are not guaranteed to be so tidy. The heads' outputs are joined and passed through one more linear layer.

### Step 6: position information

Attention by itself ignores order: shuffling the input shuffles the output the same way. Transformers therefore add position information, either as position embeddings added to the input or by rotating queries and keys according to position, as in rotary embeddings.

### Cost

Because every query meets every key, compute and memory for the score matrix grow with the square of the sequence length. Doubling the context roughly quadruples that cost. During text generation, a [KV cache](/tools/kv-cache) stores past keys and values so they are not recomputed for every new token.

## Key concepts and terms

| Term | Meaning |
|---|---|
| Query | Vector describing what a position is looking for |
| Key | Vector used to match against queries |
| Value | Vector carrying the information that is mixed into the output |
| Attention weights | Softmax-normalised relevance scores |
| Self-attention | Attention within a single sequence |
| Cross-attention | Queries from one source, keys and values from another |
| Causal mask | Prevents positions from attending to later positions |
| Multi-head attention | Several attention operations in parallel |
| Scaled dot-product | Dot-product scores divided by the square root of the dimension |
| Context length | Number of positions attention can cover at once |
| Grouped-query attention | Several query heads share one set of keys and values |
| KV cache | Stored keys and values reused during generation |

## A worked example

This example uses made-up numbers to show the mechanics. Take the sentence "The cat sat because it was tired" and focus on the token "it".

1. The model projects "it" into a query vector. It also projects every token into a key and a value.
2. It computes the dot product of the "it" query with each key. Suppose the raw scores, after scaling, are: The 0.1, cat 2.3, sat 0.4, because 0.2, it 0.9, was 0.3, tired 0.6.
3. In a decoder with a causal mask, "it" may only see tokens up to itself, so "was" and "tired" are masked out.
4. A softmax over the remaining five scores gives weights of roughly: The 0.07, cat 0.61, sat 0.09, because 0.08, it 0.15. They sum to one, allowing for rounding.
5. The output for "it" is 0.61 times the value of "cat", plus smaller shares of the other values. The representation of "it" now carries information about the cat.
6. In a real model this happens in every head of every layer, and later layers build on earlier ones.

The point is that the weights were not programmed. Training adjusted the projection matrices until weights like these helped the model predict text well.

## Strengths and limitations

**Strengths:**

- Connects any two positions directly, however far apart.
- Computes all positions in parallel, which suits GPUs.
- Weights adapt to each input rather than being fixed.
- Works for text, images, audio, graphs and tables with the same operation.
- Attention weights offer a partial window into what the model uses.

**Limitations:**

- Cost grows with the square of sequence length.
- Generation needs a growing KV cache, which consumes memory.
- Order must be added separately through position encodings.
- Attention weights are not a full explanation of a model's reasoning.
- Very long contexts can still dilute focus on the relevant parts.

## How it compares with alternatives

| Approach | How it mixes information | Cost with length | Notes |
|---|---|---|---|
| Full attention | Every position with every position | Quadratic | Most flexible; the transformer default |
| [Recurrent networks](/tools/lstm) | Passes a hidden state step by step | Linear, but sequential | Hard to parallelise, struggles with long range |
| [Convolutions](/tools/cnn) | Fixed local windows | Linear | Strong for local patterns, needs depth for range |
| [State space models such as Mamba](/tools/mamba) | Learned recurrence with selective updates | Linear | Efficient for long sequences |
| [Axial attention](/tools/axial-attention) | Attention along one axis at a time | Below quadratic | Used for images and grids |

## Practical guidance and best practices

1. Use your framework's fused attention function, such as PyTorch's `scaled_dot_product_attention`, rather than writing the softmax by hand.
2. Check masks carefully; a wrong causal or padding mask silently leaks information or wastes capacity.
3. Choose head dimension and head count together so their product equals the model width.
4. Use grouped-query attention for models that must generate long outputs efficiently.
5. Plan memory for the KV cache when estimating serving costs.
6. Keep computation in a stable precision for the softmax, or use kernels that handle this for you.
7. Treat attention maps as a debugging aid, not as proof of why a model made a decision.

## Common mistakes

- Forgetting the scaling factor, which makes training unstable for large head sizes.
- Applying a mask after the softmax instead of before it.
- Assuming attention understands order without any position encoding.
- Reading a single head's attention map as the model's full explanation.
- Underestimating memory when doubling context length.

## Frequently asked questions

### What are queries, keys and values in plain words?

A query is a question each position asks, a key is a label each position wears, and a value is the content each position can share. The model compares the question with every label, and the better the match, the more of that position's content it takes.

### What is the difference between self-attention and cross-attention?

Self-attention mixes information within one sequence, so words in a sentence look at each other. Cross-attention lets one sequence look at another, such as a decoder reading an encoded source, or an image generator reading a text prompt.

### Why divide by the square root of the dimension?

Dot products of long vectors tend to be large. Large scores make the softmax put almost all weight on one position, which gives tiny gradients and slows learning. Dividing by the square root of the dimension keeps scores in a workable range.

### Why does attention get expensive for long inputs?

Every query is compared with every key, so the number of comparisons is the sequence length squared. Techniques such as FlashAttention reduce memory traffic, and alternatives such as state space models change the scaling, but exact full attention remains quadratic in compute.

### Is attention only used in language models?

No. Vision transformers apply it to image patches, speech models to audio frames, graph models to neighbouring nodes and tabular models to rows and columns. Diffusion models use cross-attention to follow text prompts.

## Glossary

- Alignment: a learned correspondence between positions in two sequences.
- Embedding: a vector that represents a token or element.
- Head: one of several parallel attention operations in a layer.
- Masking: blocking certain positions from receiving attention.
- Positional encoding: information added so the model knows element order.
- Softmax: a function that turns scores into positive weights that sum to one.
- Token: a unit of input, such as a word piece.
- Transformer: an architecture built from attention and feed-forward layers.

## Summary

Attention lets a network build each output as a weighted mix of its inputs, with weights computed from how well queries match keys. It began as a fix for long sentences in translation in 2014 and became the heart of the transformer in 2017, with scaled dot-product and multi-head attention. It is flexible and parallel but costs grow with the square of length, which drives ongoing work on efficient kernels, shared keys and values, and alternative sequence models.