PyTorch Reading Attention Weights
Reading attention weights means checking a weight grid against the words that produced it. The grid is a claim about where the model looks, and the words are the evidence for that claim.
Print a grid with labels
Rows name the word doing the looking, and columns name the word being looked at. The grid below is fixed by hand, so the reading steps stay easy to follow.
words = ["cat", "sat", "mat"]
weights = [[1.00, 0.00, 0.00],
[0.30, 0.70, 0.00],
[0.10, 0.60, 0.30]]
print(" " + " ".join(f"{w:>5}" for w in words))
for word, row in zip(words, weights):
print(f"{word:>4} " + " ".join(f"{v:5.2f}" for v in row))Rows add up to one
Softmax made each row of any real weight grid, so each row must sum to 1. The sums of the hand-fixed grid print as three ones.
import torch
w = torch.tensor(weights)
print("Row sums:", w.sum(dim=1))Find where each row leans
argmax reads the position of the biggest weight in each row. mat leans on sat, while cat and sat lean on themselves.
best = w.argmax(dim=1)
for word, idx in zip(words, best.tolist()):
print(f"{word} looks most at {words[idx]}")Example
import torch
torch.manual_seed(0)
words = ["the", "cat", "sat", "on", "a", "mat"]
seq = torch.tensor([0, 1, 2, 3, 4, 5])
target = torch.tensor([1, 2, 3, 4, 5, 5])
emb = torch.nn.Embedding(6, 8)
wq = torch.nn.Linear(8, 8, bias=False)
wk = torch.nn.Linear(8, 8, bias=False)
wv = torch.nn.Linear(8, 8, bias=False)
head = torch.nn.Linear(8, 6)
params = (list(emb.parameters()) + list(wq.parameters()) + list(wk.parameters())
+ list(wv.parameters()) + list(head.parameters()))
opt = torch.optim.Adam(params, lr=0.02)
loss_fn = torch.nn.CrossEntropyLoss()
mask = torch.triu(torch.full((6, 6), float("-inf")), diagonal=1)
for step in range(100):
x = emb(seq).unsqueeze(0)
Q, K, V = wq(x), wk(x), wv(x)
scores = Q @ K.transpose(-2, -1) / (8 ** 0.5)
weights = torch.softmax(scores + mask, dim=-1)
mixed = weights @ V
logits = head(mixed).squeeze(0)
loss = loss_fn(logits, target)
opt.zero_grad()
loss.backward()
opt.step()
print(f"Final loss: {loss.item():.4f}")
with torch.no_grad():
x = emb(seq).unsqueeze(0)
Q, K, V = wq(x), wk(x), wv(x)
scores = Q @ K.transpose(-2, -1) / (8 ** 0.5)
weights = torch.softmax(scores + mask, dim=-1).squeeze(0)
print(" " + " ".join(f"{word:>5}" for word in words))
for word, row in zip(words, weights):
print(f"{word:>4} " + " ".join(f"{v:5.2f}" for v in row.tolist()))Key points
- Weight grid: one row per word, holding the weights that word hands out to the others
- Each row adds up to 1, because each row comes from one softmax
- The biggest entry in a row names the word that row leaned on most
- A grid can show where the weight went, but never why it went there
- In short, read the grid against the input and report only the pattern it actually shows
In this example, a tiny attention model trains to predict the next word of one sentence. The sentence is the cat sat on a mat, and the target for the last word repeats mat. Embeddings give each word a vector, and one causal attention head mixes the vectors. After 100 steps the final loss prints as 0.0006. The grid then prints the trained weights with the words as labels. The rows for a and mat put all of their weight on sat. The sat row spreads its weight across the, cat, and sat. The grid shows a model leaning on sat to finish the sentence, and it shows nothing about grammar.
Note
This lesson's grid comes from one short sentence, so its patterns describe this model on this sentence only
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
In the trained grid, which word do the rows for a and mat lean on?