Short Version
Why do we cache the K (key) and V (value) matrices, but not he Q (query) matrix?
Long Version
Given a toy set of 2-dimensional embedding vectors:
| Token | Embedding Vector |
|---|---|
| quick | [0.27 0.78] |
| lazy | [0.38 0.58] |
| brown | [0.50 0.83] |
| jumps | [0.20 0.53] |
| over | [0.46 0.59] |
| the | [0.45 0.55] |
| fox | [0.19 0.69] |
| dog | [0.51 0.47] |
Meaning the input tokens the quick brown fox would be represented by their embeddings vectors in the X (input) matrix:
X ╭ ╮
│0.45 0.55│ ;the
│0.27 0.78│ ;quick
│0.50 0.83│ ;brown
│0.19 0.69│ ;fox
╰ ╯
And we have a simple set of toy Wq, Wk, Wv weight matrices. Our toy embedding vectors are 2-dimensional. That then requires the weight matrices to have 2 rows; but they can have an arbitrarily decided number of columns; we choose 3 columns. Here are our toy weight matrices:
Wq (query)
╭ ╮
│0.6 -0.4 0.1│
│0.3 0.9 -0.2│
╰ ╯
Wk (key)
╭ ╮
│ 0.5 0.2 -0.3│
│-0.1 0.8 0.4│
╰ ╯
Wv (value)
╭ ╮
│0.7 -0.3 -0.2│
│0.4 0.5 -0.1│
╰ ╯
Then we multiply the input matrix by the three WQ, WK, and WV weight matrices:
Query Matrix:
Q = X*Wq
╭ ╮ ╭ ╮
│0.45 0.55│ ╭ ╮ │ 0.43 0.32 -0.06│
│0.27 0.78│ │ 0.6 -0.4 0.1│ = │ 0.40 0.59 -0.13│
│0.50 0.83│ │ 0.3 0.9 -0.2│ │ 0.55 0.55 -0.12│
│0.19 0.69│ ╰ ╯ │ 0.32 0.54 -0.12│
╰ ╯ ╰ ╯
Key Matrix:
K = X*Wk
╭ ╮ ╭ ╮
│0.45 0.55│ ╭ ╮ │ 0.17 0.53 0.08│
│0.27 0.78│ │ 0.5 0.2 -0.3│ = │ 0.06 0.68 0.23│
│0.50 0.83│ │-0.1 0.8 0.4│ │ 0.17 0.76 0.18│
│0.19 0.69│ ╰ ╯ │ 0.03 0.59 0.22│
╰ ╯ ╰ ╯
Value Matrix:
V = X*Wv
╭ ╮ ╭ ╮
│0.45 0.55│ ╭ ╮ │ 0.54 0.14 -0.14│
│0.27 0.78│ │0.7 -0.3 -0.2│ = │ 0.50 0.31 -0.13│
│0.50 0.83│ │0.4 0.5 -0.1│ │ 0.68 0.26 -0.18│
│0.19 0.69│ ╰ ╯ │ 0.41 0.29 -0.11│
╰ ╯ ╰ ╯
Note: The Q, K, and V do undergo further processing, but that isn't really important here, as it's not relevant to the question. But it's essentially:
A = softmax(mask(Q * transpose(K) / sqrt(2)); 2=number of dimensions in embedding space
O = A xor V
Adding a new token to the input
Let's say after the quick brown fox, the next predicted token was jumps. So now we append that embedding to our X input matrix, giving us a new input matrix X:
X ╭ ╮
│0.45 0.55│ ;the
│0.27 0.78│ ;quick
│0.50 0.83│ ;brown
│0.19 0.69│ ;fox
│0.20 0.53│ ;jumps
╰ ╯
And again we compute the Q, K, V matrices for this new set of input tokens:
Query Matrix:
Q = X*Wq
╭ ╮ ╭ ╮
│0.45 0.55│ ╭ ╮ │0.43 0.32 -0.06│
│0.27 0.78│ │ 0.6 -0.4 0.1│ = │0.40 0.59 -0.13│
│0.50 0.83│ │ 0.3 0.9 -0.2│ │0.55 0.55 -0.12│
│0.19 0.69│ ╰ ╯ │0.32 0.54 -0.12│
│0.20 0.53│ │0.28 0.40 -0.09│
╰ ╯ ╰ ╯
Key Matrix:
K = X*Wk
╭ ╮ ╭ ╮
│0.45 0.55│ ╭ ╮ │0.17 0.53 0.08│
│0.27 0.78│ │ 0.5 0.2 -0.3│ = │0.06 0.68 0.23│
│0.50 0.83│ │-0.1 0.8 0.4│ │0.17 0.76 0.18│
│0.19 0.69│ ╰ ╯ │0.03 0.59 0.22│
│0.20 0.53│ │0.05 0.46 0.15│
╰ ╯ ╰ ╯
Value Matrix:
V = X*Wv
╭ ╮ ╭ ╮
│0.45 0.55│ │ 0.54 0.14 -0.14 │
│0.27 0.78│ ╭ ╮ │ 0.50 0.31 -0.13 │
│0.50 0.83│ │0.7 -0.3 -0.2│ = │ 0.68 0.26 -0.18 │
│0.19 0.69│ │0.4 0.5 -0.1│ │ 0.41 0.29 -0.11 │
│0.20 0.53│ ╰ ╯ │ 0.35 0.21 -0.09 │
╰ ╯ ╰ ╯
Except the Q, K, V matrices are nearly identical, so cache them
You can see that the first four rows of Q, K, and V are the same as last time. That's because we just added a new token on the end; all the previous rows in X are the same. So rather than recompute the entire Q, K, and V matrix, we can keep them in memory, and just compute the final additional row:
Query Matrix:
Q = X*Wq
╭ ╮ ╭ ╮
│░░░░ ░░░░│ │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ ╭ ╮ │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ │0.6 -0.4 0.1│ = │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ │0.3 0.9 -0.2│ │░░░░ ░░░░ ░░░░│
jumps│0.20 0.53│ ╰ ╯ │0.28 0.40 -0.09│
╰ ╯ ╰ ╯
Key Matrix:
K = X*Wk
╭ ╮ ╭ ╮
│░░░░ ░░░░│ │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ ╭ ╮ │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ │ 0.5 0.2 -0.3│ = │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ │-0.1 0.8 0.4│ │░░░░ ░░░░ ░░░░│
jumps│0.20 0.53│ ╰ ╯ │0.05 0.46 0.15│
╰ ╯ ╰ ╯
Value Matrix:
V = X*Wv
╭ ╮ ╭ ╮
│░░░░ ░░░░│ │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ ╭ ╮ │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ │0.7 -0.3 -0.2│ = │░░░░ ░░░░ ░░░░│
│░░░░ ░░░░│ │0.4 0.5 -0.1│ │░░░░ ░░░░ ░░░░│
jumps│0.20 0.53│ ╰ ╯ │0.35 0.21 -0.09│
╰ ╯ ╰ ╯
By caching the Q, K, V, we save ourselves 72 multiplications, and 36 additions, at the cost of storing 36 floats.
Except, apparently we don't cache Q.
We only cache K and V.
Why?
Why not also cache Q?
I've asked ChatGPT over and over to try to explain it to me. It keeps saying:
We don't need the Query (Q) matrix after it is first used, so it is recomputed every time.
Which I do not understand; because it is used again after the first time, and it doesn't have to be recomputed: because we could cache it. I'm certain it's for a reason i'll never actually understand, but obviously i'll accept the answer that everyone else agrees is the correct answer. I was just also hoping to understand it myself: so i thought i would try asking.