What are query, key, and value in attention?

Attention is the mechanism that lets a model decide, for each token, which other tokens to draw information from. It is the core of the transformer, and once you see how it works, most of modern language modeling clicks into place.

Each token is projected, through learned weight matrices, into three vectors. The query represents what the token is looking for. The key represents what a token offers as a match. The value is the content actually passed along. The original paper defines attention as “mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors.”^[1]

The output for a token is a weighted sum of the values, where, in the authors’ words, “the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.”^[1] A strong query-key match gives that value a large weight, so the token pulls in information from the positions most relevant to it.

Three separate projections exist instead of one shared vector for a specific reason. Each projection is its own learned linear transform of the same input, so the model can give a single token three different jobs at once: what it is searching for, what it advertises as a match, and what it actually hands over once chosen. Collapsing those into one shared vector would force a token’s search criteria and its content to be the same thing, which throws away a degree of freedom the model would otherwise have. In the original transformer, the 512-dimensional model is split across 8 heads, so each head’s query, key, and value vectors are 64-dimensional.^[1] The paper is explicit that this is not a cost concession: “due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.”^[1]