Documentation

LeanMlir.Proofs.Nets.Small.Cifar8BbnWholeBackB

The batch-statistics cifar8-BN backward chain IS the certified whole-net VJP #

Cifar8BbnStepTieB.cifar8Bbn_train_step_tiedB ties the thirty-eight gradient NODES of the cifar8wb_bbn_* train steps: each denotes the layer's Jacobian contracted with the cotangent the rendered backward chain hands it. That theorem takes the cotangents as the chain delivers them. This file is the other half — that the chain delivers the right ones: the backward chain those artifacts render, read in the BackLinks vocabulary the node tie uses (reluMaskB, bnBackB, cInB, the per-example pool scatter, rowDenseBackFlat), IS the backward of the certified VJP of the whole batched forward, relu ∘ dense ∘ relu ∘ dense ∘ dense ∘ pool ∘ (conv → BN → relu)² ∘ …, at every batch size N. ResNet-50's r50InputGradB_eq_r34B_full_vjp is the statement this one mirrors; what ResNet-50 does with sixteen opaque bottleneck witnesses, this net does with seventeen concrete CertLayers, because every stage of it is a layer the suite already certifies.

How it is built #

Scope #

A smooth-point statement at one precision: the conv nodes are the f32 kinds, as in ResNet-50's whole-net tie (the bf16 artifacts' gradient nodes are Cifar8BbnStepTieB's, on the renderers' bf16 switch at the identity rounding). One device: cifar8wb_bbn has no data-parallel render. It is about the INPUT cotangent chain; the thirty-eight parameter gradients are the node tie's, and that each of them is the loss's derivative in its parameter is Cifar8BbnParamGrad.lean (cifar8Bbn_net_lossGrad).

def Proofs.StableHLO.MaxPool2SmoothB (N c h w : ℕ) (x : Vec (N * (c * (2 * h) * (2 * w)))) :

Per-example smoothness of a batched pool input: every example's 2×2 windows attain their maximum at exactly one cell (MaxPool2Smooth, read on that example's [C, 2h, 2w] slice).

Equations
Instances For
    theorem Proofs.StableHLO.maxPoolFlat_differentiableAt_row (N c h w : ℕ) (x : Vec (N * (c * (2 * h) * (2 * w)))) (hx : MaxPool2SmoothB N c h w x) (r : Fin N) :
    noncomputable def Proofs.StableHLO.maxPool2BHasVJPAt (N c h w : ℕ) (x : Vec (N * (c * (2 * h) * (2 * w)))) (hx : MaxPool2SmoothB N c h w x) :

    The batched 2×2 pool's VJP at a smooth point — batchMapHasVJPAt over each example's own maxPoolFlatHasVJPAt'. Built field by field (no transport), so its .backward reduces.

    Equations
    Instances For
      theorem Proofs.StableHLO.maxPool2BHasVJPAt_backward (N c h w : ℕ) (x : Vec (N * (c * (2 * h) * (2 * w)))) (hx : MaxPool2SmoothB N c h w x) :

      Its backward is the render's scatter: maxPoolBackFlat per example on that example's own saved activation, which is den (.maxPoolBackB …).

      noncomputable def Proofs.StableHLO.maxPool2Layer (N c h w : ℕ) :
      CertLayer (N * (c * (2 * h) * (2 * w))) (N * (c * h * w))

      Batched 2×2/s2 max pool as a CertLayer, certified where every example's windows are smooth; its backward graph is the render's .maxPoolBackB at the saved pool input.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        theorem Proofs.StableHLO.maxPool2Layer_fwd_apply (N c h w : ℕ) (v : Vec (N * (c * (2 * h) * (2 * w)))) :
        (maxPool2Layer N c h w).fwd v = batchMap N (maxPoolFlat c h w) v
        theorem Proofs.StableHLO.CertLayer.comp_vjp_backward {m n p : ℕ} (L₁ : CertLayer m n) (L₂ : CertLayer n p) (x : Vec m) (hx : (L₁.comp L₂).ok x) (dy : Vec p) :
        ((L₁.comp L₂).vjp x hx).backward dy = (L₁.vjp x ⋯).backward ((L₂.vjp (L₁.fwd x) ⋯).backward dy)

        comp's VJP backward, peeled: L₂'s then L₁'s. Definitional (vjpCompAt_backward); stated so a whole-net chain unwinds by rw rather than by simp on the structure.

        theorem Proofs.StableHLO.cbReluB_back_eq (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x : Vec (N * (ic * h * w))) (hs : ∀ (k : Fin (N * (oc * h * w))), bnBatchLA N oc h w ε γ β (batchMap N (flatConv W b) x) k ≠ 0) (dy : Vec (N * (oc * h * w))) :
        (cbReluBHasVJPAt N W b ε hε γ β x hs).backward dy = BackLinks.cInB N W b (BackLinks.bnBackB N oc h w ε hε γ β (batchMap N (flatConv W b) x) (BackLinks.reluMaskB (N * (oc * h * w)) (bnBatchLA N oc h w ε γ β (batchMap N (flatConv W b) x)) dy))

        The conv → BN → relu stage's VJP backward, in the BackLinks vocabulary: the relu mask at the pre-activation, the true-BN backward at the conv output, the conv input-VJP. The relu peer of cbsB_back_eq.

        theorem Proofs.StableHLO.cbReluLayer_vjp_backward (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x : Vec (N * (ic * h * w))) (hx : (cbReluLayer N W b ε hε γ β).ok x) (dy : Vec (N * (oc * h * w))) :
        ((cbReluLayer N W b ε hε γ β).vjp x hx).backward dy = BackLinks.cInB N W b (BackLinks.bnBackB N oc h w ε hε γ β (batchMap N (flatConv W b) x) (BackLinks.reluMaskB (N * (oc * h * w)) (bnBatchLA N oc h w ε γ β (batchMap N (flatConv W b) x)) dy))

        cbReluLayer's VJP backward, in the BackLinks vocabulary.

        theorem Proofs.StableHLO.maxPool2Layer_vjp_backward (N c h w : ℕ) (x : Vec (N * (c * (2 * h) * (2 * w)))) (hx : (maxPool2Layer N c h w).ok x) (dy : Vec (N * (c * h * w))) :
        ((maxPool2Layer N c h w).vjp x hx).backward dy = batchMapAux N (maxPoolBackFlat c h w) x dy

        maxPool2Layer's VJP backward is the per-example scatter.

        theorem Proofs.StableHLO.denseLayer_vjp_backward (N : ℕ) {a nC : ℕ} (W : Mat a nC) (b : Vec nC) (x : Vec (N * a)) (hx : (denseLayer N W b).ok x) (dy : Vec (N * nC)) :
        ((denseLayer N W b).vjp x hx).backward dy = rowDenseBackFlat N a nC W dy

        denseLayer's VJP backward is rowDenseBackFlat (den of .denseRowBack).

        theorem Proofs.StableHLO.reluOut_vjp_backward (n : ℕ) (x : Vec n) (hx : (CertLayer.reluOut n).ok x) (dy : Vec n) :

        reluOut's VJP backward is the relu mask.

        structure Proofs.C8BbnW (ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ) :

        The thirty-eight tensors of the eight-conv CIFAR net with BatchNorm, as one record: eight (kernel, bias, γ, β) and the three dense layers. One ε is shared by every BatchNorm, as the render bakes it.

        Instances For

          The running activations, one def each, in forward order. Stage k runs at spatial 2^(5-k)·h, written as the explicit 2 * (…) nest the render's types use.

          noncomputable def Proofs.c8A1 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
          Vec (N * (c1 * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))

          conv₁ → BN → relu, at 32².

          Equations
          Instances For
            noncomputable def Proofs.c8A2 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
            Vec (N * (c1 * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))

            conv₂ → BN → relu, at 32².

            Equations
            Instances For
              noncomputable def Proofs.c8P1 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
              Vec (N * (c1 * (2 * (2 * (2 * h))) * (2 * (2 * (2 * w)))))

              pool₁: 32² → 16².

              Equations
              Instances For
                noncomputable def Proofs.c8A3 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                Vec (N * (c2 * (2 * (2 * (2 * h))) * (2 * (2 * (2 * w)))))

                conv₃ → BN → relu, at 16².

                Equations
                Instances For
                  noncomputable def Proofs.c8A4 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                  Vec (N * (c2 * (2 * (2 * (2 * h))) * (2 * (2 * (2 * w)))))

                  conv₄ → BN → relu, at 16².

                  Equations
                  Instances For
                    noncomputable def Proofs.c8P2 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                    Vec (N * (c2 * (2 * (2 * h)) * (2 * (2 * w))))

                    pool₂: 16² → 8².

                    Equations
                    Instances For
                      noncomputable def Proofs.c8A5 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                      Vec (N * (c3 * (2 * (2 * h)) * (2 * (2 * w))))

                      conv₅ → BN → relu, at 8².

                      Equations
                      Instances For
                        noncomputable def Proofs.c8A6 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                        Vec (N * (c3 * (2 * (2 * h)) * (2 * (2 * w))))

                        conv₆ → BN → relu, at 8².

                        Equations
                        Instances For
                          noncomputable def Proofs.c8P3 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                          Vec (N * (c3 * (2 * h) * (2 * w)))

                          pool₃: 8² → 4².

                          Equations
                          Instances For
                            noncomputable def Proofs.c8A7 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                            Vec (N * (c4 * (2 * h) * (2 * w)))

                            conv₇ → BN → relu, at 4².

                            Equations
                            Instances For
                              noncomputable def Proofs.c8A8 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                              Vec (N * (c4 * (2 * h) * (2 * w)))

                              conv₈ → BN → relu, at 4².

                              Equations
                              Instances For
                                noncomputable def Proofs.c8P4 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                                Vec (N * (c4 * h * w))

                                pool₄: 4² → 2², the flatten the head reads (c4·h·w wide).

                                Equations
                                Instances For
                                  noncomputable def Proofs.c8H9 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                                  Vec (N * d1)

                                  dense₉ pre-activation.

                                  Equations
                                  Instances For
                                    noncomputable def Proofs.c8R9 {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                                    Vec (N * d1)

                                    relu₉.

                                    Equations
                                    Instances For
                                      noncomputable def Proofs.c8HA {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                                      Vec (N * d1)

                                      denseₐ pre-activation.

                                      Equations
                                      Instances For
                                        noncomputable def Proofs.c8RA {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                                        Vec (N * d1)

                                        reluₐ.

                                        Equations
                                        Instances For
                                          noncomputable def Proofs.cifar8BbnForwardB {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                                          Vec (N * nC)

                                          The batched eight-conv CIFAR forward with BatchNorm at batch statistics, to the logits: the composite dense_b ∘ relu ∘ dense_a ∘ relu ∘ dense₉ ∘ pool ∘ (conv → BN → relu)² ∘ … at batch N, through the named running activations above.

                                          Equations
                                          Instances For
                                            noncomputable def Proofs.cifar8BbnLayer {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) :
                                            StableHLO.CertLayer (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w)))))) (N * nC)

                                            The net as one CertLayer: seventeen certified layers composed.

                                            Equations
                                            • One or more equations did not get rendered due to their size.
                                            Instances For
                                              theorem Proofs.cifar8BbnLayer_fwd {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) :

                                              The shape check: the composed layer's forward IS cifar8BbnForwardB.

                                              structure Proofs.Cifar8BbnSmoothAt {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :

                                              Where the net is certified: each of the eight pre-relu activations off its kink, each of the four pool inputs with unique window maxima per example, the two head pre-activations off zero — at the running activations. Everything else in the net (conv, BN at 0 < ε, dense) is smooth everywhere.

                                              Instances For
                                                theorem Proofs.cifar8BbnSmoothAt_ok {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (hs : Cifar8BbnSmoothAt N p ε x) :
                                                (cifar8BbnLayer N p ε hε).ok x

                                                Cifar8BbnSmoothAt is the layer's .ok: one comp_ok_of per layer, each naming its input.

                                                noncomputable def Proofs.cifar8BbnBHasVJPAt {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (hs : Cifar8BbnSmoothAt N p ε x) :

                                                The certified whole-net VJP of cifar8BbnForwardB at a smooth point.

                                                Equations
                                                Instances For
                                                  noncomputable def Proofs.cifar8BbnInputGradB {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (dy : Vec (N * nC)) :
                                                  Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))

                                                  The rendered backward chain, read in the BackLinks vocabulary, from the logits cotangent dy to the input cotangent: three rowDenseBackFlats with the two head relu masks between them, then per stage the pool scatter, and twice reluMaskB → bnBackB → cInB. Every step is the den of the node the artifact emits (den_convBackBatched_eq_cInB, den_maxPoolBackB, den_selectPosB, den of .denseRowBack), the BatchNorm link through bnInB_eq_bnBackB.

                                                  Equations
                                                  • One or more equations did not get rendered due to their size.
                                                  Instances For
                                                    theorem Proofs.cifar8BbnLayer_vjp_backward {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (hok : (cifar8BbnLayer N p ε hε).ok x) (dy : Vec (N * nC)) :
                                                    ((cifar8BbnLayer N p ε hε).vjp x hok).backward dy = cifar8BbnInputGradB N p ε hε x dy

                                                    The chain IS the certified whole-net backward. The layer's VJP backward, peeled one comp at a time, is each stage's backward in turn, and each of those is its BackLinks form: cbReluB_back_eq at the eight conv stages, maxPool2BHasVJPAt_backward at the four pools, and rfl at the dense and relu layers.

                                                    theorem Proofs.cifar8BbnBHasVJPAt_backward {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (hs : Cifar8BbnSmoothAt N p ε x) (dy : Vec (N * nC)) :
                                                    (cifar8BbnBHasVJPAt N p ε hε x hs).backward dy = cifar8BbnInputGradB N p ε hε x dy

                                                    …and so is the certified whole-net VJP's.

                                                    theorem Proofs.cifar8BbnInputGradB_correct {ic c1 c2 c3 c4 h w d1 nC kH kW : ℕ} (N : ℕ) (p : C8BbnW ic c1 c2 c3 c4 h w d1 nC kH kW) (ε : ℝ) (hε : 0 < ε) (x : Vec (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (hs : Cifar8BbnSmoothAt N p ε x) (dy : Vec (N * nC)) (i : Fin (N * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) :
                                                    cifar8BbnInputGradB N p ε hε x dy i = ∑ j : Fin (N * nC), pdiv (cifar8BbnForwardB N p ε) x i j * dy j

                                                    The chain IS the pdiv-contracted Jacobian of the whole net — at every batch size, every logits cotangent and every input coordinate, wherever the eight relus and the two head relus are off their kinks and the four pools have unique window maxima (Cifar8BbnSmoothAt), with 0 < ε.