L4.1 · LESSON 1 OF 6

Queries, keys and values

Calculate query-key scores, scaled softmax weights and the weighted value output for one attention head.

THE IDEA

Three views of the same token

Attention lets each token gather information from other tokens. To do that it needs to decide who is relevant and what to take from them. The model learns three separate projections of every hidden vector:

a query Q: what this position is looking for,
a key K: what each position offers to be matched against a query,
a value V: the content that is passed along if that position is chosen.

hidden x[0]1hidden x[1]2

projection	matrix (rows)	result for x = [1, 2]	role
query q = x·Wq	[[1,0],[1,1]]	[3, 2]	"what am I looking for?"
key k = x·Wk	[[0,1],[1,0]]	[2, 1]	"what do I offer to be matched?"
value v = x·Wv	[[2,0],[0,2]]	[2, 4]	"what do I pass along if chosen?"

The same hidden vector is projected three ways by three learned matrices. A query is compared with every key; the scores become weights; the output is the weighted average of the values.

One hidden vector, three learned matrices, three different results.

Predict first. Which vectors are averaged to produce the attention output?

NumPy

Q, K, V = X @ Wq, X @ Wk, X @ Wv     # X is (tokens, d_model)

Chapter context

Attention computes a data-dependent mixture of value vectors. A query scores candidate keys, a normalization turns scores into weights, and those weights combine values. Keys decide how to look; values supply what is gathered.

TRY IT YOURSELF

Part 1. Do it hands-on

Compute the three projections of a hidden vector by hand, then write the projection step with shape checks in Python.

hidden x[0]1hidden x[1]2

projection	matrix (rows)	result for x = [1, 2]	role
query q = x·Wq	[[1,0],[1,1]]	[3, 2]	"what am I looking for?"
key k = x·Wk	[[0,1],[1,0]]	[2, 1]	"what do I offer to be matched?"
value v = x·Wv	[[2,0],[0,2]]	[2, 4]	"what do I pass along if chosen?"

The same hidden vector is projected three ways by three learned matrices. A query is compared with every key; the scores become weights; the output is the weighted average of the values.

For x = [1, 2] (the sliders above are free to explore):

q =

k =

v =

letter (Q, K or V) that gets averaged

To do: A token's hidden vector is x = [1, 2]. Its query is x @ Wq with Wq = [[1, 0], [1, 1]] (rows). Type q
To do: The key is x @ Wk with Wk = [[0, 1], [1, 0]]. Type k
To do: The value is x @ Wv with Wv = [[2, 0], [0, 2]]. Type v
To do: Which of Q, K, V are averaged (weighted) to produce the attention output? Type the letter

Part 2. Do it in code

project_qkv(X, Wq, Wk, Wv): X has shape (n, d_model). Return (Q, K, V) = (X @ Wq, X @ Wk, X @ Wv). Raise ValueError if the weight matrices do not have d_model rows or if Wq and Wk have a different number of columns.
qkv_shapes(n, d_model, d_k, d_v): the shapes of Q, K and V as a dict with keys "Q", "K", "V".
output_shape(n_q, n_k, d_v): the shape of the attention output. (It has one row per query and d_v columns.)
param_count(d_model, d_k, d_v): the number of weights in Wq, Wk and Wv together (no biases).

import numpy as np

def project_qkv(X, Wq, Wk, Wv):
    pass

def qkv_shapes(n, d_model, d_k, d_v):
    pass

def output_shape(n_q, n_k, d_v):
    pass

def param_count(d_model, d_k, d_v):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Which vectors are averaged to produce the attention output?

AThe values: queries and keys only decide the weights, and the output is a weighted average of the value vectors.BThe keys, because they are matched against each other.CThe queries, because they carry what each position is looking for.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

Q\in\mathbb R^{n_q\times d_k},\ K\in\mathbb R^{n_k\times d_k},\ V\in\mathbb R^{n_k\times d_v}.

S=QK^T/\sqrt{d_k}+M,\quad A=\operatorname{softmax}_{rows}(S),\quad O=AV.

M is zero at allowed positions and conceptually negative infinity at excluded positions. Every query needs a defined policy if all its keys are masked.

Worked example

Use q=[1,0], K=[[1,0],[0,1]], V=[[2,0],[0,2]], and d_k=2. Dot scores [1,0] become [1/sqrt(2),0]. Exponentiate and normalize to approximately [0.669762,0.330238]. The output is [2×0.669762,2×0.330238]≈[1.339523,0.660477]. If only the first key is causally allowed, the output becomes [2,0].

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import sqrt,exp,isclose
q=[1,0]; keys=[[1,0],[0,1]]; values=[[2,0],[0,2]]
scores=[sum(a*b for a,b in zip(q,k))/sqrt(2) for k in keys]
weights=[exp(s-max(scores)) for s in scores]
weights=[w/sum(weights) for w in weights]
out=[sum(weights[i]*values[i][j] for i in range(2)) for j in range(2)]
assert isclose(sum(weights),1)
assert isclose(weights[0],.6697615493266569)
assert isclose(out[0],1.3395230986533138)
assert isclose(out[1],.6604769013466862)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Attention scores are not weights before normalization. Causal masking happens before softmax. Attention weights alone are not a full causal explanation of a model decision.