Documentation

LeanMlir.Proofs.Nets.ViT.ViTVecLNBackCertifiedTie

§B: the ViT encoder-block backward tie at the VECTOR LayerNorm the net runs #

The shipped depth-12 net is vitForwardKV, whose blocks are transformerBlockV at γ β : Vec D, the form ViTRender.lean emits. This file states the block backward tie there, so the whole-net fold in ViTWholeBackCertifiedTie.lean is about the blocks the artifact contains. It is package 3.1's move for ConvNeXt-T, one architecture over.

⭐ There is no new analysis, and the reason is that ConvNeXt already built ViT's LayerNorm backward. rowLNVecFlatBack (ChannelLNBack.lean) is perRowFlatPR of bn_grad_input c ε 1 (X r) ∘ diagBack γ — the γ scale in front of the unit-γ input gradient at each row's own saved activation — and its own header says it is "literally ViT's per-token LN with 'token' read as 'spatial position'". rowLNVecFlat_has_vjp_backward_eq already pins it to layerNormVec_per_token_has_vjp_mat. So the vector-LN seam (rowLNVecFlatBack_eq_vecLN_vjp) is that lemma read at a flat saved input, and everything else in the block — mhsaBackFlat, the dense Wᵀ 0 input-VJPs, the diagBack GELU derivative, the perRowFlatPR residual seams — is LayerNorm-agnostic and is reused from ViTMhsaBackCertifiedTie.lean verbatim.

⚠ Both sublayer decompositions and the block unfold are rfl at the vector LN: transformerBlockV_has_vjp_mat is a vjpMat_comp / biPathMat_has_vjp assembly, so the projections reduce. 3-axiom-clean.

source
theorem Proofs.rowLNVecFlatBack_eq_vecLN_vjp (n D : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec D) (X dy : Vec (n * D)) :
rowLNVecFlatBack n D ε γ X dy = ((layerNormVec_per_token_has_vjp_mat n D ε γ β hε).backward (Mat.unflatten X) (Mat.unflatten dy)).flatten

The vector-LN backward at a flat saved input IS the certified per-token vector-LN VJP, flattened. rowLNVecFlat_has_vjp_backward_eq (ConvNeXt's) read through hasVJPMat_to_hasVJP's projection. The only LayerNorm-specific step in this file.

source
theorem Proofs.transformerMlpSublayerV_backward_decomp {h N dh : ℕ} (dff : ℕ) (ε : ℝ) (hε : 0 < ε) (γ2 β2 : Vec (h * dh)) (Wfc1 : Mat (h * dh) dff) (bfc1 : Vec dff) (Wfc2 : Mat dff (h * dh)) (bfc2 : Vec (h * dh)) (hM dz : Mat N (h * dh)) :
(transformerMlpSublayerV_has_vjp_mat N h dh dff ε γ2 β2 hε Wfc1 bfc1 Wfc2 bfc2).backward hM dz = fun (i : Fin N) (j : Fin (h * dh)) => dz i j + (layerNormVec_per_token_has_vjp_mat N (h * dh) ε γ2 β2 hε).backward hM ((transformerMlp_has_vjp_mat N (h * dh) dff Wfc1 bfc1 Wfc2 bfc2).backward (fun (n : Fin N) => layerNormVec (h * dh) ε γ2 β2 (hM n)) dz) i j

The vector-LN MLP sublayer's VJP backward decomposes — the MLP peer, also rfl.

source
theorem Proofs.attnSubFlatTieV {h N dh : ℕ} (ε : ℝ) (hε : 0 < ε) (γ1 β1 : Vec (h * dh)) (Wq Wk Wv Wo : Mat (h * dh) (h * dh)) (bq bk bv bo : Vec (h * dh)) (A : Mat N (h * dh)) (w : Vec (N * (h * dh))) :
residual (rowLNVecFlatBack N (h * dh) ε γ1 A.flatten ∘ mhsaBackFlat Wq Wk Wv Wo (fun (r : Fin N) => dense Wq bq (layerNormVec (h * dh) ε γ1 β1 (A r))) (fun (r : Fin N) => dense Wk bk (layerNormVec (h * dh) ε γ1 β1 (A r))) fun (r : Fin N) => dense Wv bv (layerNormVec (h * dh) ε γ1 β1 (A r))) w = ((transformerAttnSublayerV_has_vjp_mat N h dh ε γ1 β1 hε Wq Wk Wv Wo bq bk bv bo).backward A (Mat.unflatten w)).flatten

The attention-sublayer backward float-half IS the certified sublayer VJP, flat. residual (rowLNVecFlatBack ∘ mhsaBackFlat) with Q/K/V pinned at LNᵥ₁(A) and the LN backward at the block's own saved input. The sdpa leaf is mhsaBackFlat_eq_mhsa_vjp, which never mentions a LayerNorm.

source
theorem Proofs.mlpSubFlatTieV {h N dh : ℕ} (dff : ℕ) (ε : ℝ) (hε : 0 < ε) (γ2 β2 : Vec (h * dh)) (Wfc1 : Mat (h * dh) dff) (bfc1 : Vec dff) (Wfc2 : Mat dff (h * dh)) (bfc2 : Vec (h * dh)) (hM : Mat N (h * dh)) (v : Vec (N * (h * dh))) :
residual (rowLNVecFlatBack N (h * dh) ε γ2 hM.flatten ∘ perRowFlatPR N (h * dh) fun (r : Fin N) => dense Wfc1.transpose 0 ∘ (diagBack fun (c : Fin dff) => geluScalarDeriv (dense Wfc1 bfc1 (layerNormVec (h * dh) ε γ2 β2 (hM r)) c)) ∘ dense Wfc2.transpose 0) v = ((transformerMlpSublayerV_has_vjp_mat N h dh dff ε γ2 β2 hε Wfc1 bfc1 Wfc2 bfc2).backward hM (Mat.unflatten v)).flatten

The MLP-sublayer backward float-half IS the certified sublayer VJP, flat. The residual is lifted out of the per-token fold. The per-token body tie (transformerMlp_back_flat_eq_perRowFlatPR) is LayerNorm-agnostic and reused verbatim; only the LN₂-back seam changes.

source
theorem Proofs.vitBlockBackV_eq_transformerBlockV_vjp {h N dh : ℕ} (dff : ℕ) (ε : ℝ) (hε : 0 < ε) (γ1 β1 γ2 β2 : Vec (h * dh)) (Wq Wk Wv Wo : Mat (h * dh) (h * dh)) (bq bk bv bo : Vec (h * dh)) (Wfc1 : Mat (h * dh) dff) (bfc1 : Vec dff) (Wfc2 : Mat dff (h * dh)) (bfc2 : Vec (h * dh)) (A : Mat N (h * dh)) :
vitBlockBackV Wq Wk Wv Wo (fun (r : Fin N) => dense Wq bq (layerNormVec (h * dh) ε γ1 β1 (A r))) (fun (r : Fin N) => dense Wk bk (layerNormVec (h * dh) ε γ1 β1 (A r))) (fun (r : Fin N) => dense Wv bv (layerNormVec (h * dh) ε γ1 β1 (A r))) ε γ1 A.flatten Wfc1 Wfc2 (fun (r : Fin N) (c : Fin dff) => geluScalarDeriv (dense Wfc1 bfc1 (layerNormVec (h * dh) ε γ2 β2 (transformerAttnSublayerV N h dh ε γ1 β1 Wq Wk Wv Wo bq bk bv bo A r)) c)) γ2 (transformerAttnSublayerV N h dh ε γ1 β1 Wq Wk Wv Wo bq bk bv bo A).flatten = fun (dY : Vec (N * (h * dh))) => ((transformerBlockV_has_vjp_mat N h dh dff ε γ1 β1 hε Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2).backward A (Mat.unflatten dY)).flatten

⭐⭐ THE VECTOR-LN ViT BLOCK §B TIE. vitBlockBackV, with every saved activation pinned to the real forward (Q/K/V at LNᵥ₁ A; the LN₁ backward at the block's own input A and the LN₂ backward at the attention sublayer's output; the GELU derivative at dense₁(LNᵥ₂(attn A))), IS the certified transformerBlockV input-gradient VJP, flattened.

Assembled from the block unfold and the two sublayer flat ties; no new analysis, and general in the head count.

source
theorem Proofs.vitBlockBackVAt_eq_vjp (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (hε : 0 < ε) (p : BlockParamsV (heads * d_head) mlpDim) (v : Vec (Np1 * (heads * d_head))) :
vitBlockBackVAt Np1 heads d_head mlpDim ε p v = (hasVJPMat_to_hasVJP (transformerBlockV_has_vjp_mat Np1 heads d_head mlpDim ε p.γ1 p.β1 hε p.Wq p.Wk p.Wv p.Wo p.bq p.bk p.bv p.bo p.γ2 p.β2 p.Wfc1 p.bfc1 p.Wfc2 p.bfc2)).backward v

The block tie in the form the tower recursion needs — vitBlockBackVAt at a FLAT saved input v is the flat block's VJP backward at v, i.e. exactly the HasVJP that vitBodyKVFlat_has_vjp's chain step consumes. Mat.flatten_unflatten is the only step.