Self-attention from scratch in PyTorch

Transformer on-ramp · Part 5 of 8

Reading about attention is one thing. Typing it is another. Let us build it from scratch in PyTorch, in about ten lines, and then check it against the library.

From the formula to a function

Self-attention looks abstract until you write it. The whole mechanism fits in a few lines of PyTorch. Here is a function that takes queries, keys and values and returns the attended output together with the weights, so we can inspect them:

import torch

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Q, K, V: (batch, heads, seq, d_k)
    scores = Q @ K.transpose(-2, -1) / (K.shape[-1] ** 0.5)   # (batch, heads, seq, seq)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))  # e.g. a causal mask
    weights = scores.softmax(dim=-1)                           # each row sums to 1
    return weights @ V, weights

What each line does

1. Scores. Q @ K.transpose(-2, -1) compares every query with every key at once. The result has one score for each pair of tokens, so it is a seq × seq table.
2. Scale. Dividing by the square root of d_k is the fix from the last post.
3. Mask (optional). Setting a score to −∞ makes its softmax weight exactly 0. We use it to stop a token from reading the future.
4. Softmax and mix. Softmax turns each row of scores into weights that add to 1; multiplying by V blends the values.

The causal mask

A model that writes text one word at a time must not peek at words that have not been written yet. The mask is a lower-triangular table: token i may read tokens up to and including i, and never later ones. After softmax the masked positions have weight exactly 0. All the language models you meet later are built this way.

Check that it behaves

B, H, T, dk = 2, 4, 5, 8
Q, K, V = (torch.randn(B, H, T, dk) for _ in range(3))
mask = torch.tril(torch.ones(T, T))                 # 1 where reading is allowed
out, w = scaled_dot_product_attention(Q, K, V, mask)

print(out.shape)                                    # torch.Size([2, 4, 5, 8])
print(torch.allclose(w.sum(-1), torch.ones(B, H, T)))   # True: every row sums to 1
print((w.triu(1) == 0).all().item())                # True: nothing leaks from the future
ref = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=True)
print(torch.allclose(out, ref, atol=1e-5))          # True: matches PyTorch's own version

The last line is the reassuring one: our short version gives the same numbers as the optimised function that ships inside PyTorch.

Open in Colab Run it, draw the mask, and compare with PyTorch’s own version. Free in your browser.

Takeaways

• Attention is one matrix multiply for scores, a softmax, and one more multiply.
• Masking with −∞ before softmax gives exact zero weights.
• A causal mask lets each token read only the past: the shape of every LLM.
• Our version matches PyTorch’s own implementation to numerical precision.

Next in the series (Part 6): Multi-head attention: many ways of looking at one sentence