Self-Attention

Explore transformer self-attention with interactive query, key and value matrices. Follow attention scores and see how tokens combine information.

query of “it”a thing? alive?catsatonmat…ittiredkey:thing, animalkey:thing

1Every word asks and offers

Each word makes a query (what am I looking for?) and a key (what can I be found by?). “it” is looking for a thing, ideally a living one.

catsatonmat…ittired49%16%

2Score each word, then softmax

Query · key gives a score per word. Softmax turns the scores into shares that add up to 100%. Thicker arcs here mean more attention.

catmatothers0.65thing0.49animal0.00person0.07refers0.07actionnew “it”

3Blend the values into new meaning

Mix every word's value by its share. The new “it” is mostly “cat”, so it now carries “animal”: the model knows what “it” refers to.

Step through it

Watch one word attend

Pick a word and step through how it scores every other word, turns the scores into shares, and blends their values. This toy has five named features and a hand-set W_Q; to keep it readable, W_K and W_V are the identity, so each key and value is just the word's features.

Sentence

Toy model: five named features and a hand-set query matrix, so you can follow the arithmetic. The computation is the real one.

Click any word to see where it looks.

Query of “it”appears at step 2
WordEmbedding (thing, animal, person, refers back, action)Score÷ √5Attention
the–···
cat–···
sat–···
on–···
the–···
mat–···
because–···
it–···
was–···
tired–···
New “it” (blended values)appears at the last step

Every word at once: row = who is looking, column = who they look at

the
cat
sat
on
the
mat
because
it
was
tired

Query matrix Wq (hand-set): row = feature the word has, column = feature it looks for

thing00003
animal00000
person00000
refers back32.4-400
action13300
Step 0 / 7
1400 ms

Press Step to compute attention for “it” one stage at a time.

The three roles

Query, key and value in plain words

The paper defines attention as “mapping a query and a set of key-value pairs to an output”, where the output is a weighted sum of the values. In self-attention, every word plays all three parts at once.

Think of a library search

  • Query: what you type into the search box. For a word, it describes what information would help it. “it” looks for the thing it refers to.
  • Key: the label on each book's spine that your search is matched against. For a word, it advertises what the word is: “I'm an animal”, “I'm a verb”.
  • Value: what's inside the book, the thing you actually take away. For a word, it's the information it passes on to whoever attends to it.

The twist: you don't take one book. You take a little from every book, more from the ones whose labels match your search best. That “soft” lookup is what softmax does.

embedding xthe word itself× W_Qquery qwhat am I looking for?× W_Kkey kwhat can I be found by?× W_Vvalue vwhat do I hand over?

Three different learned matrices turn the same word into three different vectors.

Why not one vector?

Looking and being found are different jobs. In the worked example below, “cats” gives 67% of its attention to “chase”, but “chase” gives only 45% back. If query and key were the same vector, the score from cats to chase would always equal the score from chase to cats.

Why is the value separate?

What a word is matched on doesn't have to be what it hands over. In the example, keys ignore the “small” feature, but values carry it, so “mice” can pass on “small” to “chase” even though size played no part in the match.

Real sizes

In the paper's base model each word is a vector of d_model = 512 numbers. Each attention head projects it down to a query and key of d_k = 64 and a value of d_v = 64. The matrices W_Q, W_K, W_V start random and are learned during training; nobody sets them by hand.

Equation (1), by hand

The maths, step by step

The whole of attention is one line from the paper: Attention(Q, K, V) = softmax(QKᵀ / √d_k) V. Here it is on three words, “cats chase mice”, with tiny matrices you can check with a pencil. The numbers in the weight matrices are chosen to be easy, not learned. Every result is calculated by the same code the tests check.

  1. 1

    Each word is a row of numbers

    Stack the words' embeddings into a matrix X: one row per word. These three features have names so you can read them; a real model learns hundreds of unnamed ones (512 in the paper's base model).

    X (3 words × 3 features)

    animalactionsmall
    cats100
    chase010
    mice101
  2. 2

    Make Q, K and V with three weight matrices

    Multiply X by three different matrices. Each one asks a different question of the same word. In a real model these matrices are learned during training.

    W_Q

    animal02
    action20
    small00

    W_K

    animal10
    action01
    small00

    W_V

    animal10
    action01
    small01

    Q = X · W_Q

    cats02
    chase20
    mice02

    K = X · W_K

    cats10
    chase01
    mice10

    V = X · W_V

    cats10
    chase01
    mice11

    One row worked out: “chase” is [0, 1, 0], so its query is 0 × (row 1) + 1 × (row 2) + 0 × (row 3) of W_Q = [2, 0]. Its query points at column 1 of the keys, the “I'm an animal” column: an action looks for who's involved.

  3. 3

    Score every query against every key: QKᵀ

    Row i, column j is query i · key j: multiply matching numbers and add. A big score means word j has what word i is looking for.

    QKᵀ (rows: who looks, columns: who is looked at)

    catschasemice
    cats020
    chase202
    mice020

    For example, cats → chase = q_cats · k_chase = 0 × 0 + 2 × 1 = 2. Notice the matrix isn't symmetric in general: what cats looks for and what chase looks for are different questions.

  4. 4

    Divide by √d_k = √2

    d_k is the length of each key (here 2). Dividing keeps the scores in a range where softmax still spreads attention. The section below shows why this matters when d_k is large.

    QKᵀ ÷ √2 (√2 ≈ 1.4142)

    catschasemice
    cats01.41420
    chase1.414201.4142
    mice01.41420
  5. 5

    Softmax each row

    Turn each row into shares that are all positive and add up to 1: raise e to each number, then divide by the row's total. These are the attention weights.

    row “cats”: e^0 = 1, e^1.4142 = 4.1133, e^0 = 1

    total = 6.1133, so weights = 1 ÷ 6.1133, 4.1133 ÷ 6.1133, 1 ÷ 6.1133 = [0.1636, 0.6728, 0.1636]

    softmax(QKᵀ ÷ √d_k): the attention weights

    catschasemice
    cats0.16360.67280.1636
    chase0.44580.10840.4458
    mice0.16360.67280.1636

    “cats” gives 67% of its attention to “chase”, but “chase” gives only 45% back to “cats” (it splits between cats and mice). Separate queries and keys are what let attention be one-sided like this.

  6. 6

    Blend the values: weights × V

    Each word's output is the value rows mixed by its weights. That's the word's new vector, now carrying information from the words it attended to.

    out_cats = 0.1636 × [1, 0] + 0.6728 × [0, 1] + 0.1636 × [1, 1]

    = [0.3272, 0.8364]

    Output = weights · V

    cats0.32720.8364
    chase0.89160.5542
    mice0.32720.8364

    That's Equation (1), done. Every word is handled at once with matrix multiplication, instead of one word after another as in recurrent networks. That's why the paper calls the Transformer “more parallelizable” and quick to train.

The √d_k in Equation (1)

Why divide by √d_k?

The paper's reason: “for large values of d_k, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients”. (The authors say they “suspect” this; it's their explanation, not a proof.) Here it is, measured: random queries and keys whose numbers have mean 0 and spread 1, exactly as in the paper's footnote.

One random query against 10 random keys, d_k = 64. Bars are the softmax weights.

Without scaling

100%10 keys

Divided by √64

100%10 keys

Without scaling, one key usually grabs nearly all the attention, and the rest get almost nothing.

Key length d_k

Averages over 400 random queries:

Spread of q·k
8.07
… after ÷ √d_k
1.01
Top weight
87%
… after ÷ √d_k
33%
Its slope a(1−a)
0.09
… after ÷ √d_k
0.21

  1. 1. Scores grow with length

    q·k adds up d_k products. The paper's footnote: if each number has mean 0 and variance 1, the sum has mean 0 and variance d_k, so its spread is √d_k. Measured above: about 8.07 for d_k = 64 (√64 = 8.00).

  2. 2. Big scores make softmax pick one

    Softmax uses e^score, so a gap of 10 between two scores means a ratio of e¹⁰ ≈ 22,000. Spread-out scores turn into one weight near 100%.

  3. 3. Then learning stalls

    A softmax weight a changes with its own score at rate a(1 − a). Near a = 1 that's almost 0, so training gets almost no signal to adjust the scores. Dividing by √d_k brings the spread back to about 1, whatever d_k is.

No peeking

Masking: hiding the future

A model that writes text one word at a time (the paper's decoder, and chat models like GPT and Claude) mustn't see words that haven't been written yet. The paper does this inside the attention formula by “masking out (setting to −∞) all values in the input of the softmax which correspond to illegal connections”. Since e^−∞ = 0, those words get exactly 0% attention.

Attention weights (%) for the toy sentence. Row = the word looking, column = the word looked at.

thecatsatonthematbecauseitwastired
the100
cat5050
sat137513
on25252525
the2020202020
mat111143111111
because14141414141414
it5535551855
was111111111111111111
tired921999119999

Everything above the diagonal is hidden: each word sees only itself and the words before it. “it” can still find “cat” (53%), because the cat came earlier, but “the” at the start can only look at itself.

Section 3.2.2

Multi-head attention: several lookups at once

One attention lookup blends everything into a single average, so a word can't easily look for two different things (say, its subject and its object) at the same time. The paper's fix is to run several attention “heads” side by side, each with its own W_Q, W_K and W_V.

each word512 numbershead 1: own W_Q, W_K, W_V → 64head 2: own W_Q, W_K, W_V → 64head 3: own W_Q, W_K, W_V → 64head 4: own W_Q, W_K, W_V → 64head 5: own W_Q, W_K, W_V → 64head 6: own W_Q, W_K, W_V → 64head 7: own W_Q, W_K, W_V → 64head 8: own W_Q, W_K, W_V → 64concat8 × 64× W_O512

head_i = Attention(Q·W_i^Q, K·W_i^K, V·W_i^V)

MultiHead(Q, K, V) = Concat(head_1, …, head_h) · W^O

h = 8, d_model = 512, d_k = d_v = d_model / h = 64

  • In the paper's words, multiple heads let the model “jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.”
  • Each head works with 64 numbers instead of 512, so 8 heads cost about the same as one full-size head.
  • The paper tested this (Table 3): a single head was 0.9 BLEU worse than the best setting, but quality also dropped with too many heads. More isn't always better.

Section 3.5

Attention can't see word order, so we add it

Attention only compares vectors; nothing in Equation (1) knows which word came first. Shuffle the words and each word's output is exactly the same. The paper: “Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens”.

Same words, opposite meaning, same answer

“maya chased the dog”: chased → [0.86, 0.43, 0.43, 0.00, 0.07]

“the dog chased maya”: chased → [0.86, 0.43, 0.43, 0.00, 0.07]

Using this page's toy model, “chased” comes out identical in both sentences, even though who chased whom is reversed. Without position information, attention treats a sentence as a bag of words.

The paper's fix: add a wave pattern for each position

PE(pos, 2i) = sin(pos / 10000^(2i / d_model))

PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

023dimension 0 (fast wave)dimension 31 (slow wave)

Each position gets its own pattern of sines and cosines (rows above, d_model = 32 to fit), added to the word's embedding before attention. The authors also tried learned position vectors and got “nearly identical results”; they chose the waves because they “may allow the model to extrapolate to sequence lengths longer than the ones encountered during training”.

Section 3.2.3

Three places attention appears in the Transformer

Encoder self-attention

Q, K and V all come from the input sentence. Every word can look at every other word (no mask).

Decoder self-attention

Q, K and V come from the output written so far, with the mask, so each word sees only earlier words.

Encoder–decoder attention

Queries come from the decoder; keys and values come from the encoder's output. This is how the translation looks back at the source sentence. It's attention, but not self-attention.

Section 4, Table 1

What it costs, and why it won

n is the number of words and d the vector size. Self-attention compares every word with every word, so its work grows with n². In exchange, every word reaches every other word in one step, and nothing has to wait for the previous word.

Layer typeWork per layerSteps that must wait in lineLongest path between two words
Self-attentionO(n² · d)O(1)O(1)
Recurrent (RNN)O(n · d²)O(n)O(n)
ConvolutionalO(k · n · d²)O(1)O(log_k n)

The paper notes self-attention is cheaper than a recurrent layer when n is smaller than d, which was usual for sentences in translation. For very long texts the n² term dominates, which is why long context windows are expensive.

Read the original

The paper

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser and Illia Polosukhin. “Attention Is All You Need.” Advances in Neural Information Processing Systems 30 (NIPS 2017), pages 5998–6008.

It introduced the Transformer, a translation model built from attention alone, with no recurrence or convolution. Its big model scored 28.4 BLEU on WMT 2014 English-to-German translation, beating previous results including ensembles by over 2 BLEU, after 3.5 days of training on 8 NVIDIA P100 GPUs. The same building block now sits inside GPT, Claude, BERT and vision transformers.

arXiv:1706.03762NeurIPS proceedings

Sections used on this page: 3.2 (attention), 3.2.1 (Equation 1 and the √d_k footnote), 3.2.2 (multi-head), 3.2.3 (masking and the three uses), 3.5 (positional encoding), 4 (Table 1), 6.2 (Table 3).

How it works: Self-Attention

Step by step

  • The word doing the looking
  • Attention weight (brighter is more)
  • Query
  • Keys

Press Start or Step: the step the animation is on lights up below.

  1. 1

    Start from each word's embedding

    A vector of numbers standing for the word. Here its five features are named so you can read them.

    e = embedding[word]
  2. 2

    Make the query

    Multiply the embedding by Wq to get what this word is looking for. “it” looks for a thing, ideally an animal, not a person.

    q = e · Wq
  3. 3

    Every word offers a key

    Keys describe what each word can be found by. In this toy, Wk is the identity, so a key is just the word's features.

    k_j = e_j · Wk
  4. 4

    Score every word

    The dot product of the query with each key. It's large when they point the same way.

    s_j = q · k_j
  5. 5

    Scale down by √d

    Long vectors give big dot products, which would make softmax pick one word and ignore the rest. Dividing by √d keeps scores in a sensible range.

    s_j = s_j / √d
  6. 6

    Turn scores into shares

    Softmax makes every score positive and the shares add up to 1. These are the attention weights.

    a = softmax(s)
  7. 7

    Blend the values

    The word's new vector is the values mixed by the weights. “it” now carries “animal” from “cat”.

    out = Σ a_j · v_j

In “The cat sat on the mat because it was tired”, you know “it” means the cat. A language model has to work that out too. Self-attention is how each word looks at the other words and pulls in the ones that matter for its meaning.

Every word makes three vectors from its embedding. A query: what am I looking for? A key: what can I be found by? A value: what do I pass on if chosen? A word compares its query with every key; a good match gives a high score.

Softmax turns the scores into shares that add up to 100%, and the word's new meaning is the values blended by those shares. Transformers (GPT, Claude, BERT) stack many of these layers, each with several attention “heads” running side by side. The recipe comes from “Attention Is All You Need” (Vaswani et al., NIPS 2017): Attention(Q, K, V) = softmax(QKᵀ / √d_k) V.

The numbers on this page are a toy: each word has five named features and the query matrix is set by hand, so you can follow the arithmetic. A real model learns hundreds of unnamed dimensions, but the computation is the same.

Where you'll meet it

  • Inside every transformer: language models, translation, and vision transformers that treat image patches like words.
  • Interview classics: “what are Q, K and V?”, “why divide by √d?”, “why is attention O(n²)?”.
  • Attention maps (the grid here) are a common, if rough, way to peek at what a model connects.
time O(n² · d)space O(n²)

Every word scores every other word, so doubling the text length quadruples the work. That's why long context windows are expensive.