Attention mechanism: the query, key and value diagram

In this article

1. Before you start
2. Why interviewers ask this
3. The attention mechanism from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

Why interviewers ask this

Attention is the operation inside every transformer, so a candidate who cannot draw it cannot explain a modern language model. Interviewers ask for the picture because it exposes whether you know what each of the three inputs does. A common opening is "walk me through attention", and the follow-ups are why the weights come from a softmax, what the scaling is for, and where the queries, keys and values come from.

A weak answer says "the model focuses on the important words". A strong answer says that each output is a weighted average of values, that the weights come from query-key comparisons, and that the whole thing is differentiable (a small change in a score gives a small change in the output), so the network learns what to look for. The history helps too: attention was invented to fix a specific bottleneck, which is the first step below.

The attention mechanism from first principles

In plain words

Attention mechanism

An attention mechanism is a function that takes one query and a set of key-value pairs and returns a weighted average of the values, where each value's weight reflects how well its key matches the query. It lets a network read different parts of its input for different outputs.

Soft lookup

A soft lookup is a retrieval step that returns a blend of every stored item, weighted by how well each one matches the request, instead of returning only the single best match. Its weights change smoothly when the match scores change.

Think of a library with a forgiving librarian. You say what you are after. The librarian compares your request with the label on every shelf. Instead of one book, you get a few pages from every shelf, thickest from the best matches. The analogy breaks in one place: the librarian uses fixed judgement, while a network learns what a good match is from its training data.

In one precise sentence, attention is a soft dictionary lookup. Compare the query with every key, turn the scores into weights, and average the values. Nothing else happens inside it: no rules, no search loop, only arithmetic that a training run can adjust.

Step 1: A query asks, keys advertise, values carry the content

Attention splits each stored item into a key that is compared and a value that is returned, and adds a query that does the asking. The three are the query, key and value vectors of a transformer; all are lists of numbers, and each is a projection of the input. A projection is the input multiplied by a learned matrix. The key advertises what an item contains; the value is what it hands over.

It was invented to fix a bottleneck. Sutskever et al. (2014) encoded a source sentence into a vector of fixed dimensionality, and a second network decoded the translation from it. An encoder is the network that reads the input and turns it into vectors, and a decoder is the network that writes the output from them. Bahdanau et al. (2014, v7 2016) conjectured that the fixed-length vector "is a bottleneck" and let the decoder "(soft-)search" the source for relevant parts. In their model the decoder's state is the query, and the encoder's outputs, which they call annotations, serve as both keys and values.

Step 2: Score each key against the query

Each key gets one number saying how well it matches the query, computed by a score function. Bahdanau et al. parametrised the score as a small feedforward network (layers that pass numbers forward only) trained with the rest of the model. In their appendix it is a single layer that applies a hyperbolic tangent (a function that squashes any number into the range -1 to 1) to the decoder state and annotation combined, called additive attention.

Luong et al. (2015) compared three content-based scores: the plain dot product of the decoder state and a source state, a "general" form with a weight matrix in the middle, and a concatenation form, which is the additive score Bahdanau et al. used. The first two are cheaper.

Vaswani et al. (2017, v7 2023) chose the dot product, divided by the square root of the key size, and wrote that it "is much faster and more space-efficient in practice". In the toy below the query is (1.0, 0.4, 0.1) and the key for "cat" is (1.5, 0.3, 0.2), so the score is 1.5 + 0.12 + 0.02 = 1.64¹.

Step 3: Softmax turns scores into weights, so the lookup is soft

Softmax converts the scores into weights that are positive and add up to one, and this makes the lookup soft: every value keeps some weight, and the weights move smoothly when the scores move. In Bahdanau et al. the weight is ${\alpha }_{ij}=\exp (e_{ij})/{\sum }_k\exp (e_{ik})$, which they read as the probability that output word $i$ is aligned to source word $j$.

The reason to use softmax instead of picking the best key is training. Nudge the score of "cat" up by 0.1 in the toy and the soft output moves by (0.0774, 0.0951), while a hard pick of the highest score stays at (4.0, 4.0) and moves by exactly 0¹. A hard pick is flat almost everywhere, so a small change in a score gives training no gradient, no signal about which way to adjust. Bahdanau et al. make the same point in their own words: the soft alignment "allows the gradient of the cost function to be backpropagated through", where backpropagated means the error is traced backward through each step to learn how each number should change.

Step 4: The weights mix the values into the context vector

The output is the weighted sum of the value vectors, the context vector, and it is computed fresh for every query. With weights of 0.1297, 0.1934, 0.5473 and 0.1297 over the toy values, the context vector is (2.5780, 2.2528). "cat" carries just over half the weight (0.5473), and because its value is the largest it supplies most of the vector: 2.1890 of 2.5780 and 2.1890 of 2.2528¹. A word's share of the weight is not its share of the output. Bahdanau et al. describe this as an expected annotation, the average of the annotations under the alignment probabilities.

Where the three inputs come from decides the name. When the queries come from a decoder and the keys and values from an encoder, as here, it is cross-attention. When all three come from the same sequence, it is self-attention, covered in self-attention explained; the difference is laid out in cross-attention versus self-attention. Running several of these lookups side by side is multi-head attention.

The formula, and what the scaling is for

Stacking all queries into a matrix Q, all keys into K and all values into V, the whole mechanism in the form Vaswani et al. print is:

Here Q times the transpose of K holds every query-key dot product, d_k is the key size, softmax is applied to each row, and the product with $V$ is the weighted average of the values. The division by $\sqrt{d_k}$ has a stated reason. Vaswani et al. suspected that for large key sizes the dot products grow large in magnitude, pushing the softmax into regions where it has extremely small gradients, and they scale the dot products by one over the square root of d_k to counteract that. The toy in this article leaves the scores unscaled for clarity; the self-attention page works through why the scale matters (self-attention explained).

1 · Query, keys, values2 · Score3 · Softmax4 · MixReplay

Step-by-step description

1. A query asks; keys advertise; values carry content.
2. Each key gets one score against the query.
3. Scores become soft weights summing to one.
4. The weights average the values.

• The query is compared with every key to give one score per key.
• Softmax turns the scores into weights that add up to one.
• The output is the weighted average of the values.

Worked example

A four-word source phrase makes the diagram concrete. All numbers are illustrative values chosen for the walkthrough, not taken from a trained model, and the scores are left unscaled for clarity¹. A decoder is about to write the word for "cat", and its query is (1.0, 0.4, 0.1).

The weights add up to 1.0, so the context vector is 0.1297 × (1, 0) + 0.1934 × (0, 1) + 0.5473 × (4, 4) + 0.1297 × (2, −1) = (2.5780, 2.2528)¹. A hard pick would return (4.0, 4.0) and discard the information in the other three.

Text description of the diagram

• The weights over the toy words are 0.1297, 0.1934, 0.5473 and 0.1297, and they add up to 1.0.
• The best match gets the largest weight, but every word keeps some.

The nudge test from step 3 makes the training argument measurable. Raising the "cat" score by 0.1 changes the soft output by (0.0774, 0.0951) and the hard output by (0, 0). Raising the score of "sat" instead, which is not the winner, changes the soft output by (−0.0078, −0.0438) and the hard output by (0, 0) again¹.

Text description of the diagram

• Raising the cat score by 0.1 changes the soft output by 0.0774 and 0.0951.
• The hard pick changes by exactly 0.

Cost grows quickly because every output position scores every input position. Bahdanau et al. note that the alignment model is evaluated $T_x\times T_y$ times for a sentence pair, so a source and target of 10 words need 100 scores, 50 words need 2,500 and 500 words need 250,000¹. Doubling both lengths from 10 to 20 words gives 400 scores, four times the work, while doubling only the source gives 200¹.

Trade-offs and edge cases

Attention buys flexible reading of the input at a price in compute and memory, and the right variant depends on length and hardware. The main cases:

• Additive versus dot-product. Additive attention uses a small network and is more expensive in practice; dot-product attention is a matrix multiplication and is much faster, though Vaswani et al. call the two similar in theoretical complexity. They report that for small key sizes the two perform similarly, and that additive attention does better than unscaled dot-product attention for larger ones, which is the evidence that led them to suspect large dot products and add the scaling factor.
• Global versus local. Luong et al. compared a global form that attends to all source words with a local form that looks at a subset of source words at a time, trading some context for lower cost.
• Quadratic cost. The score count is the product of the two lengths, so doubling both lengths, as in self-attention over one sequence, quadruples the work; the FlashAttention page covers the memory-traffic fix.
• Memory at generation time. A decoder that reuses past keys and values stores them, which is the key-value (KV) cache.
• Weights are not an explanation. A large weight shows where the lookup concentrated. Whether that also explains why the model answered as it did is debated, so in an interview claim only what the weights show.

How to answer it in an interview

The spoken structure is: define it as a soft lookup, walk compare, normalise, mix, justify softness, then name the origin and the cost.

Common mistakes

• "Attention picks the most important word" (wrong). It averages all values with weights, and the best match only gets the largest share (right: say every position contributes, in proportion to its weight).
• "Keys and values are the same thing" (wrong). They can be the same vectors, as in the original translation model's annotations, but in a transformer they are separate projections (the input times a learned matrix), one used for matching and one for content (right: the key decides how much, the value decides what).
• "The weights are probabilities of correctness" (wrong). They are normalised match scores, and they say nothing about whether the output is right (right: call them attention weights).
• "Attention was introduced with the transformer" (wrong). Bahdanau et al. introduced it for translation in 2014, and Vaswani et al. in 2017 built a model out of attention alone (right: name both papers).
• "Softmax is just there to make the weights add to one" (wrong). It also makes the lookup differentiable, which is what lets the model learn what to match (right: say both reasons).

Follow-up questions

What is the attention mechanism?

It is a way for a network to read the parts of its input that matter for the current output. It compares a query with a key for every item, turns the matches into weights, and averages the item values with those weights.

What is the attention mechanism formula?

Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V. Q, K and V stack the queries, keys and values; the softmax turns each row of scaled scores into weights, and the product with V averages the values. Vaswani et al. scale by the square root of the key size because large dot products push the softmax into regions with extremely small gradients.

When was the attention mechanism invented?

Bahdanau et al. introduced it for translation in 2014 (arXiv v1 of 1409.0473, 1 September 2014), and Vaswani et al. built a model from attention alone in 2017. It was invented because early encoder-decoder translation squeezed a whole sentence into one fixed-length vector, which Bahdanau et al. conjectured was a bottleneck; attention lets the decoder look back at every source position instead.

Why does attention use a softmax instead of picking the best match?

A hard pick has a zero gradient almost everywhere, so training cannot improve the matching. In the toy above a 0.1 nudge moves the soft output by about 0.08 and 0.10 and the hard pick by exactly 0.

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

Additive attention scores a query and key with a small feedforward network, and dot-product attention multiplies them directly. Vaswani et al. say the dot product is much faster and more space-efficient in practice, and add a scaling factor for large key sizes.

What happens to the cost if the sentence doubles?

The work is one score per query-key pair. Doubling both lengths, as in self-attention, takes 100 scores to 400, four times the work; doubling only the source in cross-attention takes 100 to 200. In cross-attention the queries come from the decoder and the keys and values from the encoder; in self-attention all three come from one sequence, as in self-attention explained and the transformer architecture diagram.

Sources

1. Neural Machine Translation by Jointly Learning to Align and Translate (Bahdanau, Cho and Bengio), arXiv 1409.0473 v7 — arXiv (v1 submitted 2014-09-01, v7 revised 2016-05-19), 1 Sep 2014
2. Effective Approaches to Attention-based Neural Machine Translation (Luong, Pham and Manning), arXiv 1508.04025 v5 — arXiv (v5 2015-09-20; EMNLP 2015), 17 Aug 2015
3. Attention Is All You Need (Vaswani et al.), arXiv 1706.03762 v7 — arXiv (v7 revised 2023-08-02), 12 Jun 2017
4. Sequence to Sequence Learning with Neural Networks (Sutskever, Vinyals and Le), arXiv 1409.3215 — arXiv (last revised 2014-12-14), 10 Sep 2014