scaled dot-product attention
In plain wordsThe standard formula for attention in transformers: score how well each question matches each label, scale the scores down, turn them into percentages, and use those to blend the stored information.
The attention function of the transformertransformer, n.: The design behind today's chatbots and much of modern AI. It handles text by letting every word look at every other word and decide which of them matter for understanding it.: every query is compared with every key by a dot productdot product, n.: A score of how well two lists of numbers line up, found by multiplying matching entries and adding the products. It is large when the lists point the same way and zero when they are at right angles., each score is divided by √d_k (the square root of the key dimension), a softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score. over the keys converts the scores into weights, and the output is the correspondingly weighted sum of the values. The division matters because the dot product of two d_k-dimensional vectors with independent, unit-variance components has variance d_k: left unscaled, wide heads would produce scores so spread out that the softmax puts almost all weight on one key and passes back almost no gradient. Scoring by dot products rather than with a small neural network (additive attention) reduces the whole computation to matrix multiplications, the operation GPUs run fastest. Masks add −∞ to disallowed scores before the softmax. Fused kernels such as FlashAttentionFlashAttention, prop. n.: A faster way for AI chips to compute attention, the step in which a language model relates every word to every other. It gets exactly the same answer while moving far less data around the chip. compute exactly the same function without storing the full score matrix.
Swap the hand-written attention for PyTorch's fused SDPA call and memory use drops sharply.
Origin
Named and defined in 'Attention Is All You Need' (Vaswani et al., 2017). Unscaled multiplicative ('dot' and 'general') scoring had been proposed in 2015 by Minh-Thang Luong, Hieu Pham and Christopher Manning; the 2017 paper added the 1/√d_k factor.
Usage note
PyTorch exposes it as torch.nn.functional.scaled_dot_product_attention, hence the abbreviation SDPA. Some models use a different scale or normalize queries and keys ('QK-norm') to keep scores stable.
Related
- attentionattention, n.: A way for an AI model to choose what to focus on. For each word it processes, the model scores how relevant every other word is and draws information mostly from the relevant ones.
- query, key and valuequery, key and value, n.: The three roles each word plays inside attention: the query is what it is looking for, the key is what it can be found by, and the value is the information it hands over when chosen.
- softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score.
- multi-head attentionmulti-head attention, n.: Running several attention processes side by side, each free to track a different kind of relationship between words — grammar, topic, who did what — and then combining their results.
- attention maskattention mask, n.: A grid of yes-or-no switches telling an attention layer which words each word may look at — for example, ignore blank padding, or do not peek at words that come later.
- FlashAttentionFlashAttention, prop. n.: A faster way for AI chips to compute attention, the step in which a language model relates every word to every other. It gets exactly the same answer while moving far less data around the chip.
Sources
- Vaswani, A. et al. (2017). Attention Is All You Need. NIPS 2017. arxiv.org/abs/1706.03762
- Luong, M.-T., Pham, H. & Manning, C. D. (2015). Effective Approaches to Attention-based Neural Machine Translation. EMNLP 2015. arxiv.org/abs/1508.04025
Mentioned in
attentionattention, n.: A way for an AI model to choose what to focus on. For each word it processes, the model scores how relevant every other word is and draws information mostly from the relevant ones., attention maskattention mask, n.: A grid of yes-or-no switches telling an attention layer which words each word may look at — for example, ignore blank padding, or do not peek at words that come later., dot productdot product, n.: A score of how well two lists of numbers line up, found by multiplying matching entries and adding the products. It is large when the lists point the same way and zero when they are at right angles., FlashAttentionFlashAttention, prop. n.: A faster way for AI chips to compute attention, the step in which a language model relates every word to every other. It gets exactly the same answer while moving far less data around the chip., multi-head attentionmulti-head attention, n.: Running several attention processes side by side, each free to track a different kind of relationship between words — grammar, topic, who did what — and then combining their results., Noam ShazeerNoam Shazeer, prop. n.: An American computer scientist who co-invented the transformer and several techniques found in nearly every large language model, co-founded Character.AI, and later returned to Google to co-lead Gemini., query, key and valuequery, key and value, n.: The three roles each word plays inside attention: the query is what it is looking for, the key is what it can be found by, and the value is the information it hands over when chosen., self-attentionself-attention, n.: Attention in which a passage looks at itself: every word checks every other word in the same text to work out its meaning in context, such as which noun 'it' refers to., linear attentionlinear attention, n.: A family of faster substitutes for attention whose cost grows in step with the length of the text rather than with its square, by keeping a running summary instead of comparing every pair of words., QK-normQK-norm, n.: A small change to how a language model pays attention — rescaling two sets of internal numbers before comparing them — that prevents a common way large models' training goes haywire.