Query, Key and Value in transformers

Module 05 · Transformers · 19 concepts

In depth / 15 lines

Query, Key and Value, in more detail

Attention is easiest to follow as a lookup in a library.

The query is what you are looking for right now.

The key is the label on the spine of each book, advertising what it holds.

The value is the content you actually take away if you choose that book.

Every token produces all three, by multiplying its vector by three learned matrices.

Those matrices are the parameters, and they are what training adjusts.

A token's query is compared against every token's key to produce a relevance score.

High score means that token is worth listening to for this particular query.

The scores become weights, and the output is the weighted mixture of the values.

Splitting key from value is the part people find strange at first.

It lets a token advertise one thing and contribute another.

A pronoun can advertise I am a reference and contribute the identity it refers to.

A single vector doing both jobs would be forced to compromise between them.

In self-attention all three come from the same sequence.

In cross-attention the query comes from one sequence and the key and value from another.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Query, Key and Value

Three projections of the same thing

Attention needs three different views of every position, and it gets them by multiplying the input by three learned matrices.

Q = X W_Q \qquad K = X W_K \qquad V = X W_V

$X$ : the input, $n \times d_{\text{model}}$ : one row per position.
$W_Q, W_K$ : projections into the comparison space, $d_{\text{model}} \times d_k$ .
$W_V$ : projection into the output space, $d_{\text{model}} \times d_v$ .

All three start from the identical input. What makes them different is only that the three matrices are trained separately.

What each one is for

The query is what this position is looking for. The key is what a position advertises about itself. The value is what it actually hands over if selected.

The comparison happens between queries and keys, and the thing returned is a mixture of values.

\text{score}_{ij} = q_i \cdot k_j \qquad \text{out}_i = \sum_j \alpha_{ij} v_j

That split is the reason there are three and not two: what a position offers as a match key need not be what it contributes as content.

A pronoun might key on "I am a noun phrase" while its value carries the identity of who it refers to.

The shapes

Q: n \times d_k \qquad K: n \times d_k \qquad V: n \times d_v

Queries and keys must share a width, because they are dotted together. Values need not, because they are only ever summed.

In practice $d_k = d_v = d_{\text{model}} / h$ , where $h$ is the number of heads, so that concatenating the heads returns to the original width.

One layer, in order

Take the input, one vector per position.
Multiply by $W_Q$ to get a query for every position.
Multiply the same input by $W_K$ to get a key for every position.
And by $W_V$ to get a value.
Hand all three to the attention formula.

Where they come from decides the layer's name

If all three are projections of the same sequence, it is self-attention.

If the queries come from one sequence and the keys and values from another, it is cross-attention.

Nothing else changes. The formula, the matrices and the shapes are identical, and only the source of $X$ differs.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

What role does the value vector play?

Transformers

19 concepts in this module, then a graded test with every answer explained.

Study this module →