Documentation

LeanMlir.Proofs.Nets.ViT.ViTMhsaBackCertifiedTie

§B: the ViT MHSA backward chain IS the certified VJP (the sdpa adjoint) #

The substantive vit-specific §B leaf: the hand-composed multi-head self-attention backward mhsaBackFlat (ViTBackChains.lean) IS the certified MHSA input-gradient VJP mhsa_has_vjp_mat (Attention.lean), flattened — the attention analogue of the depthwise/conv adjoint gates.

Unlike the CNN convFlatBack (a free reversed-kernel conv that needed a gate), the ViT sdpa cores are ALREADY certified-sdpa_back-grounded by construction (coreQFlat = flatten ∘ mhsaSdpaBackQ ∘ unflatten, mhsaSdpaBackQ = sdpa_back_Q per mhSlab head). What this file closes is the assembly reconciliation: mhsaBackFlat is a flat per-head fan-in with SEPARATE dense Wᵀq/Wᵀk/Wᵀv projection-backwards (perRowFlat), while the certified mhsa_has_vjp_mat.backward is a Mat-space VJP over the qkv-MERGED projection. ViTBackB0's mhsa_backward_collapseMH already collapses the certified Mat backward to the clean per-head merged sum mhsaBackCollapsedMH = ∑ₕ (Σⱼ Wq c (h,j)·dQ + Σⱼ Wk·dK + Σⱼ Wv·dV); this file shows mhsaBackFlat (Q/K/V pinned to the actual projections dense W· bq (X·)) equals that, coordinatewise: dense Wᵀ 0 = Mat.mulVec W, the Σ k over h·dh reindexes to Σₕ Σⱼ, and the chain's separate projBack sums regroup into the certified ∑ₕ(Q+K+V) by Finset.sum_add_distrib. So mhsaBackFlat IS the certified attention gradient — the genuinely-new (sdpa) half of the ViT block §B tie.

The block tie that wraps this in the per-token LN/dense/gelu sublayer reconciliations and the residual fan-in is vitBlockBackV_eq_transformerBlockV_vjp (ViTVecLNBackCertifiedTie.lean). 3-axiom-clean.

source
theorem Proofs.projBack_core_coord {h N dh : ℕ} (W : Mat (h * dh) (h * dh)) (S : Mat N (h * dh)) (idx : Fin (N * (h * dh))) :
perRowFlat N (h * dh) (dense W.transpose 0) S.flatten idx = ∑ k : Fin (h * dh), W (finProdFinEquiv.symm idx).2 k * S (finProdFinEquiv.symm idx).1 k

The projection-back leaf coordinate. A per-token dense (Wᵀ) 0 projection-backward (perRowFlat) applied to a flattened saved cotangent Mat.flatten S reads, at output coordinate (r, c), as Σ k, W c k · S r k — i.e. Mat.mulVec W of row r. The dense_transpose = mulVec fact lifted through the perRowFlat / flatten-unflatten reindex. The shared leaf for all four projections (Wq/Wk/Wv contracted against the cores, Wo against the block cotangent).

source
theorem Proofs.woback_unflatten {h N dh : ℕ} (Wo : Mat (h * dh) (h * dh)) (dconcat : Vec (N * (h * dh))) :
Mat.unflatten (perRowFlat N (h * dh) (dense Wo.transpose 0) dconcat) = fun (i : Fin N) (c : Fin (h * dh)) => Wo.mulVec (Mat.unflatten dconcat i) c

The Wo-back, unflattened. unflatten (perRowFlat (dense Wᵀo 0) dconcat) = fun i c => mulVec Wo (unflatten dconcat i) c — the block cotangent run through the output-projection backward, in Mat form. This is the per-head slab the cores read (mhSlab h (unflatten woflat) = dAttg h).

source
theorem Proofs.mhsaBackFlat_eq_mhsa_vjp {h N dh : ℕ} (Wq Wk Wv Wo : Mat (h * dh) (h * dh)) (bq bk bv bo : Vec (h * dh)) (X : Mat N (h * dh)) :
(mhsaBackFlat Wq Wk Wv Wo (fun (r : Fin N) => dense Wq bq (X r)) (fun (r : Fin N) => dense Wk bk (X r)) fun (r : Fin N) => dense Wv bv (X r)) = fun (dconcat : Vec (N * (h * dh))) => ((mhsa_has_vjp_mat N h dh Wq Wk Wv Wo bq bk bv bo).backward X (Mat.unflatten dconcat)).flatten

THE ViT MHSA BACKWARD §B TIE. The MHSA backward chain mhsaBackFlat, with its saved Q/K/V projections pinned to the actual dense W· b· (X·) projections at the saved block input X, IS the certified MHSA input-gradient VJP (mhsa_has_vjp_mat …).backward X, flattened. So the attention backward the ViT chain is spelled in IS the certified attention gradient, not a look-alike. Closes under [propext, Classical.choice, Quot.sound].

source
theorem Proofs.dense_transpose_eq_mulVec {m n : ℕ} (W : Mat m n) :
dense W.transpose 0 = W.mulVec

The chain's dense input-VJP dense (Wᵀ) 0 IS the certified contraction Mat.mulVec W (the certified dense_has_vjp.backward, which ignores its affine activation); mul_comm per term. The function-level form (no x arg) the simp matches against.

source
theorem Proofs.diagBack_eq_gelu_vjp {n : ℕ} (s : Vec n) :
(diagBack fun (c : Fin n) => geluScalarDeriv (s c)) = (gelu_has_vjp n).backward s

The chain's GELU backward diagBack (act'(s)) IS the certified gelu_has_vjp.backward at the saved pre-activation s (the elementwise derivative scaling — gelu_has_vjp.backward s dy i = dy i · geluScalarDeriv (s i), diagBack is the same scaling, mul_comm).

source
theorem Proofs.transformerMlp_backward_pertoken (N D dff : ℕ) (Wfc1 : Mat D dff) (bfc1 : Vec dff) (Wfc2 : Mat dff D) (bfc2 : Vec D) (Y dz : Mat N D) :
(transformerMlp_has_vjp_mat N D dff Wfc1 bfc1 Wfc2 bfc2).backward Y dz = fun (r : Fin N) => Wfc1.mulVec ((gelu_has_vjp dff).backward (dense Wfc1 bfc1 (Y r)) (Wfc2.mulVec (dz r)))

The transformerMlp backward in explicit per-token form. The nested vjpMat_comp (dense₂ ∘ gelu ∘ dense₁, per token) reduces to: each token's dz r runs mulVec Wfc2, the GELU backward at the saved pre-activation dense₁(Y r), then mulVec Wfc1. Pure rfl (the per-token VJPs are rowwise/vjpMat_comp structure projections).

source
theorem Proofs.transformerMlp_back_flat_eq_perRowFlatPR (N D dff : ℕ) (Wfc1 : Mat D dff) (bfc1 : Vec dff) (Wfc2 : Mat dff D) (bfc2 : Vec D) (Y : Mat N D) (v : Vec (N * D)) :
((transformerMlp_has_vjp_mat N D dff Wfc1 bfc1 Wfc2 bfc2).backward Y (Mat.unflatten v)).flatten = perRowFlatPR N D (fun (r : Fin N) => dense Wfc1.transpose 0 ∘ (diagBack fun (c : Fin dff) => geluScalarDeriv (dense Wfc1 bfc1 (Y r) c)) ∘ dense Wfc2.transpose 0) v

L2 — the transformerMlp backward, flattened, IS perRowFlatPR of the flat chain. The certified per-token MLP-body backward (mulVec Wfc1 ∘ gelu-back ∘ mulVec Wfc2) equals the chain's dense Wᵀ₁ 0 ∘ diagBack(act'(dense₁ Y)) ∘ dense Wᵀ₂ 0, row by row.