Documentation

LeanMlir.Proofs.Nets.ViT.ViTMhsaBackCertifiedTie

The ViT MHSA backward chain IS the certified VJP (the sdpa adjoint) #

The ViT-specific backward leaf: the hand-composed multi-head self-attention backward mhsaBackFlat (ViTBackChains.lean) IS the certified MHSA input-gradient VJP mhsaHasVJPMat (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-sdpaBack*-grounded by construction (coreQFlat = flatten ∘ mhsaSdpaBackQ ∘ unflatten, mhsaSdpaBackQ = sdpaBackQ per headSliceMat 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 mhsaHasVJPMat.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 backward 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).

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 (headSliceMat _ _ _ 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))) => ((mhsaHasVJPMat N h dh Wq Wk Wv Wo bq bk bv bo).backward X (Mat.unflatten dconcat)).flatten

The ViT MHSA backward 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 (mhsaHasVJPMat …).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.transformerMlp_back_flat_eq_perRowFlatPR {gf : GeluForm} (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)) :
((transformerMlpHasVJPMat gf 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) => gf.scalarDeriv (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.