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:

Piece	Chapter	VJP move
Q = X Wq	MLP.lean	dense backward
K = X Wk	MLP.lean	dense backward
V = X Wv	MLP.lean	dense backward
Q * K^T	(matmul = dense)	chain rule
/ sqrt(d)	(scalar)	chain rule + scale
softmax(...)	this file	closed-form collapse
... * V	(matmul = dense)	chain rule
three-way fan-in at X	Residual.lean	biPath_has_vjp

So the only genuinely new ingredient in attention is the standalone softmax VJP (previously we only had it bundled inside CE loss). Once that's in hand, everything else is composition via tools we built in earlier chapters.

Structure of this file

1. Standalone softmax VJP — the last closed-form trick.
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.

theorem Proofs.matmul_right_const_flat_diff {m p q : ℕ} (D : Mat p q) : Differentiable ℝ fun (v : Vec (m * p)) => ((Mat.unflatten v).mul D).flatten

theorem Proofs.matmul_left_const_flat_diff {m p q : ℕ} (C : Mat m p) : Differentiable ℝ fun (v : Vec (p * q)) => (C.mul (Mat.unflatten v)).flatten

theorem Proofs.scalarScale_flat_diff {m n : ℕ} (s : ℝ) : Differentiable ℝ fun (v : Vec (m * n)) => Mat.flatten fun (r : Fin m) (c : Fin n) => s * Mat.unflatten v r c

theorem Proofs.transpose_flat_diff {m n : ℕ} : Differentiable ℝ fun (v : Vec (m * n)) => (Mat.unflatten v).transpose.flatten

theorem Proofs.dense_per_token_flat_diff {N inD outD : ℕ} (W : Mat inD outD) (b : Vec outD) : Differentiable ℝ fun (v : Vec (N * inD)) => Mat.flatten ((fun (X : Mat N inD) (n : Fin N) => dense W b (X n)) (Mat.unflatten v))

Differentiability of the flattened per-token dense map. fun X => fun n => dense W b (X n) is linear in X, so the flattened version is Differentiable everywhere.

theorem Proofs.gelu_per_token_flat_diff (N D : ℕ) : Differentiable ℝ fun (v : Vec (N * D)) => Mat.flatten ((fun (X : Mat N D) (n : Fin N) => gelu D (X n)) (Mat.unflatten v))

Differentiability of the flattened per-token GELU map. geluScalar = 0.5 · x · (1 + tanh(√(2/π)(x + 0.044715·x³))). With Real.differentiable_tanh available to fun_prop, the proof discharges automatically.

theorem Proofs.dense_diff {m n : ℕ} (W : Mat m n) (b : Vec n) : Differentiable ℝ (dense W b)

Differentiability of dense W b as a function of the input.

theorem Proofs.softmax_diff (c : ℕ) : Differentiable ℝ (softmax c)

Differentiability of softmax c — same recipe as rowSoftmax_flat_diff, but on the unflattened vector.

theorem Proofs.layerNorm_diff (D : ℕ) (ε γ β : ℝ) (hε : 0 < ε) : Differentiable ℝ (layerNormForward D ε γ β)

Differentiability of layerNormForward D ε γ β. layerNormForward = bnForward = bnAffine ∘ bnNormalize (definitionally), where the chain is differentiable when ε > 0.

theorem Proofs.layerNorm_per_token_flat_diff (N D : ℕ) (ε γ β : ℝ) (hε : 0 < ε) : Differentiable ℝ fun (v : Vec (N * D)) => Mat.flatten ((fun (X : Mat N D) (n : Fin N) => layerNormForward D ε γ β (X n)) (Mat.unflatten v))

Differentiability of the flattened per-token LayerNorm map. Now a theorem (was an axiom): each output coord projects through a row-projection CLM into layerNorm_diff at that row.

theorem Proofs.identity_mat_flat_diff (a b : ℕ) : Differentiable ℝ fun (v : Vec (a * b)) => ((fun (X : Mat a b) => X) (Mat.unflatten v)).flatten

Differentiability of the flattened identity matrix map. Mat.flatten ∘ id ∘ Mat.unflatten = id on Vec (a*b).

The softmax Jacobian

For p = softmax(z) with p_j = exp(z_j) / sum_k exp(z_k), the quotient rule gives:

dp_j/dz_i = p_j * (delta_{ij} - p_i)

This is the famous "diag minus outer product" form:

J = diag(p) - p * p^T

Dense (every output depends on every input), but rank-1 correction to a diagonal — which means the VJP has a closed-form collapse, just like BatchNorm did.

theorem Proofs.pdiv_softmax (c : ℕ) (z : Vec c) (i j : Fin c) : pdiv (softmax c) z i j = softmax c z j * ((if i = j then 1 else 0) - softmax c z i)

Partial derivative of softmax (quotient rule on the exponentials).

d(softmax(z))_j/dz_i = softmax(z)_j * (delta_{ij} - softmax(z)_i)

Proved (was an axiom). The j-th coord of softmax c z is Real.exp (z j) / S with S := Σ_k Real.exp (z k) > 0, so the j-th output coord function z' ↦ exp(z' j) * (Σ_k exp(z' k))⁻¹ has HasFDerivAt derivative built from HasFDerivAt.exp, HasFDerivAt.fun_sum, (hasDerivAt_inv ·).comp_hasFDerivAt, and HasFDerivAt.mul. Evaluating that CLM at basisVec i and collapsing Σ_k exp(z k) · δ_{ki} = exp(z i) gives the formula.

noncomputable def Proofs.softmax_has_vjp (c : ℕ) : HasVJP (softmax c)

Softmax VJP — the closed-form collapse.

back(z, dy)_i = p_i * (dy_i - <p, dy>)

where p = softmax(z) and <p, dy> = sum_j p_j * dy_j is one scalar.

Read this carefully. The naive VJP would be: dz_i = sum_j J_{ji} * dy_j = sum_j (p_j * (delta_{ij} - p_i)) * dy_j

That's O(c) per entry, O(c^2) total. But expanding: dz_i = p_i * dy_i - p_i * sum_j p_j * dy_j = p_i * (dy_i - <p, dy>)

The rank-1 correction lets you precompute one scalar (<p, dy>) and apply it to every entry. Total work: O(c). Same optimization pattern as BN (one reduction + a broadcast) and max-pool (one comparison + a select).

Interpretation. Softmax outputs a probability distribution. Its backward subtracts the "weighted average of the incoming gradient under that distribution" from each entry, then scales by the entry's probability. Entries with low probability get small gradients (because the softmax flattened them in the forward); entries with high probability get gradients proportional to how much they deviate from the weighted-average cotangent.

This is the one place where "softmax means softly select one thing" maps directly to "softmax backward selectively amplifies the gradient for the winning class."

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.softmaxCE_grad (c : ℕ) (logits : Vec c) (label j : Fin c) : pdiv (fun (z : Vec c) (x : Fin 1) => crossEntropy c z label) logits j 0 = softmax c logits j - oneHot c label j

Softmax cross-entropy scalar gradient — proved (was an axiom in MLP.lean; relocated here to use pdiv_softmax).

∂(-log softmax(z)[label])/∂z_j = softmax(z)_j - onehot(label)_j

Stated using pdiv on a Vec 1-valued wrapper (cross-entropy is naturally scalar, but pdiv is defined for Vec → Vec; we just take the only output index). Proof: fderiv_apply extracts the only coord, then HasFDerivAt.log (with softmax z label > 0) composed with softmax_diff gives the derivative of the inner Real.log. Negating and evaluating at basisVec j reduces via pdiv_softmax to the expected formula.

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.

noncomputable def Proofs.rowSoftmax {m n : ℕ} (A : Mat m n) : Mat m n

Row-wise softmax of a matrix.

Equations

• Proofs.rowSoftmax A i = Proofs.softmax n (A i)

theorem Proofs.rowSoftmax_flat_diff (m n : ℕ) : Differentiable ℝ fun (v : Vec (m * n)) => (rowSoftmax (Mat.unflatten v)).flatten

Smoothness of rowSoftmax — proved from Mathlib calculus (planning/archive/VJP.md follow-up B).

rowSoftmax M r c = exp(M r c) / Σⱼ exp(M r j). The denominator is everywhere positive (sum of Real.exp_pos terms over a nonempty index set when n ≥ 1), so the function is C^∞ via Real.exp, Finset.sum, and div with positivity. The n = 0 case is trivial because Vec (m * 0) = Vec 0 is 0-dimensional.

noncomputable def Proofs.rowSoftmax_has_vjp_mat {m n : ℕ} : HasVJPMat fun (A : Mat m n) (r : Fin m) => softmax n (A r)

Row-wise softmax VJP — proved, no sorry.

Rows are independent, so the Jacobian is block-diagonal with the standalone softmax Jacobian in each block. The backward just applies softmax_has_vjp per row.

Equations

• Proofs.rowSoftmax_has_vjp_mat = { backward := fun (A dY : Proofs.Mat m n) (r : Fin m) (c : Fin n) => (Proofs.softmax_has_vjp n).backward (A r) (dY r) c, correct := ⋯ }

noncomputable def Proofs.rowSoftmax_has_vjp_mat' (m n : ℕ) : HasVJPMat rowSoftmax

Alias so rowSoftmax_has_vjp_mat types against the actual rowSoftmax definition (definitionally equal, but lets Lean unify on the name).

Equations

• Proofs.rowSoftmax_has_vjp_mat' m n = Proofs.rowSoftmax_has_vjp_mat

noncomputable def Proofs.sdpa (n d : ℕ) (Q K V : Mat n d) : Mat n d

Scaled dot-product attention, for a single sequence and a single head. Q K V : Mat n d.

sdpa Q K V = softmax_row(Q * K^T / sqrt(d)) * V

MLIR (emitMHSAForward, lines 754-781): %mh_sc = dot_general %mh_q, %mh_k, contracting_dims = [3] x [3] %mh_ss = multiply %mh_sc, broadcast(1/sqrt(d)) %mh_sm = softmax(%mh_ss) -- via reduce max, shift, exp, reduce sum, divide %mh_av = dot_general %mh_sm, %mh_v, contracting_dims = [3] x [2]

Equations

• Proofs.sdpa n d Q K V = (Proofs.rowSoftmax fun (i j : Fin n) => 1 / √↑d * Q.mul K.transpose i j).mul V

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 dense_has_vjp, 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 softmax_has_vjp 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 dense_has_vjp:

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 biPath_has_vjp 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

Earlier drafts of this section ended with a single axiom sdpa_has_vjp whose type was just (... functions) × (... functions) × (... functions). That was vacuous as a correctness claim — a triple of zero functions satisfies it. The current state is:

1. Concrete definitions of sdpa_back_Q, sdpa_back_K, sdpa_back_V transcribed from the step-by-step derivation above.
2. Honest correctness theorems stated in terms of pdivMat (the matrix-level partial derivative primitive from Tensor.lean), proved compositionally via the four-step vjpMat_comp chain described in §Q-correctness below.

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

noncomputable def Proofs.sdpa_scale (d : ℕ) : ℝ

1 / sqrt(d), the SDPA scale factor.

Equations

• Proofs.sdpa_scale d = 1 / √↑d

noncomputable def Proofs.sdpa_weights (n d : ℕ) (Q K : Mat n d) : Mat n n

Softmax-weights under the SDPA scale, reused by all three backwards.

Equations

• Proofs.sdpa_weights n d Q K = Proofs.rowSoftmax fun (i j : Fin n) => Proofs.sdpa_scale d * Q.mul K.transpose i j

noncomputable def Proofs.sdpa_dWeights {n d : ℕ} (V dOut : Mat n d) : Mat n n

Gradient flowing into weights from the final matmul out = weights · V.

Equations

• Proofs.sdpa_dWeights V dOut = dOut.mul V.transpose

noncomputable def Proofs.sdpa_dScaled (n d : ℕ) (Q K V dOut : Mat n d) : Mat n n

Per-row softmax VJP: p_i * (dw_i - <p_i, dw_i>).

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.sdpa_dScores (n d : ℕ) (Q K V dOut : Mat n d) : Mat n n

Gradient w.r.t. the pre-softmax scores, after undoing the / sqrt(d) scale.

Equations

• Proofs.sdpa_dScores n d Q K V dOut i j = Proofs.sdpa_scale d * Proofs.sdpa_dScaled n d Q K V dOut i j

noncomputable def Proofs.sdpa_back_Q (n d : ℕ) (Q K V dOut : Mat n d) : Mat n d

Backward w.r.t. Q: dQ = dScores · K.

Equations

• Proofs.sdpa_back_Q n d Q K V dOut = (Proofs.sdpa_dScores n d Q K V dOut).mul K

noncomputable def Proofs.sdpa_back_K (n d : ℕ) (Q K V dOut : Mat n d) : Mat n d

Backward w.r.t. K: dK = dScores^T · Q.

Equations

• Proofs.sdpa_back_K n d Q K V dOut = (Proofs.sdpa_dScores n d Q K V dOut).transpose.mul Q

noncomputable def Proofs.sdpa_back_V (n d : ℕ) (Q K _V dOut : Mat n d) : Mat n d

Backward w.r.t. V: dV = weights^T · dOut. (V does not appear on the RHS: V's gradient flows only through the final matmul, not through weights.)

Equations

• Proofs.sdpa_back_V n d Q K _V dOut = (Proofs.sdpa_weights n d Q K).transpose.mul dOut

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. matmul_right_const_has_vjp (Mat.transpose K) — ∂(Q · K^T)/∂Q
2. scalarScale_has_vjp (sdpa_scale d) — ∂(scale · scores)/∂scores
3. rowSoftmax_has_vjp_mat — ∂(rowSoftmax scaled)/∂scaled
4. matmul_right_const_has_vjp V — ∂(weights · V)/∂weights

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

noncomputable def Proofs.sdpa_Q_chain (n d : ℕ) (K V : Mat n d) : Mat n d → Mat n d

Explicit 4-composition forward for SDPA, varying Q with K, V fixed.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.sdpa_Q_chain_eq (n d : ℕ) (Q K V : Mat n d) : sdpa_Q_chain n d K V Q = sdpa n d Q K V

noncomputable def Proofs.sdpa_Q_chain_has_vjp (n d : ℕ) (K V : Mat n d) : HasVJPMat (sdpa_Q_chain n d K V)

HasVJPMat for the chain — built by nesting vjpMat_comp thrice.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.sdpa_back_Q_correct (n d : ℕ) (Q K V dOut : Mat n d) (i : Fin n) (j : Fin d) : sdpa_back_Q n d Q K V dOut i j = ∑ k : Fin n, ∑ l : Fin d, pdivMat (fun (Q' : Mat n d) => sdpa n d Q' K V) Q i j k l * dOut k l

Correctness of sdpa_back_Q — proved, no sorry.

Two moves: (1) replace fun Q' => sdpa n d Q' K V by the chain via sdpa_Q_chain_eq; (2) apply the chain's .correct and verify that the chain's backward reduces to sdpa_back_Q (pure unfolding).

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).

noncomputable def Proofs.sdpa_K_chain (n d : ℕ) (Q V : Mat n d) : Mat n d → Mat n d

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.sdpa_K_chain_eq (n d : ℕ) (Q K V : Mat n d) : sdpa_K_chain n d Q V K = sdpa n d Q K V

noncomputable def Proofs.sdpa_K_chain_has_vjp (n d : ℕ) (Q V : Mat n d) : HasVJPMat (sdpa_K_chain n d Q V)

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.sdpa_back_K_correct (n d : ℕ) (Q K V dOut : Mat n d) (i : Fin n) (j : Fin d) : sdpa_back_K n d Q K V dOut i j = ∑ k : Fin n, ∑ l : Fin d, pdivMat (fun (K' : Mat n d) => sdpa n d Q K' V) K i j k l * dOut k l

Correctness of sdpa_back_K — proved, no sorry.

Same shape as Q, but the chain goes through a leading transpose step. The resulting backward computes ∑ k, Q k j * dScores k i whereas sdpa_back_K is Mat.mul (Mat.transpose dScores) Q, which expands to ∑ k, dScores k i * Q k j. Equal by mul_comm at the summand level.

theorem Proofs.sdpa_eq_mul_weights (n d : ℕ) (Q K V : Mat n d) : sdpa n d Q K V = (sdpa_weights n d Q K).mul V

The final matmul in SDPA: for fixed Q, K, the function V' ↦ sdpa Q K V' is V' ↦ W · V' where W = sdpa_weights Q K. Pure rewrite; definitional.

theorem Proofs.sdpa_back_V_correct (n d : ℕ) (Q K V dOut : Mat n d) (i : Fin n) (j : Fin d) : sdpa_back_V n d Q K V dOut i j = ∑ k : Fin n, ∑ l : Fin d, pdivMat (fun (V' : Mat n d) => sdpa n d Q K V') V i j k l * dOut k l

Correctness of sdpa_back_V — proved, no sorry.

The V-path is the simplest case: V' only enters through the final matmul out = weights · V'. So fun V' => sdpa n d Q K V' is just fun V' => Mat.mul W V' where W is fixed (= sdpa_weights n d Q K), and the VJP comes directly from matmul_left_const_has_vjp.

noncomputable def Proofs.sdpa_has_vjp_mat3 (n d : ℕ) : HasVJPMat3 (sdpa n d)

Bundled SDPA ternary VJP. Packages sdpa_back_{Q, K, V}_correct into a single HasVJPMat3 instance. The backward triple (sdpa_back_Q, sdpa_back_K, sdpa_back_V) gives per-input gradients; correctness is the three existing per-input theorems in one structure.

Equations

• One or more equations did not get rendered due to their size.

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)

Earlier this section just narrated "no new VJP math" and moved on. The current state proves mhsa_has_vjp_mat end-to-end: we define mhsa_layer concretely in Lean (Q/K/V projections → per-head slice → sdpa-per-head → concat → Wo projection), then prove its HasVJPMat via the new pdivMat_colIndep + colSlabwise_has_vjp_mat framework (Phase 3, Apr 2026), which lifts the per-head SDPA backward over the head axis. Formula remains numerically gradient-checked in check_jacobians.py for cross-validation.

noncomputable def Proofs.mhsa_layer (N heads d_head : ℕ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (X : Mat N (heads * d_head)) : Mat N (heads * d_head)

Multi-head SDPA on a single sequence: Mat N (heads·d_head) → Mat N (heads·d_head).

Concretely defined (not opaque):

1. Q, K, V projections (each a per-token dense with its own Wq/Wk/Wv).
2. For each head h : Fin heads, extract the (N, d_head) slice of Q/K/V by indexing finProdFinEquiv (h, k) in the combined axis.
3. Run sdpa on each slice.
4. Concatenate the head outputs back along the feature axis.
5. Output projection Wo · concat + bo (per-token dense).

The bundled VJP theorem below packages the correctness of this whole thing — composing dense Jacobians, the per-head SDPA jacobians (we already proved sdpa_back_{Q,K,V}_correct), and the reshape/unreshape pdiv_reindex facts, with the per-head independence handled by Phase 3's column-stacking framework (pdivMat_colIndep + colSlabwise_has_vjp_mat).

Equations

• One or more equations did not get rendered due to their size.

Phase 3: Column-stacked SDPA — the bridge from HasVJPMat3 to multi-head.

The two `mhsa_*` former axioms below were the project floor for two reasons:
(1) joint differentiability of `(Q, K, V) ↦ sdpa Q K V`, which doesn't
follow from the existing per-input `_flat_diff` lemmas; (2) the per-head
"vmap" structure, which `colSlabwise_has_vjp_mat` (Phase 1) 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 `mhsa_g : 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 `mhsa_g_flat_diff` and `mhsa_g_has_vjp_mat` are then mechanical
composition of existing pieces: the joint `_flat_diff` factors through
`rowSoftmax_flat_diff` after stage-by-stage chaining, and the VJP comes
from `sdpa_has_vjp_mat3` plus a "column-third projection" argument that
matches the `(c : Fin 3)` index of the slab to the Q/K/V partial. 

noncomputable def Proofs.mhsa_g (n d : ℕ) (slab : Mat n (3 * d)) : Mat n d

Column-stacked SDPA: takes a slab Mat n (3 * d_head) whose columns encode (c : Fin 3, j : Fin d_head) via finProdFinEquiv, with c = 0 being the Q-third, c = 1 the K-third, c = 2 the V-third. Returns sdpa applied to those three thirds.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.mhsa_pre_weights (n d : ℕ) (slab : Mat n (3 * d)) : Mat n n

Pre-softmax matrix in mhsa_g: scale · Q · K^T as a function of slab. Each entry is a polynomial in the slab's coords (linear projections times each other), so fun_prop discharges flat-diff after unfolding.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.mhsa_pre_weights_flat_diff (n d : ℕ) : Differentiable ℝ fun (v : Vec (n * (3 * d))) => (mhsa_pre_weights n d (Mat.unflatten v)).flatten

noncomputable def Proofs.mhsa_weights (n d : ℕ) (slab : Mat n (3 * d)) : Mat n n

Post-softmax weights in mhsa_g: rowSoftmax(scale · Q · K^T).

Equations

• Proofs.mhsa_weights n d slab = Proofs.rowSoftmax (Proofs.mhsa_pre_weights n d slab)

theorem Proofs.mhsa_weights_flat_diff (n d : ℕ) : Differentiable ℝ fun (v : Vec (n * (3 * d))) => (mhsa_weights n d (Mat.unflatten v)).flatten

theorem Proofs.mhsa_g_flat_diff (n d : ℕ) : Differentiable ℝ fun (v : Vec (n * (3 * d))) => (mhsa_g n d (Mat.unflatten v)).flatten

Joint flat-diff of column-stacked SDPA.

The blocker for Phase 3 (per planning/archive/mhsa.md): joint diff in (Q, K, V) doesn't follow from the existing per-input _flat_diff lemmas. Here we prove it by treating the qkv-slab as the variable, factoring SDPA as Mat.mul ∘ rowSoftmax ∘ scaled-matmul, and chaining: pre-softmax is fun_prop-able (polynomial in slab coords), rowSoftmax composes via rowSoftmax_flat_diff, final matmul-with-V splits per output coord into a sum of products of two diff scalars.

Column-stacked SDPA VJP

`HasVJPMat (mhsa_g n d)`: the backward column-stacks
`(sdpa_back_Q, sdpa_back_K, sdpa_back_V)` according to the c-third
of the slab column index. Correctness reduces to `sdpa_has_vjp_mat3`
after observing that perturbing the c-th third of the slab only
perturbs the c-th input of SDPA. 

noncomputable def Proofs.mhsa_proj_c {n d : ℕ} (c : Fin 3) (slab : Mat n (3 * d)) : Mat n d

Column projection slab ↦ slab^[c] for a fixed c : Fin 3. Linear, so its flat form is a reindexCLM.

Equations

• Proofs.mhsa_proj_c c slab r j = slab r (finProdFinEquiv (c, j))

theorem Proofs.mhsa_proj_c_flat_diff (n d : ℕ) (c : Fin 3) : Differentiable ℝ fun (v : Vec (n * (3 * d))) => (mhsa_proj_c c (Mat.unflatten v)).flatten

noncomputable def Proofs.mhsa_proj_c_CLM (n d : ℕ) (c : Fin 3) : Vec (n * (3 * d)) →L[ℝ] Vec (n * d)

The flat form of the column projection mhsa_proj_c c is exactly the reindexCLM σ_c idx = fPF((decode idx).1, fPF(c, (decode idx).2)).

Equations

• Proofs.mhsa_proj_c_CLM n d c = Proofs.reindexCLM fun (idx : Fin (n * d)) => finProdFinEquiv ((finProdFinEquiv.symm idx).1, finProdFinEquiv (c, (finProdFinEquiv.symm idx).2))

theorem Proofs.mhsa_proj_c_eq_CLM (n d : ℕ) (c : Fin 3) (v : Vec (n * (3 * d))) : (mhsa_proj_c c (Mat.unflatten v)).flatten = (mhsa_proj_c_CLM n d c) v

noncomputable def Proofs.mhsa_lift_c_CLM (n d : ℕ) (c : Fin 3) : Vec (n * d) →L[ℝ] Vec (n * (3 * d))

"Lift to slab third c": embeds Vec (n * d) into Vec (n * (3 * d)) by placing u in the c-th column third and zero elsewhere. Linear, hence a CLM. The dual of mhsa_proj_c_CLM. Constructed from per-coord CLMs via ContinuousLinearMap.pi: each output coord is either a projection (if the index is in the c-third) or zero.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.mhsa_lift_c_CLM_apply (n d : ℕ) (c : Fin 3) (u : Vec (n * d)) (idx : Fin (n * (3 * d))) : (mhsa_lift_c_CLM n d c) u idx = if (finProdFinEquiv.symm (finProdFinEquiv.symm idx).2).1 = c then u (finProdFinEquiv ((finProdFinEquiv.symm idx).1, (finProdFinEquiv.symm (finProdFinEquiv.symm idx).2).2)) else 0

noncomputable def Proofs.mhsa_embed_c (n d : ℕ) (c : Fin 3) (slab : Mat n (3 * d)) (u : Vec (n * d)) : Vec (n * (3 * d))

"Embed Q' into slab at the c-th third, keep other thirds at slab's values." Affine function: mhsa_lift_c_CLM c · u + (slab with c-th third zeroed).

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.mhsa_embed_c_eq (n d : ℕ) (c : Fin 3) (slab : Mat n (3 * d)) (u : Vec (n * d)) : mhsa_embed_c n d c slab u = (mhsa_lift_c_CLM n d c) u + fun (idx : Fin (n * (3 * d))) => if (finProdFinEquiv.symm (finProdFinEquiv.symm idx).2).1 = c then 0 else slab.flatten idx

theorem Proofs.mhsa_embed_c_hasFDerivAt (n d : ℕ) (c : Fin 3) (slab : Mat n (3 * d)) (u₀ : Vec (n * d)) : HasFDerivAt (mhsa_embed_c n d c slab) (mhsa_lift_c_CLM n d c) u₀

theorem Proofs.mhsa_g_comp_embed (n d : ℕ) (c : Fin 3) (slab : Mat n (3 * d)) (u : Vec (n * d)) : (mhsa_g n d (Mat.unflatten (mhsa_embed_c n d c slab u))).flatten = if c = 0 then (sdpa n d (Mat.unflatten u) (mhsa_proj_c 1 slab) (mhsa_proj_c 2 slab)).flatten else if c = 1 then (sdpa n d (mhsa_proj_c 0 slab) (Mat.unflatten u) (mhsa_proj_c 2 slab)).flatten else (sdpa n d (mhsa_proj_c 0 slab) (mhsa_proj_c 1 slab) (Mat.unflatten u)).flatten

The composition mhsa_g ∘ mhsa_embed_c c slab equals "SDPA with the c-th argument variable, the other two fixed at slab's projections". This is the freezing identity.

theorem Proofs.mhsa_embed_c_at_proj (n d : ℕ) (c : Fin 3) (slab : Mat n (3 * d)) : mhsa_embed_c n d c slab (mhsa_proj_c c slab).flatten = slab.flatten

theorem Proofs.pdivMat_mhsa_g_split_chain (n d : ℕ) (slab : Mat n (3 * d)) (c : Fin 3) (freeze_fn : Mat n d → Mat n d) (h_g_freeze_eq : ∀ (u : Vec (n * d)), (mhsa_g n d (Mat.unflatten (mhsa_embed_c n d c slab u))).flatten = (freeze_fn (Mat.unflatten u)).flatten) : fderiv ℝ (fun (v : Vec (n * (3 * d))) => (mhsa_g n d (Mat.unflatten v)).flatten) slab.flatten ∘SL mhsa_lift_c_CLM n d c = fderiv ℝ (fun (u : Vec (n * d)) => (freeze_fn (Mat.unflatten u)).flatten) (mhsa_proj_c c slab).flatten

Helper for pdivMat_mhsa_g_split (per-c chain rule). For each c : Fin 3, the chain rule gives: fderiv flat_g flat_slab ∘L mhsa_lift_c_CLM = fderiv flat_freeze_c flat_proj_c_slab. Used in pdivMat_mhsa_g_split after the basis-vector lift identity.

theorem Proofs.pdivMat_mhsa_g_split (n d : ℕ) (slab : Mat n (3 * d)) (i : Fin n) (c : Fin 3) (j : Fin d) (k : Fin n) (l : Fin d) : pdivMat (mhsa_g n d) slab i (finProdFinEquiv (c, j)) k l = if c = 0 then pdivMat (fun (Q' : Mat n d) => sdpa n d Q' (mhsa_proj_c 1 slab) (mhsa_proj_c 2 slab)) (mhsa_proj_c 0 slab) i j k l else if c = 1 then pdivMat (fun (K' : Mat n d) => sdpa n d (mhsa_proj_c 0 slab) K' (mhsa_proj_c 2 slab)) (mhsa_proj_c 1 slab) i j k l else pdivMat (fun (V' : Mat n d) => sdpa n d (mhsa_proj_c 0 slab) (mhsa_proj_c 1 slab) V') (mhsa_proj_c 2 slab) i j k l

pdivMat of mhsa_g splits per-c into the corresponding pdivMat of SDPA against its c-th argument. The freezing lemma: changes in the c-th column third of the slab only perturb the c-th input of SDPA. Proved via the chain rule mhsa_g ∘ mhsa_embed_c = freeze_c.

noncomputable def Proofs.mhsa_g_has_vjp_mat (n d : ℕ) : HasVJPMat (mhsa_g n d)

HasVJPMat for column-stacked SDPA. Backward column-stacks the three sdpa_back_* outputs by their c : Fin 3 slot. Correctness comes from pdivMat_mhsa_g_split (case-splits on c into the corresponding one-input SDPA pdivMat) and sdpa_has_vjp_mat3.correct_*.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.colSlabApply_flat_diff {n heads d_in d_out : ℕ} (g : Mat n d_in → Mat n d_out) (hg_diff : Differentiable ℝ fun (v : Vec (n * d_in)) => (g (Mat.unflatten v)).flatten) : Differentiable ℝ fun (v : Vec (n * (heads * d_in))) => (colSlabApply g (Mat.unflatten v)).flatten

Flat-diff for colSlabApply g: each output coord is (g (slab h ·)) [n, j_out], factoring through the slab-projection CLM (linear) and g (flat-diff).

noncomputable def Proofs.mhsa_qkv_W (heads d_head : ℕ) (Wq Wk Wv : Mat (heads * d_head) (heads * d_head)) : Mat (heads * d_head) (heads * (3 * d_head))

Combined Q/K/V weight matrix: stack Wq | Wk | Wv with the per-head interleave layout. Output column (h, c, j) ↦ (Wq | Wk | Wv)[k, fPF(h, j)] based on c : Fin 3. Used to express the three Q/K/V projections as a single per-token dense, enabling clean composition with colSlabApply mhsa_g.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.mhsa_qkv_b (heads d_head : ℕ) (bq bk bv : Vec (heads * d_head)) : Vec (heads * (3 * d_head))

Equations

• One or more equations did not get rendered due to their size.

@[simp]

theorem Proofs.mhsa_qkv_W_eq0 (heads d_head : ℕ) (Wq Wk Wv : Mat (heads * d_head) (heads * d_head)) (k : Fin (heads * d_head)) (h : Fin heads) (j : Fin d_head) : mhsa_qkv_W heads d_head Wq Wk Wv k (finProdFinEquiv (h, finProdFinEquiv (0, j))) = Wq k (finProdFinEquiv (h, j))

@[simp]

theorem Proofs.mhsa_qkv_W_eq1 (heads d_head : ℕ) (Wq Wk Wv : Mat (heads * d_head) (heads * d_head)) (k : Fin (heads * d_head)) (h : Fin heads) (j : Fin d_head) : mhsa_qkv_W heads d_head Wq Wk Wv k (finProdFinEquiv (h, finProdFinEquiv (1, j))) = Wk k (finProdFinEquiv (h, j))

@[simp]

theorem Proofs.mhsa_qkv_W_eq2 (heads d_head : ℕ) (Wq Wk Wv : Mat (heads * d_head) (heads * d_head)) (k : Fin (heads * d_head)) (h : Fin heads) (j : Fin d_head) : mhsa_qkv_W heads d_head Wq Wk Wv k (finProdFinEquiv (h, finProdFinEquiv (2, j))) = Wv k (finProdFinEquiv (h, j))

@[simp]

theorem Proofs.mhsa_qkv_b_eq0 (heads d_head : ℕ) (bq bk bv : Vec (heads * d_head)) (h : Fin heads) (j : Fin d_head) : mhsa_qkv_b heads d_head bq bk bv (finProdFinEquiv (h, finProdFinEquiv (0, j))) = bq (finProdFinEquiv (h, j))

@[simp]

theorem Proofs.mhsa_qkv_b_eq1 (heads d_head : ℕ) (bq bk bv : Vec (heads * d_head)) (h : Fin heads) (j : Fin d_head) : mhsa_qkv_b heads d_head bq bk bv (finProdFinEquiv (h, finProdFinEquiv (1, j))) = bk (finProdFinEquiv (h, j))

@[simp]

theorem Proofs.mhsa_qkv_b_eq2 (heads d_head : ℕ) (bq bk bv : Vec (heads * d_head)) (h : Fin heads) (j : Fin d_head) : mhsa_qkv_b heads d_head bq bk bv (finProdFinEquiv (h, finProdFinEquiv (2, j))) = bv (finProdFinEquiv (h, j))

theorem Proofs.mhsa_layer_eq_compose (N heads d_head : ℕ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (X : Mat N (heads * d_head)) : mhsa_layer N heads d_head Wq Wk Wv Wo bq bk bv bo X = (fun (M : Mat N (heads * d_head)) (n : Fin N) => dense Wo bo (M n)) (colSlabApply (mhsa_g N d_head) ((fun (X' : Mat N (heads * d_head)) (n : Fin N) => dense (mhsa_qkv_W heads d_head Wq Wk Wv) (mhsa_qkv_b heads d_head bq bk bv) (X' n)) X))

The mhsa_layer factorization: it equals output_dense ∘ colSlabApply mhsa_g ∘ qkv_stack_dense.

All three pieces have HasVJPMat and flat-diff:

• qkv_stack_dense uses mhsa_qkv_W, mhsa_qkv_b as a single per-token dense.
• colSlabApply mhsa_g lifts mhsa_g_has_vjp_mat per-head.
• output_dense is the standard per-token dense for Wo, bo.

noncomputable def Proofs.mhsa_composed_has_vjp_mat (N heads d_head : ℕ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : HasVJPMat ((fun (M : Mat N (heads * d_head)) (n : Fin N) => dense Wo bo (M n)) ∘ colSlabApply (mhsa_g N d_head) ∘ fun (X' : Mat N (heads * d_head)) (n : Fin N) => dense (mhsa_qkv_W heads d_head Wq Wk Wv) (mhsa_qkv_b heads d_head bq bk bv) (X' n))

The composed MHSA VJP — Wo-dense ∘ colSlabApply mhsa_g ∘ qkv-dense, stated on the explicit composition (no mhsa_layer_eq_compose transport). This is the substantive witness; mhsa_has_vjp_mat below re-types it to mhsa_layer with the cast confined to the correct field.

Why the split (kernel-cost lesson, 2026-07): the previous mhsa_has_vjp_mat := by rw [show mhsa_layer = …]; exact vjpMat_comp … made the constant's VALUE an Eq.mpr cast around the structure. Any kernel defeq that whnf'd (mhsa_has_vjp_mat …).backward had to replay the whole mhsa_layer rewrite — ~200s of kernel type-checking PER downstream declaration that forced it (no cross-declaration whnf cache), which was almost all of ViTBackB0's (3×) and ViTMhsaBackCertifiedTie's (1×) build time. With backward a direct field, the projection whnfs in one hop.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.mhsa_has_vjp_mat (N heads d_head : ℕ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : HasVJPMat (mhsa_layer N heads d_head Wq Wk Wv Wo bq bk bv bo)

Multi-head SDPA VJP (Phase 8). Now a theorem (was an axiom), composed from mhsa_g_has_vjp_mat, colSlabwise_has_vjp_mat, and the per-token dense framework. backward is the composed witness's field DIRECTLY (kernel-cheap projection); the mhsa_layer_eq_compose transport lives only in correct — a Prop the kernel never reduces. See mhsa_composed_has_vjp_mat's docstring for why.

Equations

• Proofs.mhsa_has_vjp_mat N heads d_head Wq Wk Wv Wo bq bk bv bo = { backward := (Proofs.mhsa_composed_has_vjp_mat N heads d_head Wq Wk Wv Wo bq bk bv bo).backward, correct := ⋯ }

theorem Proofs.mhsa_layer_flat_diff (N heads d_head : ℕ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (mhsa_layer N heads d_head Wq Wk Wv Wo bq bk bv bo (Mat.unflatten v)).flatten

Differentiability of the flattened multi-head SDPA layer — theorem (was an axiom). Composition of three _flat_diff lemmas.

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 rowwise_has_vjp_mat (Tensor.lean). These are theorems — no new axioms.

noncomputable def Proofs.layerNorm_per_token_has_vjp_mat (N D : ℕ) (ε γ β : ℝ) (hε : 0 < ε) : HasVJPMat fun (X : Mat N D) (n : Fin N) => layerNormForward D ε γ β (X n)

Per-token layer norm across a sequence. Applies layerNormForward to each row of the (N, D) input; the backward is block-diagonal.

Equations

• Proofs.layerNorm_per_token_has_vjp_mat N D ε γ β hε = Proofs.rowwise_has_vjp_mat (Proofs.layerNorm_has_vjp D ε γ β hε) ⋯

noncomputable def Proofs.dense_per_token_has_vjp_mat (N inD outD : ℕ) (W : Mat inD outD) (b : Vec outD) : HasVJPMat fun (X : Mat N inD) (n : Fin N) => dense W b (X n)

Per-token dense projection across a sequence. Q = X · W + b, row-by-row dense with shared weights.

Equations

• Proofs.dense_per_token_has_vjp_mat N inD outD W b = Proofs.rowwise_has_vjp_mat (Proofs.dense_has_vjp W b) ⋯

noncomputable def Proofs.gelu_per_token_has_vjp_mat (N D : ℕ) : HasVJPMat fun (X : Mat N D) (n : Fin N) => gelu D (X n)

Per-token GELU across a sequence. Elementwise activation, so diagonal Jacobian both across rows and within a row.

Equations

• Proofs.gelu_per_token_has_vjp_mat N D = Proofs.rowwise_has_vjp_mat (Proofs.gelu_has_vjp D) ⋯

A transformer encoder block

From emitTransformerBlockForward (line 796 of MlirCodegen.lean):

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 — layerNorm_per_token_has_vjp_mat (theorem via rowwise_has_vjp_mat)
• MHSA — mhsa_has_vjp_mat (bundled HasVJPMat def — Phase 8)
• MLP — two dense_per_token_has_vjp_mat + one gelu_per_token_has_vjp_mat, glued with vjpMat_comp
• + residuals — biPathMat_has_vjp (theorem, Tensor.lean) with identity

The transformer block theorem below glues these with vjpMat_comp and biPathMat_has_vjp. Zero new axioms — every piece is a theorem.

noncomputable def Proofs.transformerMlp (N D mlpDim : ℕ) (Wfc1 : Mat D mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim D) (bfc2 : Vec D) : Mat N D → Mat N D

MLP sublayer of a transformer block: dense ∘ GELU ∘ dense applied per-token.

Concretely: MLP(z) = Wfc2 · gelu(Wfc1 · z + bfc1) + bfc2, applied row-wise.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.transformerMlp_flat_diff (N D mlpDim : ℕ) (Wfc1 : Mat D mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim D) (bfc2 : Vec D) : Differentiable ℝ fun (v : Vec (N * D)) => (transformerMlp N D mlpDim Wfc1 bfc1 Wfc2 bfc2 (Mat.unflatten v)).flatten

Differentiability of the flattened transformerMlp — composition of dense ∘ gelu ∘ dense per-token. Built from the three per-token-flat Diff helpers via Differentiable.comp, with the usual Mat.unflatten_flatten rewrite to push the bijection through ∘.

noncomputable def Proofs.transformerMlp_has_vjp_mat (N D mlpDim : ℕ) (Wfc1 : Mat D mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim D) (bfc2 : Vec D) : HasVJPMat (transformerMlp N D mlpDim Wfc1 bfc1 Wfc2 bfc2)

HasVJPMat for the MLP sublayer — chain of two vjpMat_comp steps over per-token liftings (dense ∘ gelu ∘ dense). Theorem, no longer axiom: every Diff hypothesis is discharged by the per-token-flat helpers above.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.transformerAttnSublayer (N heads d_head : ℕ) (ε γ1 β1 : ℝ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : Mat N (heads * d_head) → Mat N (heads * d_head)

Attention sublayer: X ↦ X + MHSA(LN1(X)). Top-level composition; the biPathMat skip-adds identity to the MHSA∘LN1 branch.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.transformerMlpSublayer (N heads d_head mlpDim : ℕ) (ε γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Mat N (heads * d_head) → Mat N (heads * d_head)

MLP sublayer: h ↦ h + MLP(LN2(h)). Same biPathMat structure.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.transformerBlock (N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Mat N (heads * d_head) → Mat N (heads * d_head)

Transformer encoder block forward: MLP-sublayer ∘ attention-sublayer. Signature matches the codegen: Mat N (heads·d_head) → Mat N (heads·d_head).

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.transformerAttnSublayer_inner_flat_diff (N heads d_head : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => ((mhsa_layer N heads d_head Wq Wk Wv Wo bq bk bv bo ∘ fun (X : Mat N (heads * d_head)) (n : Fin N) => layerNormForward (heads * d_head) ε γ1 β1 (X n)) (Mat.unflatten v)).flatten

Differentiability of the flattened attention sublayer's non-trivial arm (mhsa ∘ LN1). Used by both the sublayer VJP proof and any downstream composition that needs Diff for the sublayer's arm.

theorem Proofs.transformerAttnSublayer_flat_diff (N heads d_head : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerAttnSublayer N heads d_head ε γ1 β1 Wq Wk Wv Wo bq bk bv bo (Mat.unflatten v)).flatten

Differentiability of the flattened attention sublayer. biPathMat (id) (mhsa ∘ LN1) flattens to a sum, both arms Differentiable.

noncomputable def Proofs.transformerAttnSublayer_has_vjp_mat (N heads d_head : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : HasVJPMat (transformerAttnSublayer N heads d_head ε γ1 β1 Wq Wk Wv Wo bq bk bv bo)

Attention sublayer VJP: biPathMat of identity and mhsa ∘ LN1. Theorem, no longer axiom: discharges the Differentiable hypotheses using identity_mat_flat_diff for the skip arm and the inner Diff helper above for the mhsa ∘ LN1 arm.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.transformerMlpSublayer_inner_flat_diff (N heads d_head mlpDim : ℕ) (ε γ2 β2 : ℝ) (hε : 0 < ε) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => ((transformerMlp N (heads * d_head) mlpDim Wfc1 bfc1 Wfc2 bfc2 ∘ fun (X : Mat N (heads * d_head)) (n : Fin N) => layerNormForward (heads * d_head) ε γ2 β2 (X n)) (Mat.unflatten v)).flatten

Differentiability of the MLP sublayer's non-trivial arm (transformerMlp ∘ LN2). Composition of transformerMlp_flat_diff and layerNorm_per_token_flat_diff.

theorem Proofs.transformerMlpSublayer_flat_diff (N heads d_head mlpDim : ℕ) (ε γ2 β2 : ℝ) (hε : 0 < ε) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerMlpSublayer N heads d_head mlpDim ε γ2 β2 Wfc1 bfc1 Wfc2 bfc2 (Mat.unflatten v)).flatten

Differentiability of the flattened MLP sublayer. biPathMat (id) (transformerMlp ∘ LN2) flattens to a sum, both arms Differentiable.

noncomputable def Proofs.transformerMlpSublayer_has_vjp_mat (N heads d_head mlpDim : ℕ) (ε γ2 β2 : ℝ) (hε : 0 < ε) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : HasVJPMat (transformerMlpSublayer N heads d_head mlpDim ε γ2 β2 Wfc1 bfc1 Wfc2 bfc2)

MLP sublayer VJP: biPathMat of identity and transformerMlp ∘ LN2. Theorem, no longer axiom: same recipe as transformerAttnSublayer_has_vjp_mat.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.transformerBlock_flat_diff (N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerBlock N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 (Mat.unflatten v)).flatten

Differentiability of the flattened transformer block. MlpSublayer ∘ AttnSublayer; both sublayers' flat Diff are theorems above.

noncomputable def Proofs.transformerBlock_has_vjp_mat (N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : HasVJPMat (transformerBlock N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2)

Transformer block VJP — composition of attn + mlp sublayers. Theorem, no longer axiom: a single vjpMat_comp of the two sublayer theorems with their Diff helpers.

Equations

• One or more equations did not get rendered due to their size.

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 vjpMat_comp and induction on k, if each block has a HasVJPMat then so does the stack — for any depth.

For the formal theorem we use a single shared parameter tuple across blocks (a mild simplification; in practice every block has its own weights). The Jacobian-composition structure doesn't change — the theorem generalizes trivially to per-block parameters by replacing the Nat induction with a Fin k parameter function, which is mechanical once the single-shared-param case is proved.

noncomputable def Proofs.transformerTower (k N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Mat N (heads * d_head) → Mat N (heads * d_head)

k-fold iterated transformer block, sharing parameters across all k layers. Defined by Nat.rec so the HasVJPMat proof is a straightforward induction on k.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.transformerTower_flat_diff (k N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerTower k N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 (Mat.unflatten v)).flatten

Differentiability of the flattened k-fold transformer tower. Induction on k: zero case is identity, successor case is block ∘ tower(k) composed via Differentiable.comp.

noncomputable def Proofs.transformerTower_has_vjp_mat (k N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : HasVJPMat (transformerTower k N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2)

Transformer tower VJP — k-fold composition. Theorem, no longer axiom: induction on k via vjpMat_comp and transformerBlock_has_vjp_mat.

Equations

• One or more equations did not get rendered due to their size.

ViT body: tower + final LN

vit_body 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
vit_body(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 vjpMat_comp glues the two VJPs.

The patch-embedding and classifier-head steps exit Mat-land (they change type to/from Tensor3 and Vec respectively). Both are trivial compositions of already-proved theorems (conv2d_has_vjp3 / pdiv_reindex for patch embed, pdiv_reindex / dense_has_vjp / softmaxCE_grad for the classifier) but they don't fit in the uniform HasVJPMat frame. We mark them as future work; closing this would require a unified rank-polymorphic VJP framework that's not needed for the pedagogy.

noncomputable def Proofs.vit_body (k N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) (γF βF : ℝ) : Mat N (heads * d_head) → Mat N (heads * d_head)

ViT body — transformer tower followed by final per-token LayerNorm.

Composition is finalLN ∘ transformerTower; matches the codegen's emitForwardBody ordering for a .transformerEncoder followed by the implicit final LN block.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.vit_body_flat_diff (k N heads d_head mlpDim : ℕ) (ε : ℝ) (hε : 0 < ε) (γ1 β1 : ℝ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) (γF βF : ℝ) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (vit_body k N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 γF βF (Mat.unflatten v)).flatten

Differentiability of the flattened ViT body. finalLN ∘ transformerTower — both have flat Diff theorems above.

noncomputable def Proofs.vit_body_has_vjp_mat (k N heads d_head mlpDim : ℕ) (ε : ℝ) (hε : 0 < ε) (γ1 β1 : ℝ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) (γF βF : ℝ) : HasVJPMat (vit_body k N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 γF βF)

The ViT body VJP — finalLN ∘ transformerTower. Theorem, no longer axiom: a single vjpMat_comp of the tower + final LN with their Diff helpers.

Conceptually still the punchline: a depth-k ViT backbone has a correct VJP, composed entirely from proved building blocks. With Phase 3's column-stacking framework, even mhsa_has_vjp_mat and its flat-diff sibling are theorems now, so this whole chain is pure-Mathlib closure with no project axioms.

Equations

• One or more equations did not get rendered due to their size.

What we've proved (and what's left)

Proved (zero sorry's, machine-checked):

• Dense, ReLU (MLP.lean)
• Softmax cross-entropy loss gradient (MLP.lean)
• Conv2d, MaxPool, Flatten (CNN.lean)
• BatchNorm closed-form backward (BatchNorm.lean)
• Residual / biPath fan-in (Residual.lean)
• Depthwise conv (Depthwise.lean)
• Squeeze-and-Excitation / elementwise product VJP (SE.lean)
• LayerNorm, GELU (LayerNorm.lean)
• Standalone softmax VJP (this file)
• Scaled dot-product attention backwards sdpa_back_{Q,K,V} — proved via vjpMat_comp composition of four matrix-level VJP building blocks (matmul, scalarScale, rowSoftmax, matmul). Formulas are also numerically gradient-checked as a belt-and-braces check.

Three calculus rules do all the structural work (now theorems proved from Mathlib's fderiv, formerly axioms):

pdiv_comp   (chain rule — functions compose, derivatives compose)
pdiv_add    (linearity — derivatives of sums are sums of derivatives)
pdiv_mul    (product rule — derivatives of elementwise products)

Five closed-form "Jacobian-structure tricks" handle the layers whose Jacobians are dense but exploitable:

1. Diagonal (activations) — collapse the sum_j to one term.
2. Sparse toeplitz (conv, depthwise) — reversed/transposed kernels.
3. Binary selection (max-pool) — route gradients to argmax cells.
4. Rank-1 correction to diagonal (softmax, BN, LN, IN, GN) — one extra scalar reduction, everything else is pointwise.
5. Outer product + reductions (dense, matmul) — rank-1 update accumulation.

That is the complete taxonomy. I've thought hard about this and cannot find a sixth trick or a fourth calculus rule anywhere in the modern architecture zoo. Every paper, every block, every optimization is a rearrangement of these ten things.

What this means for the reader

If you've read this far, you have a complete decoder for the architecture-of-the-month. Pick any paper — Swin, ConvNeXt, CLIP, Mamba, anything — and walk through the forward pass. For each operation, ask:

1. Is it composition of known ops? -> chain rule.
2. Is it a sum of branches? -> fan-in add.
3. Is it an elementwise / scalar product of branches? -> fan-in mul.
4. Is it an activation? -> diagonal Jacobian template.
5. Is it a normalization? -> closed-form three-term formula.
6. Is it a convolution or linear map? -> the structured-matmul machinery.
7. Is it an attention or softmax-based selection? -> the closed-form rank-1 collapse.

If the answer is "none of the above" — which it won't be — then you've found the first genuinely new layer of the decade, and you get to write the next chapter of this book.

Until then, welcome to the end of the road.

Bridging ranks: from Mat-land back to Vec-land

vit_body_has_vjp_mat lives in HasVJPMat territory. The pieces at the boundaries — patch embedding (image → tokens) and classifier head (tokens → logits) — change tensor rank. Rather than invent new mixed-rank VJP frameworks, we flatten everything to Vec at the interfaces and compose via plain HasVJP, glued by vjp_comp.

Two ingredients needed:

• hasVJPMat_to_hasVJP (Tensor.lean, Phase 10) — bridges any HasVJPMat to HasVJP on the flattened endpoints. One theorem, no new axioms.
• cls_slice_flat_has_vjp — gathers row 0 of a flattened Mat (N+1) D. Derivable from pdiv_reindex.

noncomputable def Proofs.cls_slice_flat (N D : ℕ) : Vec ((N + 1) * D) → Vec D

CLS token extraction, stated on the flattened matrix. Row 0 of a Mat (N+1) D is a Vec D; on the flattened Vec ((N+1)*D) this is the gather v ↦ fun k => v (fPF (0, k)).

Equations

• Proofs.cls_slice_flat N D v k = v (finProdFinEquiv (0, k))

noncomputable def Proofs.cls_slice_flat_has_vjp (N D : ℕ) : HasVJP (cls_slice_flat N D)

CLS slice VJP — gather-style; backward scatters dy to row 0. Derived from pdiv_reindex.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.classifier_flat (N D nClasses : ℕ) (Wcls : Mat D nClasses) (bcls : Vec nClasses) : Vec ((N + 1) * D) → Vec nClasses

Classifier head: flattened CLS slice + dense projection to Vec nClasses.

fun v : Vec ((N+1)*D) => dense W_cls b_cls (cls_slice_flat v)

Equations

• Proofs.classifier_flat N D nClasses Wcls bcls = Proofs.dense Wcls bcls ∘ Proofs.cls_slice_flat N D

theorem Proofs.cls_slice_flat_diff (N D : ℕ) : Differentiable ℝ (cls_slice_flat N D)

Differentiability of cls_slice_flat — linear reindex.

theorem Proofs.dense_input_diff {m n : ℕ} (W : Mat m n) (b : Vec n) : Differentiable ℝ (dense W b)

Differentiability of dense W b as a function of the input vector — linear.

noncomputable def Proofs.classifier_flat_has_vjp (N D nClasses : ℕ) (Wcls : Mat D nClasses) (bcls : Vec nClasses) : HasVJP (classifier_flat N D nClasses Wcls bcls)

Classifier head VJP — composition via vjp_comp. Theorem, no longer axiom: cls_slice_flat and dense are both linear, so their Diff hypotheses discharge by fun_prop.

Equations

• One or more equations did not get rendered due to their size.

Patch embedding — Phase 6 (de-opaqued, no longer axiomatic)

The patch embedding takes a flattened image Vec (ic*H*W) and produces a flattened Vec ((N+1)*D) interpreted as Mat (N+1) D:

1. Conv projection with stride = patchSize: per-patch dense projection W_conv : Kernel4 D ic patchSize patchSize + bias b_conv : Vec D.
2. Reshape spatial (D, H', W') to tokens (N, D) — pure permutation.
3. Prepend learnable CLS token at row 0 → (N+1, D).
4. Add learnable positional embedding matrix → (N+1, D).

This was previously an opaque definition + two bundled axioms (patchEmbed_flat_has_vjp / patchEmbed_flat_diff). Phase 6 (Apr 2026) de-opaques: the forward is a concrete def and both axioms become theorems via foundation rules.

The N parameter is independent of (H, W, patchSize) — out-of-range patches contribute zero (via a hpad guard on the image read), so the API does not require N = (H/patchSize) * (W/patchSize).

noncomputable def Proofs.patchEmbed_flat (ic H W patchSize N D : ℕ) (W_conv : Kernel4 D ic patchSize patchSize) (b_conv cls_token : Vec D) (pos_embed : Mat (N + 1) D) : Vec (ic * H * W) → Vec ((N + 1) * D)

Patch embedding forward on flattened endpoints.

Output at flat-index idx_out = finProdFinEquiv (n, d):

• n = 0: cls_token d + pos_embed 0 d.
• n > 0 (let p := n - 1, h' := p / (W/patchSize), w' := p % (W/patchSize)): `b_conv d + Σ c kh kw, W_conv d c kh kw * img(c, h'*P+kh, w'*P+kw)
  • pos_embed n d, where the image read is guarded by h'*P+kh < H ∧ w'*P+kw < W` (returns 0 if out of range).

This handles arbitrary N cleanly: for n whose decoded patch (h', w') falls outside the image grid, the inner sum is identically zero.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.patchEmbed_flat_diff (ic H W patchSize N D : ℕ) (W_conv : Kernel4 D ic patchSize patchSize) (b_conv cls_token : Vec D) (pos_embed : Mat (N + 1) D) : Differentiable ℝ (patchEmbed_flat ic H W patchSize N D W_conv b_conv cls_token pos_embed)

Patch embedding differentiability — proved from foundation rules. patchEmbed_flat is linear in img plus constants (pos_embed, cls_token, b_conv); the only non-trivial part is the dependent if hpad : ... then img(σ hpad) else 0 pattern handled by differentiableAt_pad_eval.

noncomputable def Proofs.patchEmbed_input_grad_formula (ic H W patchSize N D : ℕ) (W_conv : Kernel4 D ic patchSize patchSize) (dy : Vec ((N + 1) * D)) : Vec (ic * H * W)

Closed-form input gradient for patchEmbed_flat — direct formula, written as a sum over patches p : Fin N with reconstructed kernel offsets (kh, kw) matching the input position (hh, ww) decoded from idx_in. Equivalent (under the patch-row decomposition h' := p/(W/P), w' := p%(W/P)) to the standard "deconvolution" formula for patchEmbed.

The CLS row (n = 0) does not appear here — idx_in only flows through the conv-projection branch (n > 0), so the gradient sums over p : Fin N (corresponding to output rows n = p+1).

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.patchEmbed_flat_has_vjp (ic H W patchSize N D : ℕ) (W_conv : Kernel4 D ic patchSize patchSize) (b_conv cls_token : Vec D) (pos_embed : Mat (N + 1) D) : HasVJP (patchEmbed_flat ic H W patchSize N D W_conv b_conv cls_token pos_embed)

Patch embedding VJP — proved from foundation rules.

Phase 6b (Apr 2026): the forward (de-opaqued in Phase 6a) is linear in img plus constants; the only non-trivial pattern is the dependent if hpad : ... then img(σ) else 0 handled by pdiv_const_mul_pi_pad_eval (already used for conv2d/depthwise input-VJPs). Closing collapse mirrors conv2d_has_vjp3, with one new wrinkle: split Σ n : Fin (N+1) into n = 0 (CLS row, contributes 0 to img-grad) + Σ p : Fin N (n = p+1) via Fin.sum_univ_succ.

Backward: patchEmbed_input_grad_formula W_conv dy.

Equations

• One or more equations did not get rendered due to their size.

The full ViT theorem

Compose patch embed + ViT body (via the Mat→Vec bridge) + classifier. All three are HasVJPs on Vec, so vjp_comp chains them directly.

noncomputable def Proofs.vit_full (ic H W patchSize N mlpDim heads d_head kBlocks nClasses : ℕ) (W_conv : Kernel4 (heads * d_head) ic patchSize patchSize) (b_conv cls_token : Vec (heads * d_head)) (pos_embed : Mat (N + 1) (heads * d_head)) (ε γ1 β1 : ℝ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) (γF βF : ℝ) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) : Vec (ic * H * W) → Vec nClasses

vit_full — full ViT forward from flattened image pixels to logits.

Vec (ic*H*W) → Vec nClasses

Composition: patchEmbed → (flatten ∘ vit_body ∘ unflatten) → classifier. Uses D := heads * d_head directly (no separate D parameter) so the type-level reinterpretation at the body is a no-op.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.classifier_flat_diff (N D nClasses : ℕ) (Wcls : Mat D nClasses) (bcls : Vec nClasses) : Differentiable ℝ (classifier_flat N D nClasses Wcls bcls)

Differentiability of the classifier head — composition of linear ops.

noncomputable def Proofs.vit_full_has_vjp (ic H W patchSize N mlpDim heads d_head kBlocks nClasses : ℕ) (W_conv : Kernel4 (heads * d_head) ic patchSize patchSize) (b_conv cls_token : Vec (heads * d_head)) (pos_embed : Mat (N + 1) (heads * d_head)) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) (γF βF : ℝ) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) : HasVJP (vit_full ic H W patchSize N mlpDim heads d_head kBlocks nClasses W_conv b_conv cls_token pos_embed ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 γF βF Wcls bcls)

vit_full VJP — the grand finale. Theorem, no longer axiom: three vjp_comp steps glueing patchEmbed_flat_has_vjp, hasVJPMat_to_hasVJP (vit_body_has_vjp_mat ...), and classifier_flat_has_vjp. Each vjp_comp's Diff hypotheses are discharged by the per-stage Diff theorems above.

Equations

• One or more equations did not get rendered due to their size.

Public correctness theorems for the attention defs

The _has_vjp / _has_vjp_mat defs above bundle a backward function with a .correct field; these _correct theorems expose that field as a top-level proposition so consumers can refer to the contract directly without reaching into record internals.

theorem Proofs.mhsa_has_vjp_mat_correct (N heads d_head : ℕ) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (X dY : Mat N (heads * d_head)) (i : Fin N) (j : Fin (heads * d_head)) : (mhsa_has_vjp_mat N heads d_head Wq Wk Wv Wo bq bk bv bo).backward X dY i j = ∑ k : Fin N, ∑ l : Fin (heads * d_head), pdivMat (mhsa_layer N heads d_head Wq Wk Wv Wo bq bk bv bo) X i j k l * dY k l

Public correctness theorem for mhsa_has_vjp_mat: multi-head SDPA's backward equals the pdivMat-contracted Jacobian. Phase 3's column-stacking proof closes this without any project axiom.

theorem Proofs.transformerBlock_has_vjp_mat_correct (N heads d_head mlpDim : ℕ) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) (X dY : Mat N (heads * d_head)) (i : Fin N) (j : Fin (heads * d_head)) : (transformerBlock_has_vjp_mat N heads d_head mlpDim ε γ1 β1 hε Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2).backward X dY i j = ∑ k : Fin N, ∑ l : Fin (heads * d_head), pdivMat (transformerBlock N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2) X i j k l * dY k l

Public correctness theorem for transformerBlock_has_vjp_mat: the full transformer block backward (attention sublayer + MLP sublayer glued by vjpMat_comp) equals the pdivMat-contracted Jacobian.

theorem Proofs.vit_full_has_vjp_correct (ic H W patchSize N mlpDim heads d_head kBlocks nClasses : ℕ) (W_conv : Kernel4 (heads * d_head) ic patchSize patchSize) (b_conv cls_token : Vec (heads * d_head)) (pos_embed : Mat (N + 1) (heads * d_head)) (ε γ1 β1 : ℝ) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) (γ2 β2 : ℝ) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) (γF βF : ℝ) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) (dy : Vec nClasses) (i : Fin (ic * H * W)) : (vit_full_has_vjp ic H W patchSize N mlpDim heads d_head kBlocks nClasses W_conv b_conv cls_token pos_embed ε γ1 β1 hε Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 γF βF Wcls bcls).backward x dy i = ∑ j : Fin nClasses, pdiv (vit_full ic H W patchSize N mlpDim heads d_head kBlocks nClasses W_conv b_conv cls_token pos_embed ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 γF βF Wcls bcls) x i j * dy j

Public correctness theorem for vit_full_has_vjp: the full ViT's backward equals the pdiv-contracted Jacobian (Jacobian-transpose applied to the cotangent). Exposes the witness's .correct field as a top-level proposition so consumers (and #print axioms audits) can cite the apex contract directly instead of reaching into the record. The long signature is just the full ViT hyperparameter set; the proof is the witness field.