Lensa ML

How Self-Attention Works

Step 1 of 6: What Is Self-Attention?

Each word attends to the whole sentence

self-attention: same sentence attends to itselfThe20%cat22%sat16%on10%the19%mat13%ThecatsatonthematKeysQuery (click a word)full self-attention matrix (row = query)ThecatsatonthematThecatsatonthemat

Query

"cat"

Attends most to

"cat"

Top weight

22%

"cat" attends most to itself (22%). In self-attention, each word in a sentence scores how relevant every word is to its own meaning, itself included.

Click each word in the top row to make it the query. With these 3-dimensional toy vectors the weights are fairly flat: for "cat" they range from about 10% to 22%, close to the uniform 1/6 ≈ 17%. Trained models with 64-dimensional heads produce much sharper patterns. Self-attention is often described as a soft dictionary lookup. A query is compared with every key, and instead of returning the single matching value, the lookup returns a blend of all values weighted by how well each key matched.

Step 2 of 6: Query, Key, Value

Projecting input into Q, K, V matrices

creating the query, key, and value vectorsInput MatrixThecatsatonthemat0.80.20.50.30.90.1-0.50.70.4-0.8-0.40.60.70.10.40.2-0.60.8W_Q0.6-0.30.10.40.8-0.20.10.50.7Query Matrix0.60.20.40.60.7-0.10.00.90.1-0.60.20.40.50.10.3-0.0-0.10.7×=all three projectionsQuery0.60.20.40.60.7-0.10.00.90.1-0.60.20.40.50.10.3-0.0-0.10.7Key0.70.20.10.10.90.3-0.30.50.7-0.3-0.60.40.60.10.10.4-0.70.2Value0.40.50.60.4-0.20.8-0.1-0.50.7-0.6-0.1-0.20.30.50.4-0.20.6-0.0

Projection

Query

Input shape

6 × 3

Query shape

6 × 3

The Query matrix is computed by multiplying the input embeddings by W_Q. Each row becomes a query vector: it asks "what am I looking for?"

Click Q, K, and V to highlight each projection. Each weight matrix here is 3x3, so the 6 × 3 input X becomes three 6 × 3 matrices with one multiplication each: Q = X · W_Q, K = X · W_K, V = X · W_V. In a real model X is 512 to several thousand dimensions wide, and the three projections are often computed together as one larger matrix multiply. Q and K must share a dimension d_k for the dot product to work, but V can have a different size d_v.

Step 3 of 6: Q·K Dot Product

Computing similarity between queries and keys

computing Q·K dot product for "cat"Q("cat")[0.55, 0.68, -0.08]0K("The")(0.55)(0.69)+(0.68)(0.24)+(-0.08)(0.14)0.53K("cat")(0.55)(0.15)+(0.68)(0.85)+(-0.08)(0.33)0.63K("sat")(0.55)(-0.30)+(0.68)(0.45)+(-0.08)(0.67)0.09K("on")(0.55)(-0.34)+(0.68)(-0.64)+(-0.08)(0.44)-0.66K("the")(0.55)(0.60)+(0.68)(0.15)+(-0.08)(0.07)0.43K("mat")(0.55)(0.44)+(0.68)(-0.66)+(-0.08)(0.18)-0.22

Query

"cat"

Highest score

0.63

Best match

"cat"

Score("cat", "cat") = 0.63 is the highest. The attention score is the dot product Q·K: multiply each dimension pairwise and sum. Higher scores mean the query and key are more aligned.

Click each word button to change the query and watch the six scores update. The dot product q · k = |q| |k| cos θ grows with both alignment and vector length, so a key with a large norm can score highly against many queries. The score matrix is not symmetric. Score("cat", "sat") uses cat's query and sat's key, while score("sat", "cat") uses sat's query and cat's key: after scaling, they are 0.05 and 0.47 here. All n² scores come from a single matrix product, Q · Kᵀ.

Step 4 of 6: Attention Score

Scale, softmax, and attention weights

attention score pipeline for Q("cat")ThecatsatonthematQ·K0.530.630.09-0.660.43-0.22÷ √3 = 1.73scaled0.310.370.05-0.380.25-0.13softmaxweight20.3%21.5%15.7%10.2%19.1%13.1%Σ = 1.00 (weights sum to 1)

Query

"cat"

1/√d_k

0.58

Top weight

21.5%

Attends to

"cat"

Raw Q·K scores are divided by √d_k = 1.73 to prevent large values from pushing softmax into extremes. Then softmax turns the scores into a probability distribution: "cat" puts 21.5% of its attention on "cat".

Click each word button to follow its scores through scaling and softmax. The √d_k factor comes from variance. If the entries of q and k are independent with mean 0 and variance 1, their dot product has variance d_k, so dividing by √d_k brings it back to 1. Without it, a 64-dimensional head would produce scores about 8 times larger, softmax would put nearly all weight on one key, and its gradient would almost vanish. Implementations subtract the row maximum before exponentiating, which leaves the result unchanged but avoids overflow.

Step 5 of 6: Causal Self-Attention

Masking future tokens in decoders

bidirectional self-attention: all tokens visibleKeysThecatsatonthematQueriesThecatsatonthemat0.3019%0.2118%0.0916%-0.0813%0.2418%0.1316%0.3120%0.3722%0.0516%-0.3810%0.2519%-0.1313%0.1417%0.4724%0.2720%-0.3211%0.0917%-0.3311%-0.1714%0.1419%0.3222%0.1419%-0.1614%-0.1913%0.2419%0.1417%0.0616%-0.0414%0.1918%0.1317%0.0215%0.0616%0.2419%0.2419%0.0015%0.1217%Click a row to select query. Top: scaled score, bottom: attention weight.

Query

"sat"

Visible tokens

6

Top weight

24.2%

In bidirectional self-attention, "sat" can attend to all 6 tokens. This is used in encoders (like BERT) where the full context is available. Compare with causal masking to see the difference.

Switch between Bidirectional and Causal (Masked), then click rows of the matrix. Masking is done by adding a matrix with 0 on and below the diagonal and −∞ above it, because e^(−∞) = 0. Only n(n + 1)/2 of the n² entries survive, 21 of 36 here. The mask is what makes training for text generation efficient: one forward pass computes a prediction at every position, and none of them can see the token it has to predict. Padding masks work the same way and hide padding tokens in batches of uneven length.

Step 6 of 6: Full Pipeline

End-to-end attention in one view

Input Embedding"cat" → [0.3, 0.9, 0.1]Linear ProjectionsQ[0.55, 0.68, -0.08]K[0.15, 0.85, 0.33]V[0.41, -0.15, 0.83]Q · KᵀThe0.53cat0.63sat0.09on-0.66the0.43mat-0.22scores÷ √d_k (1.73)scaled0.310.370.05-0.380.25-0.13softmaxweights20.3%21.5%15.7%10.2%19.1%13.1%× V → OutputΣ (weight × V) for each dimension20%×V(The)22%×V(cat)16%×V(sat)10%×V(on)19%×V(the)13%×V(mat)Output("cat")[0.14, 0.17, 0.45]Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V

Query

"cat"

Attends to

"cat"

Output

0.14,0.17,0.45

This is the full attention pipeline for "cat": embed → project to Q,K,V → compute Q·K scores → scale by 1/√3 → softmax → weighted sum of V vectors. The output [0.14, 0.17, 0.45] is a context-aware representation that blends information from all words, weighted by relevance.

Click each word button to see its whole path from embedding to output. The output for each word is a weighted average of the value vectors with positive weights that sum to 1, so it always lies inside the region the values span. Attention alone cannot push a value outside that range, which is one reason Transformer blocks add an output projection and a feed-forward layer after it. Standard attention stores the full n × n weight matrix. FlashAttention (Dao et al., 2022) computes exactly the same output in tiles without storing that matrix.

What is self-attention in neural networks?

Self-attention is an operation that updates each element of a sequence by taking a weighted average of all the elements, where the weights measure how relevant each one is. Every token produces a query, a key, and a value vector. The query is compared with every key by a dot product, the scores are scaled and passed through softmax, and the resulting weights mix the value vectors into a new, context-aware representation of the token.

Self-attention is the core operation of Transformers, so it runs inside every large language model, BERT-style encoder, vision transformer, and speech model such as Whisper. It lets a model link a pronoun to its noun, a verb to its subject, or an image patch to a related patch anywhere in the input, in a single step. Its masked (causal) form is what allows language models to be trained on all positions of a text at once.

Key ideas

  • Each token is projected into a query, a key, and a value vector.
  • Attention scores are dot products between one token's query and every token's key.
  • Scores are divided by √d_k so softmax does not saturate as dimensions grow.
  • The output for a token is the softmax-weighted average of all value vectors.
  • A causal mask sets future positions to −∞ so a token sees only itself and earlier tokens.
  • The score matrix is not symmetric, because queries and keys use different weights.

Short takes on this topic

Common questions

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

In self-attention the queries, keys, and values all come from the same sequence, so a sentence attends to itself. In cross-attention the queries come from one sequence and the keys and values from another, for example a translation decoder attending to the encoded source sentence. The math is identical; only the inputs differ.

What is the difference between additive and dot-product attention?

Additive attention, used by Bahdanau et al., scores a query and key with a small feed-forward network: v · tanh(W_1 q + W_2 k). Dot-product attention, used by Luong et al. and in Transformers, scores them with q · k. The dot-product form is much faster because all scores come from one matrix multiplication, and scaling by √d_k fixes its problem with large dimensions.

Do attention weights explain a model's predictions?

Only partly. Attention maps show where information was gathered in one layer, but later layers, value vectors, and residual connections also shape the output. Studies such as Jain and Wallace (2019) found that very different attention patterns can lead to the same prediction, so attention is a useful hint, not a full explanation.

Is self-attention only used for text?

No. Vision transformers apply it to image patches, speech models to audio frames, and graph attention networks to the neighbors of each node. Any input that can be split into a set of vectors can use it, as long as position or structure is supplied separately.

Related explainers

References

  1. Vaswani et al. (2017), "Attention Is All You Need", NeurIPS
  2. Luong, Pham & Manning (2015), "Effective Approaches to Attention-based Neural Machine Translation"
  3. Dao et al. (2022), "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
  4. Jain & Wallace (2019), "Attention is not Explanation"

Last updated