ViT scaling pass — multi-head rendering + faithfulness

The representative ViT close (Items A–D) rendered at heads = 1, where MHSA collapses to three matmuls + a row-softmax (mhsa_layer_one_head). The MATH was always general in heads (mhsa_has_vjp_mat, transformerBlock(V)_has_vjp_mat, vitForward2(V)_has_vjp); what was missing is RENDERING + faithfulness at heads > 1. This file closes that:

1. mhsa_layer_spelled — the load-bearing tie, the general-heads successor of mhsa_layer_one_head: mhsa_layer N heads d IS, per head, slice → matmul-spelled SDPA → pad-scatter, summed over heads. The concat is spelled as Σ_h headPadMat h ∘ (per-head SDPA) — every output column receives exactly one head's value, and the sum stays at the single index N·(heads·d) (no (N·a)+(N·b) Nat-cast trouble a binary concat would hit).

2. vitBlockGraphMH(V) over the two new ch10 tokens headSliceF/headPadF (+ headsSumG, a left-assoc addV fold), with vitFwdGraphMH(V)_faithful: the multi-head forward graphs denote the proven vitForward2/vitForward2V at heads := hm1 + 1 — faithfulness against mhsa_layer N heads d DIRECTLY, not a 1-head specialization.

The graph layer is stated at heads = hm1 + 1 (the head fold needs a first head); the Mat-level spelling is fully general in heads.

noncomputable def Proofs.headSliceMat (N heads d : ℕ) (h : Fin heads) (A : Mat N (heads * d)) : Mat N d

Head h's [N,d] column block of an [N,heads·d] matrix — the finProdFinEquiv (h, ·) column gather mhsa_layer feeds each head's SDPA.

Equations

• Proofs.headSliceMat N heads d h A r j = A r (finProdFinEquiv (h, j))

noncomputable def Proofs.headPadMat (N heads d : ℕ) (h : Fin heads) (A : Mat N d) : Mat N (heads * d)

Scatter an [N,d] head block into head h's columns of a zero [N,heads·d]. Summed over heads this is mhsa_layer's concat; it is also the slice's VJP.

Equations

• Proofs.headPadMat N heads d h A r hj = if (finProdFinEquiv.symm hj).1 = h then A r (finProdFinEquiv.symm hj).2 else 0

theorem Proofs.sum_headPadMat_apply {N heads d : ℕ} (G : Fin heads → Mat N d) (n : Fin N) (hj : Fin (heads * d)) : (∑ h : Fin heads, headPadMat N heads d h (G h)) n hj = G (finProdFinEquiv.symm hj).1 n (finProdFinEquiv.symm hj).2

The pad-sum IS the head concat: every column hj lands in exactly one head's block, so the sum over heads of pad-scatters reads off head (symm hj).1 at column (symm hj).2 — mhsa_layer's concat indexing.

theorem Proofs.mhsa_layer_spelled (Np1 heads d : ℕ) (Wq Wk Wv Wo : Mat (heads * d) (heads * d)) (bq bk bv bo : Vec (heads * d)) (X : Mat Np1 (heads * d)) : mhsa_layer Np1 heads d Wq Wk Wv Wo bq bk bv bo X = fun (n : Fin Np1) => dense Wo bo ((∑ h : Fin heads, headPadMat Np1 heads d h ((rowSoftmax fun (i j : Fin Np1) => sdpa_scale d * (headSliceMat Np1 heads d h fun (r : Fin Np1) (c : Fin (heads * d)) => dense Wq bq (X r) c).mul (headSliceMat Np1 heads d h fun (r : Fin Np1) (c : Fin (heads * d)) => dense Wk bk (X r) c).transpose i j).mul (headSliceMat Np1 heads d h fun (r : Fin Np1) (c : Fin (heads * d)) => dense Wv bv (X r) c))) n)

MHSA at general heads is per-head slice → matmul-spelled SDPA → pad-scatter, summed over heads. The general-heads successor of mhsa_layer_one_head and the load-bearing tie for the multi-head graph faithfulness: each head's SDPA is exactly the ch10 token spelling (Q_h·K_hᵀ → ·1/√d → row-softmax → P_h·V_h) on the sliced Q/K/V, and the concat is the pad-sum (sum_headPadMat_apply).

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

The ch10 spelled pre-norm transformer block at general heads (Mat level) — the exact op sequence vitBlockGraphMH denotes: LN₁ → Q/K/V per-token dense → per head (slice → Q_h·K_hᵀ → ·1/√d → row-softmax → P_h·V_h → pad) → Σ heads → output dense → +res → LN₂ → fc1 → GELU → fc2 → +res.

Equations

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

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

The spelled multi-head block IS transformerBlock at general heads. The sublayer/residual structure matches definitionally once mhsa_layer_spelled turns the per-head plumbing into the pad-sum.

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

The spelled multi-head block at vector-[D] LN — each LN site decomposed as the graph (and ViTRender) emit it: pure normalize (scalar-LN at 1,0) → per-channel scale → per-channel bias; attention as vitBlockSpelledMH.

Equations

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

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

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

def Proofs.StableHLO.headsSumG {n hm1 : ℕ} : (Fin (hm1 + 1) → SHlo n) → SHlo n

Left-assoc addV fold of one graph per head — the token-level Σ over heads (heads = hm1 + 1: the fold needs a first head).

Equations

• Proofs.StableHLO.headsSumG f = f 0
• Proofs.StableHLO.headsSumG f = (Proofs.StableHLO.headsSumG fun (i : Fin (hm1 + 1)) => f i.castSucc).addV (f (Fin.last (hm1 + 1)))

theorem Proofs.StableHLO.den_headsSumG {n hm1 : ℕ} (f : Fin (hm1 + 1) → SHlo n) : den (headsSumG f) = fun (j : Fin n) => ∑ h : Fin (hm1 + 1), den (f h) j

The head fold denotes the pointwise sum over heads.

theorem Proofs.StableHLO.headSliceFlat_flat {N heads d : ℕ} (h : Fin heads) (A : Mat N (heads * d)) : headSliceFlat N heads d h A.flatten = (headSliceMat N heads d h A).flatten

Per-head slice commutation bridge.

theorem Proofs.StableHLO.headPadFlat_flat {N heads d : ℕ} (h : Fin heads) (A : Mat N d) : headPadFlat N heads d h A.flatten = (headPadMat N heads d h A).flatten

Per-head pad commutation bridge.

theorem Proofs.StableHLO.flatten_sum {m n H : ℕ} (G : Fin H → Mat m n) : (fun (j : Fin (m * n)) => ∑ h : Fin H, (G h).flatten j) = (∑ h : Fin H, G h).flatten

Pointwise sums over heads commute with Mat.flatten.

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

One spelled pre-norm transformer block over the ch10 tokens at heads = hm1 + 1: lnRowF → Q/K/V denseRowF → per head (headSliceF → matmulF(Q_h, transposeF K_h) → scaleF → softmaxRowF → matmulF(P_h, V_h) → headPadF) → headsSumG → output denseRowF → addV residual → lnRowF → fc1 → geluF → fc2 → addV residual.

Equations

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

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

The vector-LN multi-head block over the tokens: each LN site is lnRowF(1,0) → rowScaleF γ → rowBiasF β (the ViTRender decomposition); attention as vitBlockGraphMH.

Equations

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

theorem Proofs.StableHLO.vitBlockGraphMHV_den_aux {Np1 hm1 d mlpDim : ℕ} (pfx epsStr sStr oneStr zeroStr : String) (ε : ℝ) (γ1 β1 : Vec ((hm1 + 1) * d)) (Wq Wk Wv Wo : Mat ((hm1 + 1) * d) ((hm1 + 1) * d)) (bq bk bv bo γ2 β2 : Vec ((hm1 + 1) * d)) (Wfc1 : Mat ((hm1 + 1) * d) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim ((hm1 + 1) * d)) (bfc2 : Vec ((hm1 + 1) * d)) (e : SHlo (Np1 * ((hm1 + 1) * d))) (A : Mat Np1 ((hm1 + 1) * d)) (hA : den e = A.flatten) : den (vitBlockGraphMHV pfx epsStr sStr oneStr zeroStr ε (sdpa_scale d) γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 e) = (vitBlockSpelledMHV Np1 (hm1 + 1) d mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 A).flatten

Multi-head vector-LN block-graph denotation. Public — the depth-k faithfulness induction (ViTDepthK.lean) chains it per block.

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

Whole multi-head ViT forward graph (scalar LN): patch embed → 2 spelled multi-head blocks (distinct params) → final per-token LN → CLS slice → dense head, at heads = hm1 + 1, D = (hm1+1)·d.

Equations

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

theorem Proofs.StableHLO.vitFwdGraphMH_faithful (ic H W patchSize N hm1 d mlpDim nClasses : ℕ) (epsStr sStr : String) (Wc : Kernel4 ((hm1 + 1) * d) ic patchSize patchSize) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (ε γ1₁ β1₁ : ℝ) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat ((hm1 + 1) * d) ((hm1 + 1) * d)) (bq₁ bk₁ bv₁ bo₁ : Vec ((hm1 + 1) * d)) (γ2₁ β2₁ : ℝ) (Wfc1₁ : Mat ((hm1 + 1) * d) mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim ((hm1 + 1) * d)) (bfc2₁ : Vec ((hm1 + 1) * d)) (γ1₂ β1₂ : ℝ) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat ((hm1 + 1) * d) ((hm1 + 1) * d)) (bq₂ bk₂ bv₂ bo₂ : Vec ((hm1 + 1) * d)) (γ2₂ β2₂ : ℝ) (Wfc1₂ : Mat ((hm1 + 1) * d) mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim ((hm1 + 1) * d)) (bfc2₂ : Vec ((hm1 + 1) * d)) (γF βF : ℝ) (Wcls : Mat ((hm1 + 1) * d) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) : den (vitFwdGraphMH epsStr sStr ε (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) = vitForward2 ic H W patchSize N mlpDim (hm1 + 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

Multi-head ViT forward faithfulness — the scaling-pass apex: the multi-head forward graph denotes the proven distinct-param 2-block vitForward2 at heads := hm1 + 1 — against mhsa_layer directly (per-block vitBlockGraphMH_den_aux + vitBlockSpelledMH_eq, mhsa_layer_spelled under the hood).

def Proofs.StableHLO.vitFwdGraphMHV {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)) (γ1₁ β1₁ : Vec ((hm1 + 1) * d)) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat ((hm1 + 1) * d) ((hm1 + 1) * d)) (bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ : Vec ((hm1 + 1) * d)) (Wfc1₁ : Mat ((hm1 + 1) * d) mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim ((hm1 + 1) * d)) (bfc2₁ γ1₂ β1₂ : Vec ((hm1 + 1) * d)) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat ((hm1 + 1) * d) ((hm1 + 1) * d)) (bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ : Vec ((hm1 + 1) * d)) (Wfc1₂ : Mat ((hm1 + 1) * d) mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim ((hm1 + 1) * d)) (bfc2₂ γF βF : Vec ((hm1 + 1) * d)) (Wcls : Mat ((hm1 + 1) * d) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) : SHlo nClasses

Whole multi-head vector-LN ViT forward graph: patch embed → 2 spelled multi-head vector-LN blocks (distinct params) → final vector-LN (the same three-token decomposition) → CLS slice → dense head. The production ViTRender shape at the graph layer.

Equations

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

theorem Proofs.StableHLO.vitFwdGraphMHV_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)) (ε : ℝ) (γ1₁ β1₁ : Vec ((hm1 + 1) * d)) (Wq₁ Wk₁ Wv₁ Wo₁ : Mat ((hm1 + 1) * d) ((hm1 + 1) * d)) (bq₁ bk₁ bv₁ bo₁ γ2₁ β2₁ : Vec ((hm1 + 1) * d)) (Wfc1₁ : Mat ((hm1 + 1) * d) mlpDim) (bfc1₁ : Vec mlpDim) (Wfc2₁ : Mat mlpDim ((hm1 + 1) * d)) (bfc2₁ γ1₂ β1₂ : Vec ((hm1 + 1) * d)) (Wq₂ Wk₂ Wv₂ Wo₂ : Mat ((hm1 + 1) * d) ((hm1 + 1) * d)) (bq₂ bk₂ bv₂ bo₂ γ2₂ β2₂ : Vec ((hm1 + 1) * d)) (Wfc1₂ : Mat ((hm1 + 1) * d) mlpDim) (bfc1₂ : Vec mlpDim) (Wfc2₂ : Mat mlpDim ((hm1 + 1) * d)) (bfc2₂ γF βF : Vec ((hm1 + 1) * d)) (Wcls : Mat ((hm1 + 1) * d) nClasses) (bcls : Vec nClasses) (x : Vec (ic * H * W)) : den (vitFwdGraphMHV 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 (hm1 + 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

Multi-head vector-LN ViT forward faithfulness — the multi-head graph at the production LN form denotes vitForward2V at heads := hm1 + 1.