Q/K/V Math — Scaled Dot-Product Attention Derived

Why this session exists

Scaled dot-product attention combines a learned compatibility score, a constrained normalization and a value aggregation. Knowing the formula is not enough: shape-compatible mistakes in normalization, masking and backward reduction can silently change the computation. We will derive a soft lookup objective, state the assumptions behind scaling, correct a hand-worked example, derive every backward path and test against finite differences and PyTorch's CPU SDPA.

Prerequisites

S107 attention intuition, S098 backpropagation, matrix multiplication and softmax derivatives. We use row-batch arrays and mathematical projection weights (input_width,output_width); PyTorch Linear stores their transpose.

(a) Intuition · 5 min

Think of retrieving a weighted blend of records rather than literally averaging web pages. Matching features and returned content serve different purposes, so separate query/key/value projections are useful. That separation is structural, not evidence that an individual learned Q coordinate means a human question. In cross-attention, queries and source records may even come from different modalities and widths.

The four questions this session answers

Why separate projections? They allow asymmetric matching and independently transformed content. Why square-root scaling? It controls variance under a stated random-component model. Why softmax? It is one smooth normalized choice, derivable from entropy-regularized score maximization. How does backward work? Two matrix-product paths and a row-local softmax vector–Jacobian product, followed by summation of shared-input paths.

The full formula, dissected

For query input $X_q\in\mathbb R^{B\times L\times d_q^{in}}$ and source $X_s\in\mathbb R^{B\times S\times d_s^{in}}$:

Here Q is (B,L,d_k), K (B,S,d_k), V (B,S,d_v), Z/A (B,L,S) and O (B,L,d_v). Transpose only K's last two axes. Q/K score widths must match; K/V lengths must match. Value width is independent. GQA changes numbers of K/V heads, not inherently the equality of query/key/value widths. A rectangular score matrix is a useful test case even though different sequence sources can have equal lengths.

Why sqrt(d_k) — the variance derivation

Assume independent coordinates, independent Q and K, zero mean and unit variance. Then $E[q_i k_i]=0$ and $E[q_i^2k_i^2]=1$. Cross-coordinate covariance terms vanish, so

Dividing by $d_k$ instead gives variance $1/d_k$ under these assumptions; leaving raw scores gives variance $d_k$. Neither alternative is mathematically illegal or guaranteed to diverge. Learned temperatures/norms change behavior. LayerNorm and Xavier initialization do not guarantee independence of projected queries and keys. In particular, setting Q=K=X yields diagonal self-scores $\|x_i\|^2$ with mean $d_k$; even after scaling their mean grows like $\sqrt{d_k}$. That is why the old tied-Q/K demo did not test its claimed unit-variance argument.

For large score gaps, softmax can become sharply concentrated and its Jacobian small. This is a conditioning risk, not proof every omitted scale freezes training. Increasing sequence length changes score count and extreme-value behavior, not necessarily per-score variance.

Why softmax and not linear normalisation

For legal keys, choose a distribution maximizing expected score plus entropy:

The Lagrangian derivative gives $s_j-\tau(\log a_j+1)+\lambda=0$, hence $a_j\propto\exp(s_j/\tau)$. Normalizing gives softmax. It is the solution to this objective, not the only conceivable attention mechanism. Dividing raw signed scores by their sum can produce negative weights or division by zero; positive-score normalization is differentiable away from a zero denominator and can define another model.

Softmax is shift invariant: subtract each finite row maximum before exponentiation. That prevents overflow in exponentials but cannot repair an already infinite dot product. An all-masked -inf row has no probability distribution and naive max subtraction makes NaNs; our reference rejects it. Dtype, accumulation and backend behavior must be tested explicitly.

Hand trace: a 3-token, 2-dim toy

Let Q=K=V=X with identity projections and

Divide every score by $\sqrt2$. Row 0 has unnormalized weights $(e^{1/\sqrt2},1,e^{1/\sqrt2})$. Correct rounded results are

The previous first rows (0.39,0.22,0.39) were not accurate rounding. Row 2's larger self score is a dot product, not normalized cosine. Symmetric scores need not produce a symmetric attention matrix because rows have different normalizers. Under a causal mask row 0 must be (1,0,0); row 1 mixes only the first two values.

Backward pass — the gradient anatomy

Let upstream $G=\partial L/\partial O$. For each batch:

The rowsum is broadcast over keys. The full row Jacobian is $\operatorname{diag}(a)-aa^T$, but applying it this way costs O(S) per row, not O(S²) per row. For fixed excluded logits, masked-score derivatives are zero. Derivatives of learned biases/masks, if any, require separate handling.

For self-attention with shared X:

Do not omit two of the input paths. In cross-attention, query/source gradients are separate, and the source receives both key and value contributions. If the scalar loss is averaged, its normalization enters G; do not divide by batch/length a second time.

Two circuits: where to look and what to bring back

With linear, position-free self-attention, scores depend on $W_QW_K^T$. For invertible R, $W_Q'=W_QR$ and $W_K'=W_KR^{-T}$ preserve that product. For orthogonal R, both can multiply by the same R. Calling those 'inverse orthogonal rotations' while multiplying both identically obscures the convention. When an output projection exists, a value/output basis change can likewise preserve $W_VW_O$. This non-identifiability is qualified by biases, positional transformations such as RoPE, normalization and architectural constraints; Q/K/V are not independently semantically identifiable from scores alone.

Additive versus dot-product scoring

Additive scoring uses $e_{ij}=v_a^T\tanh(W_qx_i+W_ky_j+b)$; the projection widths can match through a chosen hidden scoring width even when input widths differ. The nonlinear pair interaction differs from a bilinear dot-product score. Tanh bounds hidden activations, but the output vector's norm and hidden width still affect score scale—there is no unconditional immunity to saturation. Dot products map naturally to dense matrix multiplication, while additive scores can also be batched but usually need a larger pairwise nonlinear intermediate. Choose by task evidence and actual kernels; neither a universal expressiveness ordering nor a universal speed winner follows. Multi-head/deep compositions further change the comparison.

Masking, order and cached decoding

Masking forbidden logits before softmax both excludes their values and excludes their scores from the denominator. Merely zeroing weights after softmax leaves legal weights dependent on future scores: future leakage can remain even if future values have zero weight. Renormalizing the retained weights can be algebraically equivalent for finite well-conditioned rows, so pre-softmax placement is not logically the only method; it is the clearer numerical contract.

A finite negative sentinel can underflow to exact zero or leave nonzero mass depending on dtype/scale. It does not generically 'round back into range in fp16'. Explicit masks and prefix-perturbation tests are better than folklore. A forced first causal row is not a sufficient explanation of learned attention sinks in later rows.

For a cached prefix of length P and a new query block of length L, query row i has absolute position P+i, so allow key j when $j\le P+i$. A naive L×(P+L) top-left lower triangle discards valid cached keys. PyTorch 2.9 SDPA's non-square causal convention is upper-left aligned; pass an explicit allowed mask for this case. SDPA Boolean True means allowed; nn.MultiheadAttention Boolean masks use True to block. Set dropout_p=0.0 explicitly for SDPA evaluation—it does not infer module eval state.

(c) Hands-on · 25 min

The standalone float64 CPU program tests rectangular, causal and cached-prefix cases, exact analytical backward and three-epsilon finite differences. It also rejects wrong-axis softmax, zero-after-softmax leakage and tied-Q/K variance reasoning. Timing is a bounded CPU fixture, not FlashAttention proof.

import math,time,statistics
import torch
from torch.nn import functional as F
 
torch.set_num_threads(1);torch.manual_seed(19)
def attention(q,k,v,allowed=None):
    if q.shape[-1]!=k.shape[-1] or k.shape[-2]!=v.shape[-2]:raise ValueError('incompatible roles')
    if q.shape[-1]==0 or k.shape[-2]==0:raise ValueError('empty keys/width')
    z=q@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
    if allowed is not None:
        allowed=torch.broadcast_to(allowed,z.shape)
        if not allowed.any(-1).all():raise ValueError('all-masked row')
        z=z.masked_fill(~allowed,-torch.inf)
    a=(z-z.max(-1,keepdim=True).values).softmax(-1)
    return a@v,a
 
def backward(q,k,v,a,g):
    dv=a.transpose(-2,-1)@g;da=g@v.transpose(-2,-1)
    dz=a*(da-(a*da).sum(-1,keepdim=True))
    return dz@k/math.sqrt(q.shape[-1]),dz.transpose(-2,-1)@q/math.sqrt(q.shape[-1]),dv
 
for L,S,causal in [(3,5,False),(4,4,True),(2,5,True)]:
    q=torch.randn(1,L,3,dtype=torch.double,requires_grad=True)
    k=torch.randn(1,S,3,dtype=torch.double,requires_grad=True)
    v=torch.randn(1,S,2,dtype=torch.double,requires_grad=True)
    # For 2x5, three keys are a cached prefix.
    mask=torch.arange(S)[None,:]<=torch.arange(S-L,S)[:,None] if causal else torch.ones(L,S,dtype=torch.bool)
    out,a=attention(q,k,v,mask);g=torch.randn_like(out)
    auto=torch.autograd.grad((out*g).sum(),(q,k,v))
    manual=backward(q,k,v,a,g)
    for x,y in zip(auto,manual):assert torch.allclose(x,y,atol=1e-12)
    builtin=F.scaled_dot_product_attention(q,k,v,attn_mask=mask,dropout_p=0.)
    assert torch.allclose(out,builtin,atol=1e-12)
    assert torch.allclose(a.sum(-1),torch.ones_like(a.sum(-1)))
    # Independent scalar finite differences for every input coordinate.
    values=[q.detach().clone(),k.detach().clone(),v.detach().clone()]
    for eps in [1e-4,1e-5,1e-6]:
        for tensor,expected in zip(values,auto):
            for idx in range(tensor.numel()):
                flat=tensor.view(-1);old=flat[idx].item()
                flat[idx]=old+eps;plus=(attention(*values,mask)[0]*g).sum().item()
                flat[idx]=old-eps;minus=(attention(*values,mask)[0]*g).sum().item();flat[idx]=old
                assert abs((plus-minus)/(2*eps)-expected.reshape(-1)[idx].item())<1e-7
    if L==2:assert not torch.allclose(out,F.scaled_dot_product_attention(q,k,v,is_causal=True,dropout_p=0.))
# Complete self-attention projection gradients, including all three X paths.
x=torch.randn(1,3,4,dtype=torch.double,requires_grad=True)
w=[torch.randn(4,d,dtype=torch.double,requires_grad=True) for d in [2,2,3]]
q,k,v=[x@p for p in w];out,a=attention(q,k,v);g=torch.randn_like(out)
auto=torch.autograd.grad((out*g).sum(),[x]+w);dq,dk,dv=backward(q,k,v,a,g)
dx=sum(d@p.T for d,p in zip([dq,dk,dv],w))
assert torch.allclose(dx,auto[0],atol=1e-12)
for d,expected in zip([dq,dk,dv],auto[1:]):assert torch.allclose((x.transpose(-2,-1)@d).sum(0),expected,atol=1e-12)
assert not torch.allclose(dq@w[0].T,auto[0])
# Hand trace and wrong-axis control.
x=torch.tensor([[1.,0.],[0.,1.],[1.,1.]],dtype=torch.double)
out,a=attention(x,x,x)
assert torch.allclose(a[0],torch.tensor([.401112,.197776,.401112],dtype=torch.double),atol=1e-6)
wrong=(x@x.T/math.sqrt(2)).softmax(-2)
assert not torch.allclose(wrong.sum(-1),torch.ones(3,dtype=torch.double))
# Post-softmax zeroing still depends on forbidden scores.
z=torch.tensor([[1.,2.,3.]]);legal=torch.tensor([[True,False,False]])
bad=z.softmax(-1)*legal;z2=z.clone();z2[0,2]=30
assert not torch.equal(bad,z2.softmax(-1)*legal)
assert torch.equal(z.masked_fill(~legal,-torch.inf).softmax(-1),z2.masked_fill(~legal,-torch.inf).softmax(-1))
# Orthogonal basis invariance.
r,_=torch.linalg.qr(torch.randn(2,2,dtype=torch.double))
assert torch.allclose(attention(x@r,x@r,x)[1],a,atol=1e-12)
# Correct variance experiment: independent q and k, not Q=K.
for d in [8,32,128]:
    q=torch.randn(5000,d);k=torch.randn_like(q);scaled=(q*k).sum(-1)/math.sqrt(d)
    assert .85<scaled.var().item()<1.15
    print('d',d,'scaled score variance',scaled.var().item())
# Deterministic cyclic shift: K_j = X_(j-1), so query i matches key i+1, not i-1.
x=4*torch.eye(6,dtype=torch.double);_,a=attention(x,torch.roll(x,1,0),x)
assert torch.equal(a.argmax(-1),(torch.arange(6)+1)%6)
try:attention(x,x,x,torch.zeros(6,6,dtype=torch.bool));raise AssertionError('all-masked accepted')
except ValueError:pass
q=torch.randn(1,16,8,dtype=torch.double);k=torch.randn_like(q);v=torch.randn_like(q)
for label,fn in [('scratch',lambda:attention(q,k,v)[0]),('sdpa',lambda:F.scaled_dot_product_attention(q,k,v,dropout_p=0.))]:
    for _ in range(3):fn()
    ts=[]
    for _ in range(5):
        start=time.perf_counter();fn();ts.append(time.perf_counter()-start)
    print(label,'CPU fixture median seconds',statistics.median(ts))
print('forward, manual backward, finite differences and negative controls passed')

Anatomy of the script

The same function handles differing lengths/value widths. Masks are Boolean allowed positions and reject empty support. Float64 and small shapes make tight CPU parity meaningful; other dtypes/backends need appropriate tolerances. SDPA selects eligible kernels by device/dtype/layout and settings; a CPU parity run does not establish a GPU FlashAttention path. There is no enable_math argument on the SDPA function in the cited API; backend selection uses documented context/configuration APIs separately.

The shift oracle uses an orthogonal basis so the intended argmax is guaranteed; random rows could have an accidental larger cross-dot. With roll(X,+1), the strong match is the cyclic super-diagonal (i+1), not the claimed sub-diagonal. Incorrect softmax may define some other trainable operation; the test rejects it because it violates this attention contract, not because every such model necessarily learns nonsense.

(d) Production reality · 15 min

FlashAttention: maintain running row maximum m, exponential sum l and weighted accumulator u over key tiles. For a new tile maximum $m'$ (including the old m), rescale old accumulators by $\exp(m-m')$, add new exponentials $\exp(s_j-m')$ and their weighted values, then output u/l. This online softmax is algebraically equivalent to dense normalization while avoiding the full score matrix in HBM; floating operation order can change rounding. Arithmetic remains quadratic for dense attention. No universal speedup, closed-model adoption or GPU benchmark is asserted.

Numerical precision: stable softmax is an established identity, not a verified T5 incident. Bfloat16's wide exponent range differs from fp16; exponentiation and dot-product overflow are separate failure points. Higher-precision accumulation can help but cannot guarantee finite values under all inputs. The execution tests cover float64 CPU only.

KV cache: for B sequences, N layers, Hkv key/value heads, T positions, width d and b bytes per element, bytes = $2BNH_{kv}Tdb$. For B=1, N=80, Hkv=64, T=32,768, d=128, b=2, this is 85,899,345,920 bytes = 80 GiB = 85.899 GB. Hkv=8 gives 10 GiB = 10.737 GB, an eightfold reduction in this cache component. At 4,096 tokens those values are 10 GiB and 1.25 GiB, not 40 GB per request. These are explicitly hypothetical full-MHA versus grouped configurations, not a claim that a shipped 70B model has 64 KV heads. Weights, allocator overhead and other state are additional.

Where this shows up in the rest of the plan

S109 adds multiple routing distributions and grouped caches. Position embeddings/RoPE modify score inputs; sinusoidal dot products oscillate rather than decrease monotonically with distance, and RoPE combines relative offsets with content rather than guaranteeing extrapolation. Sparse visibility, IO-aware kernels and cache layout change different aspects of the computation. Original CLIP should not be mislabeled a token-level cross-attention architecture.

(e) Recall + stretch · 10 min

1. Formula? $O=\operatorname{softmax}(QK^T/\sqrt{d_k}+M)V$ over legal keys.
2. Three projections? Distinct structural matching/content roles; Q/K factorization has a change-of-basis symmetry under the stated assumptions.
3. Scaling proof? Sum independent unit-variance products, then divide by their standard deviation; not a universal uniqueness theorem.
4. Shapes? Q (B,L,dk), K (B,S,dk), V (B,S,dv), scores (B,L,S), output (B,L,dv).
5. Wrong axis? Columns rather than rows normalize; reject with a rectangular case and row-sum/output oracle.
6. Softmax backward? $a\odot(g-\sum a g)$ per row, linear work in its key count.
7. Permutation property? Position-free unmasked self-attention is equivariant; fixed causal masks break arbitrary reorderings.
8. Flash insight? Tile and maintain online normalization/weighted sums; do not store the quadratic intermediate globally.
9. Different sources? Cross-attention; equal or unequal lengths are possible.
10. 32K cache? Use KV heads and explicit bytes; the worked 80-layer example is 80 GiB MHA or 10 GiB with eight KV heads, before other memory.

Explain-out-loud test

Differentiate every path to X, explain why masked future scores can leak through a denominator, and reproduce the corrected 3×2 hand trace without confusing dot product and cosine.

What comes next

Multi-head attention turns these single-head contracts into a shape- and cache-aware module. Learning-plan hub.

Primary references

• Transformer §3.2 and scaling footnote.
• PyTorch 2.9 SDPA: masks, dropout, scale and backend eligibility.
• FlashAttention, GQA: different IO/cache mechanisms.