Documentation

LeanMlir.Proofs.Nets.ViT.ViTMultiHead

ViT scaling pass — multi-head rendering + faithfulness #

The MATH is general in heads (mhsaHasVJPMat, transformerBlockVHasVJPMat); this file supplies the RENDERING + faithfulness at heads > 1:

  1. mhsaLayer_spelled — the general-heads tie, on which the rest of the file builds: mhsaLayer 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. vitBlockGraphMHV over the two new ch10 tokens headSliceF/headPadF (+ headsSumG, a left-assoc addV fold), with the block denotation vitBlockGraphMHV_den_aux that ViTDepthK's vitFwdGraphKMHV_faithful chains per block at heads := hm1 + 1 — faithfulness against mhsaLayer 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.

References #

source
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 mhsaLayer feeds each head's SDPA.

Equations
Instances For
    source
    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 mhsaLayer's concat; it is also the slice's VJP.

    Equations
    Instances For
      source
      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 — mhsaLayer's concat indexing.

      source
      theorem Proofs.mhsaLayer_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)) :
      mhsaLayer 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) => sdpaScale 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 tie the multi-head graph faithfulness rests on: 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).

      source
      noncomputable def Proofs.vitBlockSpelledMHV (gf : GeluForm) (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 spelled per head (mhsaLayer_spelled).

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        source
        theorem Proofs.vitBlockSpelledMHV_eq {gf : GeluForm} (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 gf Np1 heads d mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 X = transformerBlockV gf 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 mhsaLayer_spelled.

        source
        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
        Instances For
          source
          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.

          source
          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.

          source
          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.

          source
          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.

          source
          def Proofs.StableHLO.vitBlockGraphMHV (gf : GeluForm) {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 per head: headSliceF → SDPA → headPadF, folded by headsSumG.

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            source
            theorem Proofs.StableHLO.vitBlockGraphMHV_den_aux {gf : GeluForm} {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 gf pfx epsStr sStr oneStr zeroStr ε (sdpaScale d) γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 e) = (vitBlockSpelledMHV gf 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.