Attention — the Capstone #

The fanciest architectural primitive in modern vision and language models, formalized in one file. If you're reading the book straight through, this is the chapter where everything you've learned clicks together and you realize there's nothing left to learn.

The cast of characters #

Scaled dot-product attention:

out = softmax((Q * K^T) / sqrt(d)) * V

where Q = X Wq, K = X Wk, V = X Wv — three dense projections of the same input X. Every piece is something we already have:

So the only genuinely new ingredient in attention is the standalone softmax VJP (softmaxHasVJP, separate from the CE-loss gradient). Once that's in hand, everything else is composition via tools we built in earlier chapters.

Structure of this file #

0. The multi-head matrix kit — column-slab independence (pdivMat_colIndep, colSlabwiseHasVJPMat) and the ternary matrix VJP HasVJPMat3; then the differentiability helpers.
1. The standalone softmax VJP is in Softmax.lean (imported here).
2. Scaled dot-product attention — SDPA as a composition.
3. Multi-head wrapper — reshape/transpose boilerplate, no new math.
4. Transformer block — LN -> MHSA -> + -> LN -> MLP -> +, pure composition.
5. Final commentary — why the taxonomy is complete.
6. Bridging ranks — patch embedding, classifier head, and the weight-tied whole-ViT witness vitFullHasVJP (vitFullHasVJP_correct).

References #

• Vaswani et al. 2017, Attention Is All You Need. https://arxiv.org/abs/1706.03762
• Dosovitskiy et al. 2021, An Image is Worth 16x16 Words (ViT). https://arxiv.org/abs/2010.11929

Per-head / per-slab column independence #

Multi-head attention applies the same per-head function to each of heads column slabs of width d_in from a Mat n (heads * d_in) input. The column-slab analog of rowwiseHasVJPMat factors that vmap-over-heads structure: each head's output depends only on its own slab of the input, so the matrix Jacobian is block-diagonal across the head axis. The general form lets each slab have its own map (colSlabApplyH); the shared-map colSlabApply is its constant-family case.

Ternary matrix VJP #

For ternary-input functions like SDPA (Q, K, V) ↦ out, package the three per-input VJPs as a single structure analogous to HasVJPMat. The backward returns the triple of per-input gradients; correctness holds independently for each input (with the others fixed).

Attention as a composition #

For a single sequence of n tokens, each with feature dim d, let X : Mat n d be the input. Attention produces out : Mat n d via:

Q = X * Wq        -- (n x d), dense projection
K = X * Wk        -- (n x d)
V = X * Wv        -- (n x d)
scores = Q * K^T   -- (n x n)
scaled = scores / sqrt(d)
weights = softmax_row(scaled)   -- softmax applied per row
out = weights * V               -- (n x d)

Because the input X is a matrix, we need matrix-level types. We work with Mat n d throughout this section (already defined in Tensor.lean).

Row-wise softmax is just "apply the 1D softmax to each row independently." Its VJP is just "apply the 1D softmax VJP to each row independently." No new derivation; the fan-out structure is trivially parallel.

The backward pass through SDPA (by hand, then compositionally) #

Working backward from d_out : Mat n d, four steps:

Step 1. Through the final matmul out = weights * V. By the dense layer VJP generalized to matrices (same derivation as denseHasVJP, just with a batch dimension):

d_V       = weights^T * d_out     -- (n x d)
d_weights = d_out * V^T           -- (n x n)

Step 2. Through the per-row softmax. Each row is independent, so we apply softmaxHasVJP row-by-row:

d_scaled_i = weights_i * (d_weights_i - <weights_i, d_weights_i> * 1)

Step 3. Through the scalar scale scaled = scores / sqrt(d). Just divide the incoming gradient by sqrt(d):

d_scores = d_scaled / sqrt(d)

Step 4. Through scores = Q * K^T. Same matrix-matmul VJP as step 1, but now Q and K both flow back:

d_Q = d_scores * K                       -- (n x d)
d_K = d_scores^T * Q                     -- (n x d)

Step 5. Three parallel dense backwards from Q, K, V back to X. Each uses denseHasVJP:

d_X_via_Q = d_Q * Wq^T
d_X_via_K = d_K * Wk^T
d_X_via_V = d_V * Wv^T

Step 6. Fan-in at X — the three paths add:

d_X = d_X_via_Q + d_X_via_K + d_X_via_V

This is biPathHasVJP from Residual.lean, applied twice (to combine three paths). The three-way fan-in is the attention backward pass at the input. Q, K, V are parallel branches reading from X, so their gradients accumulate at X.

And the parameter gradients (for W_q, W_k, W_v, W_o) are collected at each dense layer along the way — exactly as with any other dense layer in the book.

There is no novel structural move in attention. It's three dense layers, two matmuls, one row-softmax, one scale, and a three-way fan-in. Every piece has been proved. The composition is mechanical.

The backward, concretely #

This section gives:

1. Concrete definitions of sdpaBackQ, sdpaBackK, sdpaBackV transcribed from the step-by-step derivation above.
2. Correctness theorems sdpaBackQ_correct, sdpaBackK_correct, sdpaBackV_correct, stated in terms of pdivMat (the matrix-level partial derivative primitive from Tensor.lean), each proved compositionally as a vjpMatComp chain (four steps for Q).

The concrete formulas are also numerically gradient-checked in check_jacobians.py (test_sdpaBackQ/K/V) for cross-validation.

Q and K correctness via compositional SDPA forward chain #

For Q (with K, V fixed), sdpa n d · K V is the composition:

Q ↦ Q · K^T   ↦   scale * _   ↦   rowSoftmax _   ↦   _ · V

Four steps, four already-proved HasVJPMat building blocks:

1. matmulRightConstHasVJP (Mat.transpose K) — ∂(Q · K^T)/∂Q
2. scalarScaleHasVJP (sdpaScale d) — ∂(scale · scores)/∂scores
3. rowSoftmaxHasVJPMat — ∂(rowSoftmax scaled)/∂scaled
4. matmulRightConstHasVJP V — ∂(weights · V)/∂weights

Chain them with vjpMatComp thrice → a HasVJPMat for the full Q-path. Then show the chain's backward function equals sdpaBackQ pointwise (trivial — the chain's backward literally computes the same nested formula) and invoke its .correct to discharge the goal.

K case #

K enters through a transpose before the first matmul. One extra step in the chain: K ↦ K^T, then follow the Q chain (but with the matmul being "left factor constant" this time because Q is fixed and K^T is on the right).

Multi-head: parallelism over a partition #

Multi-head attention is:

1. Project X : Mat N D three ways: Q = X·Wq + bq, K = X·Wk + bk, V = X·Wv + bv.
2. Reshape each projection (N, D) → (N, heads, d_head) by slicing the feature axis.
3. Run SDPA independently on each of the heads slices.
4. Concatenate the head outputs back to (N, D).
5. Apply the output projection Y = concat · Wo + bo.

In the MLIR (emitMHSAForward): reshape (B, N, D) -> (B, N, H, D_h) transpose -> (B, H, N, D_h) [SDPA per head, using batching_dims = [0, 1]] transpose -> (B, N, H, D_h) reshape -> (B, N, D) dense projection (the "output projection" Wo)

This section proves mhsaHasVJPMat end-to-end: we define mhsaLayer concretely in Lean (Q/K/V projections → per-head slice → sdpa-per-head → concat → Wo projection), then prove its HasVJPMat via the pdivMat_colIndep + colSlabwiseHasVJPMat framework, which lifts the per-head SDPA backward over the head axis. Formula remains numerically gradient-checked in check_jacobians.py for cross-validation.

Column-stacked SDPA — the bridge from HasVJPMat3 to multi-head. #

Two facts are needed for multi-head attention:
(1) joint differentiability of `(Q, K, V) ↦ sdpa Q K V`, which doesn't
follow from the per-input `_flat_diff` lemmas; (2) the per-head
"vmap" structure, which `colSlabwiseHasVJPMat` handles for
*unary* per-slab functions but SDPA is naturally ternary.

The fix: column-stack `(Q | K | V)` into a single `Mat n (3 * d_head)`
"qkv slab", define `mhsaG : Mat n (3 * d_head) → Mat n d_head` as the
unary view of SDPA on this slab, and lift via the existing framework.

Both `mhsaG_flat_differentiable` and `mhsaGHasVJPMat` are then mechanical
composition of existing pieces: the joint `_flat_diff` factors through
`rowSoftmax_flat_differentiable` after stage-by-stage chaining, and the VJP comes
from `sdpaHasVJPMat3` plus a "column-third projection" argument that
matches the `(c : Fin 3)` index of the slab to the Q/K/V partial. 

Column-stacked SDPA VJP #

`HasVJPMat (mhsaG n d)`: the backward column-stacks
`(sdpaBackQ, sdpaBackK, sdpaBackV)` according to the c-third
of the slab column index. Correctness reduces to `sdpaHasVJPMat3`
after observing that perturbing the c-th third of the slab only
perturbs the c-th input of SDPA. 

Per-token liftings (theorems) #

Every per-token operation in a transformer (LN, dense, GELU) lifts from HasVJP on Vec D to HasVJPMat on Mat N D via the single helper rowwiseHasVJPMat (Tensor.lean).

A transformer encoder block #

From MlirCodegen.emitTransformerBlockForward:

block(x) = h1 + MLP(LN2(h1))       where h1 = x + MHSA(LN1(x))

Expanding:

h1 = x + MHSA(LN1(x))       -- attention sub-layer with residual
out = h1 + MLP(LN2(h1))     -- MLP sub-layer with residual

where MLP(z) = dense(Wfc2, bfc2, gelu(dense(Wfc1, bfc1, z))).

Every piece is now a HasVJPMat on Mat N D:

• LN1, LN2 — layerNormPerTokenHasVJPMat (theorem via rowwiseHasVJPMat)
• MHSA — mhsaHasVJPMat (bundled HasVJPMat def)
• MLP — two densePerTokenHasVJPMat + one geluPerTokenHasVJPMat, glued with vjpMatComp
• + residuals — biPathMatHasVJP (theorem, Tensor.lean) with identity

The transformer block theorem below glues these with vjpMatComp and biPathMatHasVJP.

Stacking transformer blocks #

ViT-Tiny has 12 transformer blocks; ViT-Base has 12, ViT-Large has 24. The stack is just k-fold composition of individual blocks. By vjpMatComp and induction on k, if each block has a HasVJPMat then so does the stack — for any depth.

The formal tower transformerTower uses a single shared parameter tuple across all k blocks (Nat.rec over one tuple), with scalar-affine LayerNorms; in practice every block has its own weights and a per-feature LN. Per-block parameters are not formalized in this file.

ViT body: tower + final LN #

vitBody is the ViT backbone operating on a single (N, D) sequence, after the patch embedding produced a Mat N D input and before the classifier head slices the CLS token and runs dense+softmax CE.

patch_embed(X : Tensor3 ic h w) : Mat N D   ← outside Mat-land
vitBody(M : Mat N D) : Mat N D             ← the backbone (this file)
classifier(M) = dense(W_cls, b_cls, M[0])   ← Mat → Vec, then softmax CE loss

The backbone is finalLN ∘ transformerTower. Both sides are Mat N D, so vjpMatComp glues the two VJPs.

The patch-embedding and classifier-head steps exit Mat-land (they change type to/from Tensor3 and Vec respectively), so they don't fit in the uniform HasVJPMat frame. They are bridged below (§ Bridging ranks): the body is flattened with HasVJPMat.toHasVJP and composed with patchEmbedFlatHasVJP and classifierFlatHasVJP in vitFullHasVJP.