Documentation

LeanMlir.Proofs.Nets.ViT.ViTVecLNBackCertifiedTie

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 the same move as ConvNeXtBackCertifiedTie.lean makes for ConvNeXt-T.

No new analysis: ConvNeXt's LayerNorm backward is ViT's. rowLNVecFlatBack (ChannelLNBack.lean) is perRowFlatPR of bnGradInput 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'". rowLNVecFlatHasVJP_backward_eq already pins it to layerNormVecPerTokenHasVJPMat. 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: transformerBlockVHasVJPMat is a vjpMatComp / biPathMatHasVJP assembly, so the projections reduce.

source
theorem Proofs.rowLNVecFlatBack_eq_vecLN_vjp (n D : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec D) (X dy : Vec (n * D)) :
rowLNVecFlatBack n D ε γ X dy = ((layerNormVecPerTokenHasVJPMat 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. rowLNVecFlatHasVJP_backward_eq (ConvNeXt's) read through HasVJPMat.toHasVJP's projection. The only LayerNorm-specific step in this file.

source
theorem Proofs.attnSubFlat_tie_v {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 = ((transformerAttnSublayerVHasVJPMat 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.mlpSubFlat_tie_v {h N dh : ℕ} {gf : GeluForm} (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) => gf.scalarDeriv (dense Wfc1 bfc1 (layerNormVec (h * dh) ε γ2 β2 (hM r)) c)) ∘ dense Wfc2.transpose 0) v = ((transformerMlpSublayerVHasVJPMat gf 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 : ℕ} {gf : GeluForm} (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) => gf.scalarDeriv (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))) => ((transformerBlockVHasVJPMat gf 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 backward 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 {gf : GeluForm} (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (hε : 0 < ε) (p : BlockParamsV (heads * d_head) mlpDim) (v : Vec (Np1 * (heads * d_head))) :
vitBlockBackVAt gf Np1 heads d_head mlpDim ε p v = (transformerBlockVHasVJPMat gf 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).toHasVJP.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 vitBodyKVFlatHasVJP's chain step consumes. Mat.flatten_unflatten is the only step.