Documentation

LeanMlir.Proofs.Nets.Small.MnistCNN

Chapter 3: MNIST 2D CNN (no BatchNorm) — whole-network VJP #

The Chapter-3 demo model mnistCnnNoBn:

conv2d 1→c (relu) → conv2d c→c (relu) → maxPool 2×2 → flatten → dense (relu) → dense (relu) → dense (identity)

mnistCnnNoBnHasVJPAt is the structural whole-network VJP: the composed backward equals the pdiv-Jacobian VJP of the full forward pass, conditional on smoothness hypotheses (no ReLU kink / MaxPool tie at the running activations). The Chapter-3 sibling of cnnHasVJPAt, minus BN and residual blocks. TrainedCnn.trainedCnnHasVJPAt discharges every hypothesis on a reduced instance of this forward — input a 24×24 MNIST crop 4×4-average-pooled to 6×6, conv 1→2→2 (3×3), dense 18→8→8→10, /128-rational trained weights (the max-pool no-tie hypothesis needed a pool-tie regulariser during training) — at one pooled MNIST test digit.

noncomputable def Proofs.convReluHasVJPAt {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (v : Vec (ic * h * w)) (h_smooth : ∀ (k : Fin (oc * h * w)), flatConv W b v k ≠ 0) :
HasVJPAt (relu (oc * h * w) ∘ flatConv W b) v

conv → relu block VJP at a smooth point (no BatchNorm). relu ∘ flatConv W b. The plain-conv analogue of convBnReluHasVJPAt — conv is linear (global VJP via the HasVJP3 bridge), relu carries the smoothness hypothesis.

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    theorem Proofs.convRelu_differentiableAt {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (v : Vec (ic * h * w)) (h_smooth : ∀ (k : Fin (oc * h * w)), flatConv W b v k ≠ 0) :

    relu ∘ flatConv W b is differentiable at a smooth point.

    noncomputable def Proofs.denseReluHasVJPAt {m n : ℕ} (W : Mat m n) (b : Vec n) (v : Vec m) (h_smooth : ∀ (k : Fin n), dense W b v k ≠ 0) :

    dense → relu block VJP at a smooth point. relu ∘ dense W b.

    Equations
    Instances For
      theorem Proofs.denseRelu_differentiableAt {m n : ℕ} (W : Mat m n) (b : Vec n) (v : Vec m) (h_smooth : ∀ (k : Fin n), dense W b v k ≠ 0) :

      relu ∘ dense W b is differentiable at a smooth point.

      noncomputable def Proofs.mnistCnnNoBnForward {ic c h w d1 nClasses kH kW : ℕ} (W₁ : Kernel4 c ic kH kW) (b₁ : Vec c) (W₂ : Kernel4 c c kH kW) (b₂ : Vec c) (W₃ : Mat (c * h * w) d1) (b₃ : Vec d1) (W₄ : Mat d1 d1) (b₄ : Vec d1) (W₅ : Mat d1 nClasses) (b₅ : Vec nClasses) :
      Vec (ic * (2 * h) * (2 * w)) → Vec nClasses

      The Chapter-3 mnistCnnNoBn forward, in flattened Vec space. Conv stage runs at spatial (2*h, 2*w); the maxPool halves it to (h, w); then three dense layers (two with ReLU).

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        noncomputable def Proofs.mnistCnnNoBnHasVJPAt {ic c h w d1 nClasses kH kW : ℕ} (W₁ : Kernel4 c ic kH kW) (b₁ : Vec c) (W₂ : Kernel4 c c kH kW) (b₂ : Vec c) (W₃ : Mat (c * h * w) d1) (b₃ : Vec d1) (W₄ : Mat d1 d1) (b₄ : Vec d1) (W₅ : Mat d1 nClasses) (b₅ : Vec nClasses) (x : Vec (ic * (2 * h) * (2 * w))) (h1 : ∀ (k : Fin (c * (2 * h) * (2 * w))), flatConv W₁ b₁ x k ≠ 0) (h2 : ∀ (k : Fin (c * (2 * h) * (2 * w))), flatConv W₂ b₂ ((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x) k ≠ 0) (h_mp : MaxPool2Smooth (Tensor3.unflatten (((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x))) (h3 : ∀ (k : Fin d1), dense W₃ b₃ (maxPoolFlat c h w (((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x)) k ≠ 0) (h4 : ∀ (k : Fin d1), dense W₄ b₄ ((relu d1 ∘ dense W₃ b₃) (maxPoolFlat c h w (((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x))) k ≠ 0) :
        HasVJPAt (mnistCnnNoBnForward W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅) x

        MNIST 2D CNN (no BN) whole-network VJP at a smooth point.

        The composed backward of the full Chapter-3 forward equals the pdiv-contracted Jacobian (Jacobian-transpose applied to the cotangent), conditional on smoothness at the four ReLU kinks and the one MaxPool. Built by vjpCompAt through convRelu → convRelu → maxPool → denseRelu → denseRelu → dense. The Chapter-3 sibling of cnnHasVJPAt (BN-free, no resblocks).

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For
          theorem Proofs.mnistCnnNoBnHasVJPAt_correct {ic c h w d1 nClasses kH kW : ℕ} (W₁ : Kernel4 c ic kH kW) (b₁ : Vec c) (W₂ : Kernel4 c c kH kW) (b₂ : Vec c) (W₃ : Mat (c * h * w) d1) (b₃ : Vec d1) (W₄ : Mat d1 d1) (b₄ : Vec d1) (W₅ : Mat d1 nClasses) (b₅ : Vec nClasses) (x : Vec (ic * (2 * h) * (2 * w))) (h1 : ∀ (k : Fin (c * (2 * h) * (2 * w))), flatConv W₁ b₁ x k ≠ 0) (h2 : ∀ (k : Fin (c * (2 * h) * (2 * w))), flatConv W₂ b₂ ((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x) k ≠ 0) (h_mp : MaxPool2Smooth (Tensor3.unflatten (((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x))) (h3 : ∀ (k : Fin d1), dense W₃ b₃ (maxPoolFlat c h w (((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x)) k ≠ 0) (h4 : ∀ (k : Fin d1), dense W₄ b₄ ((relu d1 ∘ dense W₃ b₃) (maxPoolFlat c h w (((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x))) k ≠ 0) (dy : Vec nClasses) (i : Fin (ic * (2 * h) * (2 * w))) :
          (mnistCnnNoBnHasVJPAt W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ x h1 h2 h_mp h3 h4).backward dy i = ∑ j : Fin nClasses, pdiv (mnistCnnNoBnForward W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅) x i j * dy j

          Public correctness theorem for mnistCnnNoBnHasVJPAt — the Chapter-3 CNN's backward equals the pdiv-contracted Jacobian.

          theorem Proofs.maxPool2Smooth_of_injective {c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) (hinj : ∀ (ci : Fin c) (r r' : Fin (2 * h)) (s s' : Fin (2 * w)), x ci r s = x ci r' s' → r = r' ∧ s = s') :

          Positional injectivity ⇒ MaxPool2Smooth. If, on each channel, the position map (r, s) ↦ x ci r s is injective, then every 2×2 window has pairwise-distinct values, so its max is attained once. One injectivity argument in place of 6·c·h·w per-window decides.

          theorem Proofs.flatten_pos_of_pos {c h w : ℕ} {T : Tensor3 c h w} (hT : ∀ (ci : Fin c) (hi : Fin h) (wi : Fin w), 0 < T ci hi wi) (k : Fin (c * h * w)) :
          0 < T.flatten k

          A tensor that is positive everywhere flattens to a positive vector (flatten T k just reads T at the decoded index). Discharges the ReLU ∀ k, … ≠ 0 conditions once the layer is shown positive.

          theorem Proofs.maxPool2_pos {c h w : ℕ} {x : Tensor3 c (2 * h) (2 * w)} (hx : ∀ (ci : Fin c) (r : Fin (2 * h)) (s : Fin (2 * w)), 0 < x ci r s) (ci : Fin c) (hi : Fin h) (wi : Fin w) :
          0 < maxPool2 x ci hi wi

          A 2×2 max-pool of an everywhere-positive tensor is positive (the max dominates the top-left cell).

          theorem Proofs.dense_pos_of_nonneg {m n : ℕ} {W : Mat m n} {b : Vec n} {u : Vec m} (hW : ∀ (i : Fin m) (j : Fin n), 0 ≤ W i j) (hb : ∀ (j : Fin n), 0 < b j) (hu : ∀ (i : Fin m), 0 ≤ u i) (j : Fin n) :
          0 < dense W b u j

          A dense layer with nonnegative weights, a strictly-positive bias, and a nonnegative input is strictly positive — the propagating positivity invariant that discharges the dense ReLU ≠ 0 conditions without per-coordinate case analysis.

          theorem Proofs.conv2d_1x1 {ic oc h w : ℕ} (W : Kernel4 oc ic 1 1) (b : Vec oc) (t : Tensor3 ic h w) (o : Fin oc) (hi : Fin h) (wi : Fin w) :
          conv2d W b t o hi wi = b o + ∑ c : Fin ic, W o c 0 0 * t c hi wi

          1×1 conv collapses to a per-pixel channel mix. With a 1×1 kernel (SAME padding is a no-op), conv2d at each pixel is just the bias plus a channel-weighted sum of that same pixel — the closed form the 1×1-kernel witnesses compute their forward pass with.

          noncomputable def Proofs.MlpConcrete.W₀ :
          Mat 2 2

          A concrete 3-layer MLP (dense → relu → dense → relu → dense) with all-ones weights/biases and a positive input. Every ReLU pre-activation is then strictly positive (hence ≠ 0), so both smoothness hypotheses of mlpHasVJPAt discharge. The net is non-constant, so this is a live witness (non-trivial Jacobian), not a degenerate one.

          Equations
          Instances For
            noncomputable def Proofs.MlpConcrete.b₀ :
            Vec 2
            Equations
            Instances For
              noncomputable def Proofs.MlpConcrete.W₁ :
              Mat 2 2
              Equations
              Instances For
                noncomputable def Proofs.MlpConcrete.b₁ :
                Vec 2
                Equations
                Instances For
                  noncomputable def Proofs.MlpConcrete.W₂ :
                  Mat 2 2
                  Equations
                  Instances For
                    noncomputable def Proofs.MlpConcrete.b₂ :
                    Vec 2
                    Equations
                    Instances For
                      noncomputable def Proofs.MlpConcrete.x :
                      Vec 2
                      Equations
                      Instances For

                        Unconditional whole-network VJP for a concrete 3-layer MLP. Both ReLU ≠ 0 hypotheses are discharged via dense_pos_of_nonneg (positive bias + nonnegative weights/input propagate strict positivity), with relu_id_of_pos collapsing the inner ReLU.

                        Equations
                        • One or more equations did not get rendered due to their size.
                        Instances For
                          noncomputable def Proofs.CnnConcrete.Ws :
                          Kernel4 1 1 1 1
                          Equations
                          Instances For
                            noncomputable def Proofs.CnnConcrete.bs :
                            Vec 1
                            Equations
                            Instances For
                              noncomputable def Proofs.CnnConcrete.X :
                              Vec (1 * (2 * 1) * (2 * 1))
                              Equations
                              Instances For
                                noncomputable def Proofs.CnnConcrete.W₁ :
                                Kernel4 1 1 1 1
                                Equations
                                Instances For
                                  noncomputable def Proofs.CnnConcrete.b₁ :
                                  Vec 1
                                  Equations
                                  Instances For
                                    noncomputable def Proofs.CnnConcrete.W₂ :
                                    Kernel4 1 1 1 1
                                    Equations
                                    Instances For
                                      noncomputable def Proofs.CnnConcrete.b₂ :
                                      Vec 1
                                      Equations
                                      Instances For
                                        noncomputable def Proofs.CnnConcrete.W₁' :
                                        Kernel4 1 1 1 1
                                        Equations
                                        Instances For
                                          noncomputable def Proofs.CnnConcrete.b₁' :
                                          Vec 1
                                          Equations
                                          Instances For
                                            noncomputable def Proofs.CnnConcrete.W₂' :
                                            Kernel4 1 1 1 1
                                            Equations
                                            Instances For
                                              noncomputable def Proofs.CnnConcrete.b₂' :
                                              Vec 1
                                              Equations
                                              Instances For
                                                noncomputable def Proofs.CnnConcrete.Wp :
                                                Kernel4 1 1 1 1
                                                Equations
                                                Instances For
                                                  noncomputable def Proofs.CnnConcrete.bp :
                                                  Vec 1
                                                  Equations
                                                  Instances For
                                                    noncomputable def Proofs.CnnConcrete.Wd :
                                                    Mat 1 2
                                                    Equations
                                                    Instances For
                                                      noncomputable def Proofs.CnnConcrete.bd :
                                                      Vec 2
                                                      Equations
                                                      Instances For
                                                        theorem Proofs.CnnConcrete.cbr_x :
                                                        cbr Ws bs (11 / 4) 1 10 X = bnForward (1 * (2 * 1) * (2 * 1)) (11 / 4) 1 10 X

                                                        cbr X collapses to bn X (identity conv, then relu of a positive).

                                                        noncomputable def Proofs.CnnConcrete.cnnConcreteHasVJPAt :
                                                        HasVJPAt (cnnForward Ws bs (11 / 4) 1 10 W₁ b₁ W₂ b₂ 1 0 1 1 0 1 W₁' b₁' W₂' b₂' Wp bp 1 0 1 1 0 1 1 0 1 Wd bd) X

                                                        Whole-network VJP for a concrete ResNet-style CNN with BatchNorm — every smoothness hypothesis discharged: the stem produces distinct positive BN outputs (so maxpool has no ties and bn ≠ 0), and the resblock BNs use γ=0 (constant).

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