Documentation

LeanMlir.Proofs.Nets.ViT.ViTFoldGB

T3 §1 fold for ViT-Tiny at the BATCHED index — the op set 4c leg 4 renders #

ViTFoldG.lean folds the ten gradient nodes of the PER-EXAMPLE traversal (ViTRender.vitBackAll at adam := true). This is its batched peer, at ViTRenderB.vitBackAllB's constructors, and it exists because 4c leg 4 moves every committed ViT artifact onto that traversal.

The bytes do not move and the denotation does. Measured 2026-09-07: all nineteen drop-free ViT artifacts — vit_fwd, vitin_fwd and the seventeen AdamW/EMA train steps — re-render byte-identically off vitBackAllB, because every batched form was built to emit its per-example peer's text and tests/TestBatchedEmitTie.lean pins each one individually. So this file is not about different bytes; it is about the AST those bytes are pretty of.

⭐⭐ And on one parameter the AST is genuinely better, which is the whole point of the leg. The CLS token is ONE shared [192] vector, so its gradient is the sum of every example's CLS-row cotangent. The per-example render emits denseBiasGradB (N := 1) — "sum one thing", correct there because pretty B performed the batch lift OUTSIDE the AST — where the batched one emits denseBiasGradB (N := vbB) and the sum is inside den. ViTRenderB.lean flags that line as "THE ONE LINE WHERE N := 1 → N := vbB CHANGES THE FUNCTION", and clsGrad_denB below is the statement the per-example clsGrad_den could not make. Same emitted text either way, which is why the byte tie cannot see it and den_rowDenseBiasGradB_at_one exists to argue the point.

The op table of every committed ViT train step, after leg 4 #

emitted nodelemmaper-example peer (fused ViTPoC op unless noted)
veclnGammaGradB (25 LN γ: LN1/LN2 × 12 + final)veclnGammaGradB_denViTPoC.veclnGammaSgd_den
rowDenseBiasGradB (25 LN β)rowDenseBiasGradB_den_lnbetaViTPoC.rowDenseBiasSgd_den_lnbeta
rowDenseWeightGradB (Wq/Wk/Wv/Wo/Wfc1/Wfc2 × 12)rowDenseWeightGradB_denViTPoC.rowDenseWeightSgd_den
rowDenseBiasGradB (bq/bk/bv/bo/bfc1/bfc2 × 12)rowDenseBiasGradB_denViTPoC.rowDenseBiasSgd_den
patchEmbedWeightGradB / patchEmbedBiasGradBpatchEmbedWeightGradB_den / patchEmbedBiasGradB_denViTPoC.patchEmbedWeightSgd_den / patchEmbedBiasSgd_den
posEmbedGradBposEmbedGradB_denViTPoCG.posEmbedGrad_den (un-fused)
denseBiasGradB at N = B (the CLS token)clsGrad_denBViTPoCG.clsGrad_den (un-fused), at N = 1
weightGradB / biasGradB (the classifier)headWGradB_den / headBGradB_denViTPoC.headW_den / headB_den

No new mathematics: every proof is Finset.sum_congr rfl over the batch and then the per-example bridge at batchSlice n. That is ResNet34FoldB.denseWGradB_den's shape, and it is available because each batched den arm is literally ∑_batch of the per-example one — the constructors were written that way (StableHLO.lean's own comment on veclnGammaGradB: "TWO-LEVEL: the outer Σ_n is the batch, the inner Σ_r the rows within one example").

One lemma per op kind certifies every optimizer tail at once — AdamW, the wx/clip variants, the EMA shadow and the 4× accumulation all consume the same *GradB node, and the vitin_adamdp128x4wxclipdrop artifact whose accuracy the book quotes is one of them.

The LayerNorm form is the VECTOR one (γ β : Vec D), which is what the shipped vitForwardKV runs; the scalar-affine spelling this cone was caught on three times is nowhere here.

Honest residual #

theorem Proofs.ViTPoCGB.veclnGammaGradB_den {N R D : } (xN epsStr cotN : String) (ε : ) (βv : Vec D) (x : Vec (N * (R * D))) (γ : Vec D) (dy : Vec (N * (R * D))) (k : Fin D) :
StableHLO.den (StableHLO.SHlo.veclnGammaGradB xN epsStr ε x (StableHLO.SHlo.operand cotN dy)) k = n : Fin N, o : Fin (R * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin R) => layerNormVec D ε gv βv (Mat.unflatten (StableHLO.batchSlice N (R * D) x n) r)) γ k o * StableHLO.batchSlice N (R * D) dy n o

Batched vector-LN γ GRADIENT denotes the certified Σ_n γ gradient. Two levels: the outer sum is the batch, the inner one the tokens within one example. All 25 sites.

theorem Proofs.ViTPoCGB.rowDenseBiasGradB_den_lnbeta {N R D : } (cotN : String) (ε : ) (γv : Vec D) (X : Fin NMat R D) (β : Vec D) (dy : Vec (N * (R * D))) (i : Fin D) :
StableHLO.den (StableHLO.SHlo.operand cotN dy).rowDenseBiasGradB i = n : Fin N, o : Fin (R * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin R) => layerNormVec D ε γv bv (X n r)) β i o * StableHLO.batchSlice N (R * D) dy n o

The SAME row-reduce op, certified against the vector-LN β forward. An LN β gradient and a dense bias gradient are the identical two-level reduce, so this constructor appears twice in the table against two different certified Jacobians — as it does per example. All 25 β sites.

theorem Proofs.ViTPoCGB.rowDenseWeightGradB_den {N tk a c : } (xN cotN : String) (bb : Vec c) (x : Vec (N * (tk * a))) (W : Mat a c) (dy : Vec (N * (tk * c))) (i : Fin a) (j : Fin c) :
StableHLO.den (StableHLO.SHlo.rowDenseWeightGradB xN x (StableHLO.SHlo.operand cotN dy)) (finProdFinEquiv (i, j)) = n : Fin N, o : Fin (tk * c), pdiv (fun (v : Vec (a * c)) => Mat.flatten fun (r : Fin tk) => dense (Mat.unflatten v) bb (Mat.unflatten (StableHLO.batchSlice N (tk * a) x n) r)) W.flatten (finProdFinEquiv (i, j)) o * StableHLO.batchSlice N (tk * c) dy n o

Batched per-token dense weight GRADIENT denotes the certified Σ_n Σ_tokens x⊗dy. All six denses in each of the 12 blocks. ⚠ The emitted dot_general contracts batch AND token in ONE op; the two sums here are that contraction read apart.

theorem Proofs.ViTPoCGB.rowDenseBiasGradB_den {N tk a c : } (cotN : String) (W : Mat a c) (X : Fin NMat tk a) (b : Vec c) (dy : Vec (N * (tk * c))) (i : Fin c) :
StableHLO.den (StableHLO.SHlo.operand cotN dy).rowDenseBiasGradB i = n : Fin N, o : Fin (tk * c), pdiv (fun (b' : Vec c) => Mat.flatten fun (r : Fin tk) => dense W b' (X n r)) b i o * StableHLO.batchSlice N (tk * c) dy n o

Batched per-token dense bias GRADIENT denotes the certified Σ_n Σ_tokens dy. bq/bk/bv/bo/bfc1/bfc2.

theorem Proofs.ViTPoCGB.patchEmbedWeightGradB_den {ic H W P tk D N : } (xN cotN : String) (bc cls : Vec D) (pos : Mat (tk + 1) D) (img : Vec (N * (ic * H * W))) (Wp : Kernel4 D ic P P) (dy : Vec (N * ((tk + 1) * D))) (d : Fin D) (c : Fin ic) (kh kw : Fin P) :
StableHLO.den (StableHLO.SHlo.patchEmbedWeightGradB xN img (StableHLO.SHlo.operand cotN dy)) (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) = n : Fin N, o : Fin ((tk + 1) * D), pdiv (fun (v : Vec (D * ic * P * P)) => patchEmbed_flat ic H W P tk D (Kernel4.unflatten v) bc cls pos (StableHLO.batchSlice N (ic * H * W) img n)) Wp.flatten (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) o * StableHLO.batchSlice N ((tk + 1) * D) dy n o

Batched patch-embed conv weight GRADIENT denotes the certified Σ_n patchify weight gradient. ⚠ The emitted convolution contracts the batch axis itself, so the outer sum here is inside one op rather than across N of them.

theorem Proofs.ViTPoCGB.patchEmbedBiasGradB_den {ic H W P tk D N : } (cotN : String) (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (tk + 1) D) (img : Vec (N * (ic * H * W))) (dy : Vec (N * ((tk + 1) * D))) (i : Fin D) :
StableHLO.den (StableHLO.SHlo.operand cotN dy).patchEmbedBiasGradB i = n : Fin N, o : Fin ((tk + 1) * D), pdiv (fun (b' : Vec D) => patchEmbed_flat ic H W P tk D Wc b' cls pos (StableHLO.batchSlice N (ic * H * W) img n)) bc i o * StableHLO.batchSlice N ((tk + 1) * D) dy n o

Batched patch-embed conv bias GRADIENT denotes the certified Σ_n bias gradient (Σ_patches dy, the CLS row excluded — the inner p.succ).

theorem Proofs.ViTPoCGB.posEmbedGradB_den {ic H W P tk D N : } (cotN : String) (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (tk + 1) D) (img : Vec (N * (ic * H * W))) (dy : Vec (N * ((tk + 1) * D))) (i : Fin ((tk + 1) * D)) :
StableHLO.den (StableHLO.SHlo.operand cotN dy).posEmbedGradB i = n : Fin N, o : Fin ((tk + 1) * D), pdiv (fun (p : Vec ((tk + 1) * D)) => patchEmbed_flat ic H W P tk D Wc bc cls (Mat.unflatten p) (StableHLO.batchSlice N (ic * H * W) img n)) pos.flatten i o * StableHLO.batchSlice N ((tk + 1) * D) dy n o

Batched positional-embed GRADIENT denotes the certified Σ_n gradient — the summed cotangent, since the positional table is added to every token and its Jacobian is the identity. ⚠⚠ den_patchEmbedBiasGradB's neighbour warns that this batch sum is INVISIBLE at N = 1: a render that dropped it type-checks and emits the same bytes.

theorem Proofs.ViTPoCGB.clsGrad_denB {N : } (cotN : String) (Wc : Kernel4 192 3 16 16) (bc cls : Vec 192) (pos : Mat 197 192) (img : Vec (N * (3 * 224 * 224))) (dyEmbed : Vec (N * (197 * 192))) (i : Fin 192) :
StableHLO.den (StableHLO.SHlo.operand cotN (StableHLO.batchMap N (StableHLO.clsSliceFlat 196 192) dyEmbed)).denseBiasGradB i = n : Fin N, j : Fin (197 * 192), pdiv (fun (cl : Vec 192) => patchEmbed_flat 3 224 224 16 196 192 Wc bc cl pos (StableHLO.batchSlice N (3 * 224 * 224) img n)) cls i j * StableHLO.batchSlice N (197 * 192) dyEmbed n j

The CLS-token GRADIENT, and this is the one statement the per-example fold could not make. The render slices row 0 of every example's embed cotangent (clsSlice, lifted by batchOp) and reduces the result as an [N, D] batch, so the op is denseBiasGradB at N = B and its den sums over the batch — which is what a shared [192] parameter's gradient IS.

ViTPoCG.clsGrad_den is the same theorem at N = 1, where the batch lift lived in pretty B outside the AST. The bytes are identical (biasGrad's emitted reduce takes the B axis either way) and the functions are not; den_rowDenseBiasGradB_at_one is the general form of the trap.

⚠ Stated at the committed ViT-Tiny dims rather than generically, for ViTPoCG.clsGrad_den's reason: the CLS operand's type is Vec (N * (1 * D)), which reduces to Vec (N * D) only at a literal D.

theorem Proofs.ViTPoCGB.headWGradB_den {N D nC : } (aN cotN : String) (a : Vec (N * D)) (Wc : Mat D nC) (bc : Vec nC) (cot : Vec (N * nC)) (i : Fin D) (j : Fin nC) :
StableHLO.den (StableHLO.SHlo.weightGradB aN a (StableHLO.SHlo.operand cotN cot)) (finProdFinEquiv (i, j)) = n : Fin N, k : Fin nC, pdiv (fun (v : Vec (D * nC)) => dense (Mat.unflatten v) bc (StableHLO.batchSlice N D a n)) Wc.flatten (finProdFinEquiv (i, j)) k * StableHLO.batchSlice N nC cot n k

Batched classifier weight GRADIENT denotes the certified Σ_n outer product.

theorem Proofs.ViTPoCGB.headBGradB_den {N D nC : } (cotN : String) (Wc : Mat D nC) (a : Vec D) (bc : Vec nC) (cot : Vec (N * nC)) (n : Fin N) (i : Fin nC) :
StableHLO.batchSlice N nC (StableHLO.den (StableHLO.SHlo.operand cotN cot).biasGradB) n i = j : Fin nC, pdiv (fun (b' : Vec nC) => dense Wc b' a) bc i j * StableHLO.batchSlice N nC cot n j

Batched classifier bias GRADIENT denotes the certified cotangent, PER EXAMPLE.

biasGradB is the identity on its operand — the reduce over the batch is in the emitted text, outside the AST — so the statement this node supports is the per-example one at every batchSlice n, and it is the per-example biasGrad carve-out carried over rather than a new one. StableHLO.lean's constructor comment records the same thing on the emitter side.

def Proofs.ViTPoCGB.RowDenseWTiedB (N tk : ) {a c : } (xN cotN : String) (bb : Vec c) (x : Vec (N * (tk * a))) (W : Mat a c) (dy : Vec (N * (tk * c))) :

A batched per-token dense weight gradient node, tied (rowDenseWeightGradB_den).

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    def Proofs.ViTPoCGB.RowDenseBTiedB (N tk : ) {a c : } (cotN : String) (W : Mat a c) (x : Vec (N * (tk * a))) (b : Vec c) (dy : Vec (N * (tk * c))) :

    A batched per-token dense bias gradient node, tied (rowDenseBiasGradB_den).

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      def Proofs.ViTPoCGB.VecLNGammaTiedB (N R : ) {D : } (xN epsStr cotN : String) (ε : ) (βv : Vec D) (x : Vec (N * (R * D))) (γ : Vec D) (dy : Vec (N * (R * D))) :

      A batched vector-LN γ gradient node, tied (veclnGammaGradB_den).

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        def Proofs.ViTPoCGB.VecLNBetaTiedB (N R : ) {D : } (cotN : String) (ε : ) (γv : Vec D) (x : Vec (N * (R * D))) (β : Vec D) (dy : Vec (N * (R * D))) :

        A batched vector-LN β gradient node, tied (rowDenseBiasGradB_den_lnbeta).

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