Query, Key, and Value

Think of a library. You walk in with a question. That's your query. Every book has a label on its spine, which is the key. You compare your question with each label. And when you find a match, you read the book's contents, which is the value. Attention uses exactly these three roles.

• A question, a label, the contents
• Query, key, and value
• Three learned roles

Learned attention

In the simple version, each word used the same vector for everything. Now we let the network learn three different views of each word, using three linear layers. Queries are compared to keys, which gives the scores. We divide by the square root of the dimension, to keep the numbers stable. Softmax gives the weights, and those weights blend the values. Run it. There are five words, so the weights form a five by five grid, and each row adds up to one. This one formula is the heart of the transformer.

import torch
from torch import nn

torch.manual_seed(0)
dim = 8
x = torch.rand(1, 5, dim)

query = nn.Linear(dim, dim, bias=False)
key = nn.Linear(dim, dim, bias=False)
value = nn.Linear(dim, dim, bias=False)

q, k, v = query(x), key(x), value(x)
scores = q @ k.transpose(1, 2) / dim ** 0.5
weights = scores.softmax(dim=-1)
out = weights @ v

print("Attention weights:", list(weights.shape))
print("Each row adds to 1:", weights[0].sum(dim=-1).round(decimals=2).tolist())
print("Output:", list(out.shape))

Attention weights: [1, 5, 5]
Each row adds to 1: [1.0, 1.0, 1.0, 1.0, 1.0]
Output: [1, 5, 8]

Do not read the future

When a model is generating text, it must not peek at words that haven't been written yet. So we use a mask. Torch tril gives a triangle, where each word can only see itself, and the words before it. We set the hidden scores to minus infinity, so that softmax turns them into exactly zero. Run it. Look at the mask, and see the triangle. Word two can only attend to words zero, one, and two. This one line is what makes a language model possible.

import torch

scores = torch.rand(5, 5)

mask = torch.tril(torch.ones(5, 5))
print("What each word may see:\n", mask.int())

hidden = scores.masked_fill(mask == 0, float("-inf"))
weights = hidden.softmax(dim=-1)
print("Word 2 attends to:", weights[2].round(decimals=2).tolist())

What each word may see:
 tensor([[1, 0, 0, 0, 0],
        [1, 1, 0, 0, 0],
        [1, 1, 1, 0, 0],
        [1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1]], dtype=torch.int32)
Word 2 attends to: [0.5099999904632568, 0.23000000417232513, 0.25999999046325684, 0.0, 0.0]

Summary