Documentation

LeanMlir.Proofs.Nets.ViT.ViTBackNet

ViT folded — the last net onto CertLayer, and the first one folded END TO END #

BackNetFolds.lean deliberately left ViT out ("its blocks are per-token Mat-shaped with a different backward vocabulary … a separate sitting"). This file is that sitting.

⚠⚠ FIRST, A CORRECTION TO THE LEDGER — ViT was never the LEAST folded net, it was the MOST #

CertifiedChain.lean's header says "Measured before writing this file: nothing in LeanMlir/Proofs/ folds those blocks into a stage or a net". That is wrong, and ViT is the counterexample: ViTBackB0.lean has carried the depth-k tower backward graph vitBodyBackGraphKMHV and the whole-net graph vitNetBackGraph (patchEmbed → tower → final vec-LN → classifier, at every depth) since before either file existed; this file's vitNetBackGraph_faithful, which ties the whole-net graph to the whole-net VJP, is in tests/AuditAxioms.lean.

So the accurate statement of what the other six nets have is block capstones; ViT alone had a concrete whole-net backward graph tied to the whole-net VJP, stem and head included. What ViT lacked was not the fold — it was the generic fold: its tower was a bespoke induction that no other net could reuse and that reused nothing.

⭐ This file closes that, and the closing is a THEOREM, not a re-implementation. vitTrunkV_graph proves the generic CertLayer chain produces the hand-written vitBodyBackGraphKMHV term for term, and vitTrunkV_fwd proves its forward is the shipped vitBodyKVFlat. So the bespoke induction is not replaced and not trusted alongside the generic one — it is derived from it.

The one piece of per-net work, exactly as the recipe predicted #

Making the blocks pluggable. ViT's capstones took the incoming cotangent as dY : Vec n and wrapped it internally as .operand "%dz" / .operand "%dh", so a block could only ever be the LAST thing in a graph. ViTBackB0.lean now threads ecot : SHlo n through the vec-LN production chain (transformerMlpBackGraph → mlpSublayerV* → attnSublayerV* → transformerBlockV* → the tower), each faithfulness statement carrying den ecot = Mat.flatten dz. Strictly more general: every old statement is the new one at ecot := .operand "%d…" (Mat.flatten dz).

⭐ That also deleted a real seam inside the block: transformerBlockVBackGraphMH used to feed the attention sublayer den (mlpSublayerVBackGraph …) — the MLP arm's value, re-embedded as a constant. The two sublayers now compose as subgraphs, and so do successive blocks in the tower. ⚠ One seam remains and is NOT this refactor's: 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 — but it is the next thing to generalize if the graphs are ever to be emitted rather than only denoted.

⭐⭐ AND THE FOLD RUNS IMAGE → LOGITS — the first one in the repo that does #

No conv net's fold runs that far. §8b records why for R50: the stem is blocked on a proof gap (no den-level faithfulness for the batched maxPool3s2BackB graph — a pool is kinked wherever a window ties), and the head is simply unbuilt. ViT has neither obstacle, because its 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 — including that the fold's VJP is the shipped vitForwardKV_has_vjp, not merely another VJP of the same map. That last step is HasVJPAt.backward_unique_of_eq along the forward equation vitNetLayer_fwd: the two witnesses are VJPs of propositionally equal maps, so both backwards are the same pdiv contraction.

The tier: ok = True, and that is the STRONGER certificate #

GELU and LayerNorm are smooth everywhere, so a ViT block has a global HasVJPMat, lifted pointwise by .toHasVJPAt. ViT joins enet (swish) and convnext (gelu) in the unconditional tier; r34/r50 (relu), mnv2 (relu6) and mnv4 (relu) carry _at hypotheses because those activations genuinely have no derivative at their kinks.

⚠ Not tied to the committed artifact. Same status as every other net's fold: this is a certified composition, not a proof that verified_mlir/vit_train_step.mlir IS this graph.

source
noncomputable def Proofs.StableHLO.vitBlockVLayer {Np1 hm1 d mlpDim : ℕ} (ε : ℝ) (hε : 0 < ε) (p : BlockParamsV ((hm1 + 1) * d) mlpDim) :
CertLayer (Np1 * ((hm1 + 1) * d)) (Np1 * ((hm1 + 1) * d))

A vec-LN multi-head transformer block as a CertLayer, at the flat index.

The bridge from ViT's per-token Mat world to CertLayer's Vec → Vec one is hasVJPMat_to_hasVJP, whose statement HasVJP (fun v => Mat.flatten (f (Mat.unflatten v))) is definitionally blockVFlat — so the lift costs nothing. The backward graph is the committed transformerBlockVBackGraphMHP at the unflattened saved activation.

⭐ ok = True: GELU and LayerNorm are smooth, so the graph denotes the VJP at every input, with no side condition to discharge.

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    noncomputable def Proofs.StableHLO.vitTrunkV {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 {Np1 hm1 d mlpDim : ℕ} (ε : ℝ) (hε : 0 < ε) (k : ℕ) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (v : Vec (Np1 * ((hm1 + 1) * d))) :
      (vitTrunkV ε hε k ps).fwd v = vitBodyKVFlat 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 {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 ε hε k ps).graph A.flatten e = vitBodyBackGraphKMHV ε k ps A e

      ⭐⭐ THE PAYOFF: 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.

      source
      theorem Proofs.StableHLO.vitTrunkV_ok {Np1 hm1 d mlpDim : ℕ} (ε : ℝ) (hε : 0 < ε) (k : ℕ) (ps : Fin k → BlockParamsV ((hm1 + 1) * d) mlpDim) (v : Vec (Np1 * ((hm1 + 1) * d))) :
      (vitTrunkV ε hε k ps).ok v

      ⭐ A ViT trunk's ok is True at every depth — the smooth tier's payoff. CertLayer.comp conjoins preconditions, so for a relu net this would be a deepening stack of side conditions (r50's is 3 clauses per block × 16 blocks); for ViT the conjunction collapses and the whole depth-k trunk is certified unconditionally.

      ⚠ Not a weaker statement than the _at nets' — a stronger one.

      No other net has this: §8b records R50's stem as blocked on a proof gap (no den-level faithfulness for the batched maxPool3s2BackB graph — a pool is kinked wherever a window ties), and the head is simply unbuilt. ViT has neither obstacle: its 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. That is exactly why ViT's stem is free where R50's is blocked on a maxpool.

      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 (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. The first net in the repo whose fold is the entire network.

            Equations
            • One or more equations did not get rendered due to their size.
            Instances For
              source
              theorem Proofs.StableHLO.vitNetLayer_fwd (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 ic H W patchSize N mlpDim hm1 d nClasses k ε hε Wc bc cls pos ps γF βF Wcls bcls).fwd x = vitForwardKV 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 (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 ic H W patchSize N mlpDim hm1 d nClasses k ε hε Wc bc cls pos ps γF βF Wcls bcls).graph x e = vitNetBackGraph ic H W patchSize N mlpDim hm1 d nClasses k ε Wc ps γF Wcls (Mat.unflatten (patchEmbed_flat ic H W patchSize N ((hm1 + 1) * d) Wc bc cls pos x)) (Mat.unflatten (vitBodyKVFlat (N + 1) (hm1 + 1) d mlpDim ε k ps (patchEmbed_flat 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 (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 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 (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 ic H W patchSize N mlpDim hm1 d nClasses k ε Wc ps γF Wcls (Mat.unflatten (patchEmbed_flat ic H W patchSize N ((hm1 + 1) * d) Wc bc cls pos x)) (Mat.unflatten (vitBodyKVFlat (N + 1) (hm1 + 1) d mlpDim ε k ps (patchEmbed_flat ic H W patchSize N ((hm1 + 1) * d) Wc bc cls pos x))) e) = (vitForwardKV_has_vjp 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 vitForwardKV_has_vjp.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 for all seven nets, 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).