What is attention in neural networks?

What does attention compute?

Attention is a mechanism that lets a neural network weigh which parts of its input matter most for each piece of output. Instead of squashing a whole sentence into one fixed vector, the network keeps a representation for every input position. At each decoding step it decides how much to draw on each one.

Bahdanau, Cho and Bengio introduced this idea in 2014 for machine translation. Their encoder read a source sentence into a sequence of vectors rather than a single one. A decoder learned to search that sequence for the positions most relevant to the word it was about to generate. They called this soft alignment: a weighted sum over source positions, rather than a hard choice of one word to translate.

The weights came from an alignment model, a small feedforward network scoring how well each source position matched the decoder's current state. Bahdanau et al. reported that this approach, called RNNsearch, matched a conventional phrase-based translation system on English-to-French translation. It also held up far better than a fixed-vector encoder as sentences grew longer.

What are queries, keys and values?

The transformer paper by Vaswani et al. (2017) recast attention in a more general form. An attention function maps a query and a set of key-value pairs to an output, all of them vectors. The output is a weighted sum of the values, where each weight comes from a compatibility score between the query and the matching key.

Vaswani et al. used scaled dot-product attention. Queries and keys are compared by a dot product, divided by the square root of the key dimension, then passed through a softmax to produce weights on the values. The scaling matters: the paper notes that for large key dimensions, unscaled dot products grow large and push the softmax into regions with vanishingly small gradients. Dividing by that square root keeps training stable.

In practice, queries, keys and values are packed into matrices and processed together. That is what lets attention run as fast matrix multiplication on a GPU rather than step by step.

What is multi-head attention?

Multi-head attention runs several smaller attention computations in parallel instead of one large one. Vaswani et al. linearly project the queries, keys and values into several lower-dimensional subspaces. They run scaled dot-product attention on each in parallel, then concatenate the results and project them once more.

The point is representational variety. A single attention head produces one weighted average per query, and averaging limits what it can represent. Multiple heads let the model attend to different kinds of relationships at different positions at once. The original transformer used 8 heads, each working in 64 dimensions, which the authors note keeps total computation close to that of one full-size head.

Their ablations found single-head attention scored 0.9 BLEU points worse than their best configuration on English-to-German translation. Too many heads also hurt quality. There was a sweet spot, not a simple dial toward more.

Why is attention quadratic, and why does that matter?

Self-attention compares every position in a sequence with every other position. Its computational and memory cost grows with the square of sequence length. Vaswani et al.'s own complexity table puts the per-layer cost of self-attention at order n-squared times the representation dimension, against order n times dimension-squared for a recurrent layer.

That quadratic term is cheap to state and expensive to live with. The FlashAttention paper by Dao et al. (2022) points out that standard implementations also materialize the full n-by-n attention matrix in GPU high-bandwidth memory, which is slow to read and write even when the arithmetic itself is modest. Dao et al. found that most attention operations are bottlenecked by these memory transfers, not by raw compute. That is why naive quadratic cost translates so directly into slow, memory-hungry training on long sequences.

FlashAttention vs. standard attention

FlashAttention, from Dao et al. (2022), computes the exact same attention output without ever writing the full n-by-n score matrix to slow GPU memory. It splits queries, keys and values into blocks, loads each block into fast on-chip SRAM, and incrementally updates a running softmax. The final result is identical to standard attention.

The paper reports up to 7.6 times faster attention computation on GPT-2 and up to 9 times fewer memory reads and writes than a standard implementation, for typical head dimensions. That translated into a 15% faster BERT-large training time than the MLPerf 1.1 record, a 3 times speedup training GPT-2 against HuggingFace's implementation, and a 2.4 times speedup on the Long-Range Arena benchmark, all reported by Dao et al. in 2022.

Because it uses less memory, FlashAttention also let the authors train on longer sequences than before. They report the first transformer to beat chance on the Path-X task, a 16,000-token sequence, with 61.4% accuracy. A block-sparse variant did the same on Path-256, a 64,000-token sequence, with 63.1% accuracy.

Approximate attention: trading exactness for speed

Vaswani et al. already flagged restricting self-attention to a neighbourhood around each position as a way to cut cost for long sequences. That raises the longest path between positions from a constant to order n divided by the neighbourhood size. Later work built on that idea with sparse and low-rank approximations such as Linformer, Performer and Reformer, cited in the FlashAttention paper.

Dao et al. found these approximate methods often cut floating-point operations without delivering real wall-clock speedups, because they still ignore memory-access overhead. Their block-sparse FlashAttention instead skips whole blocks of the attention matrix according to a sparsity mask. The paper reports it running faster than every approximate attention method they tested on the Long-Range Arena benchmark, while matching standard attention's accuracy.