PyTorch Queries, Keys and Values
A projection is a Linear layer that rewrites each token vector in a new role. One attention head makes three projections of every token, one for seeking, one for offering, and one for carrying.
One projection shrinks the vector
A Linear layer from n_embd inputs to head_size outputs rewrites each 128-number token as a 32-number query. The batch and time axes pass through untouched.
import torch
from torch import nn
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
head_size = n_embd // n_head
torch.manual_seed(42)
x = torch.randn(1, 4, n_embd)
query = nn.Linear(n_embd, head_size, bias=False)
q = query(x)
print("x shape:", tuple(x.shape))
print("q shape:", tuple(q.shape))Same token, two different roles
The query layer and the key layer hold different weights, so one token produces two different vectors. The printed openings of the two rows disagree, which is the whole point of separate roles.
import torch
from torch import nn
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
head_size = n_embd // n_head
torch.manual_seed(42)
x = torch.randn(1, 4, n_embd)
query = nn.Linear(n_embd, head_size, bias=False)
key = nn.Linear(n_embd, head_size, bias=False)
q = query(x)
k = key(x)
print("same input, different output:", not torch.equal(q, k))
print("q row 0, first 3:", [round(v, 2) for v in q[0, 0, :3].tolist()])
print("k row 0, first 3:", [round(v, 2) for v in k[0, 0, :3].tolist()])Example
import torch
from torch import nn
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
head_size = n_embd // n_head
torch.manual_seed(42)
# One window of four tokens, already embedded as vectors.
x = torch.randn(1, 4, n_embd)
query = nn.Linear(n_embd, head_size, bias=False)
key = nn.Linear(n_embd, head_size, bias=False)
value = nn.Linear(n_embd, head_size, bias=False)
q = query(x)
k = key(x)
v = value(x)
print("x shape:", tuple(x.shape))
print("q shape:", tuple(q.shape))
print("k shape:", tuple(k.shape))
print("v shape:", tuple(v.shape))
print("q row 0, first 4:", [round(n, 2) for n in q[0, 0, :4].tolist()])
print("k row 0, first 4:", [round(n, 2) for n in k[0, 0, :4].tolist()])
print("v row 0, first 4:", [round(n, 2) for n in v[0, 0, :4].tolist()])Key points
- Query: the projection that says what a position is looking for
- Key: the projection that says what a position offers to others
- Value: the projection that carries the content a position hands over
- Head size: the length of each projected vector, here 128 divided by 4 heads
- The three layers start from the same token and still produce three different vectors
- A projection without bias only stretches and rotates the input vector
In this example, a seeded random table stands in for four embedded tokens in one window. Three Linear layers with no bias project the same table into queries, keys, and values. Every projection keeps the batch and time shape. It shrinks the last axis from 128 to the head size of 32. The printed openings show three different rows from the same first token, one per role. The next lesson compares queries against keys to score how much each position should read from each earlier one.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
What does the value projection of a token carry?