Documentation

LeanMlir.Proofs.Nets.ViT.ViTBackChains

The ViT-Tiny backward chains — the ℝ maps the ViT ties are about #

The hand-composed reverse of the committed ViT-Tiny forward, as plain defs on the cotangent, from the attention core outwards:

Convention table for the net vitInputGradK reverses (vitForwardKV at ViT-Tiny, read from apps/baselines/MainVitTrain.lean and LeanMlir/Proofs/Codegen/ViTRender.lean): depth 12 with distinct per-block parameters, D = 192 = 3 heads × 64, MLP dim 768, 197 tokens (196 patches + CLS), 16×16/s16 patchify (no conv2d, so no padding phase and none of the even-kernel question EvenKernelConvBack.lean found for ConvNeXt), vector-[D] LayerNorm at all 25 sites, GELU, ε = 1e-5. No BatchNorm anywhere. ⚠ N throughout is the TOKEN count, not a batch; the batch is B, the binder of vitInputGradKB below — the per-example chain lifted stage by stage over B examples, the way the batched T3 tie (ViTStepTieGB) lifts every activation and cotangent — tied by vitInputGradKB_eq_batchMap_vitForwardKV_vjp (ViTWholeBackCertifiedTieB.lean).

Moved here from the three float bridges that defined them beside their float twins on 2026-09-08 (planning/archive/float_second_pass.md); no number is stated about any of these chains.

noncomputable def Proofs.mhSlab {n h dh : } (hd : Fin h) (Q : Mat n (h * dh)) :
Mat n dh

The column slab [hd·dh, (hd+1)·dh) of a Mat n (h·dh) as a Mat n dh (head hd's view) — the finProdFinEquiv (hd, ·) column restriction, matching mhsa_layer's per-head extraction.

Equations
Instances For
    noncomputable def Proofs.mhsaSdpaBackV {h N dh : } (Q K V dOut : Mat N (h * dh)) :
    Mat N (h * dh)

    Multi-head sdpa backward w.r.t. V — per head, the certified sdpa_back_V on the head slabs; concatenated by the finProdFinEquiv column layout.

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      noncomputable def Proofs.mhsaSdpaBackQ {h N dh : } (Q K V dOut : Mat N (h * dh)) :
      Mat N (h * dh)

      Multi-head sdpa backward w.r.t. Q.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        noncomputable def Proofs.mhsaSdpaBackK {h N dh : } (Q K V dOut : Mat N (h * dh)) :
        Mat N (h * dh)

        Multi-head sdpa backward w.r.t. K.

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For
          noncomputable def Proofs.coreVFlat {h N dh : } (Q K V : Mat N (h * dh)) (v : Vec (N * (h * dh))) :
          Vec (N * (h * dh))

          Flattened multi-head sdpa backward w.r.t. V (saved projections Q K V fixed).

          Equations
          Instances For
            noncomputable def Proofs.coreQFlat {h N dh : } (Q K V : Mat N (h * dh)) (v : Vec (N * (h * dh))) :
            Vec (N * (h * dh))

            Flattened multi-head sdpa backward w.r.t. Q.

            Equations
            Instances For
              noncomputable def Proofs.coreKFlat {h N dh : } (Q K V : Mat N (h * dh)) (v : Vec (N * (h * dh))) :
              Vec (N * (h * dh))

              Flattened multi-head sdpa backward w.r.t. K.

              Equations
              Instances For
                noncomputable def Proofs.mhsaBackFlat {h N dh : } (Wq Wk Wv Wo : Mat (h * dh) (h * dh)) (Q K V : Mat N (h * dh)) :
                Vec (N * (h * dh))Vec (N * (h * dh))

                The full multi-head self-attention input-gradient backward (cotangent dY ↦ dX): output-projection backward (dense Woᵀ 0, per token) → the three sdpa cores → Q/K/V projection backwards (dense Wᵀ 0, per token), fanning in at X (the three paths add). The certified MHSA backward at the input (mhsa_layer, Attention.lean) is dconcat = dY·Woᵀ, (dQ, dK, dV) = sdpa_back(dconcat) per head, dX = dQ·Wqᵀ + dK·Wkᵀ + dV·Wvᵀ; mhsaBackFlat_eq_mhsa_vjp says this chain is that.

                Equations
                • One or more equations did not get rendered due to their size.
                Instances For
                  noncomputable def Proofs.vitBlockBack {h N dh dff : } (Wq Wk Wv Wo : Mat (h * dh) (h * dh)) (Q K V : Mat N (h * dh)) (lnB₁ : Vec (h * dh)Vec (h * dh)) (W₁ : Mat (h * dh) dff) (W₂ : Mat dff (h * dh)) (sgelu : Vec dff) (lnB₂ : Vec (h * dh)Vec (h * dh)) :
                  Vec (N * (h * dh))Vec (N * (h * dh))

                  The ViT encoder-block input-gradient backward — the reverse of LN → MHSA → +x → LN → MLP → +x. The block is mlpResidual ∘ attnSub (forward), so the backward is attnSubBack ∘ mlpResidualBack:

                  • MLP-residual backward (per token): residual (LN₂-back ∘ dense W₁ᵀ ∘ geluBack ∘ dense W₂ᵀ) — the reverse of dense W₂ ∘ gelu ∘ dense W₁ ∘ LN₂, lifted over the sequence (perRowFlat);
                  • attention-sublayer backward: residual (LN₁-back ∘ mhsaBackFlat) — the residual skip's cotangent flows both through the MHSA backward and directly to x.

                  The LN backwards (lnB₁/lnB₂) are supplied as one shared map per site; geluBack is the saved-derivative diagBack.

                  Equations
                  • One or more equations did not get rendered due to their size.
                  Instances For
                    noncomputable def Proofs.vitBlockBackV {h N dh dff : } (Wq Wk Wv Wo : Mat (h * dh) (h * dh)) (Q K V : Mat N (h * dh)) (ε : ) (γ1 : Vec (h * dh)) (X1 : Vec (N * (h * dh))) (W₁ : Mat (h * dh) dff) (W₂ : Mat dff (h * dh)) (sgelu : Fin NVec dff) (γ2 : Vec (h * dh)) (X2 : Vec (N * (h * dh))) :
                    Vec (N * (h * dh))Vec (N * (h * dh))

                    The vector-LayerNorm ViT encoder-block input-gradient backward. rowLNVecFlatBack in both LN slots at the site's flat saved input, per-token GELU slopes in the MLP sublayer's per-row fold (perRowFlatPR), and each sublayer's residual outside its fold.

                    Reading right to left, this is the reverse of transformerBlockV: the MLP sublayer's dense W₂ ∘ gelu ∘ dense W₁ ∘ LN₂ backward under a residual, then the attention sublayer's mhsa ∘ LN₁ backward under a residual. The shipped depth-12 net is vitForwardKV, whose blocks are transformerBlockV at γ β : Vec D; this is the spelling the whole-net tie is stated at.

                    Equations
                    • One or more equations did not get rendered due to their size.
                    Instances For
                      noncomputable def Proofs.vitAttnOutAt (Np1 heads d_head mlpDim : ) (ε : ) (p : BlockParamsV (heads * d_head) mlpDim) (v : Vec (Np1 * (heads * d_head))) :
                      Mat Np1 (heads * d_head)

                      The block's attention-sublayer output at a flat saved input — the LN₂ site's saved activation, and the point the MLP sublayer's backward is taken at. Named because it appears three times in vitBlockBackVAt and once more in every tie about it.

                      Equations
                      Instances For
                        noncomputable def Proofs.vitBlockBackVAt (Np1 heads d_head mlpDim : ) (ε : ) (p : BlockParamsV (heads * d_head) mlpDim) (v : Vec (Np1 * (heads * d_head))) :
                        Vec (Np1 * (heads * d_head))Vec (Np1 * (heads * d_head))

                        vitBlockBackV with every saved slot pinned to the real forward at the block's own input. The Q/K/V projections at LN₁(A), the LN₁ backward at A and the LN₂ backward at the attention sublayer's output, the GELU derivative at dense₁(LN₂(attn A)) — one function of the block's flat input v, so the tower recursion can be written down. cnxBlockChBackAt's shape.

                        Equations
                        • One or more equations did not get rendered due to their size.
                        Instances For
                          noncomputable def Proofs.vitTowerBackK (Np1 heads d_head mlpDim : ) (ε : ) (k : ) (ps : Fin kBlockParamsV (heads * d_head) mlpDim) :
                          Vec (Np1 * (heads * d_head))Vec (Np1 * (heads * d_head))Vec (Np1 * (heads * d_head))

                          The depth-k encoder-tower backward at a saved tower input v.

                          Head-first, like the forward it reverses. vitBodyKVFlat (k+1) ps = vitBodyKVFlat k (ps ∘ succ) ∘ blockVFlat (ps 0) runs block 0 FIRST, so the backward applies block 0's reverse LAST, and the tail's saved input is block 0's forward OUTPUT. cnxStageChKBack's recursion verbatim, one architecture over. Writing the fold as its own recursion (rather than a list fold) is what makes the saved-activation thread visible, and the thread is the content.

                          Equations
                          • One or more equations did not get rendered due to their size.
                          • Proofs.vitTowerBackK Np1 heads d_head mlpDim ε 0 x_3 x✝ = id
                          Instances For
                            noncomputable def Proofs.clsScatter (N D : ) (dy : Vec D) :
                            Vec ((N + 1) * D)

                            The CLS-slice backward — the adjoint of cls_slice_flat (gather row 0 of the (N+1)×D sequence): scatter the head cotangent dy back to row 0 (the CLS token), zero on the patch rows. The certified cls_slice_flat_has_vjp.backward.

                            Equations
                            Instances For
                              noncomputable def Proofs.vitSavedPE (ic H W patchSize N heads d_head : ) (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)) (x : Vec (ic * H * W)) :
                              Vec ((N + 1) * (heads * d_head))

                              The patch-embed output — the encoder tower's saved input. Named as a FUNCTION of the image so that the same constant is both the activation the tower's slots are saved at and the f argument of the chain's first vjp_comp (cnxSavedA0 … cnxSavedA10's reason).

                              Equations
                              Instances For
                                noncomputable def Proofs.vitSavedBody (ic H W patchSize N mlpDim heads d_head 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 kBlockParamsV (heads * d_head) mlpDim) (x : Vec (ic * H * W)) :
                                Vec ((N + 1) * (heads * d_head))

                                The encoder tower's output — the final LayerNorm's saved input.

                                Equations
                                • One or more equations did not get rendered due to their size.
                                Instances For
                                  noncomputable def Proofs.vitInputGradK (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 kBlockParamsV (heads * d_head) mlpDim) (γF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (x : Vec (ic * H * W)) :
                                  Vec nClassesVec (ic * H * W)

                                  THE WHOLE-NET ViT-TINY INPUT GRADIENT, at the depth, head count and LayerNorm spelling the net runs. The reverse of vitForwardKV = classifier_flat ∘ LNᵥ ∘ vitBodyKVFlat ∘ patchEmbed:

                                  patchEmbedBack ∘ towerBack ∘ finalLNBack ∘ clsScatter ∘ dense Wclsᵀ

                                  with every slot concrete — the patch-embed backward is patchEmbed_input_grad_formula (which IS patchEmbed_flat_has_vjp's backward, definitionally), the head is the free dense Wᵀ 0 followed by the CLS scatter, the final LN is rowLNVecFlatBack at the tower's output, and the tower is vitTowerBackK. vitInputGradK_eq_vitForwardKV_vjp is the apex that says this is the certified whole-net gradient.

                                  Equations
                                  • One or more equations did not get rendered due to their size.
                                  Instances For
                                    noncomputable def Proofs.vitSavedPEB (B ic H W patchSize N heads d_head : ) (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)) (x : Vec (B * (ic * H * W))) :
                                    Vec (B * ((N + 1) * (heads * d_head)))

                                    The batched patch embedding: StableHLO.batchMap B of vitSavedPE, the tower's saved input at every example.

                                    Equations
                                    • One or more equations did not get rendered due to their size.
                                    Instances For
                                      noncomputable def Proofs.vitSavedBodyB (B ic H W patchSize N mlpDim heads d_head 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 kBlockParamsV (heads * d_head) mlpDim) (x : Vec (B * (ic * H * W))) :
                                      Vec (B * ((N + 1) * (heads * d_head)))

                                      The batched tower output: StableHLO.batchMap B of the tower at the batched patch embedding — saved stage by stage, not batchMap B of the composed per-example prefix (the two agree only up to batchMap_comp, ViTWholeBackCertifiedTieB.lean).

                                      Equations
                                      • One or more equations did not get rendered due to their size.
                                      Instances For
                                        noncomputable def Proofs.vitInputGradKB (B 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 kBlockParamsV (heads * d_head) mlpDim) (γF : Vec (heads * d_head)) (Wcls : Mat (heads * d_head) nClasses) (x : Vec (B * (ic * H * W))) :
                                        Vec (B * nClasses)Vec (B * (ic * H * W))

                                        THE BATCHED WHOLE-NET ViT INPUT GRADIENTvitInputGradK at each of B examples, stage by stage, as the batched render computes it. The head backward and the CLS scatter are StableHLO.batchMap B of their per-example leaves (both input-independent); the final-LN and tower backwards are StableHLO.batchMapAux B of their per-example maps, each at the batched saved activation (vitSavedBodyB, vitSavedPEB); the patch-embed backward is batchMap B of the linear formula. Every slot is a lift because no ViT op couples examples — the same honesty argument ViTStepTieGB makes for the batched T3 tie. B is a variable: this chain carries no batch numeral. vitInputGradKB_eq_batchMap_vitForwardKV_vjp (ViTWholeBackCertifiedTieB.lean) says it IS the certified gradient of batchMap B vitForwardKV.

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