Self-Attention by Hand: Queries, Keys, and Values

Where lesson one left off

Lesson one ended with a single tensor. A batch of character sequences was tokenized to integer ids of shape (B, T), a token embedding table turned each id into a vector, a position embedding table added order, and the sum was x = tok_emb + pos_emb of shape (B, T, C) — our running character-level GPT on tiny-Shakespeare, with B the batch, T the time or sequence axis, and C the channel width we call d_model, fixed at 384 for this course. Every value in that tensor describes one token in isolation. Position 5 does not yet know that position 2 was a comma or that position 0 was a capital letter. The whole purpose of attention is to fix exactly that: to let each position gather information from the other positions before it.¹

The mechanism we build is the one Vaswani and colleagues named scaled dot-product attention in "Attention Is All You Need," and the concrete code follows the causal self-attention in Andrej Karpathy's nanoGPT, the same reference implementation lesson one drew its running example from.¹² We stay with one head for the whole lesson; splitting a head into many is the next lesson's job.

Start with a plain average

Before attention, take the crudest thing that lets a token see its past: replace each token's vector with the average of itself and every token before it. It throws away word order within the window and weights every past token equally, so it is a poor model — but it is a real operation on the lesson-one tensor, and it is the seed the rest of the lesson grows from. This staged reveal, from a loop to a matrix multiply to learned weights, follows Karpathy's "Let's build GPT" walkthrough.³

import torch

B, T, C = 4, 8, 384                  # batch, time (sequence), channels (d_model)
x = torch.randn(B, T, C)             # stands in for tok_emb + pos_emb from lesson 1

# a plain causal average, written as the loop it literally is
xbow = torch.zeros(B, T, C)          # "bag of words": (B, T, C)
for b in range(B):
    for t in range(T):
        past = x[b, : t + 1]         # (t+1, C): this token and everything before it
        xbow[b, t] = past.mean(dim=0)  # (C,): mean over the time axis

The loop hides a piece of linear algebra. Averaging position t over positions 0..t is a matrix multiply by a lower-triangular matrix whose rows sum to one. Build a (T, T) matrix of ones on and below the diagonal, normalize each row by its own sum, and multiplying it into x produces the identical result — the same average, now as one batched operation instead of a double loop.³

w = torch.tril(torch.ones(T, T))       # (T, T): 1 on/below the diagonal, 0 above
w = w / w.sum(dim=1, keepdim=True)     # normalize each row so it sums to 1
xbow2 = w @ x                          # (T, T) @ (B, T, C) -> (B, T, C)

torch.allclose(xbow, xbow2)            # True: same average, one matmul

Now the whole idea of attention is visible in one sentence: those weights do not have to be uniform. If row t could hold any non-negative numbers that sum to one, the output at position t would be a weighted average of the past, and if those numbers were computed from the tokens themselves, the average would be content-dependent. Position t could lean hard on the one past token that matters and ignore the rest. That is self-attention. The rest of the lesson is the machinery that turns the fixed 1/(t+1) weights of Fig. 2 into learned, per-pair weights.

Queries, keys, and values

To make the weights content-dependent we give every token three learned roles, each a linear projection of its (C,) vector. The query is what a token is looking for. The key is what a token offers to others as a label for itself. The value is the information a token actually contributes if it is attended to. A query meets a key to decide how much; the value is the what that then gets averaged. These are not physical quantities in the data; they are useful roles, and each is produced by its own weight matrix W_q, W_k, W_v.²

In PyTorch each projection is an nn.Linear layer, which applies x @ W (we drop the bias, as attention projections conventionally do); its two arguments are the input width and the output width.⁴ The output width is the head size, written d_k in the paper — the dimension the query and key live in, and the length over which their dot product is taken. It is a hyperparameter, smaller than C; we use 16 here so the tensors stay small enough to reason about.

import torch.nn as nn

head_size = 16                         # d_k: the width of one attention head
key   = nn.Linear(C, head_size, bias=False)   # W_k: (C, head_size)
query = nn.Linear(C, head_size, bias=False)   # W_q: (C, head_size)
value = nn.Linear(C, head_size, bias=False)   # W_v: (C, head_size)

k = key(x)     # (B, T, head_size): what each token offers
q = query(x)   # (B, T, head_size): what each token wants
v = value(x)   # (B, T, head_size): what each token would contribute

The scaled dot-product

Compatibility between a query and a key is their dot product: it is large and positive when the two vectors point the same way, so it reads as "this key answers this query." Every query dotted against every key is a single matmul, q @ k transposed on its last two axes, giving a (T, T) grid of raw scores per batch element. Entry (i, j) is how much position i's query matches position j's key. These raw scores are the attention logits.²

import math

scores = q @ k.transpose(-2, -1)       # (B, T, hs) @ (B, hs, T) -> (B, T, T)
scores = scores / math.sqrt(head_size) # divide by sqrt(d_k): the 1/sqrt(d_k) scaling

A grid of logits is not yet a set of averaging weights; the weights must be non-negative and sum to one along each row. The softmax function delivers exactly that. For a row of logits z, softmax returns exp(z_j) / Σ exp(z): every output is positive because exp is positive, and the outputs sum to one because they share the denominator. PyTorch's softmax takes the dim to normalize over, and here that is the last axis, the keys.⁵ Applying it to each row of the score grid turns compatibilities into a proper weighted average over the values.

Why divide by √d_k first? Suppose the components of q and k are independent with mean 0 and variance 1. Their dot product Σ q_i k_i is a sum of d_k such terms, so it has mean 0 and variance d_k: its typical magnitude grows like √d_k. The paper states the consequence directly — "for large values of d_k, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients" — and so it scales the logits by 1/√d_k to hold their variance near one.¹ The picture is concrete: feed softmax logits like [8, 0] and it returns roughly [0.9997, 0.0003], a near one-hot vector whose gradient with respect to the logits is almost zero, and learning stalls. Dividing by √d_k keeps the logits in the range where softmax stays soft and its gradients stay usable.

Masking the future

The score grid as built is dense: position i has a weight for every position j, including positions after it. For a model trained to predict the next token, that is cheating: position i would see its own future, and the task would be trivial at train time and impossible at generation time, when the future does not exist yet. A decoder must be causal: position i may attend only to positions 0..i. This is the same lower-triangular structure as the averaging matrix in Fig. 2, now imposed on learned scores.²

The trick is to set the forbidden scores to negative infinity before the softmax. Since exp(-inf) = 0, those positions receive exactly zero weight, and the surviving weights in each row still sum to one on their own. torch.tril gives the lower-triangular mask of allowed positions, and masked_fill writes float("-inf") everywhere the mask is zero.⁶⁷

import torch.nn.functional as F

tril = torch.tril(torch.ones(T, T))                    # (T, T): 1 where j <= i
scores = scores.masked_fill(tril == 0, float("-inf"))  # (B, T, T): future -> -inf
weights = F.softmax(scores, dim=-1)                    # (B, T, T): each row sums to 1
out = weights @ v                                       # (B, T, T) @ (B, T, hs) -> (B, T, hs)

Read the first row of the weight matrix after masking and it is [1, 0, 0, ...]: the first token has only itself to attend to, so its output is its own value unchanged. The last row is a full learned distribution over the whole sequence. Between them, each row t is a content-dependent weighted average over exactly the tokens at or before t — the Fig. 2 average, but with the weights now chosen by the data.

A single head, end to end

Assemble the pieces into one module and the whole head is a dozen lines. The forward pass reads as the five moves we derived: project to q, k, v; score with the scaled dot product; mask the future; softmax to weights; take the weighted average of the values. This mirrors nanoGPT's CausalSelfAttention, reduced to a single head and written for readability.²

class Head(nn.Module):
    """One causal self-attention head."""

    def __init__(self, C, head_size, block_size):
        super().__init__()
        self.key   = nn.Linear(C, head_size, bias=False)
        self.query = nn.Linear(C, head_size, bias=False)
        self.value = nn.Linear(C, head_size, bias=False)
        # a (block_size, block_size) lower-triangular buffer, not a parameter
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):                       # x: (B, T, C)
        B, T, C = x.shape
        k = self.key(x)                          # (B, T, head_size)
        q = self.query(x)                        # (B, T, head_size)
        v = self.value(x)                        # (B, T, head_size)

        scores = q @ k.transpose(-2, -1)         # (B, T, T)
        scores = scores / k.shape[-1] ** 0.5     # scale by 1/sqrt(d_k)
        scores = scores.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)      # (B, T, T), rows sum to 1
        return weights @ v                       # (B, T, head_size)

Every tensor's shape is annotated at the point it is made, and the shapes are the proof the code is right: (B, T, C) enters, three (B, T, head_size) projections come out, the score grid is (B, T, T), and the final weighted average returns to (B, T, head_size). In production this exact computation is one fused call, F.scaled_dot_product_attention(q, k, v, is_causal=True), which applies the same scaling and causal mask and runs a memory-efficient kernel; writing it by hand first is how you earn the right to trust the one-liner.⁸

Go deeper

Read Section 3.2.1 of "Attention Is All You Need": it is the half-page that defines softmax(QK^T/√d_k)V and gives the variance argument for the scaling you just derived.¹ Then open nanoGPT's model.py and read CausalSelfAttention; every line of Fig. 6 is there, and the parts this lesson deferred — reshaping into multiple heads and the output projection — are the few lines you cannot yet fully account for.² To feel the mechanism in your hands, print the weights matrix from Fig. 5 for a real batch and confirm each row is lower-triangular and sums to one, then swap the hand-rolled block for F.scaled_dot_product_attention and check the outputs agree.⁸ For lineage, the idea of a learned, content-dependent weighting over a sequence predates the transformer: Bahdanau, Cho, and Bengio introduced an attention mechanism for machine translation in 2014, which "Attention Is All You Need" distilled into the pure dot-product form used here.⁹