ViT scaling pass — the vector-[D] LayerNorm upgrade

The representative ViT close (Items A–D) used the proof's scalar LN γ/β. The committed production render (ViTRender.lean, the GPU-trained ViT-Tiny) is MORE faithful: vector γ, β : [D] per LN site, decomposed as scalar-LN(1,0) ∘ per-channel scale γ ∘ + β. This file brings the close to that form — planning/archive/vit_close.md's top scaling-pass item:

• layerNormVec — per-token normalize (scalar-LN at γ=1, β=0) then the per-channel affine γ ⊙ · + β, with HasVJP composed from layerNorm_has_vjp (at 1,0), layerScale_has_vjp (ch9), and the bias translation.
• transformerBlockV / vitForward2V — the vector-LN block and 2-block net, with the whole-net VJP re-composed through the same sublayer recipe as the scalar one (biPathMat_has_vjp + vjpMat_comp + rowwise_has_vjp_mat). UNCONDITIONAL except 0 < ε.
• vitFwdGraphV + vitFwdGraphV_faithful — the graph spells each LN site with the new broadcast tokens: lnRowF(1,0) → rowScaleF γ → rowBiasF β (exactly the ViTRender decomposition); faithful at heads = 1.
• Vector γ/β param bridges (vit_render_vecln{gamma,beta}_certified) — the per-channel grads dγ_k = Σ_tokens dy_(r,k)·x̂_r(k), dβ_k = Σ_tokens dy_(r,k) (reduce over batch+tokens, KEEP the channel axis), certified via the masked-gather Jacobian recipe.

noncomputable def Proofs.layerNormVec (D : ℕ) (ε : ℝ) (γv βv x : Vec D) : Vec D

Vector-[D] LayerNorm: per-token normalize (the scalar LN at γ=1, β=0 — pure x̂), then the per-channel affine γ ⊙ x̂ + β. The committed ViTRender LN form.

Equations

• Proofs.layerNormVec D ε γv βv x k = γv k * Proofs.layerNormForward D ε 1 0 x k + βv k

theorem Proofs.layerNormVec_diff (D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) : Differentiable ℝ (layerNormVec D ε γv βv)

noncomputable def Proofs.biasAdd_has_vjp {n : ℕ} (βv : Vec n) : HasVJP fun (z : Vec n) (k : Fin n) => z k + βv k

The bias translation's VJP — backward is the identity (dx = dy).

Equations

• Proofs.biasAdd_has_vjp βv = { backward := fun (_z dy : Proofs.Vec n) => dy, correct := ⋯ }

noncomputable def Proofs.layerNormVec_has_vjp (D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) : HasVJP (layerNormVec D ε γv βv)

Vector-LN VJP — (+β) ∘ layerScale γ ∘ LN(1,0), three proven pieces glued by vjp_comp. Only 0 < ε.

Equations

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

noncomputable def Proofs.layerNormVec_per_token_has_vjp_mat (N D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) : HasVJPMat fun (X : Mat N D) (r : Fin N) => layerNormVec D ε γv βv (X r)

Per-token vector-LN across a sequence — the rowwise lift.

Equations

• Proofs.layerNormVec_per_token_has_vjp_mat N D ε γv βv hε = Proofs.rowwise_has_vjp_mat (Proofs.layerNormVec_has_vjp D ε γv βv hε) ⋯

theorem Proofs.rowwise_flat_diff {N D P : ℕ} (g : Vec D → Vec P) (hg : Differentiable ℝ g) : Differentiable ℝ fun (v : Vec (N * D)) => Mat.flatten ((fun (X : Mat N D) (n : Fin N) => g (X n)) (Mat.unflatten v))

Generic flat differentiability of a rowwise lift — each output coordinate projects through a row-projection CLM into the per-row map (the layerNorm_per_token_flat_diff recipe with the row map abstracted).

theorem Proofs.layerNormVec_per_token_flat_diff (N D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) : Differentiable ℝ fun (v : Vec (N * D)) => Mat.flatten ((fun (X : Mat N D) (n : Fin N) => layerNormVec D ε γv βv (X n)) (Mat.unflatten v))

noncomputable def Proofs.transformerAttnSublayerV (N heads d_head : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : Mat N (heads * d_head) → Mat N (heads * d_head)

Attention sublayer with vector-LN: X ↦ X + MHSA(LNᵥ(X)).

Equations

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

noncomputable def Proofs.transformerMlpSublayerV (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Mat N (heads * d_head) → Mat N (heads * d_head)

MLP sublayer with vector-LN: h ↦ h + MLP(LNᵥ(h)).

Equations

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

noncomputable def Proofs.transformerBlockV (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Mat N (heads * d_head) → Mat N (heads * d_head)

Vector-LN transformer block: MLPᵥ-sublayer ∘ attentionᵥ-sublayer — the ViTRender block form.

Equations

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

theorem Proofs.transformerAttnSublayerV_inner_flat_diff (N heads d_head : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => ((mhsa_layer N heads d_head Wq Wk Wv Wo bq bk bv bo ∘ fun (X : Mat N (heads * d_head)) (n : Fin N) => layerNormVec (heads * d_head) ε γ1 β1 (X n)) (Mat.unflatten v)).flatten

Flat Diff of the attentionᵥ sublayer's non-trivial arm (mhsa ∘ LNᵥ).

theorem Proofs.transformerAttnSublayerV_flat_diff (N heads d_head : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerAttnSublayerV N heads d_head ε γ1 β1 Wq Wk Wv Wo bq bk bv bo (Mat.unflatten v)).flatten

Flat Diff of the attentionᵥ sublayer.

noncomputable def Proofs.transformerAttnSublayerV_has_vjp_mat (N heads d_head : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) : HasVJPMat (transformerAttnSublayerV N heads d_head ε γ1 β1 Wq Wk Wv Wo bq bk bv bo)

Attentionᵥ sublayer VJP.

Equations

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

theorem Proofs.transformerMlpSublayerV_inner_flat_diff (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ2 β2 : Vec (heads * d_head)) (hε : 0 < ε) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => ((transformerMlp N (heads * d_head) mlpDim Wfc1 bfc1 Wfc2 bfc2 ∘ fun (X : Mat N (heads * d_head)) (n : Fin N) => layerNormVec (heads * d_head) ε γ2 β2 (X n)) (Mat.unflatten v)).flatten

Flat Diff of the MLPᵥ sublayer's non-trivial arm.

theorem Proofs.transformerMlpSublayerV_flat_diff (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ2 β2 : Vec (heads * d_head)) (hε : 0 < ε) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerMlpSublayerV N heads d_head mlpDim ε γ2 β2 Wfc1 bfc1 Wfc2 bfc2 (Mat.unflatten v)).flatten

Flat Diff of the MLPᵥ sublayer.

noncomputable def Proofs.transformerMlpSublayerV_has_vjp_mat (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ2 β2 : Vec (heads * d_head)) (hε : 0 < ε) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : HasVJPMat (transformerMlpSublayerV N heads d_head mlpDim ε γ2 β2 Wfc1 bfc1 Wfc2 bfc2)

MLPᵥ sublayer VJP.

Equations

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

theorem Proofs.transformerBlockV_flat_diff (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerBlockV N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 (Mat.unflatten v)).flatten

Flat Diff of the vector-LN block.

noncomputable def Proofs.transformerBlockV_has_vjp_mat (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) : HasVJPMat (transformerBlockV N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2)

Vector-LN block VJP — one vjpMat_comp of the two sublayer witnesses.

Equations

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

noncomputable def Proofs.vitForward2V (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)) (ε : ℝ) (γ1₁ β1₁ : Vec (heads * d_head)) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat (heads * d_head) (heads * d_head)) (bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ : Vec (heads * d_head)) (Wfc1₁ : Mat (heads * d_head) mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim (heads * d_head)) (bfc2₁ γ1₂ β1₂ : Vec (heads * d_head)) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat (heads * d_head) (heads * d_head)) (bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ : Vec (heads * d_head)) (Wfc1₂ : Mat (heads * d_head) mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim (heads * d_head)) (bfc2₂ γF βF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) : Vec (ic * H * W) → Vec nClasses

Distinct-param 2-block ViT forward at vector-[D] LN — the production ViTRender LN form at the representative architecture.

Equations

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

noncomputable def Proofs.vitForward2V_has_vjp (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)) (ε : ℝ) (hε : 0 < ε) (γ1₁ β1₁ : Vec (heads * d_head)) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat (heads * d_head) (heads * d_head)) (bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ : Vec (heads * d_head)) (Wfc1₁ : Mat (heads * d_head) mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim (heads * d_head)) (bfc2₁ γ1₂ β1₂ : Vec (heads * d_head)) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat (heads * d_head) (heads * d_head)) (bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ : Vec (heads * d_head)) (Wfc1₂ : Mat (heads * d_head) mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim (heads * d_head)) (bfc2₂ γF βF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (bcls : Vec nClasses) : HasVJP (vitForward2V ic H W patchSize N mlpDim heads d_head nClasses W_conv b_conv cls_token pos_embed ε γ1₁ β1₁ Wq₁ Wk₁ Wv₁ Wo₁ bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ Wfc1₁ bfc1₁ Wfc2₁ bfc2₁ γ1₂ β1₂ Wq₂ Wk₂ Wv₂ Wo₂ bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ Wfc1₂ bfc1₂ Wfc2₂ bfc2₂ γF βF Wcls bcls)

Whole-net VJP for the vector-LN 2-block ViT (global) — only 0 < ε.

Equations

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

theorem Proofs.vitForward2V_has_vjp_correct (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)) (ε : ℝ) (hε : 0 < ε) (γ1₁ β1₁ : Vec (heads * d_head)) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat (heads * d_head) (heads * d_head)) (bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ : Vec (heads * d_head)) (Wfc1₁ : Mat (heads * d_head) mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim (heads * d_head)) (bfc2₁ γ1₂ β1₂ : Vec (heads * d_head)) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat (heads * d_head) (heads * d_head)) (bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ : Vec (heads * d_head)) (Wfc1₂ : Mat (heads * d_head) mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim (heads * d_head)) (bfc2₂ γ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)) : (vitForward2V_has_vjp ic H W patchSize N mlpDim heads d_head nClasses W_conv b_conv cls_token pos_embed ε hε γ1₁ β1₁ Wq₁ Wk₁ Wv₁ Wo₁ bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ Wfc1₁ bfc1₁ Wfc2₁ bfc2₁ γ1₂ β1₂ Wq₂ Wk₂ Wv₂ Wo₂ bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ Wfc1₂ bfc1₂ Wfc2₂ bfc2₂ γF βF Wcls bcls).backward x dy i = ∑ j : Fin nClasses, pdiv (vitForward2V ic H W patchSize N mlpDim heads d_head nClasses W_conv b_conv cls_token pos_embed ε γ1₁ β1₁ Wq₁ Wk₁ Wv₁ Wo₁ bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ Wfc1₁ bfc1₁ Wfc2₁ bfc2₁ γ1₂ β1₂ Wq₂ Wk₂ Wv₂ Wo₂ bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ Wfc1₂ bfc1₂ Wfc2₂ bfc2₂ γF βF Wcls bcls) x i j * dy j

Public correctness theorem for vitForward2V_has_vjp — the vector-LN 2-block ViT's backward equals the pdiv-contracted Jacobian at every input.

noncomputable def Proofs.vitBlockSpelledV (Np1 d mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (1 * d)) (Wq Wk Wv Wo : Mat (1 * d) (1 * d)) (bq bk bv bo γ2 β2 : Vec (1 * d)) (Wfc1 : Mat (1 * d) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (1 * d)) (bfc2 : Vec (1 * d)) (X : Mat Np1 (1 * d)) : Mat Np1 (1 * d)

The spelled vector-LN block at heads = 1 — each LN site decomposed as the graph (and ViTRender) emit it: pure normalize (scalar-LN at 1,0) → per-channel scale → per-channel bias.

Equations

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

theorem Proofs.vitBlockSpelledV_eq (Np1 d mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (1 * d)) (Wq Wk Wv Wo : Mat (1 * d) (1 * d)) (bq bk bv bo γ2 β2 : Vec (1 * d)) (Wfc1 : Mat (1 * d) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (1 * d)) (bfc2 : Vec (1 * d)) (X : Mat Np1 (1 * d)) : vitBlockSpelledV Np1 d mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 X = transformerBlockV Np1 1 d mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 X

The spelled vector-LN block IS transformerBlockV at one head — the three-stage LN decomposition collapses to layerNormVec definitionally; the per-head plumbing collapses via mhsa_layer_one_head.

theorem Proofs.StableHLO.rowScaleFlat_flat {m n : ℕ} (γ : Vec n) (A : Mat m n) : rowScaleFlat m n γ A.flatten = Mat.flatten fun (r : Fin m) => layerScale γ (A r)

Row-broadcast scale commutation bridge.

theorem Proofs.StableHLO.rowBiasFlat_flat {m n : ℕ} (β : Vec n) (A : Mat m n) : rowBiasFlat m n β A.flatten = Mat.flatten fun (r : Fin m) (k : Fin n) => A r k + β k

Row-broadcast bias commutation bridge.

def Proofs.StableHLO.vitBlockGraphV {Np1 D mlpDim : ℕ} (pfx epsStr sStr oneStr zeroStr : String) (ε s : ℝ) (γ1 β1 : Vec D) (Wq Wk Wv Wo : Mat D D) (bq bk bv bo γ2 β2 : Vec D) (Wfc1 : Mat D mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim D) (bfc2 : Vec D) (x : SHlo (Np1 * D)) : SHlo (Np1 * D)

One spelled vector-LN block over the tokens: each LN site is lnRowF(1,0) → rowScaleF γ → rowBiasF β (the ViTRender decomposition); everything else as vitBlockGraph. oneStr/zeroStr name the rendered constant-1/0 scalars the pure-normalize sites reference.

Equations

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

def Proofs.StableHLO.vitFwdGraphV {ic H W P N D mlpDim nClasses : ℕ} (epsStr sStr oneStr zeroStr : String) (ε s : ℝ) (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (γ1₁ β1₁ : Vec D) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat D D) (bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ : Vec D) (Wfc1₁ : Mat D mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim D) (bfc2₁ γ1₂ β1₂ : Vec D) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat D D) (bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ : Vec D) (Wfc1₂ : Mat D mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim D) (bfc2₂ γF βF : Vec D) (Wcls : Mat D nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) : SHlo nClasses

Whole vector-LN ViT forward graph: patch embed → 2 spelled vector-LN blocks (distinct params) → final vector-LN (same three-token decomposition) → CLS slice → dense head.

Equations

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

theorem Proofs.StableHLO.vitFwdGraphV_faithful (ic H W patchSize N d mlpDim nClasses : ℕ) (epsStr sStr oneStr zeroStr : String) (Wc : Kernel4 (1 * d) ic patchSize patchSize) (bc cls : Vec (1 * d)) (pos : Mat (N + 1) (1 * d)) (ε : ℝ) (γ1₁ β1₁ : Vec (1 * d)) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat (1 * d) (1 * d)) (bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ : Vec (1 * d)) (Wfc1₁ : Mat (1 * d) mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim (1 * d)) (bfc2₁ γ1₂ β1₂ : Vec (1 * d)) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat (1 * d) (1 * d)) (bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ : Vec (1 * d)) (Wfc1₂ : Mat (1 * d) mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim (1 * d)) (bfc2₂ γF βF : Vec (1 * d)) (Wcls : Mat (1 * d) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) : den (vitFwdGraphV epsStr sStr oneStr zeroStr ε (sdpa_scale d) Wc bc cls pos γ1₁ β1₁ Wq₁ Wk₁ Wv₁ Wo₁ bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ Wfc1₁ bfc1₁ Wfc2₁ bfc2₁ γ1₂ β1₂ Wq₂ Wk₂ Wv₂ Wo₂ bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ Wfc1₂ bfc1₂ Wfc2₂ bfc2₂ γF βF Wcls bcls x) = vitForward2V ic H W patchSize N mlpDim 1 d nClasses Wc bc cls pos ε γ1₁ β1₁ Wq₁ Wk₁ Wv₁ Wo₁ bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ Wfc1₁ bfc1₁ Wfc2₁ bfc2₁ γ1₂ β1₂ Wq₂ Wk₂ Wv₂ Wo₂ bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ Wfc1₂ bfc1₂ Wfc2₂ bfc2₂ γF βF Wcls bcls x

Vector-LN ViT forward faithfulness — the graph denotes vitForward2V at one head. The scaling-pass peer of vitFwdGraph_faithful.

theorem Proofs.pdiv_vecLN_gamma {N D : ℕ} (ε : ℝ) (βv : Vec D) (X : Mat N D) (γ : Vec D) (i : Fin D) (o : Fin (N * D)) : pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε gv βv (X r)) γ i o = layerNormForward D ε 1 0 (X (finProdFinEquiv.symm o).1) (finProdFinEquiv.symm o).2 * if i = (finProdFinEquiv.symm o).2 then 1 else 0

Jacobian of the rowwise vector-LN site w.r.t. γv — ∂y_(r,k)/∂γv_i = δ_(i,k)·x̂_r(k).

theorem Proofs.pdiv_vecLN_beta {N D : ℕ} (ε : ℝ) (γv : Vec D) (X : Mat N D) (β : Vec D) (i : Fin D) (o : Fin (N * D)) : pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε γv bv (X r)) β i o = if i = (finProdFinEquiv.symm o).2 then 1 else 0

Jacobian of the rowwise vector-LN site w.r.t. βv — ∂y_(r,k)/∂βv_i = δ_(i,k).

noncomputable def Proofs.vecLN_grad_gamma (N D : ℕ) (ε : ℝ) (X dY : Mat N D) : Vec D

The rendered vector-LN γ gradient: per-channel, the batch+token reduce dγ_k = Σ_r dY_(r,k)·x̂_r(k) (KEEPS the channel axis — ViTRender's form).

Equations

• Proofs.vecLN_grad_gamma N D ε X dY i = ∑ r : Fin N, dY r i * Proofs.layerNormForward D ε 1 0 (X r) i

noncomputable def Proofs.vecLN_grad_beta (N D : ℕ) (dY : Mat N D) : Vec D

The rendered vector-LN β gradient: dβ_k = Σ_r dY_(r,k).

Equations

• Proofs.vecLN_grad_beta N D dY i = ∑ r : Fin N, dY r i

theorem Proofs.vit_veclnGamma_grad_bridge {N D : ℕ} (ε : ℝ) (βv γ : Vec D) (X : Mat N D) (dy : Vec (N * D)) (i : Fin D) : vecLN_grad_gamma N D ε X (Mat.unflatten dy) i = ∑ o : Fin (N * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε gv βv (X r)) γ i o * dy o

Vector-LN γ-gradient bridge.

theorem Proofs.vit_veclnBeta_grad_bridge {N D : ℕ} (ε : ℝ) (γv β : Vec D) (X : Mat N D) (dy : Vec (N * D)) (i : Fin D) : vecLN_grad_beta N D (Mat.unflatten dy) i = ∑ o : Fin (N * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε γv bv (X r)) β i o * dy o

Vector-LN β-gradient bridge.

theorem Proofs.vit_render_veclngamma_certified {N D : ℕ} (ε : ℝ) (βv γ : Vec D) (X : Mat N D) (dy : Vec (N * D)) (lr : ℝ) (i : Fin D) : γ i - lr * vecLN_grad_gamma N D ε X (Mat.unflatten dy) i = γ i - lr * ∑ o : Fin (N * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε gv βv (X r)) γ i o * dy o

Vector-LN γ output, certified. γvⁿ_k = γv_k − lr·(Σ_tokens dy·x̂)_k denotes the certified rowwise vector-LN ∂/∂γv contraction. Covers all five LN sites of the vector-LN representative (and is the ViTRender per-channel LN-γ reduce).

theorem Proofs.vit_render_veclnbeta_certified {N D : ℕ} (ε : ℝ) (γv β : Vec D) (X : Mat N D) (dy : Vec (N * D)) (lr : ℝ) (i : Fin D) : β i - lr * vecLN_grad_beta N D (Mat.unflatten dy) i = β i - lr * ∑ o : Fin (N * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε γv bv (X r)) β i o * dy o

Vector-LN β output, certified.

noncomputable def Proofs.vitCotHV {Np1 D mlpDim : ℕ} (ε : ℝ) (γ2 : Vec D) (Wfc1 : Mat D mlpDim) (Wfc2 : Mat mlpDim D) (h : Vec (Np1 * D)) (m1 : Vec (Np1 * mlpDim)) (dyOut : Vec (Np1 * D)) : Vec (Np1 * D)

Cot at the attention-sublayer output h, vector-LN form: dyOut + the decomposed LN₂ input-VJP (rowScaleFlat γ2 then rowLNBackFlat at γ=1).

Equations

• Proofs.vitCotHV ε γ2 Wfc1 Wfc2 h m1 dyOut i = dyOut i + Proofs.StableHLO.rowLNBackFlat Np1 D ε 1 h (Proofs.StableHLO.rowScaleFlat Np1 D γ2 (Proofs.vitCotLn2 Wfc1 Wfc2 m1 dyOut)) i

noncomputable def Proofs.vitCotAttV {Np1 D mlpDim : ℕ} (ε : ℝ) (γ2 : Vec D) (Wo : Mat D D) (Wfc1 : Mat D mlpDim) (Wfc2 : Mat mlpDim D) (h : Vec (Np1 * D)) (m1 : Vec (Np1 * mlpDim)) (dyOut : Vec (Np1 * D)) : Vec (Np1 * D)

Cot at the SDPA output, vector-LN form.

Equations

• Proofs.vitCotAttV ε γ2 Wo Wfc1 Wfc2 h m1 dyOut = Proofs.StableHLO.rowDenseBackFlat Np1 D D Wo (Proofs.vitCotHV ε γ2 Wfc1 Wfc2 h m1 dyOut)

noncomputable def Proofs.vitCotXinV {Np1 D : ℕ} (ε : ℝ) (γ1 : Vec D) (Wq Wk Wv : Mat D D) (xin dQ dK dV cotH : Vec (Np1 * D)) : Vec (Np1 * D)

Cot at the block input, vector-LN form: cotH + the decomposed LN₁ input-VJP of the three-way Q/K/V fan-in.

Equations

• Proofs.vitCotXinV ε γ1 Wq Wk Wv xin dQ dK dV cotH i = cotH i + Proofs.StableHLO.rowLNBackFlat Np1 D ε 1 xin (Proofs.StableHLO.rowScaleFlat Np1 D γ1 (Proofs.vitCotLn1 Wq Wk Wv dQ dK dV)) i

noncomputable def Proofs.vitCotB2outV (N D nClasses : ℕ) (ε : ℝ) (γF : Vec D) (Wcls : Mat D nClasses) (b2out : Vec ((N + 1) * D)) (dy : Vec nClasses) : Vec ((N + 1) * D)

Cot at block 2's output, vector-LN form: the decomposed final-LN input-VJP of the classifier-back row-0 scatter.

Equations

• Proofs.vitCotB2outV N D nClasses ε γF Wcls b2out dy = Proofs.StableHLO.rowLNBackFlat (N + 1) D ε 1 b2out (Proofs.StableHLO.rowScaleFlat (N + 1) D γF (Proofs.vitCotFl N D nClasses Wcls dy))

theorem Proofs.vit_render_vecln2gamma_chain_certified {Np1 D mlpDim : ℕ} (ε : ℝ) (βv γ : Vec D) (h : Vec (Np1 * D)) (Wfc1 : Mat D mlpDim) (Wfc2 : Mat mlpDim D) (m1 : Vec (Np1 * mlpDim)) (dyOut : Vec (Np1 * D)) (lr : ℝ) (i : Fin D) : γ i - lr * vecLN_grad_gamma Np1 D ε (Mat.unflatten h) (Mat.unflatten (vitCotLn2 Wfc1 Wfc2 m1 dyOut)) i = γ i - lr * ∑ o : Fin (Np1 * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin Np1) => layerNormVec D ε gv βv (Mat.unflatten h r)) γ i o * vitCotLn2 Wfc1 Wfc2 m1 dyOut o

Vector LN₂ γ, chain-certified at vitCotLn2 (the fc2-back → GELU mask → fc1-back cotangent), with the saved attn-sublayer output h as the LN input.

theorem Proofs.vit_render_vecln2beta_chain_certified {Np1 D mlpDim : ℕ} (ε : ℝ) (γv β : Vec D) (h : Vec (Np1 * D)) (Wfc1 : Mat D mlpDim) (Wfc2 : Mat mlpDim D) (m1 : Vec (Np1 * mlpDim)) (dyOut : Vec (Np1 * D)) (lr : ℝ) (i : Fin D) : β i - lr * vecLN_grad_beta Np1 D (Mat.unflatten (vitCotLn2 Wfc1 Wfc2 m1 dyOut)) i = β i - lr * ∑ o : Fin (Np1 * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin Np1) => layerNormVec D ε γv bv (Mat.unflatten h r)) β i o * vitCotLn2 Wfc1 Wfc2 m1 dyOut o

Vector LN₂ β, chain-certified.

theorem Proofs.vit_render_vecln1gamma_chain_certified {Np1 D : ℕ} (ε : ℝ) (βv γ : Vec D) (xin : Vec (Np1 * D)) (Wq Wk Wv : Mat D D) (dQ dK dV : Vec (Np1 * D)) (lr : ℝ) (i : Fin D) : γ i - lr * vecLN_grad_gamma Np1 D ε (Mat.unflatten xin) (Mat.unflatten (vitCotLn1 Wq Wk Wv dQ dK dV)) i = γ i - lr * ∑ o : Fin (Np1 * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin Np1) => layerNormVec D ε gv βv (Mat.unflatten xin r)) γ i o * vitCotLn1 Wq Wk Wv dQ dK dV o

Vector LN₁ γ, chain-certified at the three-way Q/K/V fan-in vitCotLn1, with the saved block input xin as the LN input.

theorem Proofs.vit_render_vecln1beta_chain_certified {Np1 D : ℕ} (ε : ℝ) (γv β : Vec D) (xin : Vec (Np1 * D)) (Wq Wk Wv : Mat D D) (dQ dK dV : Vec (Np1 * D)) (lr : ℝ) (i : Fin D) : β i - lr * vecLN_grad_beta Np1 D (Mat.unflatten (vitCotLn1 Wq Wk Wv dQ dK dV)) i = β i - lr * ∑ o : Fin (Np1 * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin Np1) => layerNormVec D ε γv bv (Mat.unflatten xin r)) β i o * vitCotLn1 Wq Wk Wv dQ dK dV o

Vector LN₁ β, chain-certified.

theorem Proofs.vit_render_veclnFgamma_chain_certified {N D nClasses : ℕ} (ε : ℝ) (βv γ : Vec D) (b2out : Vec ((N + 1) * D)) (Wcls : Mat D nClasses) (dy : Vec nClasses) (lr : ℝ) (i : Fin D) : γ i - lr * vecLN_grad_gamma (N + 1) D ε (Mat.unflatten b2out) (Mat.unflatten (vitCotFl N D nClasses Wcls dy)) i = γ i - lr * ∑ o : Fin ((N + 1) * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin (N + 1)) => layerNormVec D ε gv βv (Mat.unflatten b2out r)) γ i o * vitCotFl N D nClasses Wcls dy o

Final vector-LN γ, chain-certified at vitCotFl (classifier-back scattered to row 0), with the saved block-2 output as the LN input.

theorem Proofs.vit_render_veclnFbeta_chain_certified {N D nClasses : ℕ} (ε : ℝ) (γv β : Vec D) (b2out : Vec ((N + 1) * D)) (Wcls : Mat D nClasses) (dy : Vec nClasses) (lr : ℝ) (i : Fin D) : β i - lr * vecLN_grad_beta (N + 1) D (Mat.unflatten (vitCotFl N D nClasses Wcls dy)) i = β i - lr * ∑ o : Fin ((N + 1) * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin (N + 1)) => layerNormVec D ε γv bv (Mat.unflatten b2out r)) β i o * vitCotFl N D nClasses Wcls dy o

Final vector-LN β, chain-certified.