Documentation

LeanMlir.Proofs.Nets.ViT.ViTDepthK

ViT depth-k (general-depth tower, distinct per-block params) #

The proven transformerTowerHasVJPMat shares ONE param tuple across blocks. This file builds the net with distinct per-block params at every 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 vitBodyKVFlatHasVJP by induction on k (the chain step is vjpComp + the bridged transformerBlockVHasVJPMat). vitForwardKV(HasVJP[_correct]) — the whole net at depth k, UNCONDITIONAL except 0 < ε.

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

The ViT-Tiny instantiation (depth 12, P=16, D=192, heads=3) of the whole-net backward is vitTinyHasVJP_correct, in ViTWholeBackCertifiedTie.

References #

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

Instances For
    source
    noncomputable def Proofs.blockV (gf : GeluForm) (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
    Instances For
      source
      noncomputable def Proofs.blockVFlat (gf : GeluForm) (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.

      Equations
      Instances For
        source
        noncomputable def Proofs.vitBodyKV (gf : GeluForm) (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
        Instances For
          source
          noncomputable def Proofs.vitBodyKVFlat (gf : GeluForm) (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, so the VJP composes block-at-a-time.

          Equations
          • One or more equations did not get rendered due to their size.
          • Proofs.vitBodyKVFlat gf Np1 heads d_head mlpDim ε 0 x_2 = fun (v : Proofs.Vec (Np1 * (heads * d_head))) => v
          Instances For
            source
            theorem Proofs.vitBodyKVFlat_eq_flatten {gf : GeluForm} (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (k : ℕ) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) (A : Mat Np1 (heads * d_head)) :
            vitBodyKVFlat gf Np1 heads d_head mlpDim ε k ps A.flatten = (vitBodyKV gf 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).

            source
            theorem Proofs.vitBodyKVFlat_differentiable {gf : GeluForm} (Np1 heads d_head mlpDim : ℕ) (ε : ℝ) (hε : 0 < ε) (k : ℕ) (ps : Fin k → BlockParamsV (heads * d_head) mlpDim) :
            Differentiable ℝ (vitBodyKVFlat gf Np1 heads d_head mlpDim ε k ps)

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

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

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

            Equations
            Instances For
              source
              noncomputable def Proofs.vitForwardKV (gf : GeluForm) (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.

              Equations
              • One or more equations did not get rendered due to their size.
              Instances For
                source
                noncomputable def Proofs.vitForwardKVHasVJP (gf : GeluForm) (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 gf 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 vjpComp steps gluing patchEmbedFlatHasVJP, the inductive vitBodyKVFlatHasVJP, the bridged per-token vector-LN, and classifierFlatHasVJP.

                Equations
                • One or more equations did not get rendered due to their size.
                Instances For
                  source
                  theorem Proofs.vitForwardKVHasVJP_correct {gf : GeluForm} (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)) :
                  (vitForwardKVHasVJP gf 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 gf 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

                  The .correct field of vitForwardKVHasVJP, restated: the depth-k ViT's backward is the pdiv-contracted Jacobian of vitForwardKV at every input. The tie of the hand-written input-gradient chain to this backward is vitInputGradK_correct.

                  source
                  def Proofs.StableHLO.vitBlockGraphMHVP (gf : GeluForm) {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.
                  Instances For
                    source
                    def Proofs.StableHLO.vitBodyGraphKMHV (gf : GeluForm) {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
                    Instances For
                      source
                      theorem Proofs.StableHLO.vitBodyGraphKMHV_den {gf : GeluForm} {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 gf epsStr sStr oneStr zeroStr ε (sdpaScale d) base k ps e) = (vitBodyKV gf 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.

                      source
                      def Proofs.StableHLO.vitFwdGraphKMHV (gf : GeluForm) {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.
                      Instances For
                        source
                        theorem Proofs.StableHLO.vitFwdGraphKMHV_faithful {gf : GeluForm} (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 gf epsStr sStr oneStr zeroStr ε (sdpaScale d) Wc bc cls pos k ps γF βF Wcls bcls x) = vitForwardKV gf 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.