Self-Attention: Queries, Keys, and Values

A transformer block moves information between positions. Self-attention is the mechanism that decides, for each position, how much to look at every other position — including itself. This lesson builds it from scratch in PyTorch: project inputs into queries, keys, and values; compute scaled dot-product scores; normalize with stable softmax; and aggregate values. We verify the implementation against PyTorch’s own scaled_dot_product_attention.

The one-sentence version

Self-attention projects each position’s embedding into a query, key, and value; compares every query with every key via scaled dot products; normalizes with softmax; and returns a weighted sum of all values.

Everything else in this lesson makes that concrete, with shapes tracked at every step.

Lab: Run the companion CPU exercise at /labs/track1/01-attention/ (shapes, masks, PyTorch parity).

1. Q/K/V projections

A decoder-only transformer processes a sequence of embeddings X of shape [B, T, d]. Self-attention projects each position into three vectors:

Tensor	Shape	Meaning
Q = X @ W_Q	[B, T, d_k]	Query: “what should I pay attention to?”
K = X @ W_K	[B, T, d_k]	Key: “what information do I offer?”
V = X @ W_V	[B, T, d_v]	Value: “what information do I contribute?”

W_Q and W_K have shape [d, d_k]; W_V has shape [d, d_v]. In the single-head case used here, d_k = d_v = 4 while d = 8; in multi-head attention (lesson 6), d_k = d_v = d / h.

2. Scaled dot-product attention scores

Compare every query with every key:

scores = Q @ K^T / sqrt(d_k)   # shape [B, T, T]

Position (i, j) of the scores matrix measures how much position i should attend to position j. The 1/sqrt(d_k) factor prevents dot products from growing too large in magnitude — without it, large scores push softmax into regions with vanishing gradients (see the failure modes table below).

3. Stable softmax

Normalize each row of scores into a probability distribution:

attn_weights = softmax(scores, dim=-1)   # shape [B, T, T]

Naive exp(x) / sum(exp(x)) overflows when scores are large. Stable softmax subtracts the row maximum first:

scores_shifted = scores - max(scores, dim=-1, keepdim=True)
attn_weights = exp(scores_shifted) / sum(exp(scores_shifted), dim=-1, keepdim=True)

This is mathematically identical but numerically safe. Every row of attn_weights sums to 1.

Why stable softmax matters

With scores like [1000, 1001, 999, 1002], naive exp(x) overflows to inf. Stable softmax subtracts the row max first, producing finite values. The shift cancels in the division, so the result is mathematically identical but safe to compute.

4. Weighted value aggregation

The output at each position is a weighted sum of all value vectors:

output = attn_weights @ V   # shape [B, T, d_v]

Position i‘s output is a convex combination of all positions’ values, where the weights come from how much position i attends to each position j. Positions that attend more to themselves will have output close to their own value; positions that attend to earlier context will reflect that context.

5. Fully worked numeric example (verified, seed 42)

Setup: B = 1, T = 4, d = 8, d_k = 4, d_v = 4, torch.manual_seed(42), torch.set_default_dtype(torch.float32).

import torch
torch.manual_seed(42)
torch.set_default_dtype(torch.float32)

X = torch.randn(1, 4, 8)
W_Q = torch.randn(8, 4)
W_K = torch.randn(8, 4)
W_V = torch.randn(8, 4)

Q = X @ W_Q   # [1, 4, 4]
K = X @ W_K   # [1, 4, 4]
V = X @ W_V   # [1, 4, 4]

scores = Q @ K.transpose(-2, -1) / (4 ** 0.5)  # [1, 4, 4]
scores_max = scores.max(dim=-1, keepdim=True)[0]
scores_shifted = scores - scores_max
exp_scores = torch.exp(scores_shifted)
attn_weights = exp_scores / exp_scores.sum(dim=-1, keepdim=True)  # [1, 4, 4]
output = attn_weights @ V  # [1, 4, 4]

Step 1 — Projections:

Q = [[[ -0.7142, -10.7618,   4.7928,   1.6591],
      [  5.9094,  -2.2237,   2.2946,  -1.0836],
      [ -3.2318,   3.9874,  -3.4993,   4.7026],
      [ -3.5491,   3.5556,  -3.1290,   1.0415]]]

K = [[[  4.0115,   3.6513,   2.9872,   3.6108],
      [ -3.4479,  -0.5821,   1.8296,   2.9016],
      [  0.9276,  -0.0589,   6.4627,   1.9089],
      [  1.3121,  -0.1874,   0.3295,   1.9857]]]

V = [[[ -8.2489,  -0.9326,   2.3788,  -0.8355],
      [  4.1438,  -1.3904,   2.1127,  -1.3241],
      [ -0.2973,   1.8336,   2.2672,   1.4077],
      [ -0.8263,   3.7984,   3.3404,  -0.7702]]]

Step 2 — Scores Q @ K^T / √4:

scores = [[[ -10.9254,  11.1552,  17.0567,   2.9766],
           [  9.2642,  -9.0133,   9.1868,   3.3872],
           [  4.0607,   8.0325,  -8.4355,   1.5988],
           [ -3.4206,   3.7321, -10.8679,  -2.1428]]]

Step 3 — Stable softmax (row-wise):

attn_weights = [[[ 7.02e-13, 2.73e-03, 9.97e-01, 7.65e-07],
                 [ 5.19e-01, 5.98e-09, 4.80e-01, 1.45e-03],
                 [ 1.85e-02, 9.80e-01, 6.91e-08, 1.57e-03],
                 [ 7.80e-04, 9.96e-01, 4.55e-07, 2.80e-03]]]

Each row sums to 1.0 (within floating-point precision).

Step 4 — Output attn_weights @ V:

output = [[[ -0.2852,  1.8248,  2.2668,  1.4002],
           [ -4.4217,  0.4019,  2.3266,  0.2412],
           [  3.9072, -1.3738,  2.1196, -1.3142],
           [  4.1202, -1.3755,  2.1164, -1.3222]]]

Verification: Matches torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False, scale=0.5) with torch.nn.attention.sdpa_kernel(backends=[torch.nn.attention.SDPBackend.MATH]) to max_diff < 1e-6 (PyTorch 2.5, CPU, float32).

Self-attention mechanism

Q/K/V projections, scaled dot-product scores, stable softmax, and weighted value aggregation, with tensor shapes at each step.

Read this diagram as text

Text equivalent. Panel A: input X [B, T, d] projects to Q, K, V via separate weight matrices. Panel B: scores = Q @ K^T / sqrt(d_k) produces a [B, T, T] attention score matrix. Panel C: stable softmax converts scores to attention weights where each row sums to 1. Panel D: attention weights @ V produces the output [B, T, d_v] as a weighted sum of value vectors.

6. PyTorch implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k, d_v, bias=False):
        super().__init__()
        self.W_Q = nn.Parameter(torch.randn(d_model, d_k))
        self.W_K = nn.Parameter(torch.randn(d_model, d_k))
        self.W_V = nn.Parameter(torch.randn(d_model, d_v))
        self.scale = 1.0 / (d_k ** 0.5)
        # Note: some implementations (e.g., GPT-2) use bias=False for Q/K/V projections.
        # PyTorch's nn.Linear defaults to bias=True; set bias=False to match.

    def forward(self, X):
        Q = X @ self.W_Q   # [B, T, d_k]
        K = X @ self.W_K   # [B, T, d_k]
        V = X @ self.W_V   # [B, T, d_v]

        scores = Q @ K.transpose(-2, -1) * self.scale   # [B, T, T]

        # Stable softmax
        scores_max = scores.max(dim=-1, keepdim=True)[0]
        scores_shifted = scores - scores_max
        exp_scores = torch.exp(scores_shifted)
        attn_weights = exp_scores / exp_scores.sum(dim=-1, keepdim=True)

        output = attn_weights @ V   # [B, T, d_v]
        return output, attn_weights

Numerical check: Compare against F.scaled_dot_product_attention(Q, K, V, is_causal=False, scale=scale) using torch.nn.attention.sdpa_kernel(backends=[torch.nn.attention.SDPBackend.MATH]) for deterministic CPU execution. Max difference should be below 1e-6 (verified with seed 42, float32).

7. Failure modes

Failure mode	Symptom	Root cause	Test
Missing scale	One attention weight dominates (~1.0); training unstable	Q @ K^T without 1/sqrt(d_k) yields large dot products, pushing softmax to extremes	Check scores.abs().max() is reasonable (not >> 10)
Unstable softmax	NaN gradients; loss blows up	exp(scores) overflows for large scores	Verify with large-magnitude inputs; check no NaN in attn_weights
Wrong transpose	Shape mismatch or silently wrong attention	Q @ K instead of Q @ K^T, or wrong transpose dims	Assert scores.shape == [B, T, T]
Value/query mismatch	Output doesn’t reflect input content	Using Q instead of V in final matmul	Assert output.shape == [B, T, d_v]

8. What this lesson does not cover

• Causal masking (preventing position i from attending to positions > i): lesson 5, causal-attention-masks.
• Multi-head attention (splitting into h heads and merging): lesson 6, multi-head-attention.
• Dropout, training loop, optimizer: covered in later training-and-debugging lessons.
• Grouped/multi-query attention: covered in advanced modules.

9. Exercises

Exercise	Prompt
1	Given X of shape [1, 3, 6], W_Q of shape [6, 2], W_K of shape [6, 2], W_V of shape [6, 2], what are the shapes of Q, K, V, scores, attn_weights, and output?
2	If d_k = 64 and the unscaled scores have max magnitude 500, what happens to naive softmax? What does stable softmax produce instead?
3	Modify the reference implementation to return attention weights for only the first query position (row 0 of attn_weights). What does this row represent?

10. Related lessons

Related lessons

• LLM Internals & Architecture · Transformer architecture and model inputsTransformer Architecture Explained: Trace a Decoder-Only LLM14 min
• LLM Internals & Architecture · Transformer architecture and model inputsTokenisation and Embeddings: From Text to Tensors12 min
• LLM Internals & Architecture · Attention mechanisms from scratchCausal and Padding Masks: Implementation and Tests15 min

11. Sources

Source	Relevance
Vaswani et al., “Attention Is All You Need” (arXiv:1706.03762), §3.2.1	Original scaled dot-product attention: Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V
PyTorch F.scaled_dot_product_attention docs (v2.5+)	Reference API: query/key/value shapes, scale parameter, is_causal flag, SDP backends
PyTorch source torch/nn/functional.py (search scaled_dot_product_attention)	Reference implementation: attn_weight = query @ key.transpose(-2, -1) * scale_factor, then softmax, then @ value

---