MLP VJP Proofs

Formal VJP correctness for the layers of a 3-layer MLP. All definitions over ℝ, proofs use Mathlib's Finset.sum.

noncomputable def Proofs.dense {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) : Vec n

Equations

• Proofs.dense W b x j = ∑ i : Fin m, x i * W i j + b j

theorem Proofs.pdiv_dense {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (i : Fin m) (j : Fin n) : pdiv (dense W b) x i j = W i j

Dense Jacobian — ∂(W·x + b)_j/∂x_i = W_{ij}. Derived from the foundation theorems (pdiv_add, pdiv_const, pdiv_finset_sum, pdiv_mul, pdiv_reindex). The proof factors dense W b into (∑ i', x i' * W i' j) + b j, distributes pdiv over the outer sum and finset sum, applies the product rule per summand, and collapses the Kronecker δ.

theorem Proofs.pdiv_dense_W {m n : ℕ} (b : Vec n) (x : Vec m) (W : Mat m n) (i : Fin m) (j' j : Fin n) : pdiv (fun (v : Vec (m * n)) => dense (Mat.unflatten v) b x) W.flatten (finProdFinEquiv (i, j')) j = if j = j' then x i else 0

Jacobian of dense wrt W — ∂dense(W, b, x)_j/∂W_{i, j'} = x_i·δ(j, j'). Derived from foundation theorems (pdiv_add, pdiv_const, pdiv_finset_sum, pdiv_mul, pdiv_reindex) over the flatten bijection. Symmetric counterpart to pdiv_dense.

noncomputable def Proofs.dense_has_vjp {m n : ℕ} (W : Mat m n) (b : Vec n) : HasVJP (dense W b)

Dense VJP — proved.

Equations

• Proofs.dense_has_vjp W b = { backward := fun (_x : Proofs.Vec m) (dy : Proofs.Vec n) => W.mulVec dy, correct := ⋯ }

noncomputable def Proofs.mnistLinear {m n : ℕ} (W : Mat m n) (b : Vec n) : Vec m → Vec n

The Chapter-1 demo model: a linear classifier is a single dense layer.

Equations

• Proofs.mnistLinear W b = Proofs.dense W b

theorem Proofs.mnistLinear_has_vjp_correct {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (dy : Vec n) (i : Fin m) : (dense_has_vjp W b).backward x dy i = ∑ j : Fin n, pdiv (mnistLinear W b) x i j * dy j

Whole-model VJP contract for the linear classifier — the degenerate simplest case of the per-architecture *_has_vjp_correct capstones, built straight from the Chapter-1 kit.

theorem Proofs.dense_differentiable {m n : ℕ} (W : Mat m n) (b : Vec n) : Differentiable ℝ (dense W b)

Dense is everywhere differentiable. dense W b is affine in x, hence smooth; this is the underlying Differentiable ℝ statement that vjp_comp_at needs when composing through dense layers.

theorem Proofs.dense_weight_grad_correct {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (dy : Vec n) (i : Fin m) (j : Fin n) : Mat.outer x dy i j = ∑ k : Fin n, pdiv (fun (v : Vec (m * n)) => dense (Mat.unflatten v) b x) W.flatten (finProdFinEquiv (i, j)) k * dy k

Dense weight gradient is the outer product — theorem (Phase 7).

Mat.outer x dy is the cotangent-contracted Jacobian of dense(W, b, x) with respect to W, at every index. This promotes the previous vacuous rfl about Mat.outer into a real theorem connecting the outer product to the actual weight gradient of dense.

(Mat.outer x dy) i j = ∑ k, pdiv (…) (Mat.flatten W) (fPF (i, j)) k · dy k

theorem Proofs.pdiv_dense_b {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (i j : Fin n) : pdiv (fun (b' : Vec n) => dense W b' x) b i j = if i = j then 1 else 0

Dense bias gradient is identity — theorem (Phase 7).

∂ dense(W, b, x)_j / ∂ b_{j'} = δ(j, j'), so the bias backward is just dy itself. Derived from pdiv_add + pdiv_const + pdiv_id — no new axiom.

theorem Proofs.dense_bias_grad_correct {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (dy : Vec n) (i : Fin n) : dy i = ∑ j : Fin n, pdiv (fun (b' : Vec n) => dense W b' x) b i j * dy j

noncomputable def Proofs.dense_weight_grad {m n : ℕ} (x : Vec m) (dy : Vec n) : Mat m n

Dense weight backward — named accessor. dW = x ⊗ dy (outer product).

Equations

• Proofs.dense_weight_grad x dy = Proofs.Mat.outer x dy

def Proofs.dense_bias_grad {n : ℕ} (dy : Vec n) : Vec n

Dense bias backward — named accessor. db = dy.

Equations

• Proofs.dense_bias_grad dy = dy

noncomputable def Proofs.relu (n : ℕ) (x : Vec n) : Vec n

Equations

• Proofs.relu n x i = if x i > 0 then x i else 0

noncomputable def Proofs.reluLinearPart (n : ℕ) (x : Vec n) : Vec n →L[ℝ] Vec n

ReLU's local linear part at a smooth point — the diagonal indicator CLM. At each coordinate k, projects to y k if x k > 0, otherwise zero. Two smooth points with the same sign pattern share this same CLM.

Equations

• Proofs.reluLinearPart n x = ContinuousLinearMap.pi fun (k : Fin n) => if x k > 0 then ContinuousLinearMap.proj k else 0

@[simp]

theorem Proofs.reluLinearPart_apply (n : ℕ) (x y : Vec n) (k : Fin n) : (reluLinearPart n x) y k = if x k > 0 then y k else 0

theorem Proofs.relu_hasFDerivAt (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) : HasFDerivAt (relu n) (reluLinearPart n x) x

ReLU is differentiable at smooth points. Within Metric.ball x r for r := min |x k|, every coordinate keeps its sign — so relu n agrees with reluLinearPart n x on a neighborhood. EventuallyEq promotes the CLM's HasFDerivAt to ReLU's.

theorem Proofs.relu_differentiableAt_of_smooth (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) : DifferentiableAt ℝ (relu n) x

ReLU is DifferentiableAt at smooth points. Corollary of relu_hasFDerivAt; lets vjp_comp_at chain through ReLU.

theorem Proofs.pdiv_relu (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) (i j : Fin n) : pdiv (relu n) x i j = if i = j then if x i > 0 then 1 else 0 else 0

ReLU partial derivative — proved via relu_hasFDerivAt and direct evaluation at basisVec i.

noncomputable def Proofs.relu_has_vjp (n : ℕ) : HasVJP (relu n)

ReLU bundled VJP — canonical (junk-at-kink) witness.

HasVJP.correct is satisfied by the canonical pdiv-derived backward: at smooth points it is the diagonal indicator (per pdiv_relu); at points where some coordinate is zero, pdiv (relu n) x agrees with fderiv's junk default of 0, so the canonical backward is 0 there too — and correct holds by rfl.

The codegen (MlirCodegen.lean) emits the standard subgradient formula if x > 0 then dy else 0 instead, which agrees with the canonical witness at smooth points and differs at the kinks (the convention relu'(0) := 0 used by every ML framework). The smooth-point agreement is formal: see relu_codegen_matches_canonical below. The Lean-vs-codegen gap at the kinks is the codegen trust boundary — see LeanMlir/Proofs/README.md.

Equations

• Proofs.relu_has_vjp n = { backward := fun (x dy : Proofs.Vec n) (i : Fin n) => ∑ j : Fin n, Proofs.pdiv (Proofs.relu n) x i j * dy j, correct := ⋯ }

theorem Proofs.relu_codegen_matches_canonical (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) (dy : Vec n) (i : Fin n) : (relu_has_vjp n).backward x dy i = if x i > 0 then dy i else 0

Bridge: relu_has_vjp's canonical backward matches the codegen formula at smooth points.

At any point where no coordinate of x is zero, the canonical pdiv-derived backward ∑ j, pdiv (relu n) x i j * dy j collapses to the framework subgradient if x i > 0 then dy i else 0 that MlirCodegen.lean actually emits. Closes the smooth-point half of the codegen trust boundary — what's left is just the kink convention.

theorem Proofs.relu_canonical_diagonal (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) (dy : Vec n) (i : Fin n) : (relu_has_vjp n).backward x dy i = (if x i > 0 then 1 else 0) * dy i

Diagonal-indicator restatement of the smooth-point bridge. relu_has_vjp.backward x dy i = 1_{x i > 0} · dy i at smooth points — same content as relu_codegen_matches_canonical, factored as (indicator) · dy i for downstream use.

noncomputable def Proofs.relu_has_vjp_at (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) : HasVJPAt (relu n) x

ReLU pointwise VJP — no canonical-witness escape.

Constructs HasVJPAt (relu n) x at a smooth point. The backward is the codegen-shape if x i > 0 then dy i else 0 directly; the correct field is a real proof via pdiv_relu (the smooth-point Jacobian) + sum-collapse, not rfl.

Equations

• Proofs.relu_has_vjp_at n x h_smooth = { backward := fun (dy : Proofs.Vec n) (i : Fin n) => if x i > 0 then dy i else 0, correct := ⋯ }

noncomputable def Proofs.softmax (c : ℕ) (z : Vec c) : Vec c

Equations

• Proofs.softmax c z = fun (j : Fin c) => (fun (j : Fin c) => Real.exp (z j)) j / ∑ k : Fin c, (fun (j : Fin c) => Real.exp (z j)) k

noncomputable def Proofs.oneHot (c : ℕ) (label : Fin c) : Vec c

Equations

• Proofs.oneHot c label j = if j = label then 1 else 0

noncomputable def Proofs.crossEntropy (c : ℕ) (logits : Vec c) (label : Fin c) : ℝ

Equations

• Proofs.crossEntropy c logits label = -Real.log (Proofs.softmax c logits label)

noncomputable def Proofs.mlpForward {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) : Vec d₀ → Vec d₃

Equations

• Proofs.mlpForward W₀ b₀ W₁ b₁ W₂ b₂ = Proofs.dense W₂ b₂ ∘ Proofs.relu d₂ ∘ Proofs.dense W₁ b₁ ∘ Proofs.relu d₁ ∘ Proofs.dense W₀ b₀

noncomputable def Proofs.mlp_has_vjp {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) : HasVJP (mlpForward W₀ b₀ W₁ b₁ W₂ b₂)

MLP composition VJP — canonical witness.

The MLP forward composes dense W b (everywhere Differentiable) with relu (non-Differentiable at the kinks). vjp_comp would require Differentiable ℝ (relu n), which doesn't hold globally, so the chain-rule route is blocked. The canonical pdiv-derived backward inhabits HasVJP.correct directly via rfl — the codegen substitutes the subgradient formula at the kinks (see LeanMlir/Proofs/README.md for the trust-boundary discussion).

Equations

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

noncomputable def Proofs.mlp_has_vjp_at {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x : Vec d₀) (h_smooth_0 : ∀ (k : Fin d₁), dense W₀ b₀ x k ≠ 0) (h_smooth_1 : ∀ (k : Fin d₂), dense W₁ b₁ (relu d₁ (dense W₀ b₀ x)) k ≠ 0) : HasVJPAt (mlpForward W₀ b₀ W₁ b₁ W₂ b₂) x

MLP pointwise VJP — no canonical-witness escape.

Constructs HasVJPAt (mlpForward …) x by chaining vjp_comp_at through dense → relu_at → dense → relu_at → dense. Requires the intermediate pre-activations dense W₀ b₀ x and dense W₁ b₁ z₀ to avoid zero (no coordinate ties the ReLU kink) — exactly the "smooth input" condition. Replaces the vacuous mlp_has_vjp.correct := rfl with a real chain-rule proof at smooth inputs.

Equations

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

Public correctness theorems for the canonical-witness defs

Each _has_vjp def above bundles a backward function with a .correct field; these _correct theorems expose that field as a top-level proposition so consumers (downstream code, tests/comparator/, doc-gen4) can refer to the contract directly without reaching into record internals.

theorem Proofs.relu_has_vjp_correct (n : ℕ) (x dy : Vec n) (i : Fin n) : (relu_has_vjp n).backward x dy i = ∑ j : Fin n, pdiv (relu n) x i j * dy j

Public correctness theorem for relu_has_vjp: the canonical witness's backward equals the pdiv-contracted Jacobian by definition.

theorem Proofs.mlp_has_vjp_correct {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x : Vec d₀) (dy : Vec d₃) (i : Fin d₀) : (mlp_has_vjp W₀ b₀ W₁ b₁ W₂ b₂).backward x dy i = ∑ j : Fin d₃, pdiv (mlpForward W₀ b₀ W₁ b₁ W₂ b₂) x i j * dy j

Public correctness theorem for mlp_has_vjp: same pattern as relu_has_vjp_correct, lifted to the three-layer MLP forward.

theorem Proofs.relu_has_vjp_at_correct (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) (dy : Vec n) (i : Fin n) : (relu_has_vjp_at n x h_smooth).backward dy i = ∑ j : Fin n, pdiv (relu n) x i j * dy j

Public correctness theorem for relu_has_vjp_at — the pointwise (smooth-input) variant. Unlike relu_has_vjp_correct, this wrapper's underlying .correct field is a real proof (pdiv_relu + sum-collapse), not rfl; the wrapper exposes it as a top-level proposition for tests/comparator/ re-verification.

theorem Proofs.mlp_has_vjp_at_correct {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x : Vec d₀) (h_smooth_0 : ∀ (k : Fin d₁), dense W₀ b₀ x k ≠ 0) (h_smooth_1 : ∀ (k : Fin d₂), dense W₁ b₁ (relu d₁ (dense W₀ b₀ x)) k ≠ 0) (dy : Vec d₃) (i : Fin d₀) : (mlp_has_vjp_at W₀ b₀ W₁ b₁ W₂ b₂ x h_smooth_0 h_smooth_1).backward dy i = ∑ j : Fin d₃, pdiv (mlpForward W₀ b₀ W₁ b₁ W₂ b₂) x i j * dy j

Public correctness theorem for mlp_has_vjp_at — the pointwise variant composed via vjp_comp_at through dense → relu_at → dense → relu_at → dense. The underlying .correct field chains real chain-rule proofs (no rfl escape at the ReLU kinks).