Self-Attention - Built From Matmuls

Context

Self-attention replaces every token's vector with a weighted average of the vectors of every token in the same query, and each token computes its own weights: its dot product with every token, passed through a softmax so the weights sum to 1. Two matrix multiplications and one softmax do the whole averaging, for every token of every query in a batch at once.

Every transformer runs this operation in every layer. The pretrained encoder Section 5.4 fine-tunes on Banking77 reads "card" in "my card never arrived" through it, and a model that writes text runs it with the later words hidden. The function written here, attention, is the one Lesson 5.3.3 builds the transformer block on.

Lesson 5.3.1's bag of embeddings averaged a query's real tokens with equal weight, so "card" counted the same in every query whatever stood beside it, and the bag could not tell one word order from another. Attention lets the weights depend on the tokens, and it fails without a word in four places: unscaled scores turn the average into a copy of one token, unmasked padding changes a query's vectors with the number of <pad> tokens after it, a missing causal mask lets a later word reach an earlier position, and with no positions at all the layer is as blind to order as the bag.

Goal

• Build attention from q @ k.transpose(-2, -1), a softmax over each row and a weighted sum of values, and check one token's output against a loop over its weights.
• Measure the softmax saturating without the sqrt(d) scale, a largest weight of 0.9321 on average at 512 dimensions against 0.1285 with it.
• Mask <pad> keys with masked_fill(..., float("-inf")), show a real token's output moving with the padding without the mask and holding with it, and shape the mask to broadcast against [batch, seq, seq] scores.
• Build a causal mask with torch.triu, show a later token leaking into an earlier position without it, and show attention with no positions giving a reordered query the same pooled vector.
• Save attention.py, check it against torch.nn.functional.scaled_dot_product_attention with torch.allclose, and measure the seq x seq weights growing with the square of the length.

Prerequisites

None of this is required, and the lesson stands on its own. The first is the picture of an attention pattern before any code, and the second is PyTorch's own statement of the function this lesson rebuilds.

• YouTube - "Attention in transformers, step-by-step | Deep Learning Chapter 6" by 3Blue1Brown, 26 minutes. From 4:29 the attention pattern, a grid of every query against every key with each column normalized by a softmax, from 11:08 masking the later tokens, and from 13:10 the values that the pattern averages.
• Official docs - PyTorch 2.14, torch.nn.functional.scaled_dot_product_attention. The arguments, the default scale of one over the square root of the last dimension, the boolean mask in which True marks a position that takes part, and a reference implementation in about twenty lines of Python.

Table of Contents

• Attention as a Weighted Average
• Query, Key and Value
• Scaling by the Square Root of d
• The Padding Mask
• The Causal Mask
• No Order Without Positions
• Checking Against PyTorch's Own Attention
• The Cost of a Square Matrix

Theory

Every block runs in ~/dev/bootcamp, beside text.py, Lesson 5.3.1's, and image_loop.py, Lesson 4.5.2's, which only the Assignment imports. Nothing here installs or downloads. The fourth subtitle saves attention.py, which every later block imports and Lesson 5.3.3 builds on. Every weight in Theory is a seeded draw, never trained, so the numbers show the mechanism, not what a trained model attends to. A Python block that opens with its imports is a script of its own, and a block that does not continues the script above it.

Attention as a Weighted Average

What. Token i takes its dot product with every token j, Lesson 2.3.1's, which grows as two vectors point the same way and as they grow longer. The softmax, Lesson 4.3.2's, turns those scores into positive weights that sum to 1, and token i's new vector is the sum of every token's vector times its weight. For a batch the scores are one matrix product, x @ x.transpose(-2, -1), of shape [batch, seq, seq]: row i holds token i's score against every token. Why. Lesson 5.3.1's bag weighted every real token by one over the count. Attention keeps the average and lets the weights depend on the tokens, so every output is a mix tilted toward the tokens most like the one computing it. When. Inside every attention layer, here first with the raw embeddings in every role, so the mechanism shows before anything is learned.

import math
import torch
from torch import nn
from text import vocab, tokenize

tokens = tokenize("Where is my card?")
ids = torch.tensor([[vocab[t] for t in tokens]])
print(tokens, ids.shape)
torch.manual_seed(42)
embed = nn.Embedding(len(vocab), 4, padding_idx=0)
with torch.no_grad():
    x = embed(ids)
print(x.shape)
print(x[0].round(decimals=2))

['where', 'is', 'my', 'card', '?'] torch.Size([1, 5])
torch.Size([1, 5, 4])
tensor([[ 1.2600, -0.7100,  0.8500,  0.5100],
        [-1.2700,  2.1200, -1.2300, -0.4900],
        [ 1.6400, -0.1600, -0.5000,  0.4400],
        [ 0.3200, -0.4200,  0.3100, -0.7700],
        [-0.7300, -0.5600, -0.7700,  0.7600]])

The query is a card_arrival training query, and ids is a batch of one, [1, 5], so every shape below carries the batch dimension first. math is for the next subtitle. Four dimensions keep every vector readable, and the embedding comes out of torch.no_grad(), Lesson 4.1.4's, so it carries no graph. The scores:

scores = x @ x.transpose(-2, -1)
print(scores.shape)
print(scores[0].round(decimals=2))

torch.Size([1, 5, 5])
tensor([[ 3.0800, -4.4200,  1.9800,  0.5700, -0.7800],
        [-4.4200,  7.8900, -2.0300, -1.3100,  0.3200],
        [ 1.9800, -2.0300,  3.1600,  0.1000, -0.3900],
        [ 0.5700, -1.3100,  0.1000,  0.9800, -0.8200],
        [-0.7800,  0.3200, -0.3900, -0.8200,  2.0200]])

transpose(-2, -1) swaps the last two dimensions, -1 being the last and -2 the one before it, so [1, 5, 4] becomes [1, 4, 5] and the batch dimension stays where it is. @ then multiplies the two matrices of each batch entry, Lesson 4.1.2's product, [5, 4] @ [4, 5], and every entry is one dot product: row 3, column 2 is card against my, 0.10. The matrix is symmetric, because card against my and my against card are the same dot product. Each row becomes weights:

weights = scores.softmax(dim=-1)
print(weights[0].round(decimals=2))
print(weights.sum(dim=-1))
print(scores.softmax(dim=1).sum(dim=-1)[0].round(decimals=4))

tensor([[0.7000, 0.0000, 0.2300, 0.0600, 0.0100],
        [0.0000, 1.0000, 0.0000, 0.0000, 0.0000],
        [0.2200, 0.0000, 0.7200, 0.0300, 0.0200],
        [0.2800, 0.0400, 0.1800, 0.4300, 0.0700],
        [0.0400, 0.1300, 0.0600, 0.0400, 0.7200]])
tensor([[1., 1., 1., 1., 1.]])
tensor([1.2442, 1.1785, 1.1949, 0.5581, 0.8243])

softmax(dim=-1) is torch.softmax as a method, taken across each row, so row 3 is how card divides its attention. dim=1 runs down each column instead, a valid dimension, so nothing is refused, and the rows the next line averages with no longer sum to 1. The weighted sum, and card's output checked by a loop over its five weights:

out = weights @ x
print(out.shape)
card = sum(weights[0, 3, j] * x[0, j] for j in range(5))
print(out[0, 3].round(decimals=4), card.round(decimals=4))
print(torch.allclose(out[0, 3], card))
print(weights[0].argmax(dim=-1))

torch.Size([1, 5, 4])
tensor([ 0.6753, -0.3587,  0.1758, -0.0739]) tensor([ 0.6753, -0.3587,  0.1758, -0.0739])
True
tensor([0, 1, 2, 3, 4])

Row i of weights @ x is row i of the weights times every token's vector, summed, which the loop computes the slow way and torch.allclose, Lesson 4.1.2's float comparison, confirms. The last line is the defect of the embedding in every role: every row's largest weight sits on the token itself, because a vector's dot product with itself is its squared length, and is, the longest, keeps 1.0000 of its weight.

DON'T:

• DON'T write x.T for the transpose of a batch. On the [1, 5, 4] embeddings it reverses all three dimensions to [4, 5, 1] with a deprecation warning, and x @ x.T is refused, where transpose(-2, -1) swaps the last two and leaves the batch first.
• DON'T take the softmax along dim=1 of [batch, seq, seq] scores. It runs without a word, and its rows summed to anything from 0.5581 to 1.2442, so no output was an average.
• DON'T use the raw embeddings as both sides of the scores. Every token's largest weight was on itself, and is kept 1.0000 of its own.

Query, Key and Value

What. Three nn.Linear(d, d) layers, Lesson 4.2.2's, turn each token's embedding into three vectors: a query q, what the token looks for, a key k, what it offers to be found by, and a value v, what it hands over when it is chosen. The scores are q @ k.transpose(-2, -1) divided by math.sqrt(d), the scale the next subtitle measures, the softmax runs over each row as before, and the output is weights @ v. Why. With the embedding in every role, each row's largest weight sat on its own token. Separate projections let card look for one thing, be found for another, and hand over a third, and training sets all three. When. In every attention layer: the three projections are its parameters, and the rest of the layer has none.

to_q, to_k, to_v = nn.Linear(4, 4), nn.Linear(4, 4), nn.Linear(4, 4)
with torch.no_grad():
    q, k, v = to_q(x), to_k(x), to_v(x)
print(q.shape, k.shape, v.shape)
print(torch.equal(q, k), sum(p.numel() for layer in (to_q, to_k, to_v) for p in layer.parameters()))

torch.Size([1, 5, 4]) torch.Size([1, 5, 4]) torch.Size([1, 5, 4])
False 60

The three layers continue the seeded stream after the embedding, and each holds a [4, 4] weight and 4 biases, 60 parameters between them. q and k are two different views of the same five embeddings, so a token's query against its own key is no longer its length squared. The scores and the weights:

scores = q @ k.transpose(-2, -1) / math.sqrt(4)
weights = scores.softmax(dim=-1)
print(weights[0].round(decimals=2))
print(weights[0].argmax(dim=-1))
print(torch.allclose(scores[0, 3], (q[0, 3] * k[0]).sum(dim=-1) / 2))

tensor([[0.2100, 0.1800, 0.2100, 0.2100, 0.1900],
        [0.2200, 0.1100, 0.2300, 0.2600, 0.1800],
        [0.2100, 0.1500, 0.2300, 0.2400, 0.1700],
        [0.2200, 0.1300, 0.2300, 0.2500, 0.1700],
        [0.2800, 0.1200, 0.2100, 0.2500, 0.1400]])
tensor([3, 3, 3, 3, 0])
True

Four rows now weight card most and the last weights where most, so card's row is the only one still leaning on itself, 0.25 against my's 0.23. The weights are close to even because seeded projections give scores near zero. The last line rebuilds card's scores as its query times every key, summed and divided by 2, the square root of 4. The values and what the weights make of them:

out = weights @ v
print(v[0].round(decimals=2))
print(out[0].round(decimals=2))

tensor([[-0.3800, -0.0600, -0.3600,  0.4700],
        [ 1.8400, -0.1700,  1.0500, -0.0600],
        [ 0.4500, -0.2700,  0.1200, -0.0500],
        [ 0.3600, -0.5100,  0.3900,  0.1300],
        [-0.0700, -0.7300, -0.2900,  1.1800]])
tensor([[ 0.4000, -0.3500,  0.1600,  0.3300],
        [ 0.3100, -0.3600,  0.1100,  0.3300],
        [ 0.3700, -0.3500,  0.1500,  0.3200],
        [ 0.3300, -0.3500,  0.1300,  0.3200],
        [ 0.2900, -0.3200,  0.1100,  0.3100]])

card's output, row 3, is its row of weights times the five values, added entry by entry. The five outputs sit much closer together than the five values, because nearly uniform rows average the same values in nearly the same proportions.

DON'T:

• DON'T read the weights of an untrained layer as what the model attends to. These rows are near uniform because nothing is trained, and four of them pointing at card is a seeded draw, not a finding.
• DON'T reuse one projection for queries and keys to save parameters. q equal to k puts every token's own squared length back on its diagonal, the first subtitle's defect.

Scaling by the Square Root of d

What. The scores are divided by math.sqrt(d), d being the length of each query and key. When the entries of q and k are independent with a mean of 0 and a variance of 1, their dot product adds d such products, and its variance is d, so its standard deviation is sqrt(d). Dividing by sqrt(d) brings every row of scores back to a spread of about 1, whatever d is. Why. A row whose scores differ by tens gives almost all of its weight to the largest, and the weighted average becomes a copy of one value. The spread grows with d while nothing else changes, so without the scale the softmax saturates at the sizes real layers use. When. In every attention layer, and as the default of PyTorch's own function in the seventh subtitle.

import math
import torch

gen = torch.Generator().manual_seed(42)
for d in (16, 64, 512):
    q = torch.randn(256, 48, d, generator=gen)
    k = torch.randn(256, 48, d, generator=gen)
    raw = q @ k.transpose(-2, -1)
    for name, scores in (("raw", raw), ("scaled", raw / math.sqrt(d))):
        weights = scores.softmax(dim=-1)
        top = weights.max(dim=-1).values
        spread = torch.special.entr(weights).sum(dim=-1)
        print(d, name, f"{scores.std().item():.2f}", f"{top.mean().item():.4f}", f"{spread.mean().item():.4f}", f"{(top > 0.99).float().mean().item():.4f}")
print(f"{math.log(48):.4f}")

16 raw 4.01 0.5976 1.2796 0.0251
16 scaled 1.00 0.1273 3.4085 0.0000
64 raw 7.98 0.8031 0.5434 0.2128
64 scaled 1.00 0.1285 3.4068 0.0000
512 raw 22.62 0.9321 0.1690 0.6281
512 scaled 1.00 0.1285 3.4057 0.0000
3.8712

Each d draws 256 sets of 48 queries and 48 keys from the standard normal, torch.randn on a generator of its own, Lesson 4.1.1's, 48 being MAX_LEN. torch.special.entr(w) computes -w * log(w) for every weight, 0 for a weight of 0, and its sum over a row is the row's entropy (how evenly a row spreads its weight: log(48), 3.8712, for 48 equal weights, 0 for one weight of 1). The columns after the name are the scores' standard deviation, the mean largest weight of a row, the mean entropy, and the share of rows whose largest weight passes 0.99. Unscaled, the spread follows the square root of d and the rows collapse onto one key as it grows. Scaled, every d gives the same rows. What the scale leaves alone:

print(torch.equal(raw.argmax(dim=-1), (raw / math.sqrt(512)).argmax(dim=-1)))
print(f"{q.var().item():.4f}", f"{k.var().item():.4f}", f"{raw.var().item():.1f}")

True
0.9999 1.0001 511.8

A positive divisor keeps a row's order, so the scale changes how far the softmax leans toward a row's largest score, never which key it is. The second line is the reason in numbers at 512: entries of variance 0.9999 and 1.0001 give dot products of variance 511.8.

DON'T:

• DON'T drop the scale because a small layer works without it. At 16 dimensions the unscaled rows already gave their largest weight 0.5976 on average against 0.1273 scaled.
• DON'T read a saturated row as a confident one. At 512 dimensions unscaled, 62.81 percent of rows put over 0.99 on one key, from random vectors that carry no information at all.

The Padding Mask

What. encode, Lesson 5.3.1's, pads every query to MAX_LEN tokens with id 0. A padding mask is a boolean tensor, True at every <pad> key, and scores.masked_fill(mask, float("-inf")) (a copy of scores with every entry where mask is True replaced by the value) sets those scores to minus infinity. The softmax turns minus infinity into a weight of exactly 0, so no token takes anything from a pad, and the real tokens' weights still sum to 1. The mask built from ids == 0 is [batch, seq], and unsqueeze(1) makes it [batch, 1, seq], which broadcasts against [batch, seq, seq] scores by Lesson 4.1.2's rule: the same keys blocked for every query of a batch entry. Why. padding_idx=0 gives a pad the zero embedding, Lesson 5.3.1's, and the key projection adds its bias, so a pad's key is a real vector with real scores, and even a key of zeros scores 0, which a softmax weights like any other score. Unmasked, the pads take most of every real token's attention, and the outputs change with nothing but the number of pads. When. Every time padded sequences go through attention, which is every batch of Banking77 queries.

Save the function as attention.py, the three steps of the first two subtitles with the mask in between:

import math

def attention(q, k, v, mask=None):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

q.shape[-1] is d, read from the queries, so one function serves every size. mask is True wherever a query must not look and only has to broadcast against the scores. The function returns the weights beside the output, so every check below can read them. Without a mask, on "Where is my card?" padded to 48, with 16 dimensions from here on:

import torch
from torch import nn
import torch.nn.functional as F
from text import vocab, tokenize, encode, val_set
from attention import attention

torch.manual_seed(42)
embed = nn.Embedding(len(vocab), 16, padding_idx=0)
to_q, to_k, to_v = nn.Linear(16, 16), nn.Linear(16, 16), nn.Linear(16, 16)

def project(ids):
    with torch.no_grad():
        x = embed(ids)
        return to_q(x), to_k(x), to_v(x)

ids = torch.tensor([encode("Where is my card?")])
print(ids.shape, ids[0, :8])
q, k, v = project(ids)
out, weights = attention(q, k, v)
print(weights.shape)
print(weights[0, :5, 5:].sum(dim=-1).round(decimals=4))
print(torch.equal(k[0, 5], to_k.bias), torch.equal(k[0, 5], k[0, 47]))

torch.Size([1, 48]) tensor([62, 11,  4,  9,  3,  0,  0,  0])
torch.Size([1, 48, 48])
tensor([0.8994, 0.8826, 0.9006, 0.8937, 0.9059])
True True

F is torch.nn.functional, as Lesson 5.1.1 imported it, for the seventh subtitle. About nine tenths of every real token's weight goes to the 43 pads, weights[0, :5, 5:], and every pad's key is exactly the key projection's bias, the same vector for all 43. The same query padded to 10 instead of 48:

short = ids[:, :10]
qs, ks, vs = project(short)
out_short, _ = attention(qs, ks, vs)
print(round((out[0, :5] - out_short[0, :5]).abs().max().item(), 4))

0.1516

Only the number of pads after the five real tokens changed, and their outputs moved by up to 0.1516, so the same query in two batches is two different inputs to every layer above. The mask:

pad = ids == 0
mask = pad.unsqueeze(1)
print(pad.shape, mask.shape, int(pad.sum()))
out, weights = attention(q, k, v, mask)
out_short, _ = attention(qs, ks, vs, (short == 0).unsqueeze(1))
print(weights[0, :5, 5:].sum().item(), torch.allclose(out[0, :5], out_short[0, :5], atol=1e-6))
print(weights[0, :5, :5].round(decimals=2))
real, real_short = (~pad).unsqueeze(-1).float(), (short != 0).unsqueeze(-1).float()
pooled = (out * real).sum(dim=1) / real.sum(dim=1)
pooled_short = (out_short * real_short).sum(dim=1) / real_short.sum(dim=1)
print(torch.allclose(pooled, pooled_short, atol=1e-6))

torch.Size([1, 48]) torch.Size([1, 1, 48]) 43
0.0 True
tensor([[0.1900, 0.1700, 0.2600, 0.2100, 0.1600],
        [0.1500, 0.2300, 0.1700, 0.2500, 0.2000],
        [0.3200, 0.1600, 0.1400, 0.1600, 0.2300],
        [0.2500, 0.2200, 0.1600, 0.1300, 0.2200],
        [0.1600, 0.1800, 0.2300, 0.2300, 0.2000]])
True

Masked, the pads get exactly 0.0, and the outputs padded to 48 and to 10 agree within atol=1e-6, the tolerance every comparison of two differently shaped computations uses here, since float32 sums grouped differently differ in their last digits. The pooled vector, Lesson 5.3.1's mean over the real tokens, with ~pad the mask flipped (~ turns every True of a boolean tensor to False and back), is the same however the query is padded. The pad positions still get outputs, and the pooling's mask leaves them out. Minus infinity, not zero:

zeroed = (q @ k.transpose(-2, -1) / 4).masked_fill(mask, 0.0).softmax(dim=-1)
print(zeroed[0, :5, 5:].sum(dim=-1).round(decimals=4))

tensor([0.8985, 0.8732, 0.9084, 0.8935, 0.9021])

The scores are the function's, divided by 4, the square root of 16, with the pads' set to 0.0, and the pads keep about what they had unmasked, because these seeded scores all sit near 0. A whole batch, the first 32 validation queries, with the mask shaped for it, and then shaped wrong:

X = val_set.tensors[0][:32]
qb, kb, vb = project(X)
out_b, weights_b = attention(qb, kb, vb, (X == 0).unsqueeze(1))
print(qb.shape, weights_b.shape, out_b.shape)
print((weights_b * (X == 0).unsqueeze(1)).sum().item())
try:
    attention(qb, kb, vb, X == 0)
except RuntimeError as err:
    print(err)
X48 = val_set.tensors[0][:48]
q48, k48, v48 = project(X48)
right, _ = attention(q48, k48, v48, (X48 == 0).unsqueeze(1))
wrong, _ = attention(q48, k48, v48, X48 == 0)
print(wrong.shape, torch.allclose(right, wrong, atol=1e-6))

torch.Size([32, 48, 16]) torch.Size([32, 48, 48]) torch.Size([32, 48, 16])
0.0
The size of tensor a (32) must match the size of tensor b (48) at non-singleton dimension 1
torch.Size([48, 48, 16]) False

val_set.tensors[0] is the [2000, 48] tensor of validation ids inside Lesson 5.3.1's TensorDataset. The [32, 48] mask lines up on the right against [32, 48, 48] as [1, 32, 48], its 32 meets the scores' 48, and the call is refused, the lucky case. With a batch of 48 the same mistake broadcasts as [1, 48, 48]: the token at position i of every query is masked by the padding of query i in the batch, the output has the right shape, and it is wrong.

DON'T:

• DON'T trust padding_idx=0 to keep pads out of attention. It zeroes the pad's embedding, and the key projection's bias made all 43 pad keys one real vector that took about nine tenths of every real token's weight.
• DON'T mask with 0.0. A score of 0 is a score, and the pads kept from 0.8732 to 0.9084 of each real token's weight. Minus infinity is the value whose softmax weight is exactly 0.
• DON'T hand attention the [batch, seq] mask. It was refused at a batch of 32 and silently wrong at a batch of 48, and unsqueeze(1) makes the [batch, 1, seq] that blocks the same keys for every query.

The Causal Mask

What. A causal mask blocks every key that comes after its query: torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1) (triu keeps the upper triangle of a matrix, with diagonal=1 the entries strictly above the main diagonal, and sets the rest to 0, here False) is True wherever column j is later than row i. Its shape [seq, seq] broadcasts against [batch, seq, seq] scores, so one mask serves every query of a batch, and | joins it to a padding mask when both are needed. Why. A model that writes text produces one token at a time, so at each position the later tokens do not exist yet, and its output there must come from that position and the ones before it. Without the mask, a change at the end of a sequence reaches its very first position. When. In the models that generate text, which Lesson 5.4.3's map places. A classifier reading a whole Banking77 query at once uses the padding mask alone.

causal = torch.triu(torch.ones(5, 5, dtype=torch.bool), diagonal=1)
print(causal)
ask = torch.tensor([[vocab[t] for t in tokenize("Where is my card?")]])
told = torch.tensor([[vocab[t] for t in tokenize("Where is my card.")]])
print(ask, told)
for m in (None, causal):
    out_ask, w_ask = attention(*project(ask), m)
    out_told, _ = attention(*project(told), m)
    print((out_ask[0] - out_told[0]).abs().max(dim=-1).values.round(decimals=4))
print(w_ask[0].round(decimals=2))

tensor([[False,  True,  True,  True,  True],
        [False, False,  True,  True,  True],
        [False, False, False,  True,  True],
        [False, False, False, False,  True],
        [False, False, False, False, False]])
tensor([[62, 11,  4,  9,  3]]) tensor([[62, 11,  4,  9,  5]])
tensor([0.1346, 0.1630, 0.1836, 0.1856, 0.1918])
tensor([0.0000, 0.0000, 0.0000, 0.0000, 0.1918])
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3900, 0.6100, 0.0000, 0.0000, 0.0000],
        [0.5200, 0.2500, 0.2300, 0.0000, 0.0000],
        [0.3300, 0.2900, 0.2100, 0.1700, 0.0000],
        [0.1600, 0.1800, 0.2300, 0.2300, 0.2000]])

The two queries differ in their last token only, ? against ., and *project(ask) unpacks the three tensors project returns into attention's first three arguments. Unmasked, every position's output moves, so the first position knows how the query ends. Masked, only the last moves. A causal and a padding mask together, on the batch:

both = (X == 0).unsqueeze(1) | torch.triu(torch.ones(48, 48, dtype=torch.bool), diagonal=1)
print(both.shape)
out_both, w_both = attention(qb, kb, vb, both)
print(torch.isnan(out_both).any().item(), w_both[0, 0, :3])

torch.Size([32, 48, 48])
False tensor([1., 0., 0.])

[32, 1, 48] | [48, 48] broadcasts to one [32, 48, 48] mask, a key blocked when it is a pad or when it is later. No output is nan, because every row keeps key 0, a real token in every query.

DON'T:

• DON'T leave the causal mask out of a layer whose positions must not see ahead. Changing the last token moved where's output by 0.1346, and no shape, no error and no loss reports it.
• DON'T build the mask with diagonal=0. That also blocks the diagonal, so the first position has no key left, its row of minus infinities comes out of the softmax as nan, and every other position loses its own token.

No Order Without Positions

What. Self-attention with no information about positions is permutation-equivariant (reorder the inputs and the outputs come back reordered the same way, and otherwise unchanged), since every score depends on two tokens' vectors and on nothing about where they stand. A mean over every token's output is then the same for every order of the same tokens. Why. "Where is my card?" and "card my where is ?" are different queries to a person and the same set of outputs to an attention layer with no positions, so a classifier that pools them gives them the same intent. Attention does not cure the bag's blindness to order on its own. When. In every attention layer whose input carries no position, which is every layer of this lesson. Lesson 5.3.3 adds position embeddings to close the gap.

perm = torch.tensor([3, 2, 0, 1, 4])
print([tokenize("Where is my card?")[i] for i in perm])
out_ask, _ = attention(*project(ask))
out_perm, _ = attention(*project(ask[:, perm]))
print(torch.allclose(out_perm, out_ask[:, perm], atol=1e-6), torch.allclose(out_perm.mean(dim=1), out_ask.mean(dim=1), atol=1e-6))
print(out_ask.mean(dim=1)[0, :4].round(decimals=4), out_perm.mean(dim=1)[0, :4].round(decimals=4))
out_c, _ = attention(*project(ask), causal)
out_pc, _ = attention(*project(ask[:, perm]), causal)
print(torch.allclose(out_pc, out_c[:, perm], atol=1e-6), round((out_pc.mean(dim=1) - out_c.mean(dim=1)).abs().max().item(), 4))

['card', 'my', 'where', 'is', '?']
True True
tensor([ 0.0041, -0.0589,  0.0809,  0.1853]) tensor([ 0.0041, -0.0589,  0.0809,  0.1853])
False 0.8583

The reordered query's outputs are the original outputs reordered by the same perm, and the pooled vectors agree, the first four of their sixteen entries printed, so a classifier reading the pooled vector cannot tell the two orders apart. The causal mask breaks the symmetry: it is built from positions, so the same tokens in another order see other keys, and the pooled vectors differ by up to 0.8583.

DON'T:

• DON'T expect an attention layer to read word order because every output depends on every token. The reordered query's pooled vector equaled the original's within 1e-6.
• DON'T read the causal mask's order as a fix for an encoder. It changed the pooled vector by hiding from every position the tokens after it, the rest of the query a classifier needs.

Checking Against PyTorch's Own Attention

What. F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=False) returns attention's output alone. Its boolean attn_mask marks the positions that take part, True to attend, the opposite of attention.py's True to block, and is_causal=True applies the causal mask without one being built. Why. A hand-built layer that matches the library's within a tolerance, masks included, is checked at once in its formula, its scale, its softmax dimension and its mask's direction. The two mask conventions are the trap: either function runs on the other's mask without an error. When. Whenever a hand-built layer is written or replaced by the library's.

pad_b = (X == 0).unsqueeze(1)
mine, _ = attention(qb, kb, vb, pad_b)
theirs = F.scaled_dot_product_attention(qb, kb, vb, attn_mask=~pad_b)
print(torch.allclose(mine, theirs, atol=1e-6))
flipped = F.scaled_dot_product_attention(qb, kb, vb, attn_mask=pad_b)
print(torch.allclose(mine, flipped, atol=1e-6), torch.isnan(flipped).any().item())
c48 = torch.triu(torch.ones(48, 48, dtype=torch.bool), diagonal=1)
print(torch.allclose(attention(qb, kb, vb, c48)[0], F.scaled_dot_product_attention(qb, kb, vb, is_causal=True), atol=1e-6))
print(torch.allclose(out_both, F.scaled_dot_product_attention(qb, kb, vb, attn_mask=~both), atol=1e-6))

True
False False
True
True

With the mask flipped by ~, PyTorch's function agrees with attention on all 32 queries within 1e-6. Handed attention.py's mask unflipped, it attends to the pads alone: no error, no nan, since every one of these 32 queries has pads, and every output wrong. is_causal=True matches the triu mask, and the joined mask, flipped, matches out_both. attention reads only the last two dimensions of q, k and v, so one more dimension in front, as Lesson 5.3.3's heads add, changes nothing but the mask's shape:

pad4 = (X == 0)[:, None, None, :]
four, _ = attention(qb.unsqueeze(1), kb.unsqueeze(1), vb.unsqueeze(1), pad4)
print(four.shape, torch.allclose(four.squeeze(1), out_b, atol=1e-6))
grown, _ = attention(qb.unsqueeze(1), kb.unsqueeze(1), vb.unsqueeze(1), pad_b)
print(grown.shape)

torch.Size([32, 1, 48, 16]) True
torch.Size([32, 32, 48, 16])

[:, None, None, :] inserts two dimensions of size 1, so the mask is [32, 1, 1, 48] against [32, 1, 48, 48] scores, and the output is the batch's own with a dimension of 1 in it. The [32, 1, 48] mask lines up on the right as [1, 32, 1, 48], and broadcasting builds a [32, 32, 48, 48] result without a refusal, each query masked by every query's padding in turn.

DON'T:

• DON'T pass attention.py's mask to scaled_dot_product_attention unchanged. Its True means attend, and the unflipped mask made every query average its pads with no error and no nan.
• DON'T keep a padding mask's shape when the scores gain a dimension. The [batch, 1, seq] mask against four-dimensional scores grew the output to [32, 32, 48, 16], and [batch, 1, 1, seq] is the shape that matches.

The Cost of a Square Matrix

What. The weights of one sequence of n tokens are an n by n matrix. Its bytes are numel() times element_size(), Lesson 4.3.3's measure, 4 a float32 entry, so doubling the length makes it four times larger, while the output, n vectors of d, only doubles. Why. Banking77's 48 tokens make that matrix negligible, and a document of thousands of tokens makes it the largest tensor in the layer, held for every sequence of a batch in every attention layer and kept for backward() while training, where the parameters do not depend on n at all. When. When MAX_LEN and the batch size are chosen, and before any model is pointed at long documents.

import time
import torch
from attention import attention

gen = torch.Generator().manual_seed(42)
for n in (48, 512, 1024, 2048, 4096, 8192):
    x = torch.randn(1, n, 64, generator=gen)
    seconds = []
    for _ in range(5):
        start = time.perf_counter()
        out, weights = attention(x, x, x)
        seconds.append(time.perf_counter() - start)
    print(n, tuple(weights.shape), weights.numel() * weights.element_size(), out.numel() * out.element_size(), f"{min(seconds) * 1000:.1f}")

48 (1, 48, 48) 9216 12288 1.3
512 (1, 512, 512) 1048576 131072 2.0
1024 (1, 1024, 1024) 4194304 262144 4.5
2048 (1, 2048, 2048) 16777216 524288 9.2
4096 (1, 4096, 4096) 67108864 1048576 49.7
8192 (1, 8192, 8192) 268435456 2097152 203.7

The columns are the length, the weights' shape, their bytes, the output's bytes and the fastest of five calls in milliseconds, time.perf_counter(), Lesson 3.6.3's clock. Every doubling of the length multiplied the weights by four and the output by two, to a quarter of a gigabyte of weights for one sequence of 8,192 tokens in one layer. From 2,048 tokens on, every doubling multiplied the milliseconds by four or more, and below that the time grows more slowly.

DON'T:

• DON'T raise MAX_LEN without counting the weights it buys. Each doubling multiplied every sequence's weights by four, before the batch and the layers multiply them again.
• DON'T read the parameter count as the memory a long input needs. The projections are the same at 48 tokens and at 8,192, and the weights went from 9,216 bytes to 268,435,456.

Lab Exercises

Lab Exercise 1: Attention by Loops

In loops.py, tokenize the Banking77 query "My card never arrived.", seed 42, and build nn.Embedding(len(vocab), 8, padding_idx=0) and three nn.Linear(8, 8) projections, in that order. Compute q, k and v without a graph.

For the token card, compute its row of weights with Python alone: its query's dot product with each key as (q[0, i] * k[0, j]).sum().item(), divided by math.sqrt(8), passed through math.exp, then each divided by their total, and its output as the sum of each weight times that token's value. Print the tokens and card's position, the five weights to four decimals and their sum, whether the matrix products give the same row and output, and the token card weights most.

Then rebuild only the embedding and the query and key projections at 8 and at 512 dimensions, each after torch.manual_seed(42), and print the largest weight in card's row unscaled and scaled at each size.

Lab Exercise 2: Proving the Masks on a Batch

In masks.py, seed 42, build a 32-dimension embedding with padding_idx=0 and three projections, and take the first 64 validation queries from val_set. Print the batch's shape and the shortest and longest real length in it.

Run attention with the padding mask and print the total weight any query gives a pad. Run each of the 64 queries alone, cut to its own real length and with no mask, and print whether every real token's output agrees with the batch's within 1e-6, then whether the batch agrees with F.scaled_dot_product_attention given the right mask.

Join a causal mask to the padding mask, print its shape, whether every query's first output equals that position's value within 1e-6, and whether the result agrees with PyTorch's function. Finish with the bytes the batch's weights hold, and the bytes they would hold at 512 tokens a query.

import math
import torch
from torch import nn
from text import vocab, tokenize

tokens = tokenize("My card never arrived.")
ids = torch.tensor([[vocab[t] for t in tokens]])
torch.manual_seed(42)
embed = nn.Embedding(len(vocab), 8, padding_idx=0)
to_q, to_k, to_v = nn.Linear(8, 8), nn.Linear(8, 8), nn.Linear(8, 8)
with torch.no_grad():
    x = embed(ids)
    q, k, v = to_q(x), to_k(x), to_v(x)
i = tokens.index("card")
row = [math.exp((q[0, i] * k[0, j]).sum().item() / math.sqrt(8)) for j in range(len(tokens))]
total = sum(row)
row = [r / total for r in row]
mixed = sum(row[j] * v[0, j] for j in range(len(tokens)))
print(tokens, i)
print([round(r, 4) for r in row], round(sum(row), 4))
weights = (q @ k.transpose(-2, -1) / math.sqrt(8)).softmax(dim=-1)
out = weights @ v
print(torch.allclose(weights[0, i], torch.tensor(row)), torch.allclose(out[0, i], mixed, atol=1e-6))
print(tokens[weights[0, i].argmax()])

['my', 'card', 'never', 'arrived', '.'] 1
[0.2359, 0.1761, 0.1764, 0.2292, 0.1825] 1.0
True True
my

for d in (8, 512):
    torch.manual_seed(42)
    embed = nn.Embedding(len(vocab), d, padding_idx=0)
    to_q, to_k = nn.Linear(d, d), nn.Linear(d, d)
    with torch.no_grad():
        q, k = to_q(embed(ids)), to_k(embed(ids))
    raw = q @ k.transpose(-2, -1)
    for name, scores in (("raw", raw), ("scaled", raw / math.sqrt(d))):
        print(d, name, f"{scores.softmax(dim=-1)[0, i].max().item():.4f}")

8 raw 0.3046
8 scaled 0.2359
512 raw 0.9998
512 scaled 0.3291

import torch
from torch import nn
import torch.nn.functional as F
from text import vocab, MAX_LEN, val_set
from attention import attention

torch.manual_seed(42)
embed = nn.Embedding(len(vocab), 32, padding_idx=0)
to_q, to_k, to_v = nn.Linear(32, 32), nn.Linear(32, 32), nn.Linear(32, 32)

def project(ids):
    with torch.no_grad():
        x = embed(ids)
        return to_q(x), to_k(x), to_v(x)

X = val_set.tensors[0][:64]
lengths = (X != 0).sum(dim=1)
print(X.shape, lengths.min().item(), lengths.max().item())
q, k, v = project(X)
pad = (X == 0).unsqueeze(1)
out, weights = attention(q, k, v, pad)
print((weights * pad).sum().item())
worst = 0.0
for n in range(64):
    alone = X[n:n + 1, :lengths[n]]
    o, _ = attention(*project(alone))
    worst = max(worst, (out[n, :lengths[n]] - o[0]).abs().max().item())
print(worst < 1e-6)
print(torch.allclose(out, F.scaled_dot_product_attention(q, k, v, attn_mask=~pad), atol=1e-6))
causal = torch.triu(torch.ones(MAX_LEN, MAX_LEN, dtype=torch.bool), diagonal=1)
both = pad | causal
out_c, weights_c = attention(q, k, v, both)
print(both.shape, torch.allclose(out_c[:, 0], v[:, 0], atol=1e-6))
print(torch.allclose(out_c, F.scaled_dot_product_attention(q, k, v, attn_mask=~both), atol=1e-6))
print(weights.numel() * weights.element_size(), 64 * 512 * 512 * 4)

torch.Size([64, 48]) 4 48
0.0
True
True
torch.Size([64, 48, 48]) True
True
589824 67108864

Theory Assessment

Q1: The Diagonal That Wins

With the raw embeddings in every role, every row of "Where is my card?" put its largest weight on its own token, is keeping 1.0000. With three projections, only card's row did. Say why the raw version favors the diagonal, why is above all, and what the projections change.

Hint: Write a token's score against itself as a dot product, and compare what it measures with a token's score against another.

Q2: A Spread That Grows

Unscaled, the scores' standard deviation was 4.01 at 16 dimensions and 22.62 at 512, and the average row's largest weight rose from 0.5976 to 0.9321. Say where the growth comes from, what dividing by sqrt(d) restores, and what it leaves unchanged in every row.

Hint: Count the products a dot product of length d adds, and ask what a positive divisor does to the order of a row.

Q3: Zero Is a Score

A teammate masks pads with scores.masked_fill(mask, 0.0), reasoning that a pad's embedding is zero anyway. Say what weight the pads keep and why, what value the fill must take and why, and what shape the mask must have for [batch, seq, seq] scores.

Hint: Ask what a softmax gives a score of 0 next to other scores near 0, and how broadcasting lines up [batch, seq] against three dimensions.

Q4: Causal on the Right, Causal on the Left

A teammate runs attention with the causal mask alone on Banking77, padded on the right as encode pads, and argues no padding mask is needed. Another pads on the left, pads first, with the same argument. Say whether each real token's output depends on the padding in each case, and what the pooling still needs in both.

Hint: Under a causal mask a token sees only the positions before it. Ask where the pads sit relative to the real tokens in each layout.

import torch
from torch import nn
from text import vocab, encode
from attention import attention

torch.manual_seed(42)
embed = nn.Embedding(len(vocab), 16, padding_idx=0)
to_q, to_k, to_v = nn.Linear(16, 16), nn.Linear(16, 16), nn.Linear(16, 16)

def project(ids):
    with torch.no_grad():
        x = embed(ids)
        return to_q(x), to_k(x), to_v(x)

causal = torch.triu(torch.ones(48, 48, dtype=torch.bool), diagonal=1)
right = torch.tensor([encode("Where is my card?")])
left = torch.cat([right[:, 5:], right[:, :5]], dim=1)
alone, _ = attention(*project(right[:, :5]), causal[:5, :5])
on_right, _ = attention(*project(right), causal)
on_left, _ = attention(*project(left), causal)
fixed, _ = attention(*project(left), causal | (left == 0).unsqueeze(1))
print(torch.allclose(on_right[0, :5], alone[0], atol=1e-6), round((on_left[0, 43:] - alone[0]).abs().max().item(), 4))
print(torch.allclose(fixed[0, 43:], alone[0], atol=1e-6))

True 0.9619
True

Assignment

Assignment: What Attention Buys Without Positions

Write order.py, which trains three classifiers on Banking77 in one script and asks whether one attention layer before the pooling scores better than Lesson 5.3.1's bag, and whether any of the three can tell a query from the same query backwards.

The three share one class: a 64-dimension embedding with padding_idx=0, the mean of the real tokens' vectors, and an nn.Linear(64, len(intents)) head, with an optional attention layer of three nn.Linear(64, 64) projections before the mean, called through attention with the padding mask and optionally the causal mask joined to it. Train each exactly as Lesson 5.3.1 trained its bag, order and the weights reseeded with 42 before each model, 10 epochs of train_epoch with Adam at 0.01, and print the training loss and validation accuracy after epochs 1, 5 and 10.

One call this lesson did not use. input.gather(dim, index) builds a tensor of index's shape whose entry at each position is input's entry at the position index names along dim, so ids.gather(1, index) reorders every row of ids by its own row of index. Use it to reverse each validation query's real tokens with its pads left in place, and print the first validation query's first 12 ids both ways. For each trained model, print the share of the 2,000 validation queries whose predicted intent is the same forwards and backwards, and its accuracy on the reversed queries. Say in a comment what the three shares show, and what the attention layer bought on its own.

Hints:

• For a row of real length L, position p below L reads position L - 1 - p, and every other position reads itself. torch.where, Lesson 4.1.2's, builds that index for every row at once from torch.arange and each row's length.
• accuracy hands its model back in training mode, so call model.eval() before the predictions you read by hand.

import torch
from torch import nn
from text import vocab, intents, order, MAX_LEN, train_loader, val_loader, val_set
from image_loop import accuracy, train_epoch
from attention import attention

causal = torch.triu(torch.ones(MAX_LEN, MAX_LEN, dtype=torch.bool), diagonal=1)

class Classifier(nn.Module):
    def __init__(self, attend, look_back):
        super().__init__()
        self.embed = nn.Embedding(len(vocab), 64, padding_idx=0)
        self.attend, self.look_back = attend, look_back
        if attend:
            self.to_q, self.to_k, self.to_v = nn.Linear(64, 64), nn.Linear(64, 64), nn.Linear(64, 64)
        self.head = nn.Linear(64, len(intents))

    def forward(self, ids):
        x = self.embed(ids)
        pad = ids == 0
        if self.attend:
            mask = pad.unsqueeze(1)
            if self.look_back:
                mask = mask | causal
            x, _ = attention(self.to_q(x), self.to_k(x), self.to_v(x), mask)
        real = (~pad).unsqueeze(-1).float()
        return self.head((x * real).sum(dim=1) / real.sum(dim=1))

def reverse_real(ids):
    lengths = (ids != 0).sum(dim=1, keepdim=True)
    pos = torch.arange(ids.shape[1])
    index = torch.where(pos < lengths, lengths - 1 - pos, pos)
    return ids.gather(1, index)

X, Y = val_set.tensors
backward = reverse_real(X)
print(X[0, :12])
print(backward[0, :12])

tensor([ 11,  12, 141,  17,  29,   6, 107,   4,  75, 213,   3,   0])
tensor([  3, 213,  75,   4, 107,   6,  29,  17, 141,  12,  11,   0])

for name, attend, look_back in (("bag", False, False), ("attention", True, False), ("causal", True, True)):
    order.manual_seed(42)
    torch.manual_seed(42)
    model = Classifier(attend, look_back)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    for epoch in range(1, 11):
        train_loss = train_epoch(model, optimizer, train_loader)
        if epoch in (1, 5, 10):
            print(name, epoch, f"{train_loss:.4f}", f"{accuracy(model, val_loader):.4f}")
    model.eval()
    with torch.no_grad():
        forward_calls, backward_calls = model(X).argmax(dim=1), model(backward).argmax(dim=1)
    print(name, f"{(forward_calls == backward_calls).float().mean().item():.4f}", f"{(backward_calls == Y).float().mean().item():.4f}")
# The bag and the attention model call every one of the 2,000 queries the
# same backwards: neither reads order. On its own the attention layer
# bought nothing: 0.0880 under the bag, 176 queries. The causal model reads
# order through its mask and keeps 0.8660 of its calls backwards.

bag 1 2.3178 0.7565
bag 5 0.1572 0.8665
bag 10 0.0353 0.8760
bag 1.0000 0.8760
attention 1 1.9851 0.6525
attention 5 0.5816 0.7640
attention 10 0.5261 0.7880
attention 1.0000 0.7880
causal 1 2.1870 0.6425
causal 5 0.3145 0.8115
causal 10 0.1990 0.8170
causal 0.8660 0.7955

Additional Resources

Optional, and nothing later depends on either. Worth opening for the same mechanism built from loops up, or for where the scale first came from.

• YouTube - "Let's build GPT: from scratch, in code, spelled out." by Andrej Karpathy, from 42:13 to 1:19:11. Averaging the past with loops, then as a matrix product with a softmax and a lower-triangular mask, self-attention with queries, keys and values, attention as an operation over sets, and why the scores are divided by the square root of the head size.
• Paper - Vaswani et al., "Attention Is All You Need", section 3.2.1 "Scaled Dot-Product Attention" and its footnote 4. The formula as the transformer's authors wrote it, the comparison with additive attention, and the variance argument: unit-variance entries give a dot product of variance d_k.

Summary

Attention turned out to be smaller than its reputation: an average, with weights each token works out from dot products and a softmax, in five lines of attention.py. Everything hard lives around those lines: which vectors play the three roles, how large the scores may grow, and which keys a query may see.

Most of what went wrong here went wrong without a sound. The raw embeddings made every token its own favorite, the unscaled scores made an average into a copy, the pads took nine tenths of every real token's attention, a [batch, seq] mask was refused at one batch size and lied at another, and PyTorch's own function accepted the mask the wrong way round. Each was found the same way, by changing something that should not matter, the padding, the last token, the order, and measuring whether the output moved.

And the order question did not go away. A trained attention layer still gave every query read backwards the intent it gave the query read forwards, because nothing in it knows where a token stands, and its seq x seq weights cost the square of the length. Positions and the rest of the block around the layer are Lesson 5.3.3's work.

Cheat Sheet

Part A - Reference Table

Part B - Compact Code Block

import math

def attention(q, k, v, mask=None):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

pad = (ids == 0).unsqueeze(1)  # [batch, 1, seq]
causal = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
out, weights = attention(q, k, v, pad | causal)
same = F.scaled_dot_product_attention(q, k, v, attn_mask=~(pad | causal))