Verified Model Distillation / Equivalence Checking

This module implements the logic to verify that two neural networks (e.g., a massive Teacher and a compressed Student) produce outputs that are within a specific tolerance ε of each other for ALL inputs in a given domain.

Main Definitions

• SequentialNet - A network with N layers
• checkEquivalence - Computable certificate for network equivalence
• verify_equivalence - The main theorem proving |T(x) - S(x)| ≤ ε

The Math

We define the difference graph: Diff(x) = Teacher(x) - Student(x)

We compute interval bounds for Diff(InputBox). If the resulting interval is contained in [-ε, ε], the networks are ε-equivalent.

Interval Subtraction

def LeanCert.ML.Distillation.IntervalDyadic.neg (I : Core.IntervalDyadic) : Core.IntervalDyadic

Negate an interval: -I = [-hi, -lo]

Equations

• LeanCert.ML.Distillation.IntervalDyadic.neg I = { lo := I.hi.neg, hi := I.lo.neg, le := ⋯ }

theorem LeanCert.ML.Distillation.IntervalDyadic.mem_neg {x : ℝ} {I : Core.IntervalDyadic} (hx : x ∈ I) : -x ∈ I.neg

Membership in negated interval

def LeanCert.ML.Distillation.IntervalDyadic.sub (I J : Core.IntervalDyadic) : Core.IntervalDyadic

Subtract two intervals: I - J = I + (-J)

Equations

• LeanCert.ML.Distillation.IntervalDyadic.sub I J = I.add J.neg

theorem LeanCert.ML.Distillation.IntervalDyadic.mem_sub {x y : ℝ} {I J : Core.IntervalDyadic} (hx : x ∈ I) (hy : y ∈ J) : x - y ∈ I.sub J

Membership in subtracted interval

Sequential Network Infrastructure

structure LeanCert.ML.Distillation.SequentialNet : Type

A standard feedforward network is just a list of layers

• layers : List Layer

instance LeanCert.ML.Distillation.instReprSequentialNet : Repr SequentialNet

Equations

• LeanCert.ML.Distillation.instReprSequentialNet = { reprPrec := LeanCert.ML.Distillation.instReprSequentialNet.repr }

def LeanCert.ML.Distillation.instReprSequentialNet.repr : SequentialNet → ℕ → Std.Format

Equations

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

def LeanCert.ML.Distillation.SequentialNet.forwardReal (net : SequentialNet) (x : List ℝ) : List ℝ

Real-valued forward pass

Equations

• net.forwardReal x = List.foldl (fun (acc : List ℝ) (l : LeanCert.ML.Layer) => l.forwardReal acc) x net.layers

def LeanCert.ML.Distillation.SequentialNet.forwardInterval (net : SequentialNet) (x : IntervalVector) (prec : ℤ) : IntervalVector

Interval forward pass

Equations

• net.forwardInterval x prec = List.foldl (fun (acc : LeanCert.Engine.IntervalVector) (l : LeanCert.ML.Layer) => l.forwardInterval acc prec) x net.layers

def LeanCert.ML.Distillation.SequentialNet.LayersWellFormed : List Layer → ℕ → Prop

A layer chain is well-formed if dimensions align

Equations

• One or more equations did not get rendered due to their size.
• LeanCert.ML.Distillation.SequentialNet.LayersWellFormed [] x✝ = True

def LeanCert.ML.Distillation.SequentialNet.WellFormed (net : SequentialNet) (inputDim : ℕ) : Prop

Check if all layers in the sequence are well-formed and dimensions align

Equations

• net.WellFormed inputDim = LeanCert.ML.Distillation.SequentialNet.LayersWellFormed net.layers inputDim

theorem LeanCert.ML.Distillation.SequentialNet.forwardLength_aux (layers : List Layer) (xs : List ℝ) (Is : IntervalVector) (hdim : xs.length = List.length Is) (hwf : LayersWellFormed layers xs.length) (prec : ℤ) : (List.foldl (fun (acc : List ℝ) (l : Layer) => l.forwardReal acc) xs layers).length = List.length (List.foldl (fun (acc : IntervalVector) (l : Layer) => l.forwardInterval acc prec) Is layers)

Helper: Lengths of foldl outputs match

theorem LeanCert.ML.Distillation.SequentialNet.mem_forwardInterval_aux (layers : List Layer) (xs : List ℝ) (Is : IntervalVector) (hdim : xs.length = List.length Is) (hwf : LayersWellFormed layers xs.length) (hmem : ∀ (i : ℕ) (hi : i < List.length Is), xs[i] ∈ Is[i]) (prec : ℤ) (hprec : prec ≤ 0) : let outReal := List.foldl (fun (acc : List ℝ) (l : Layer) => l.forwardReal acc) xs layers; let outInt := List.foldl (fun (acc : IntervalVector) (l : Layer) => l.forwardInterval acc prec) Is layers; ∀ (i : ℕ) (hi : i < List.length outInt), outReal[i] ∈ outInt[i]

Soundness of SequentialNet forward pass

theorem LeanCert.ML.Distillation.SequentialNet.mem_forwardInterval {net : SequentialNet} {xs : List ℝ} {Is : IntervalVector} (hdim : xs.length = List.length Is) (hwf : net.WellFormed xs.length) (hmem : ∀ (i : ℕ) (hi : i < List.length Is), xs[i] ∈ Is[i]) (prec : ℤ) (hprec : prec ≤ 0 := by norm_num) : let outReal := net.forwardReal xs; let outInt := net.forwardInterval Is prec; outReal.length = List.length outInt ∧ ∀ (i : ℕ) (hi : i < List.length outInt), outReal[i] ∈ outInt[i]

Public soundness theorem

Interval Vector Subtraction

def LeanCert.ML.Distillation.subVectors (a b : IntervalVector) : IntervalVector

Subtract two interval vectors (pointwise)

Equations

• LeanCert.ML.Distillation.subVectors a b = List.zipWith LeanCert.ML.Distillation.IntervalDyadic.sub a b

theorem LeanCert.ML.Distillation.subVectors_length (a b : IntervalVector) : List.length (subVectors a b) = min (List.length a) (List.length b)

Length of subtracted vectors

theorem LeanCert.ML.Distillation.mem_subVectors {ra rb : List ℝ} {ia ib : IntervalVector} (halen : ra.length = List.length ia) (hblen : rb.length = List.length ib) (ha : ∀ (i : ℕ) (hi : i < List.length ia), ra[i] ∈ ia[i]) (hb : ∀ (i : ℕ) (hi : i < List.length ib), rb[i] ∈ ib[i]) (i : ℕ) (hi : i < List.length (subVectors ia ib)) : ra[i] - rb[i] ∈ (subVectors ia ib)[i]

Membership in subtracted vectors

The Equivalence Checker

def LeanCert.ML.Distillation.intervalBoundedBy (I : Core.IntervalDyadic) (eps : ℚ) : Bool

Check if an interval is contained within [-eps, eps]

Equations

• LeanCert.ML.Distillation.intervalBoundedBy I eps = decide (-eps ≤ I.toIntervalRat.lo ∧ I.toIntervalRat.hi ≤ eps)

def LeanCert.ML.Distillation.isBoundedBy (v : IntervalVector) (eps : ℚ) : Bool

Check if an interval vector is contained within [-eps, eps]

Equations

• LeanCert.ML.Distillation.isBoundedBy v eps = List.all v fun (x : LeanCert.Core.IntervalDyadic) => LeanCert.ML.Distillation.intervalBoundedBy x eps

theorem LeanCert.ML.Distillation.intervalBoundedBy_spec {x : ℝ} {I : Core.IntervalDyadic} {eps : ℚ} (hx : x ∈ I) (hcheck : intervalBoundedBy I eps = true) : |x| ≤ ↑eps

Soundness of intervalBoundedBy

def LeanCert.ML.Distillation.checkEquivalence (teacher student : SequentialNet) (domain : IntervalVector) (eps : ℚ) (prec : ℤ := -53) : Bool

The Distillation Certifier

Returns true if the student network is proven to be within eps of the teacher network for all inputs in the domain.

Equations

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

Correctness Proofs

theorem LeanCert.ML.Distillation.checkEquivalence_dims {teacher student : SequentialNet} {domain : IntervalVector} {eps : ℚ} {prec : ℤ} (h : checkEquivalence teacher student domain eps prec = true) : List.length (teacher.forwardInterval domain prec) = List.length (student.forwardInterval domain prec)

Helper: extract dimension equality from checkEquivalence

theorem LeanCert.ML.Distillation.verify_equivalence (teacher student : SequentialNet) (domain : IntervalVector) (eps : ℚ) (prec : ℤ) (x : List ℝ) (hprec : prec ≤ 0) (hwf_t : teacher.WellFormed x.length) (hwf_s : student.WellFormed x.length) (h_dom_len : x.length = List.length domain) (h_mem : ∀ (i : ℕ) (hi : i < List.length domain), x[i] ∈ domain[i]) (h_cert : checkEquivalence teacher student domain eps prec = true) : let t_out := teacher.forwardReal x; let s_out := student.forwardReal x; let t_int := teacher.forwardInterval domain prec; let s_int := student.forwardInterval domain prec; List.length t_int = List.length s_int ∧ ∀ (i : ℕ) (hi_t : i < List.length t_int) (hi_s : i < List.length s_int), |t_out[i] - s_out[i]| ≤ ↑eps

Golden Theorem: Verified Model Distillation

If checkEquivalence returns true, then for ALL real inputs x in the domain, the output difference |Teacher(x) - Student(x)| is at most eps for every output neuron.