ViT scaling pass — depth-k (general-depth tower, distinct per-block params)

The proven transformerTower_has_vjp_mat shares ONE param tuple across blocks; the 2-block vitForward2(V) carried distinct params but fixed the depth. This file closes general depth at the production form (vector-[D] LN + multi-head):

1. BlockParamsV — the 16-field per-block param structure, and vitBodyKVFlat — the depth-k block fold (head recursion: block 0 first), with vitBodyKVFlat_has_vjp by induction on k (the chain step is vjp_comp + the bridged transformerBlockV_has_vjp_mat, exactly vitForward2V_has_vjp's step with a Fin k param function). vitForwardKV(_has_vjp[_correct]) — the whole net at depth k, UNCONDITIONAL except 0 < ε. vitForwardKV_two_eq: at k = 2 it IS vitForward2V (definitional).

2. vitBodyGraphKMHV — the token-level fold of vitBlockGraphMHV with per-block SSA prefixes b{base+i}_, and vitFwdGraphKMHV_faithful: the depth-k multi-head vector-LN forward graph denotes vitForwardKV at heads := hm1 + 1 — by induction on k chaining vitBlockGraphMHV_den_aux + vitBlockSpelledMHV_eq per block (the per-block den_aux was designed for exactly this).

Depth-12 ViT-Tiny shapes are now a config change away (the production capstone needs only the P=16/D=192/heads=3 instantiation of these).

structure Proofs.BlockParamsV (D mlpDim : ℕ) : Type

The 16 per-block ViT parameters (vector-LN form), bundled so depth-k signatures stay sane (params : Fin k → BlockParamsV D mlpDim).

• γ1 : Vec D
• β1 : Vec D
• Wq : Mat D D
• Wk : Mat D D
• Wv : Mat D D
• Wo : Mat D D
• bq : Vec D
• bk : Vec D
• bv : Vec D
• bo : Vec D
• γ2 : Vec D
• β2 : Vec D
• Wfc1 : Mat D mlpDim
• bfc1 : Vec mlpDim
• Wfc2 : Mat mlpDim D
• bfc2 : Vec D

noncomputable def Proofs.blockV (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (p : BlockParamsV (heads * d_head) mlpDim) : Mat Np1 (heads * d_head) → Mat Np1 (heads * d_head)

transformerBlockV at a bundled param block.

Equations

• Proofs.blockV Np1 heads d_head mlpDim ε p = Proofs.transformerBlockV Np1 heads d_head mlpDim ε p.γ1 p.β1 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

noncomputable def Proofs.blockVFlat (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (p : BlockParamsV (heads * d_head) mlpDim) : Vec (Np1 * (heads * d_head)) → Vec (Np1 * (heads * d_head))

One block at the flat index (the vitForward2V per-block spelling).

Equations

• Proofs.blockVFlat Np1 heads d_head mlpDim ε p v = (Proofs.blockV Np1 heads d_head mlpDim ε p (Proofs.Mat.unflatten v)).flatten

noncomputable def Proofs.vitBodyKV (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (k : ℕ) : (Fin k → BlockParamsV (heads * d_head) mlpDim) → Mat Np1 (heads * d_head) → Mat Np1 (heads * d_head)

Depth-k block fold (Mat level, head recursion — block 0 runs first): body (k+1) ps = body k (ps ∘ succ) ∘ block (ps 0).

Equations

• Proofs.vitBodyKV Np1 heads d_head mlpDim ε 0 x_2 = fun (A : Proofs.Mat Np1 (heads * d_head)) => A
• Proofs.vitBodyKV Np1 heads d_head mlpDim ε k.succ ps = (Proofs.vitBodyKV Np1 heads d_head mlpDim ε k fun (i : Fin k) => ps i.succ) ∘ Proofs.blockV Np1 heads d_head mlpDim ε (ps 0)

noncomputable def Proofs.vitBodyKVFlat (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (k : ℕ) : (Fin k → BlockParamsV (heads * d_head) mlpDim) → Vec (Np1 * (heads * d_head)) → Vec (Np1 * (heads * d_head))

Depth-k block fold at the flat index — per-block flat stages (the vitForward2V spelling, so the VJP composes block-at-a-time).

Equations

• Proofs.vitBodyKVFlat Np1 heads d_head mlpDim ε 0 x_2 = fun (v : Proofs.Vec (Np1 * (heads * d_head))) => v
• Proofs.vitBodyKVFlat Np1 heads d_head mlpDim ε k.succ ps = (Proofs.vitBodyKVFlat Np1 heads d_head mlpDim ε k fun (i : Fin k) => ps i.succ) ∘ Proofs.blockVFlat Np1 heads d_head mlpDim ε (ps 0)

theorem Proofs.vitBodyKVFlat_eq_flatten (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (k : ℕ) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) (A : Mat Np1 (heads * d_head)) : vitBodyKVFlat Np1 heads d_head mlpDim ε k ps A.flatten = (vitBodyKV Np1 heads d_head mlpDim ε k ps A).flatten

The flat fold on a flattened input is the flatten of the Mat fold (the per-block unflatten ∘ flatten round-trips cancel, inductively).

theorem Proofs.vitBodyKVFlat_diff (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (hε : 0 < ε) (k : ℕ) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) : Differentiable ℝ (vitBodyKVFlat Np1 heads d_head mlpDim ε k ps)

Flat differentiability of the depth-k body, by induction on k.

noncomputable def Proofs.vitBodyKVFlat_has_vjp (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (hε : 0 < ε) (k : ℕ) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) : HasVJP (vitBodyKVFlat Np1 heads d_head mlpDim ε k ps)

Depth-k body VJP — the tower induction at distinct per-block params: the chain step is vjp_comp gluing the bridged transformerBlockV_has_vjp_mat onto the depth-k tail. Only 0 < ε.

Equations

• One or more equations did not get rendered due to their size.
• Proofs.vitBodyKVFlat_has_vjp Np1 heads d_head mlpDim ε hε 0 x_2 = Proofs.identity_has_vjp (Np1 * (heads * d_head))

noncomputable def Proofs.vitForwardKV (ic H W patchSize N mlpDim heads d_head nClasses k : ℕ) (W_conv : Kernel4 (heads * d_head) ic patchSize patchSize) (b_conv cls_token : Vec (heads * d_head)) (pos_embed : Mat (N + 1) (heads * d_head)) (ε : ℝ) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) (γF βF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) : Vec (ic * H * W) → Vec nClasses

Depth-k distinct-param ViT forward (vector-LN): patch embed → k blocks (Fin k → BlockParamsV) → final vector-LN → CLS slice → dense head. vitForward2V generalized over depth.

Equations

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

theorem Proofs.vitForwardKV_two_eq (ic H W patchSize N mlpDim heads d_head nClasses : ℕ) (W_conv : Kernel4 (heads * d_head) ic patchSize patchSize) (b_conv cls_token : Vec (heads * d_head)) (pos_embed : Mat (N + 1) (heads * d_head)) (ε : ℝ) (ps : Fin 2 → BlockParamsV (heads * d_head) mlpDim) (γF βF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) : vitForwardKV ic H W patchSize N mlpDim heads d_head nClasses 2 W_conv b_conv cls_token pos_embed ε ps γF βF Wcls bcls = vitForward2V ic H W patchSize N mlpDim heads d_head nClasses W_conv b_conv cls_token pos_embed ε (ps 0).γ1 (ps 0).β1 (ps 0).Wq (ps 0).Wk (ps 0).Wv (ps 0).Wo (ps 0).bq (ps 0).bk (ps 0).bv (ps 0).bo (ps 0).γ2 (ps 0).β2 (ps 0).Wfc1 (ps 0).bfc1 (ps 0).Wfc2 (ps 0).bfc2 (ps 1).γ1 (ps 1).β1 (ps 1).Wq (ps 1).Wk (ps 1).Wv (ps 1).Wo (ps 1).bq (ps 1).bk (ps 1).bv (ps 1).bo (ps 1).γ2 (ps 1).β2 (ps 1).Wfc1 (ps 1).bfc1 (ps 1).Wfc2 (ps 1).bfc2 γF βF Wcls bcls

At k = 2 the depth-k net IS vitForward2V (definitional — the fold unrolls to exactly the 2-block composition).

noncomputable def Proofs.vitForwardKV_has_vjp (ic H W patchSize N mlpDim heads d_head nClasses k : ℕ) (W_conv : Kernel4 (heads * d_head) ic patchSize patchSize) (b_conv cls_token : Vec (heads * d_head)) (pos_embed : Mat (N + 1) (heads * d_head)) (ε : ℝ) (hε : 0 < ε) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) (γF βF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) : HasVJP (vitForwardKV ic H W patchSize N mlpDim heads d_head nClasses k W_conv b_conv cls_token pos_embed ε ps γF βF Wcls bcls)

Whole-net VJP for the depth-k ViT (global). All-smooth, so the only hypothesis is 0 < ε — at EVERY depth. Three vjp_comp steps gluing patchEmbed_flat_has_vjp, the inductive vitBodyKVFlat_has_vjp, the bridged per-token vector-LN, and classifier_flat_has_vjp.

Equations

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

theorem Proofs.vitForwardKV_has_vjp_correct (ic H W patchSize N mlpDim heads d_head nClasses k : ℕ) (W_conv : Kernel4 (heads * d_head) ic patchSize patchSize) (b_conv cls_token : Vec (heads * d_head)) (pos_embed : Mat (N + 1) (heads * d_head)) (ε : ℝ) (hε : 0 < ε) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) (γF βF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) (dy : Vec nClasses) (i : Fin (ic * H * W)) : (vitForwardKV_has_vjp ic H W patchSize N mlpDim heads d_head nClasses k W_conv b_conv cls_token pos_embed ε hε ps γF βF Wcls bcls).backward x dy i = ∑ j : Fin nClasses, pdiv (vitForwardKV ic H W patchSize N mlpDim heads d_head nClasses k W_conv b_conv cls_token pos_embed ε ps γF βF Wcls bcls) x i j * dy j

Public correctness theorem for vitForwardKV_has_vjp — the depth-k ViT's backward equals the pdiv-contracted Jacobian at every input.

theorem Proofs.vitTiny_has_vjp_correct (W_conv : Kernel4 (3 * 64) 3 16 16) (b_conv cls_token : Vec (3 * 64)) (pos_embed : Mat (196 + 1) (3 * 64)) (ε : ℝ) (hε : 0 < ε) (ps : Fin 12 → BlockParamsV (3 * 64) 768) (γF βF : Vec (3 * 64)) (Wcls : Mat (3 * 64) 10) (bcls : Vec 10) (x : Vec (3 * 224 * 224)) (dy : Vec 10) (i : Fin (3 * 224 * 224)) : (vitForwardKV_has_vjp 3 224 224 16 196 768 3 64 10 12 W_conv b_conv cls_token pos_embed ε hε ps γF βF Wcls bcls).backward x dy i = ∑ j : Fin 10, pdiv (vitForwardKV 3 224 224 16 196 768 3 64 10 12 W_conv b_conv cls_token pos_embed ε ps γF βF Wcls bcls) x i j * dy j

ViT-Tiny whole-network VJP — the production capstone.

vitForwardKV_has_vjp_correct instantiated at the exact MainVitTrain.lean vitTiny spec: a 3×224×224 image, 16×16 patches (N = 196 patch tokens

• the CLS token), embedding dim D = 192 = 3 heads × 64, MLP dim 768, 12 transformer blocks with DISTINCT per-block parameters (ps : Fin 12 → BlockParamsV 192 768), and Imagenette's 10 classes.

The full 12-block / 3-head ViT-Tiny's backward pass equals its Mathlib-fderiv Jacobian-transpose contracted with the cotangent, at every input image — UNCONDITIONAL except 0 < ε (softmax / GELU / vector-LN are kink-free, so no smoothness witness is needed, and the statement is generic in the weights, so it is non-degenerate by construction). The ViT peer of convNextForwardTCh_has_vjp (18-block ConvNeXt-T) and efficientnetForwardB_full_has_vjp (16-block EfficientNet-B0): a full-spec, real-architecture whole-network backward.

def Proofs.StableHLO.vitBlockGraphMHVP {Np1 hm1 d mlpDim : ℕ} (pfx epsStr sStr oneStr zeroStr : String) (ε s : ℝ) (p : BlockParamsV ((hm1 + 1) * d) mlpDim) (x : SHlo (Np1 * ((hm1 + 1) * d))) : SHlo (Np1 * ((hm1 + 1) * d))

vitBlockGraphMHV at a bundled param block.

Equations

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

def Proofs.StableHLO.vitBodyGraphKMHV {Np1 hm1 d mlpDim : ℕ} (epsStr sStr oneStr zeroStr : String) (ε s : ℝ) (base k : ℕ) : (Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) → SHlo (Np1 * ((hm1 + 1) * d)) → SHlo (Np1 * ((hm1 + 1) * d))

Depth-k token-level block fold — block base first, SSA prefixes b{base+1}_, b{base+2}_, … (distinct per block).

Equations

• One or more equations did not get rendered due to their size.
• Proofs.StableHLO.vitBodyGraphKMHV epsStr sStr oneStr zeroStr ε s x✝¹ 0 x_5 x✝ = x✝

theorem Proofs.StableHLO.vitBodyGraphKMHV_den {Np1 hm1 d mlpDim : ℕ} (epsStr sStr oneStr zeroStr : String) (ε : ℝ) (base k : ℕ) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (e : SHlo (Np1 * ((hm1 + 1) * d))) (A : Mat Np1 ((hm1 + 1) * d)) : den e = A.flatten → den (vitBodyGraphKMHV epsStr sStr oneStr zeroStr ε (sdpa_scale d) base k ps e) = (vitBodyKV Np1 (hm1 + 1) d mlpDim ε k ps A).flatten

Depth-k body denotation — by induction on k, chaining vitBlockGraphMHV_den_aux + vitBlockSpelledMHV_eq per block: the token fold denotes the flatten of the Mat block fold at heads := hm1 + 1.

def Proofs.StableHLO.vitFwdGraphKMHV {ic H W P N hm1 d mlpDim nClasses : ℕ} (epsStr sStr oneStr zeroStr : String) (ε s : ℝ) (Wc : Kernel4 ((hm1 + 1) * d) ic P P) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (k : ℕ) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (γF βF : Vec ((hm1 + 1) * d)) (Wcls : Mat ((hm1 + 1) * d) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) : SHlo nClasses

Whole depth-k multi-head vector-LN ViT forward graph: patch embed → k spelled multi-head vector-LN blocks (b1_…b{k}_, distinct params) → final vector-LN → CLS slice → dense head.

Equations

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

theorem Proofs.StableHLO.vitFwdGraphKMHV_faithful (ic H W patchSize N hm1 d mlpDim nClasses : ℕ) (epsStr sStr oneStr zeroStr : String) (Wc : Kernel4 ((hm1 + 1) * d) ic patchSize patchSize) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (ε : ℝ) (k : ℕ) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (γF βF : Vec ((hm1 + 1) * d)) (Wcls : Mat ((hm1 + 1) * d) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) : den (vitFwdGraphKMHV epsStr sStr oneStr zeroStr ε (sdpa_scale d) Wc bc cls pos k ps γF βF Wcls bcls x) = vitForwardKV ic H W patchSize N mlpDim (hm1 + 1) d nClasses k Wc bc cls pos ε ps γF βF Wcls bcls x

Depth-k multi-head vector-LN ViT forward faithfulness — the general-depth graph denotes vitForwardKV at heads := hm1 + 1, for EVERY depth k. The depth analogue of vitFwdGraphMHV_faithful.