Distillation / Equivalence certificates (TwoLayerMLP)

This module adds a distillation-style certificate:

prove that a Student network matches a Teacher network up to ε on an input box, i.e. |T(x) - S(x)| ≤ ε componentwise for all inputs x in the domain.

The implementation is kept simple and reuses TorchLean's existing, proved IBP soundness theorem for 2-layer MLPs (NN.MLTheory.CROWN.Theorems.bound_ibp_sound).

Small vector helpers

@[simp]

theorem NN.MLTheory.CROWN.Distillation.getScalar_sub {n : ℕ} (x y : TorchLean.Tensor ℝ [n]) (i : Fin n) : (x.subSpec y).getScalar i = x.getScalar i - y.getScalar i

Interval arithmetic on output boxes

def NN.MLTheory.CROWN.Distillation.boxSub {n : ℕ} (T S : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)) : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)

Interval subtraction: boxSub T S encloses every pointwise difference $x-y$ with $x\in T$ and $y\in S$.

def NN.MLTheory.CROWN.Distillation.boxWithinAbs {n : ℕ} (B : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)) (eps : ℝ) : Prop

Predicate asserting all coordinates of B lie in [-eps, eps].

noncomputable def NN.MLTheory.CROWN.Distillation.checkBoxWithinAbs {n : ℕ} (B : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)) (eps : ℝ) : Bool

Classical decision procedure for boxWithinAbs.

Noncomputable and decide-based because the carrier here is ℝ; the point is to have a Bool-valued checker to state agreement against, not to run it. The executable float version lives in the CROWN engine.

theorem NN.MLTheory.CROWN.Distillation.checkBoxWithinAbs_spec {n : ℕ} {B : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)} {eps : ℝ} : checkBoxWithinAbs B eps = true ↔ boxWithinAbs B eps

Correctness of checkBoxWithinAbs.

theorem NN.MLTheory.CROWN.Distillation.boxSub_contains {n : ℕ} {T S : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)} {x y : TorchLean.Tensor ℝ [n]} (hx : T.contains x) (hy : S.contains y) : (boxSub T S).contains (x.subSpec y)

If x ∈ T and y ∈ S, then x - y ∈ boxSub T S.

theorem NN.MLTheory.CROWN.Distillation.boxContains_getScalar {n : ℕ} {B : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)} {x : TorchLean.Tensor ℝ [n]} (h : B.contains x) (i : Fin n) : B.lo.getScalar i ≤ x.getScalar i ∧ x.getScalar i ≤ B.hi.getScalar i

Project a Box.contains hypothesis to scalar inequalities at a single coordinate.

Distillation certificate for 2-layer MLPs

noncomputable def NN.MLTheory.CROWN.Distillation.checkEquivalenceTwoLayerMlp {inDim hidDim outDim : ℕ} (teacher student : TwoLayerMLP ℝ inDim hidDim outDim) (xB : Box ℝ (Spec.Shape.dim inDim Spec.Shape.scalar)) (eps : ℝ) : Bool

Computable checker: returns true if IBP proves the student matches the teacher up to eps (componentwise) on the given input box.

theorem NN.MLTheory.CROWN.Distillation.checkEquivalence_twoLayerMlp_sound {inDim hidDim outDim : ℕ} (teacher student : TwoLayerMLP ℝ inDim hidDim outDim) (xB : Box ℝ (Spec.Shape.dim inDim Spec.Shape.scalar)) (eps : ℝ) (hcheck : checkEquivalenceTwoLayerMlp teacher student xB eps = true) (x : TorchLean.Tensor ℝ [inDim]) : xB.contains x → ∀ (i : Fin outDim), |((forward teacher x).subSpec (forward student x)).getScalar i| ≤ eps

Soundness: if checkEquivalenceTwoLayerMlp returns true, then for all inputs x in xB, the outputs are eps-close componentwise: |T(x)_i - S(x)_i| ≤ eps.