Transformers and Attention Questions

Approach

Scaled dot-product attention takes queries q, keys k and values v (three tensors, a tensor being a multi-dimensional array of numbers) and computes

where d is the last dimension of q and k, the softmax turns each row of scores into positive weights summing to 1, and the mask removes positions that must not be attended to (padding, or future tokens in a decoder). The steps are: score, scale by 1/sqrt(d), block masked positions by setting their scores to minus infinity, softmax over the key axis, multiply by v.

Shapes

A causal mask has shape (1, Lq, Lk) and a padding mask (B, 1, Lk); combining them with a logical AND broadcasts to (B, Lq, Lk). Broadcasting means that when two arrays have different shapes, any axis of size 1 is stretched (conceptually copied) to match the other array's size on that axis, so the result has the larger size on every axis. A tiny case, printed by the program: with Lq = Lk = 3, the causal mask is [[1,0,0],[1,1,0],[1,1,1]] (query i sees keys up to i) and the padding mask is [[1,1,0]] (key 2 is padding). The AND stretches the padding row down all 3 query rows and gives [[1,0,0],[1,1,0],[1,1,0]]: the last query row loses key 2 even though causality allowed it. For multi-head use, a head axis goes in front of Lq and the same code applies to the last two axes.

Implementation, with checks

The program below defines the PyTorch version and a NumPy version, and checks both against torch.nn.functional.scaled_dot_product_attention (PyTorch's built-in reference for this exact operation; the output labels abbreviate it SDPA). It also covers boolean versus additive masks, the fully masked row, numerical stability, a tiny broadcasting case, one masked row of scores whose weights you can check by hand, the stricter fix for fully masked rows, and the mask polarity of nn.MultiheadAttention against scaled_dot_product_attention. By hand, the row of scores 2, 1, 0, -1 with only the first two allowed has weights exp(2) / (exp(2) + exp(1)) = 0.731 and 0.269, and zero on the blocked two: the blocked entries simply drop out, as if the row had two entries.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.set_num_threads(1)

# ---------- PyTorch: batched scaled dot-product attention with an optional mask ----------
def attention(q, k, v, mask=None):
    """q: (B, Lq, d), k: (B, Lk, d), v: (B, Lk, dv).
    mask: bool, broadcastable to (B, Lq, Lk); True = may attend, False = blocked.
    Returns output (B, Lq, dv) and weights (B, Lq, Lk)."""
    scores = q @ k.transpose(-2, -1) / q.shape[-1] ** 0.5          # (B, Lq, Lk)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)                         # stable: PyTorch subtracts the row max
    if mask is not None:
        # A row with no allowed key gives softmax(-inf, ..., -inf) = NaN; define it as all zeros.
        weights = torch.nan_to_num(weights, nan=0.0)
    return weights @ v, weights

# ---------- NumPy version of the same function ----------
def attention_np(q, k, v, mask=None):
    scores = q @ np.swapaxes(k, -2, -1) / np.sqrt(q.shape[-1])
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)
    row_max = np.max(scores, axis=-1, keepdims=True)
    row_max = np.where(np.isfinite(row_max), row_max, 0.0)          # fully masked row: avoid -inf - -inf
    e = np.exp(scores - row_max)
    denom = e.sum(-1, keepdims=True)
    weights = np.divide(e, denom, out=np.zeros_like(e), where=denom > 0)
    return weights @ v, weights

def r3(t):
    return [round(float(a), 3) for a in t]

torch.manual_seed(0)
B, Lq, Lk, d, dv = 2, 4, 5, 8, 6
q, k, v = torch.randn(B, Lq, d), torch.randn(B, Lk, d), torch.randn(B, Lk, dv)
causal = torch.tril(torch.ones(Lq, Lk, dtype=torch.bool))[None]     # (1, Lq, Lk)
pad = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]], dtype=torch.bool)[:, None, :]   # (B, 1, Lk)
mask = causal & pad                                                  # broadcasts to (B, Lq, Lk)

out, w = attention(q, k, v, mask)
print("shapes: out", tuple(out.shape), "weights", tuple(w.shape))
ref = F.scaled_dot_product_attention(q, k, v, attn_mask=mask.expand(B, Lq, Lk))
print("matches F.scaled_dot_product_attention:", torch.allclose(out, ref, atol=1e-6))
print("row sums, sequence 2:", r3(w.sum(-1)[1]))
print("blocked weights are exactly zero:", bool((w[~mask.expand(B, Lq, Lk)] == 0).all()))

out_np, w_np = attention_np(q.numpy(), k.numpy(), v.numpy(), mask.numpy())
print("NumPy matches PyTorch:", np.allclose(out_np, out.numpy(), atol=1e-6))

# Boolean mask vs additive mask: same result when the additive value is -inf.
add = torch.zeros(B, Lq, Lk).masked_fill(~mask.expand(B, Lq, Lk), float("-inf"))
ref_add = F.scaled_dot_product_attention(q, k, v, attn_mask=add)
print("additive -inf mask equals boolean mask:", torch.allclose(ref_add, ref, atol=1e-6))

# Fully masked row: query 0 of the second sequence may not see anything.
m2 = mask.expand(B, Lq, Lk).clone()
m2[1, 0, :] = False
o2, w2 = attention(q, k, v, m2)
print("my fully masked row weights:", w2[1, 0].tolist(), "output sum of abs:", o2[1, 0].abs().sum().item())
add2 = torch.zeros(B, Lq, Lk).masked_fill(~m2, float("-inf"))
print("torch", torch.__version__, "on CPU, SDPA fully masked row, bool mask:", r3(F.scaled_dot_product_attention(q, k, v, attn_mask=m2)[1, 0, :3]))
print("torch", torch.__version__, "on CPU, SDPA fully masked row, additive -inf mask:", r3(F.scaled_dot_product_attention(q, k, v, attn_mask=add2)[1, 0, :3]))
big = torch.zeros(B, Lq, Lk).masked_fill(~m2, -1e9)
print("finite -1e9 instead of -inf, same row:", r3(torch.softmax(q @ k.transpose(-2, -1) / d ** 0.5 + big, -1)[1, 0]))

# Stability: naive softmax overflows on large scores, shifted softmax does not.
x = np.array([1000.0, 1001.0, 1002.0], dtype=np.float32)
with np.errstate(all="ignore"):
    naive = np.exp(x) / np.exp(x).sum()
stable = np.exp(x - x.max()) / np.exp(x - x.max()).sum()
print("naive:", r3(naive), "shifted:", r3(stable))

# ---------- Broadcasting with tiny shapes you can read ----------
c3 = torch.tril(torch.ones(3, 3, dtype=torch.bool))[None]           # (1, 3, 3) causal: query i sees keys <= i
p3 = torch.tensor([[True, True, False]])[:, None, :]                # (1, 1, 3) padding: key 2 is padding
print("causal (1,3,3):", c3[0].int().tolist(), "| padding (1,1,3):", p3[0].int().tolist())
print("causal & padding -> (1,3,3):", (c3 & p3)[0].int().tolist())

# ---------- One masked row, traced by hand ----------
row = torch.tensor([2.0, 1.0, 0.0, -1.0])                           # scores for one query over 4 keys
allowed = torch.tensor([True, True, False, False])
print("row weights with mask:", r3(torch.softmax(row.masked_fill(~allowed, float("-inf")), -1)),
      "| softmax over only the 2 allowed scores:", r3(torch.softmax(row[:2], -1)))
print("all-equal scores, softmax:", r3(torch.softmax(torch.zeros(4), -1)))

# ---------- Stricter fix for a fully masked row: zero only rows that have no allowed key ----------
def attention_strict(q, k, v, mask):
    scores = (q @ k.transpose(-2, -1) / q.shape[-1] ** 0.5).masked_fill(~mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    weights = torch.where(mask.any(-1, keepdim=True), weights, torch.zeros_like(weights))  # a NaN elsewhere stays visible
    return weights @ v, weights
o3, w3 = attention_strict(q, k, v, m2)
print("strict version: fully masked row", w3[1, 0].tolist(), "| matches nan_to_num version:", torch.allclose(o3, o2))

# ---------- Mask polarity: F.scaled_dot_product_attention vs nn.MultiheadAttention ----------
allowed2 = torch.tril(torch.ones(3, 3, dtype=torch.bool))           # True = may attend
x3 = torch.randn(1, 3, 8)
sdpa = F.scaled_dot_product_attention(x3, x3, x3, attn_mask=allowed2)          # True = take part
mha = nn.MultiheadAttention(8, 1, batch_first=True).eval()
_, wm = mha(x3, x3, x3, attn_mask=~allowed2, need_weights=True)               # True = NOT allowed
print("MultiheadAttention given ~allowed: weights above the diagonal are", wm[0].triu(1).abs().max().item())
_, wm2 = mha(x3, x3, x3, attn_mask=allowed2, need_weights=True)               # same mask without inverting
print("MultiheadAttention given allowed unchanged: query 0 weight on its own (allowed) key is", wm2[0, 0, 0].item())

Output:

shapes: out (2, 4, 6) weights (2, 4, 5)
matches F.scaled_dot_product_attention: True
row sums, sequence 2: [1.0, 1.0, 1.0, 1.0]
blocked weights are exactly zero: True
NumPy matches PyTorch: True
additive -inf mask equals boolean mask: True
my fully masked row weights: [0.0, 0.0, 0.0, 0.0, 0.0] output sum of abs: 0.0
torch 2.5.1 on CPU, SDPA fully masked row, bool mask: [0.0, 0.0, 0.0]
torch 2.5.1 on CPU, SDPA fully masked row, additive -inf mask: [0.0, 0.0, 0.0]
finite -1e9 instead of -inf, same row: [0.2, 0.2, 0.2, 0.2, 0.2]
naive: [nan, nan, nan] shifted: [0.09, 0.245, 0.665]
causal (1,3,3): [[1, 0, 0], [1, 1, 0], [1, 1, 1]] | padding (1,1,3): [[1, 1, 0]]
causal & padding -> (1,3,3): [[1, 0, 0], [1, 1, 0], [1, 1, 0]]
row weights with mask: [0.731, 0.269, 0.0, 0.0] | softmax over only the 2 allowed scores: [0.731, 0.269]
all-equal scores, softmax: [0.25, 0.25, 0.25, 0.25]
strict version: fully masked row [0.0, 0.0, 0.0, 0.0, 0.0] | matches nan_to_num version: True
MultiheadAttention given ~allowed: weights above the diagonal are 0.0
MultiheadAttention given allowed unchanged: query 0 weight on its own (allowed) key is 0.0

What happens to a fully masked row

If every key in a row is blocked, all scores are minus infinity, and a softmax computes exp(-inf - max) where the max is itself -inf, which is inf minus inf: not a number (NaN). That NaN then flows into the output and, in training, into every gradient, so one bad row can ruin a batch. A row can be fully masked, for example, when a query position is padding and the mask is applied to both axes, or when left padding meets a causal mask.

In fused kernels (single GPU programs that combine the score, softmax and value steps), the handling of this case is the kernel's choice, so check it. The code above defines the result explicitly: weights of zero and therefore an output of zero. In the run shown (torch 2.5.1, CPU), PyTorch's built-in function also returns zeros for such a row with either mask form; other versions or hardware kernels may behave differently, so a function that relies on this should do the handling itself. The finite alternative, filling blocked scores with a large negative number such as -1e9, avoids NaN but gives something worse for a fully masked row: the scores are all equal, so the weights become uniform (0.2 each over 5 keys in the run) and the row silently averages tokens it was forbidden to see. For rows with at least one allowed key, the finite fill and -inf agree up to about exp(-1e9), which is zero in float32.

A caution on the nan_to_num shortcut used here: it also hides a NaN that comes from somewhere else, such as an overflow upstream. The stricter version in the program, attention_strict, computes mask.any(-1, keepdim=True) and uses torch.where to zero only the rows that have no allowed key, leaving any other NaN visible (it gives the same result here, as the printed check shows). Multiplying by the any-mask would not work, because NaN times 0 is still NaN.

Boolean versus additive masks

A boolean mask says which positions are allowed (here True means allowed; some libraries use True for blocked, so check the convention). The polarity differs even inside PyTorch: scaled_dot_product_attention takes True as take part in attention, the same as this code, while nn.MultiheadAttention takes True in attn_mask as not allowed. The last two output lines show it side by side: passing the inverted mask to nn.MultiheadAttention zeroes the weights above the diagonal, while passing the allowed mask unchanged blocks the very keys you meant to allow (query 0 then gives its own key weight 0.0). An additive mask holds numbers added to the scores: 0 where allowed and minus infinity (or a very large negative number) where blocked. The run shows that the two forms give the same output when the additive value is -inf. Additive masks can also carry soft biases (such as relative-position biases, numbers added to the score that depend on how far apart two tokens are), which a boolean mask cannot.

Numerical stability of the softmax

exp of a large number overflows. The softmax is unchanged if the same constant is subtracted from every score in a row, so the implementation subtracts the row maximum first. The last two lines of the output show scores of 1000, 1001, 1002 in float32: the naive formula gives NaN, the shifted one gives 0.09, 0.245, 0.665. torch.softmax already does the shift internally; the NumPy version does it by hand and also guards the fully masked row.

Complexity and edge cases

• Time is proportional to B x Lq x Lk x d for the score product, and the same again for the weights-times-v product. Memory is proportional to B x Lq x Lk for the weights, which is what makes long sequences expensive.
• Edge cases: a fully masked row (above); Lq different from Lk (cross-attention, where queries come from one sequence and keys and values from another, or incremental decoding: generating one new token at a time against keys saved from earlier steps, where Lq is 1); masks with the wrong convention (True meaning blocked); dividing by the model width instead of the key width d; applying the mask after the softmax instead of before it (the weights then no longer sum to 1).
• Computing in float16 (a 16-bit number format whose largest finite value is 65,504) needs care: -inf is representable but large finite fills such as -1e9 are not (converting -1e9 to float16 overflows to -inf), so the fill value must suit the dtype, the number format of the tensor.

The idea in plain words

Self-attention lets every token in a sequence build its new representation as a weighted average of information from all the tokens in the same sequence (including itself). The weights are not fixed: they are computed from the content of the tokens. Each token asks a question (its query), every token advertises what it holds (its key), and the token with a good match contributes more of its payload (its value). A token here is one unit of input text, such as a word or word piece, already turned into a vector of numbers (an embedding).

The mechanism, in order

Start with an input matrix X holding one row per token. For one sequence with n tokens and model width d_model, X has shape (n, d_model).

1. Project. Three learned weight matrices turn X into queries, keys and values: Q = X W_Q, K = X W_K, V = X W_V. W_Q and W_K have shape (d_model, d_k) and W_V has shape (d_model, d_v), so Q and K are (n, d_k) and V is (n, d_v). The matrices W_Q, W_K, W_V are learned during training. Row i of Q is the question token i asks, row j of K is how token j can be matched, row j of V is what token j hands over if it is selected.
2. Score. Take the dot product of every query with every key: scores = Q K^T, shape (n, n). Entry (i, j) says how well token j answers token i.
3. Scale. Divide by the square root of d_k (explained below).
4. Softmax. A softmax turns a row of raw numbers into positive weights that sum to 1. Applying it along each row gives the attention weights, shape (n, n).
5. Mix. Multiply the weights by V: output = weights V, shape (n, d_v). Output row i is the weighted average of the value rows, using row i of the weights.

Here Q, K, V are the matrices above, d_k is the width of one key vector, and the softmax is taken over the last axis (over the keys).

A two-token example you can do on paper

Take two tokens with one-hot inputs, queries and keys equal to the inputs, and values that stretch feature 1 by 2 and feature 2 by 4. Then Q = K = X, V = [[2, 0], [0, 4]], and the score of token 0 on itself is 1 / 1.4142 = 0.7071 while its score on token 1 is 0. The softmax of (0.7071, 0) is e^0.7071 / (e^0.7071 + 1) = 2.028 / 3.028, about 0.6698, and 0.3302 for the other token, so output row 0 is about 0.6698 x (2, 0) + 0.3302 x (0, 4) = (1.34, 1.32). The program below checks it and also shows, for a softmax over three scores, how fast the gradient dies as one score pulls ahead.

import numpy as np

# Two tokens, width 2, key width 2, hand-set matrices so every number can be checked by hand.
X = np.array([[1.0, 0.0],
              [0.0, 1.0]])
Wq = Wk = np.eye(2)                   # queries and keys are just the inputs
Wv = np.array([[2.0, 0.0],
               [0.0, 4.0]])           # values stretch the two features by 2 and 4
Q, K, V = X @ Wq, X @ Wk, X @ Wv
scores = Q @ K.T / np.sqrt(2)
e = np.exp(scores)
weights = e / e.sum(axis=1, keepdims=True)
out = weights @ V
np.set_printoptions(precision=4, suppress=True)
print("V\n", V)
print("scores\n", scores)
print("weights\n", weights)
print("output row 0:", out[0])

# Why saturation stalls learning: for softmax, d p_i / d s_i = p_i * (1 - p_i).
for gap in (0, 2, 10, 30):
    s = np.array([float(gap), 0.0, 0.0])
    p = np.exp(s - s.max()) / np.exp(s - s.max()).sum()
    print(f"score gap {gap:2d}: largest weight {p[0]:.6f}  its gradient p*(1-p) = {p[0] * (1 - p[0]):.2e}")

Output:

V
 [[2. 0.]
 [0. 4.]]
scores
 [[0.7071 0.    ]
 [0.     0.7071]]
weights
 [[0.6698 0.3302]
 [0.3302 0.6698]]
output row 0: [1.3395 1.321 ]
score gap  0: largest weight 0.333333  its gradient p*(1-p) = 2.22e-01
score gap  2: largest weight 0.786986  its gradient p*(1-p) = 1.68e-01
score gap 10: largest weight 0.999909  its gradient p*(1-p) = 9.08e-05
score gap 30: largest weight 1.000000  its gradient p*(1-p) = 1.87e-13

The last four lines are the saturation effect in numbers: with no gap the largest weight is 0.33 and nudging its score still changes it noticeably (gradient 0.22); with a gap of 10 the weight is 0.9999 and the same nudge changes it by only about 9e-5.

Worked example: three tokens, run in code

The program below uses three tokens of width 4, a key width of 2, and a fixed seed. The projection matrices are random (drawn by torch.randn and scaled by 0.7), so the printed scores cannot be derived by hand; what you can check by hand is the rule that each weights row sums to 1 and the output is the weights times V, as in the two-token example above. It prints the scores, the weights, the output and a check against PyTorch's reference function. Then it repeats the computation for a batch to show the shapes, shows why the scaling is needed, and counts the memory of the score matrix.

import torch
import torch.nn.functional as F

torch.set_num_threads(1)

torch.manual_seed(0)

# Part 1: three tokens, model width 4, one head, written out step by step.
x = torch.tensor([[1.0, 0.0, 1.0, 0.0],
                  [0.0, 2.0, 0.0, 2.0],
                  [1.0, 1.0, 1.0, 1.0]])           # (n=3, d_model=4)
Wq = torch.randn(4, 2) * 0.7                        # (d_model, d_k=2)
Wk = torch.randn(4, 2) * 0.7
Wv = torch.randn(4, 2) * 0.7
Q, K, V = x @ Wq, x @ Wk, x @ Wv                    # each (3, 2)
scores = Q @ K.T / (2 ** 0.5)                       # (3, 3), divided by sqrt(d_k)
weights = torch.softmax(scores, dim=-1)             # each row sums to 1
out = weights @ V                                   # (3, 2)
print("scores\n", scores.round(decimals=3))
print("weights\n", weights.round(decimals=3))
print("row sums", weights.sum(-1).round(decimals=3))
print("output\n", out.round(decimals=3))
ref = F.scaled_dot_product_attention(Q[None], K[None], V[None])[0]
print("matches PyTorch reference:", torch.allclose(out, ref, atol=1e-6))

# Part 2: shapes for a batch of 2 sequences, length 5, d_model 16, d_k 8.
B, n, d_model, d_k = 2, 5, 16, 8
X = torch.randn(B, n, d_model)
Wq, Wk, Wv = (torch.randn(d_model, d_k) for _ in range(3))
Q, K, V = X @ Wq, X @ Wk, X @ Wv
S = Q @ K.transpose(-2, -1) / d_k ** 0.5
A = torch.softmax(S, -1)
O = A @ V
print("X", tuple(X.shape), "Q/K/V", tuple(Q.shape), "scores", tuple(S.shape),
      "weights", tuple(A.shape), "output", tuple(O.shape))

# Part 3: why divide by sqrt(d_k). q and k have unit-variance entries.
for d in (4, 64, 1024):
    q = torch.randn(20000, d)
    k = torch.randn(20000, d)
    raw = (q * k).sum(-1)
    print(f"d_k={d:5d} std of q.k = {raw.std():7.3f}  after scaling = {(raw / d ** 0.5).std():5.3f}")
d = 1024
q, K_ = torch.randn(d), torch.randn(8, d)
for name, s in (("unscaled", K_ @ q), ("scaled", K_ @ q / d ** 0.5)):
    print(name, "largest weight over 8 keys:", round(torch.softmax(s, 0).max().item(), 3))

# Part 4: memory of the score matrix, float32 = 4 bytes, one layer.
for n_ in (1024, 4096):
    heads = 16
    print(f"n={n_}: scores per sequence = {heads * n_ * n_ * 4 / 2**20:.0f} MiB")

Output:

scores
 tensor([[-0.6580, -0.8850, -1.1000],
        [ 2.6180,  4.6980,  4.9670],
        [ 0.6510,  1.4640,  1.3830]])
weights
 tensor([[0.4100, 0.3270, 0.2630],
        [0.0510, 0.4110, 0.5380],
        [0.1870, 0.4230, 0.3900]])
row sums tensor([1., 1., 1.])
output
 tensor([[-0.0620,  0.3910],
        [-0.0350,  0.0560],
        [-0.0380,  0.1050]])
matches PyTorch reference: True
X (2, 5, 16) Q/K/V (2, 5, 8) scores (2, 5, 5) weights (2, 5, 5) output (2, 5, 8)
d_k=    4 std of q.k =   1.993  after scaling = 0.997
d_k=   64 std of q.k =   8.084  after scaling = 1.010
d_k= 1024 std of q.k =  31.864  after scaling = 0.996
unscaled largest weight over 8 keys: 1.0
scaled largest weight over 8 keys: 0.403
n=1024: scores per sequence = 64 MiB
n=4096: scores per sequence = 1024 MiB

Reading the first block: token 2 (the second row) puts a weight of 0.051 on token 1 and 0.538 on token 3, so its output is dominated by the values of tokens 2 and 3. Each row of weights sums to 1, so the output is a convex combination of the value rows (a weighted average with non-negative weights that sum to 1, so it always lies between the value rows).

Shapes for a batch of sequences

A batch adds a leading axis B, and every matrix product acts on the last two axes. In the run above: X is (2, 5, 16), Q, K and V are (2, 5, 8), the scores and weights are (2, 5, 5) (one n by n table per sequence), and the output is (2, 5, 8). With several heads there is one more axis, giving (B, heads, n, n) for the scores.

Why the scores are scaled

If the entries of a query and a key are independent with mean 0 and variance 1 (variance is the average squared distance from the mean, and the standard deviation is its square root, the typical size of the spread), their dot product is a sum of d_k such products and has variance d_k, so its standard deviation is the square root of d_k. The Part 3 loop draws 20,000 random query/key pairs for each width and measures the spread of their dot products. The printout confirms the argument: the standard deviation of q.k is about 2 at d_k=4, about 8 at d_k=64 and about 32 at d_k=1024, and dividing by the square root of d_k brings each back to about 1. Without scaling, large d_k pushes the scores far apart, and the softmax saturates: one weight is almost 1 and the rest almost 0, a row that is nearly one-hot (all the weight on a single entry). In the 8-key example (one random query of width 1024 against 8 random keys) the largest weight prints as 1.0 (to three decimals) unscaled, against 0.403 scaled. A saturated softmax has a gradient with respect to the scores (how much the weights move when a score is nudged) that is close to zero, because for the largest weight p that gradient is p(1 - p); the saturation demo above printed 9e-5 for a score gap of 10 over three scores, so training signals reaching the query and key matrices vanish and learning stalls. A row this peaked also makes the output almost a copy of one value row, so changing how strongly the model attends to that token barely changes anything: it is already at weight 1. The divisor is the key width per head, d_k, not the model width.

Why several heads, and what attention maps do and do not show

One softmax produces one set of weights per token, so a token that needs information from two different places has to blend them in a single average. Multi-head attention runs several attention computations in parallel, each in its own smaller subspace (a lower-width projection of the token vector) with its own W_Q, W_K, W_V, and concatenates the results, so each head can pick different tokens. That is the reason to have several; which relations a given head ends up capturing is learned and has to be inspected, not assumed.

Common mistakes when reading an attention map (the weights table):

• Treating a high weight as an explanation. The weight says where the value was read from, but the value vectors differ in size and content, and later layers mix everything again, so a large weight is not proof that the token caused the output.
• Reading one head or one layer as "the model's focus". Each head is one of many, and weights after the first layer apply to already mixed representations, not to the original words.
• Forgetting that rows sum to 1 over the keys, not columns: the column of a token does not sum to 1.

Time and memory in the sequence length

The score matrix has n x n entries, each a dot product of length d_k. So one layer costs on the order of n squared times d in time and stores n squared weights per head. With 16 heads in float32 (4 bytes per number), the run above gives 64 MiB per sequence per layer at n=1024 and 1024 MiB at n=4096: quadrupling n multiplies the memory by 16. That quadratic growth is what the efficient-attention variants try to avoid: methods that compute only some of the n by n scores (for example, between nearby tokens only) or avoid storing the whole table. The projections cost on the order of n times d_model squared, which dominates when n is small relative to d_model.

Trade-offs and pitfalls

• Attention has no notion of order by itself: shuffle the input rows and the output rows are shuffled the same way. Position information must therefore be put into the inputs (for example as a position vector added to each token embedding) or into the scores.
• In a decoder, a causal mask (minus infinity added to the scores of later positions, same mechanism as the padding mask) must stop tokens from reading later tokens.
• Padding tokens (filler tokens added to short sequences so that every sequence in a batch has the same length) must be masked out of the keys: adding minus infinity to their scores before the softmax makes their weights exactly 0. Otherwise they take weight.
• A missing scale on a wide head saturates the softmax, as the example shows, and a saturated softmax passes almost no gradient back to the query and key matrices.

Design

The module needs five things, in this order inside forward:

1. One fused projection. A single nn.Linear(d_model, 3 * d_model) produces q, k and v in one matrix multiplication ("fused" means the three projections are stacked into one matrix; one larger multiply is usually more efficient than three small ones). Its weight has one row per output feature, so the first d_model rows make q, the next d_model make k and the last d_model make v. That row order is also the layout PyTorch's own layer uses, which is why the weights can be copied across below. chunk(3, dim=-1) then cuts the result into three tensors of shape (B, L, d_model).
2. Split into heads without a loop. view(B, L, h, d_k).transpose(1, 2) gives (B, h, L, d_k), so every head is one slice of a batched tensor and the matrix products below run on all heads at once. Here d_k = d_model / h.
3. Scores and masks. scores = q k^T / sqrt(d_k) has shape (B, h, L, L). Build one boolean "allowed" tensor: the causal part is a lower-triangular (L, L) matrix (query i may see keys j <= i), the padding part is (B, 1, 1, L) built from key_padding_mask (True = this key is padding). Combine with AND. Broadcasting means that a tensor with an axis of length 1 is repeated along that axis to match the other tensor, so the (1, 1, L, L) causal table and the (B, 1, 1, L) padding table expand to (B, h, L, L) without being copied by hand. key_padding_mask is the per-sequence list of filler positions added to shorter sequences so a batch is rectangular; the causal table is the same for every sequence (PyTorch's layer calls it attn_mask). Blocked scores are set to minus infinity before the softmax (the function that turns scores into weights summing to 1).
4. Dropout on the attention weights. Dropout randomly zeroes a fraction p of values during training and scales the values that remain by 1/(1-p) so the expected value is unchanged; it is applied after the softmax, to the weights, and does nothing in eval mode (F.dropout(..., training=self.training)). The weights returned are the post-dropout ones in training mode.
5. Combine heads. weights @ v is (B, h, L, d_k); transpose(1, 2).reshape(B, L, d_model) puts the head axis back next to the width axis and concatenates the heads, and an output linear layer returns (B, L, d_model).

Code, with checks against PyTorch

The program copies the module's weights into nn.MultiheadAttention (PyTorch's built-in layer; its fused weight stacks the q, k and v rows in the same order as the module's fused nn.Linear) and compares outputs and weights for padding-only and padding-plus-causal masks. PyTorch's convention for both mask arguments is True = blocked, so the causal matrix is passed as the upper triangle. Rows belonging to padding queries are excluded from the comparison: their outputs carry no training signal (the loss ignores them), and with left padding plus a causal mask they are exactly the all-blocked rows where implementations are free to differ, so leaving them out is not hiding a mismatch.

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.0):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.d_k = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)      # one fused projection for q, k and v
        self.out = nn.Linear(d_model, d_model)
        self.p = dropout

    def forward(self, x, key_padding_mask=None, causal=False):
        """x: (B, L, d_model). key_padding_mask: (B, L) bool, True = padding (ignore this key).
        Returns output (B, L, d_model) and attention weights (B, h, L, L)."""
        B, L, _ = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)                              # each (B, L, d_model)
        q, k, v = (t.view(B, L, self.h, self.d_k).transpose(1, 2) for t in (q, k, v))   # (B, h, L, d_k)
        scores = q @ k.transpose(-2, -1) / self.d_k ** 0.5                  # (B, h, L, L)
        allowed = torch.ones(L, L, dtype=torch.bool, device=x.device)
        if causal:
            allowed = torch.tril(allowed)                                   # query i sees keys j <= i
        allowed = allowed[None, None]                                       # (1, 1, L, L)
        if key_padding_mask is not None:
            allowed = allowed & ~key_padding_mask[:, None, None, :]         # (B, 1, 1, L) broadcast
        scores = scores.masked_fill(~allowed, float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        weights = torch.nan_to_num(weights, nan=0.0)                        # rows with no allowed key -> 0
        weights = F.dropout(weights, p=self.p, training=self.training)      # dropout on the weights
        y = (weights @ v).transpose(1, 2).reshape(B, L, -1)                 # (B, L, d_model)
        return self.out(y), weights

torch.manual_seed(0)
torch.set_num_threads(1)
B, L, d_model, h = 3, 6, 16, 4
x = torch.randn(B, L, d_model)
pad = torch.zeros(B, L, dtype=torch.bool)
pad[1, 4:] = True                                       # sequence 1 has 2 padding positions at the end
pad[2, 2:] = True                                       # sequence 2 has 4

mine = MultiHeadAttention(d_model, h, dropout=0.2).eval()
ref = nn.MultiheadAttention(d_model, h, batch_first=True).eval()
with torch.no_grad():                                   # nn.Linear and nn.MultiheadAttention share the (q;k;v) row layout
    ref.in_proj_weight.copy_(mine.qkv.weight); ref.in_proj_bias.copy_(mine.qkv.bias)
    ref.out_proj.weight.copy_(mine.out.weight); ref.out_proj.bias.copy_(mine.out.bias)

causal_bool = torch.triu(torch.ones(L, L, dtype=torch.bool), diagonal=1)   # True = blocked, PyTorch's convention
for causal in (False, True):
    y, w = mine(x, key_padding_mask=pad, causal=causal)
    y_ref, w_ref = ref(x, x, x, key_padding_mask=pad, attn_mask=causal_bool if causal else None,
                       need_weights=True, average_attn_weights=False)
    keep = ~pad[:, :, None]                              # compare only rows of real (non-padding) queries
    same_y = torch.allclose(y * keep, y_ref * keep, atol=1e-5)
    same_w = torch.allclose(w * keep[:, None], w_ref * keep[:, None], atol=1e-6)
    print(f"causal={causal}: output shape {tuple(y.shape)}, weights shape {tuple(w.shape)}, "
          f"matches nn.MultiheadAttention on real tokens: output {same_y}, weights {same_w}")

print("weight on padded keys (sequence 2, head 0, query 5):", w[2, 0, 5].tolist()[2:])
print("future weight is zero under causal mask:", bool((w[0, 0].triu(1) == 0).all()))

mine.train()
torch.manual_seed(1)
_, wd = mine(x, causal=False)
print(f"training mode, p=0.2: fraction of weights that are exactly 0 = {(wd == 0).float().mean():.3f}")
rs = wd.sum(-1).detach()
print(f"training mode row sums: min {rs.min():.2f}, mean {rs.mean():.2f}, max {rs.max():.2f} (kept weights are scaled by 1/(1-p))")
mine.eval()
_, we = mine(x)
print("eval mode row sums:", [round(float(a), 3) for a in we.sum(-1)[0, 0, :3]])
print("parameters:", sum(p.numel() for p in mine.parameters()), "= 4*16*16 + 4*16 =", 4 * 16 * 16 + 4 * 16)

# Tiny trace by hand: L=4 tokens, equal scores, so each allowed key gets weight 1/(number of allowed keys).
causal4 = torch.tril(torch.ones(4, 4, dtype=torch.bool))            # row = query, column = key; 1 = allowed
print("causal (L, L):", causal4.int().tolist())
right = torch.tensor([False, False, False, True])                    # right padding: the LAST token is filler
left = torch.tensor([True, False, False, False])                     # left padding: the FIRST token is filler
print("padding mask (B,1,1,L) holds one row of shape", tuple(right[None, None, None, :].shape), "that broadcasts over heads and queries")
for name, pm in (("right", right), ("left", left)):
    allowed4 = causal4 & ~pm[None, :]                                # (L, L) & (1, L) broadcasts the key mask down every row
    s4 = torch.zeros(4, 4).masked_fill(~allowed4, float("-inf"))
    w4 = torch.softmax(s4, dim=-1)                                   # a row of only -inf gives NaN
    print(f"{name} padding, combined allowed:", allowed4.int().tolist())
    print(f"{name} padding, softmax weights, row 0:", [round(float(v), 2) for v in w4[0]],
          "| row 3:", [round(float(v), 2) for v in w4[3]])
    print(f"{name} padding, row 0 after nan_to_num:", torch.nan_to_num(w4[0], nan=0.0).tolist())
torch.manual_seed(3)
row = torch.tensor([0.5, 0.5, 0.0, 0.0])
for _ in range(4):
    d_row = F.dropout(row, p=0.5, training=True)
    print("dropout p=0.5 on [0.5, 0.5, 0, 0]:", d_row.tolist(), "row sum", float(d_row.sum()))

Output:

causal=False: output shape (3, 6, 16), weights shape (3, 4, 6, 6), matches nn.MultiheadAttention on real tokens: output True, weights True
causal=True: output shape (3, 6, 16), weights shape (3, 4, 6, 6), matches nn.MultiheadAttention on real tokens: output True, weights True
weight on padded keys (sequence 2, head 0, query 5): [0.0, 0.0, 0.0, 0.0]
future weight is zero under causal mask: True
training mode, p=0.2: fraction of weights that are exactly 0 = 0.220
training mode row sums: min 0.04, mean 0.96, max 1.25 (kept weights are scaled by 1/(1-p))
eval mode row sums: [1.0, 1.0, 1.0]
parameters: 1088 = 4*16*16 + 4*16 = 1088
causal (L, L): [[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]
padding mask (B,1,1,L) holds one row of shape (1, 1, 1, 4) that broadcasts over heads and queries
right padding, combined allowed: [[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 0]]
right padding, softmax weights, row 0: [1.0, 0.0, 0.0, 0.0] | row 3: [0.33, 0.33, 0.33, 0.0]
right padding, row 0 after nan_to_num: [1.0, 0.0, 0.0, 0.0]
left padding, combined allowed: [[0, 0, 0, 0], [0, 1, 0, 0], [0, 1, 1, 0], [0, 1, 1, 1]]
left padding, softmax weights, row 0: [nan, nan, nan, nan] | row 3: [0.0, 0.33, 0.33, 0.33]
left padding, row 0 after nan_to_num: [0.0, 0.0, 0.0, 0.0]
dropout p=0.5 on [0.5, 0.5, 0, 0]: [1.0, 1.0, 0.0, 0.0] row sum 2.0
dropout p=0.5 on [0.5, 0.5, 0, 0]: [0.0, 0.0, 0.0, 0.0] row sum 0.0
dropout p=0.5 on [0.5, 0.5, 0, 0]: [1.0, 0.0, 0.0, 0.0] row sum 1.0
dropout p=0.5 on [0.5, 0.5, 0, 0]: [1.0, 1.0, 0.0, 0.0] row sum 2.0

Reading the output

• Both mask settings reproduce the built-in layer on real tokens, with output (3, 6, 16) and weights (3, 4, 6, 6), the (batch, heads, query, key) layout.
• Padded keys get weight exactly 0, and the future gets weight 0 under the causal mask.
• With p=0.2, about a fifth of the weights are zeroed in training (0.220 here, a random fraction around p). Row sums then vary around 1 (mean 0.96 in this run, range 0.04 to 1.25); in eval mode they are exactly 1. Dropout on the weights is therefore only for training. Why the sums wobble is clearest on one tiny row in the last lines of the output: the row [0.5, 0.5, 0, 0] with p=0.5 keeps each entry with probability one half and doubles it (1/(1-p) = 2), so the sum is 2.0 if both survive, 1.0 if one does, and 0.0 if neither does. It averages to 1, but any single row can be far from 1; a row whose big weights all happened to be dropped can sum to nearly 0, as with the 0.04 above.
• The parameter count is 4 x d_model squared plus 4 x d_model biases: the fused projection (3 x d_model squared plus 3 x d_model) plus the output layer (d_model squared plus d_model).

Tracing the masks on 4 tokens

The last part of the program makes the masks small enough to read. With L = 4 the causal table has a 1 where a query (row) may see a key (column): row 0 sees only key 0, row 3 sees all four. A padding mask marks filler keys, and combining it with the causal table (AND) removes those columns from every row.

• Right padding puts the filler after the real tokens (last token is filler). Combined table rows: [1,0,0,0], [1,1,0,0], [1,1,1,0], [1,1,1,0]. Every row still has an allowed key, and with equal scores row 3 spreads weight 0.33 over the three real keys.
• Left padding puts the filler before the real tokens (first token is filler). Now the padded query (the row that belongs to the filler position, row 0) has no allowed key at all, its softmax is over only minus infinity, and the printed weights are nan. nan_to_num turns that row into zeros. Later rows are unaffected (row 3 gives 0.33 to each of the three real keys).

Which parts are core

The fused projection, the head split, the two masks applied before the softmax, and the head recombination are the heart of the module. Dropout on the weights and returning the weights are required by this question, but are not part of what makes attention work. nan_to_num is a safety net for one specific edge case, and the fused built-in function in the last bullet below is the production alternative, not something to implement here.

Complexity and edge cases

• Time is about B x L x d_model squared for the projections and B x h x L x L x d_k for the score and value products; memory for the weights is B x h x L x L, which dominates for long L.
• A query row with no allowed key (for example a padded query under left padding combined with a causal mask) would give NaN from the softmax (the tracing above shows exactly this); nan_to_num turns it into zeros. If real NaNs can arise elsewhere, a stricter version zeroes only the rows that have no allowed key, using torch.where on a "this row has any allowed key" test, so any other NaN stays visible. Multiplying the weights by that test would not work, because NaN times 0 is still NaN.
• Mask conventions differ: this module and PyTorch use True for padding in key_padding_mask, but code that stores True for "keep" needs inverting.
• For incremental decoding (generating one token at a time and reusing the keys and values of earlier tokens, which are stored in a key-value cache), the causal mask must be built with the cache length (the number of stored tokens) in mind: one query against L_k cached keys, not L against L. The module above assumes the same length for queries and keys.
• For production speed, torch.nn.functional.scaled_dot_product_attention is a fused kernel (one GPU routine that does the score, mask, softmax, dropout and value product together, in tiles, without writing the full score matrix to memory), but it does not return the weights, which is why the explicit version is used when weights are needed.

The short version

A Transformer encoder turns a sequence of n tokens into n context-aware vectors. Each token is embedded (looked up as a vector of d numbers), a position signal is added, and the result goes through a stack of identical layers. Each layer has two sublayers: multi-head self-attention (every token gathers information from all tokens) and a position-wise feed-forward network (a small two-layer network applied to each token separately). Each sublayer is wrapped in a residual connection and layer normalization. The cost of attention is O(n^2 d) time and O(n^2) memory per head for the score matrix, so doubling the sequence length quadruples it; that drives most serving decisions.

The walk-through, in order

1. Tokens to embeddings. Split text into tokens (subword pieces), look each up in an embedding table. Shape: n by d (n tokens, d numbers each).
2. Positional encoding. Attention by itself ignores order, so a position signal is added to each embedding. The original design adds fixed sinusoids of different frequencies, where position p and feature pair i get sin and cos of p divided by 10000 to the power 2i/d. Learned position vectors (a trainable table with one row per position), relative schemes (which encode the distance between two tokens instead of each token's absolute slot) and rotary embeddings (which rotate queries and keys inside attention rather than adding to the input) are alternatives; they change how position enters the model, not the rest of the layer.
3. Multi-head self-attention. Three learned d-by-d matrices produce queries Q, keys K and values V from the input X (each n by d). Each head works on a slice of size d_k = d/h (h heads). Scores are the dot product of every query with every key, divided by the square root of d_k, then a softmax (a function that turns a row of numbers into positive weights summing to 1) over each row gives attention weights (n by n per head), and the output is the weights times V. The h head outputs are concatenated and mixed by a final d-by-d matrix. Padding positions are masked out.
4. Residual and layer norm. The sublayer output is added to its input, then normalized per token (post-LN, meaning the normalization comes after the addition, as in the original; pre-LN, which many modern models use, normalizes the sublayer's input instead and leaves the residual path unnormalized).
5. Feed-forward block. Each token goes through a linear layer to width d_ff (commonly 4d), a nonlinearity, and a linear layer back to d. This holds a large share of the parameters.
6. Repeat for N layers. The output is n vectors of size d, which feed a task head (a small final layer for the job at hand, such as a classifier over the vectors).

A traced example: positions and shapes

For d = 4 the divisors 10000^(2i/d) are 1 (i = 0) and 100 (i = 1), printed below as their reciprocals 1.0 and 0.01, so position p gets (sin p, cos p, sin(p/100), cos(p/100)). At p = 1 that is (0.8415, 0.5403, 0.0100, 1.0000): the first pair changes quickly from one position to the next, the second pair slowly, which is how the signal separates nearby and far positions. The program prints these values and then follows four token ids through one encoder layer (d = 8, 2 heads, d_k = 4, d_ff = 32), printing the shape after each step.

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

torch.set_num_threads(1)

torch.manual_seed(0)
vocab, d, h, dff = 20, 8, 2, 32
ids = torch.tensor([5, 3, 7, 3])                       # 4 token ids (n = 4)
n = len(ids)

# Position signal for d = 4, positions 0..2, so each entry can be checked by hand.
p = torch.arange(3).unsqueeze(1).float()
div = torch.exp(torch.arange(0, 4, 2) * (-math.log(10000.0) / 4))   # 10000^(-2i/d) for i = 0, 1 -> [1.0, 0.01]
pe4 = torch.zeros(3, 4)
pe4[:, 0::2], pe4[:, 1::2] = torch.sin(p * div), torch.cos(p * div)
print("divisors 10000^(-2i/4):", [round(float(v), 4) for v in div])
print("position signal, d=4:\n", pe4.round(decimals=4))

# Shapes through one encoder layer (post-LN), d = 8, 2 heads.
emb = nn.Embedding(vocab, d)
x = emb(ids)                                           # (4, 8): one row per token id
pos = torch.zeros(n, d)
pp = torch.arange(n).unsqueeze(1).float()
dv = torch.exp(torch.arange(0, d, 2) * (-math.log(10000.0) / d))
pos[:, 0::2], pos[:, 1::2] = torch.sin(pp * dv), torch.cos(pp * dv)
x = x + pos
print("ids", tuple(ids.shape), "-> embeddings + position", tuple(x.shape))
Wq, Wk, Wv, Wo = (nn.Linear(d, d, bias=False) for _ in range(4))
q, k, v = (W(x).view(n, h, d // h).transpose(0, 1) for W in (Wq, Wk, Wv))
print("per-head q/k/v", tuple(q.shape), " (heads, tokens, d_k)")
w = F.softmax(q @ k.transpose(1, 2) / math.sqrt(d // h), dim=-1)
print("attention weights", tuple(w.shape), " rows sum to", w.sum(-1).round(decimals=4).unique().tolist())
heads = (w @ v).transpose(0, 1).reshape(n, d)
x = F.layer_norm(x + Wo(heads), (d,))
print("after attention + residual + LN", tuple(x.shape))
ff = nn.Sequential(nn.Linear(d, dff), nn.ReLU(), nn.Linear(dff, d))
print("feed-forward inner width", tuple(ff[0](x).shape), "-> back to", tuple(F.layer_norm(x + ff(x), (d,)).shape))

Output:

divisors 10000^(-2i/4): [1.0, 0.01]
position signal, d=4:
 tensor([[ 0.0000,  1.0000,  0.0000,  1.0000],
        [ 0.8415,  0.5403,  0.0100,  1.0000],
        [ 0.9093, -0.4161,  0.0200,  0.9998]])
ids (4,) -> embeddings + position (4, 8)
per-head q/k/v (2, 4, 4)  (heads, tokens, d_k)
attention weights (2, 4, 4)  rows sum to [1.0]
after attention + residual + LN (4, 8)
feed-forward inner width (4, 32) -> back to (4, 8)

Reading the shapes: four ids become a (4, 8) table, one row per token; each head sees (4, 4) queries and keys, so its weight table is 4 by 4, and the two heads give (2, 4, 4); the head outputs are concatenated back to (4, 8); the feed-forward network widens each token to 32 and narrows back to 8. The layer maps (4, 8) to (4, 8), which is why layers can be stacked.

Checking the mechanics and counting parameters and cost

The program implements steps 2 and 3 from scratch and compares attention to PyTorch's scaled_dot_product_attention. It also compares a hand formula for one layer's parameter count to a built nn.TransformerEncoderLayer (post-LN, d=16, 4 heads, d_ff=64), then prints the score-matrix size and FLOPs (floating-point operations) for an example layer with d=1024, 16 heads, d_ff = 4d and fp16 numbers (2 bytes each).

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

torch.set_num_threads(1)

torch.manual_seed(0)

def sinusoidal(n, d):
    pos = torch.arange(n).unsqueeze(1)
    div = torch.exp(torch.arange(0, d, 2) * (-math.log(10000.0) / d))
    pe = torch.zeros(n, d)
    pe[:, 0::2], pe[:, 1::2] = torch.sin(pos * div), torch.cos(pos * div)
    return pe

def mha(x, Wq, Wk, Wv, Wo, h):
    n, d = x.shape
    dk = d // h
    q, k, v = (x @ W for W in (Wq, Wk, Wv))
    q, k, v = (t.view(n, h, dk).transpose(0, 1) for t in (q, k, v))   # (h, n, dk)
    w = F.softmax(q @ k.transpose(1, 2) / math.sqrt(dk), dim=-1)      # (h, n, n)
    return (w @ v).transpose(0, 1).reshape(n, d) @ Wo

n, d, h, dff = 6, 16, 4, 64
x = torch.randn(n, d) + sinusoidal(n, d)          # token embeddings + position signal
Wq, Wk, Wv, Wo = (torch.randn(d, d) / math.sqrt(d) for _ in range(4))

# Check the head-split implementation against the library kernel.
dk = d // h
q, k, v = (x @ W for W in (Wq, Wk, Wv))
q, k, v = (t.view(n, h, dk).transpose(0, 1) for t in (q, k, v))
ref = F.scaled_dot_product_attention(q, k, v).transpose(0, 1).reshape(n, d) @ Wo
print("attention matches library kernel:", torch.allclose(mha(x, Wq, Wk, Wv, Wo, h), ref, atol=1e-5))

# Parameter count: formula vs a built post-LN encoder layer.
layer = nn.TransformerEncoderLayer(d, h, dff)
built = sum(p.numel() for p in layer.parameters())
formula = (4 * d * d + 4 * d) + (2 * d * dff + dff + d) + 2 * (2 * d)
print("params per layer: built =", built, " formula =", formula)

# Cost of one layer (batch 1). Example config: d=1024, h=16 heads, d_ff=4d, fp16 = 2 bytes.
d_, heads, bytes_per = 1024, 16, 2
print("n      score entries   score MiB   attention FLOPs  projection+FFN FLOPs  ratio")
for n_ in (1024, 2048, 4096, 8192):
    entries = heads * n_ * n_
    attn_flops = 4 * n_ * n_ * d_          # QK^T and AV: 2 matmuls, 2 FLOPs per multiply-add
    dense_flops = 24 * n_ * d_ * d_        # 4 projections (8 n d^2) + FFN with d_ff=4d (16 n d^2)
    print(f"{n_:<6} {entries:<15} {entries * bytes_per / 2**20:<11.0f} {attn_flops:<16.3e} {dense_flops:<21.3e} {attn_flops / dense_flops:.2f}")

Output:

attention matches library kernel: True
params per layer: built = 3280  formula = 3280
n      score entries   score MiB   attention FLOPs  projection+FFN FLOPs  ratio
1024   16777216        32          4.295e+09        2.577e+10             0.17
2048   67108864        128         1.718e+10        5.154e+10             0.33
4096   268435456       512         6.872e+10        1.031e+11             0.67
8192   1073741824      2048        2.749e+11        2.062e+11             1.33

The parameter formula is: attention 4d^2 + 4d (four d-by-d matrices with biases), feed-forward 2 d d_ff + d_ff + d, and two layer norms with 2d each. The FLOP model is: score and weighted-sum matmuls 4 n^2 d in total; projections 8 n d^2; feed-forward with d_ff = 4d is 16 n d^2; so the dense part is 24 n d^2. Equating gives n = 6d: for d = 1024, attention FLOPs pass the rest at n of about 6,000 tokens, as the ratio column crossing 1 between 4096 and 8192 shows.

What the costs imply in practice

• Time: O(n^2 d) for attention, O(n d^2) for projections and feed-forward. For short inputs (n well below 6d) the dense layers dominate; for long inputs attention does.
• Memory: a naive implementation stores an n by n score matrix per head: each doubling of n quadruples it (32, 128, 512 and 2048 MiB per layer at n = 1024, 2048, 4096 and 8192 here). Training also keeps activations for backpropagation.
• Consequences: cap or chunk input length; use a fused kernel (one GPU program that performs several steps without writing intermediate results to memory) that computes exact attention in tiles without materializing the score matrix (the FlashAttention idea), which cuts memory but not the O(n^2) FLOPs; use efficient-attention variants (methods that compute only some of the n by n scores, such as between nearby tokens, or avoid storing the full table) for very long inputs; batch by similar lengths to avoid paying for padding.

For a serving team, in plain terms

Cost per request grows faster than input length: a request twice as long costs twice as much for the dense layers (their cost is linear in n) and four times as much for the attention part, so the total rises by somewhere between two and four times depending on which part dominates at that length. Memory per request is the second constraint, because the score matrices and cached activations scale with length squared and length respectively, so batch size must shrink as inputs get longer. Set a maximum input length, bucket requests by length, and measure on your own traffic.

Common mistakes

• Saying the scaling uses d instead of d_k per head.
• Saying attention is always the bottleneck: below n of about 6d the dense layers cost more.
• Treating positional encoding as optional: without it the model sees a bag of tokens.
• Confusing time and memory fixes: tiling saves memory, not FLOPs.

Running the code

Save the program as demo.py and run python demo.py. The output above was produced with Python 3.12, numpy 2.1.3 and torch 2.5.1 (CPU). The seed is fixed in the program.

The setup

Single-head attention for 2 tokens with d_k = 2. Each token has a query (what it is looking for), a key (what it offers for matching) and a value (the content it passes on). Given:

$Q = \begin{bmatrix}1 & 0\\ 1 & 1\end{bmatrix},\quad K = \begin{bmatrix}1 & 0\\ 0 & 1\end{bmatrix},\quad V = \begin{bmatrix}1 & 2\\ 3 & 4\end{bmatrix}$

Row i of each matrix belongs to token i. In a real model these three matrices are not typed in: they come from the token embeddings X (one row of numbers per token) multiplied by three learned weight matrices, Q = X W_Q, K = X W_K, V = X W_V. Here they are simply handed to you so the attention step can be done on its own. The formula is

$\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

where softmax is applied separately to each row.

Step by step

1. Score matrix Q K^T (2 by 2): entry (i, j) is the dot product of query i with key j.
   Row 0: (1,0) dot (1,0) = 1 and (1,0) dot (0,1) = 0. Row 1: (1,1) dot (1,0) = 1 and (1,1) dot (0,1) = 1.
2. Scaling by the square root of d_k = 2, which is 1.4142: 1 becomes 0.7071.
3. Row-wise softmax: for row 0 the scaled scores are 0.7071 and 0; the weights are e^0.7071 over (e^0.7071 + e^0) = 0.6698 and 0.3302. For row 1 the two scores are equal, so each weight is 0.5.
4. Output = weights times V: each output row is a weighted average of the rows of V. Row 0: 0.6698 times (1,2) plus 0.3302 times (3,4) = (1.6605, 2.6605). Row 1: 0.5 times (1,2) plus 0.5 times (3,4) = (2, 3).

The program does the same computation, adds the unscaled variant and compares with the library implementation.

import numpy as np
import torch
import torch.nn.functional as F

torch.set_num_threads(1)

Q = np.array([[1.0, 0.0], [1.0, 1.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0]])
V = np.array([[1.0, 2.0], [3.0, 4.0]])
dk = Q.shape[1]

scores = Q @ K.T                      # (2 queries, 2 keys): row i, column j = q_i . k_j
scaled = scores / np.sqrt(dk)
e = np.exp(scaled - scaled.max(axis=1, keepdims=True))
weights = e / e.sum(axis=1, keepdims=True)    # softmax across each ROW
out = weights @ V

np.set_printoptions(precision=4, suppress=True)
print("scores Q K^T:\n", scores)
print("scaled by sqrt(2):\n", scaled)
print("weights (row-wise softmax):\n", weights)
eu = np.exp(scores - scores.max(axis=1, keepdims=True))
print("weights without the 1/sqrt(d_k) scaling:\n", eu / eu.sum(axis=1, keepdims=True))
print("row sums:", weights.sum(axis=1))
print("output weights @ V:\n", out)
print("check row 0 by hand-style sum:", weights[0, 0] * V[0] + weights[0, 1] * V[1])
ref = F.scaled_dot_product_attention(*(torch.tensor(a, dtype=torch.float64) for a in (Q, K, V)))
print("matches library:", np.allclose(out, ref.numpy()))

Output:

scores Q K^T:
 [[1. 0.]
 [1. 1.]]
scaled by sqrt(2):
 [[0.7071 0.    ]
 [0.7071 0.7071]]
weights (row-wise softmax):
 [[0.6698 0.3302]
 [0.5    0.5   ]]
weights without the 1/sqrt(d_k) scaling:
 [[0.7311 0.2689]
 [0.5    0.5   ]]
row sums: [1. 1.]
output weights @ V:
 [[1.6605 2.6605]
 [2.     3.    ]]
check row 0 by hand-style sum: [1.6605 2.6605]
matches library: True

What each attention row means

Row i of the weight matrix is how token i divides one unit of attention across all tokens: entry (i, j) is how much token i reads from token j. Each row sums to 1 (the printed row sums), the columns do not have to. Row 0 puts 0.6698 on token 0 because its query (1,0) matches key 0 better than key 1; row 1's query (1,1) is equally close to both keys, so it takes a plain average of the two values. The output row is the resulting blend of the value vectors; a different row of the same matrix gives a different blend, which is why each token gets its own context-aware output.

What the scaling did

Without the 1/sqrt(d_k) division, row 0's weights would be 0.7311 and 0.2689 instead of 0.6698 and 0.3302: scaling softens the distribution. With only d_k = 2 the effect is small; with d_k = 64 or more the unscaled scores become large and the softmax saturates.

Causal variant and pitfalls

• A causal mask is a rule that stops each token from attending to later positions: before the softmax, the scores of those positions are set to minus infinity, so their weights become exactly 0. With a causal mask, token 0 may not see token 1, so row 0 is forced to (1, 0) and its output is V's row 0, (1, 2). Row 1 is unchanged here because it may see both.
• Softmax across rows, not columns.
• Scale by the square root of d_k, not by d_k, and not by the number of tokens.
• Keys and values have the same number of rows (one per token), but the values can have a different width than the keys.

Running the code

Save the program as demo.py and run python demo.py. The output above was produced with Python 3.12, numpy 2.1.3 and torch 2.5.1 (CPU). The program uses no random numbers, so the output is the same on every run.