TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.MLTheory.CROWN.Proofs.Distillation

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 #

source
@[simp]
theorem NN.MLTheory.CROWN.Distillation.vecGet_sub {n : ℕ} (x y : Spec.Tensor ℝ (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin n) :
(x.subSpec y).vecGet i = x.vecGet i - y.vecGet i

Interval arithmetic on output boxes #

source
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$.

Instances For
    source
    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].

    Instances For
      source
      noncomputable def NN.MLTheory.CROWN.Distillation.checkBoxWithinAbs {n : ℕ} (B : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)) (eps : ℝ) :
      Bool
      Instances For
        source
        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.

        source
        theorem NN.MLTheory.CROWN.Distillation.boxSub_contains {n : ℕ} {T S : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)} {x y : Spec.Tensor ℝ (Spec.Shape.dim n Spec.Shape.scalar)} (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.

        source
        theorem NN.MLTheory.CROWN.Distillation.boxContains_vecGet {n : ℕ} {B : Box ℝ (Spec.Shape.dim n Spec.Shape.scalar)} {x : Spec.Tensor ℝ (Spec.Shape.dim n Spec.Shape.scalar)} (h : B.contains x) (i : Fin n) :
        B.lo.vecGet i ≤ x.vecGet i ∧ x.vecGet i ≤ B.hi.vecGet i

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

        Distillation certificate for 2-layer MLPs #

        source
        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.

        Instances For
          source
          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 : Spec.Tensor ℝ (Spec.Shape.dim inDim Spec.Shape.scalar)) :
          xB.contains x → ∀ (i : Fin outDim), |((forward teacher x).subSpec (forward student x)).vecGet 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.