Documentation

LeanMlir.Proofs.Nets.ViT.ViTBackNet

ViT on CertLayer, folded image → logits #

ViT's whole-net backward graph vitNetBackGraph (patchEmbed → tower → final vec-LN → classifier, at every depth) is derived here from the generic CertifiedChain fold: vitTrunkV_graph proves the CertLayer chain produces the hand-written vitBodyBackGraphKMHV term for term, and vitTrunkV_fwd proves its forward is the shipped vitBodyKVFlat. So the bespoke tower induction in ViTBackB0.lean is an instance of the generic one.

Pluggable blocks #

ViTBackB0.lean threads an incoming-cotangent subgraph ecot : SHlo n through the vec-LN production chain (transformerMlpBackGraph → mlpSublayerV* → attnSublayerV* → transformerBlockV* → the tower), each faithfulness statement carrying den ecot = Mat.flatten dz; every statement at a fixed Vec cotangent is the one at ecot := .operand "%d…" (Mat.flatten dz). The two sublayers of a block compose as subgraphs, and so do successive blocks in the tower. Note: attnSublayerVInnerBackGraphMH still passes den e into mhsaBackGraphMH, because that MHSA graph takes a Vec cotangent. It is inside the attention arm, not between blocks, so it does not block composition.

The fold runs image → logits #

ViT's stem is an affine patchify conv and its head is a CLS slice plus a dense — both linear, so both backward graphs are activation-independent. So vitNetLayer = stem ∘ trunk ∘ finalLN ∘ head is one CertLayer, assembled by comp alone, and vitNetBackGraph_faithful (the whole-net capstone) follows from it: the hand-written vitNetBackGraph denotes (vitForwardKVHasVJP …).backward. VJP witnesses of one map have one backward, so the last step is HasVJPAt.backward_unique_of_eq along the forward equation vitNetLayer_fwd.

ok = True #

GELU and LayerNorm are smooth everywhere, so a ViT block has a global HasVJPMat, lifted pointwise by .toHasVJPAt, and every layer's precondition is True. EfficientNet (swish) and ConvNeXt (GELU) are the same; ResNet-34/50 (ReLU), MobileNetV2 (ReLU6) and MobileNetV4 (ReLU) carry _at hypotheses because those activations have no derivative at their kinks.

Scope. This is a certified composition of backward graphs; it is not a proof that verified_mlir/vit_train_step.mlir is this graph.

source
noncomputable def Proofs.StableHLO.vitTrunkV (gf : GeluForm) {Np1 hm1 d mlpDim : ℕ} (ε : ℝ) (hε : 0 < ε) (k : ℕ) :
(Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) → CertLayer (Np1 * ((hm1 + 1) * d)) (Np1 * ((hm1 + 1) * d))

The depth-k ViT trunk as one CertLayer. Written in the shipped Fin k → BlockParamsV shape and mirroring vitBodyKVFlat's recursion (block 0 runs first), so the two can be compared term for term below. The only content is CertLayer.comp, which is already proven — depth costs nothing.

Equations
Instances For
    source
    theorem Proofs.StableHLO.vitTrunkV_fwd {gf : GeluForm} {Np1 hm1 d mlpDim : ℕ} (ε : ℝ) (hε : 0 < ε) (k : ℕ) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (v : Vec (Np1 * ((hm1 + 1) * d))) :
    (vitTrunkV gf ε hε k ps).fwd v = vitBodyKVFlat gf Np1 (hm1 + 1) d mlpDim ε k ps v

    The trunk's forward is the shipped depth-k body. Without this the fold would be a chain of blocks that merely resembles ViT's; with it, vitTrunkV is vitBodyKVFlat.

    source
    theorem Proofs.StableHLO.vitTrunkV_graph {gf : GeluForm} {Np1 hm1 d mlpDim : ℕ} (ε : ℝ) (hε : 0 < ε) (k : ℕ) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (A : Mat Np1 ((hm1 + 1) * d)) (e : SHlo (Np1 * ((hm1 + 1) * d))) :
    (vitTrunkV gf ε hε k ps).graph A.flatten e = vitBodyBackGraphKMHV gf ε k ps A e

    The generic fold reproduces the hand-written tower, term for term.

    vitBodyBackGraphKMHV is ViTBackB0's bespoke depth-k reverse fold, proven faithful there by an induction on k that re-does the chain-rule argument at every depth. This theorem says the CertLayer chain's graph — built by comp, whose faithfulness was proven ONCE and for all nets — is that same term. So the bespoke induction is not a second, parallel artifact to keep in sync; it is an instance.

    The proof is the two round-trips unflatten (flatten A) = A (the saved activation the block graph differentiates at) and blockVFlat (flatten A) = flatten (blockV A) (the activation the TAIL differentiates at) — which is exactly the fact CertLayer.comp encodes and the endo slip would break.

    ViT's stem is an affine patchify conv and its head is GAP-free (a CLS slice + dense), so both are linear and their backward graphs are activation-independent.

    source
    noncomputable def Proofs.StableHLO.vitPatchEmbedLayer (ic H W patchSize N D : ℕ) (Wc : Kernel4 D ic patchSize patchSize) (bc cls : Vec D) (pos : Mat (N + 1) D) :
    CertLayer (ic * H * W) ((N + 1) * D)

    The patch-embedding stem as a CertLayer. Its graph ignores the saved activation entirely — patchEmbed is affine, so the input-VJP is the same linear map everywhere.

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      source
      noncomputable def Proofs.StableHLO.vitFinalLNLayer (N D : ℕ) (ε : ℝ) (γF βF : Vec D) (hε : 0 < ε) :
      CertLayer ((N + 1) * D) ((N + 1) * D)

      The final (pre-head) vector-LN as a CertLayer. Smooth, so ok = True.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        source
        noncomputable def Proofs.StableHLO.vitClassifierLayer (N D nClasses : ℕ) (Wcls : Mat D nClasses) (bcls : Vec nClasses) :
        CertLayer ((N + 1) * D) nClasses

        The classifier head as a CertLayer — CLS-slice then dense, both linear, so the graph ignores its activation for the same reason the stem's does.

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For
          source
          noncomputable def Proofs.StableHLO.vitNetLayer (gf : GeluForm) (ic H W patchSize N mlpDim hm1 d nClasses k : ℕ) (ε : ℝ) (hε : 0 < ε) (Wc : Kernel4 ((hm1 + 1) * d) ic patchSize patchSize) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (γF βF : Vec ((hm1 + 1) * d)) (Wcls : Mat ((hm1 + 1) * d) nClasses) (bcls : Vec nClasses) :
          CertLayer (ic * H * W) nClasses

          The whole net as one CertLayer — stem, depth-k trunk, final LN, head, composed by comp alone. Image in, logits out, and the backward graph and its faithfulness come with it.

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            source
            theorem Proofs.StableHLO.vitNetLayer_fwd {gf : GeluForm} (ic H W patchSize N mlpDim hm1 d nClasses k : ℕ) (ε : ℝ) (hε : 0 < ε) (Wc : Kernel4 ((hm1 + 1) * d) ic patchSize patchSize) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (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)) :
            (vitNetLayer gf ic H W patchSize N mlpDim hm1 d nClasses k ε hε Wc bc cls pos ps γF βF Wcls bcls).fwd x = vitForwardKV gf ic H W patchSize N mlpDim (hm1 + 1) d nClasses k Wc bc cls pos ε ps γF βF Wcls bcls x

            The whole-net layer's forward is the shipped vitForwardKV.

            source
            theorem Proofs.StableHLO.vitNetLayer_graph {gf : GeluForm} (ic H W patchSize N mlpDim hm1 d nClasses k : ℕ) (ε : ℝ) (hε : 0 < ε) (Wc : Kernel4 ((hm1 + 1) * d) ic patchSize patchSize) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (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)) (e : SHlo nClasses) :
            (vitNetLayer gf ic H W patchSize N mlpDim hm1 d nClasses k ε hε Wc bc cls pos ps γF βF Wcls bcls).graph x e = vitNetBackGraph gf ic H W patchSize N mlpDim hm1 d nClasses k ε Wc ps γF Wcls (Mat.unflatten (patchEmbedFlat ic H W patchSize N ((hm1 + 1) * d) Wc bc cls pos x)) (Mat.unflatten (vitBodyKVFlat gf (N + 1) (hm1 + 1) d mlpDim ε k ps (patchEmbedFlat ic H W patchSize N ((hm1 + 1) * d) Wc bc cls pos x))) e

            The whole-net chain's backward graph is vitNetBackGraph. The stem/trunk/LN/head version of vitTrunkV_graph: ViTBackB0's hand-composed whole-net graph is what CertLayer.comp produces, so vitNetBackGraph_faithful is a consequence of the shared combinator rather than a parallel result. The saved activations comp threads automatically are exactly the ones that theorem pins by hand.

            source
            theorem Proofs.StableHLO.vitNetLayer_ok {gf : GeluForm} (ic H W patchSize N mlpDim hm1 d nClasses k : ℕ) (ε : ℝ) (hε : 0 < ε) (Wc : Kernel4 ((hm1 + 1) * d) ic patchSize patchSize) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (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)) :
            (vitNetLayer gf ic H W patchSize N mlpDim hm1 d nClasses k ε hε Wc bc cls pos ps γF βF Wcls bcls).ok x

            The whole net is certified at every input — ok = True end to end, because every stage is smooth (affine stem, GELU/LN blocks, LN, affine head). No side condition anywhere.

            source
            theorem Proofs.StableHLO.vitNetBackGraph_faithful {gf : GeluForm} (ic H W patchSize N mlpDim hm1 d nClasses k : ℕ) (ε : ℝ) (hε : 0 < ε) (Wc : Kernel4 ((hm1 + 1) * d) ic patchSize patchSize) (bc cls : Vec ((hm1 + 1) * d)) (pos : Mat (N + 1) ((hm1 + 1) * d)) (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)) (e : SHlo nClasses) :
            den (vitNetBackGraph gf ic H W patchSize N mlpDim hm1 d nClasses k ε Wc ps γF Wcls (Mat.unflatten (patchEmbedFlat ic H W patchSize N ((hm1 + 1) * d) Wc bc cls pos x)) (Mat.unflatten (vitBodyKVFlat gf (N + 1) (hm1 + 1) d mlpDim ε k ps (patchEmbedFlat ic H W patchSize N ((hm1 + 1) * d) Wc bc cls pos x))) e) = (vitForwardKVHasVJP gf ic H W patchSize N mlpDim (hm1 + 1) d nClasses k Wc bc cls pos ε hε ps γF βF Wcls bcls).backward x (den e)

            Whole-net backward-graph faithfulness. The reverse-composed backward graph vitNetBackGraph denotes the proven whole-net VJP vitForwardKVHasVJP.backward at every input image and output cotangent, at every depth k (multi-head, vector-LN). It falls out of CertLayer.faithful at vitNetLayer plus vitNetLayer_graph — the composition argument is comp's, proven once in CertifiedChain, and the only ViT-specific input is the forward equality.

            The peer of R34BWeights / R50BWeights / mnv4Blocks: the shipped config's shape pinned in the types instead of stated in a docstring. ViT-Tiny is 12 identical-shaped blocks at heads = 3, d_head = 64 (D = 192), mlpDim = 768, N + 1 = 197 tokens — so unlike the conv nets there is no ladder to pin, only the depth and the widths, and every block has the same type (which is why chain over a List.replicate-shaped index is the right form here and a per-row-typed parameter record buys nothing).