Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TorchLean: Formalizing Neural Networks in Lean

Robert Joseph George, Jennifer Cruden, Xiangru Zhong, Huan Zhang, Anima Anandkumar

Abstract

Neural networks are increasingly deployed in safety- and mission-critical pipelines, yet many verification and analysis results are produced outside the programming environment that defines and runs the model. This separation creates a semantic gap between the executed network and the analyzed artifact, so guarantees can hinge on implicit conventions such as operator semantics, tensor layouts, preprocessing, and floating-point corner cases. We introduce TorchLean, a framework in the Lean 4 theorem prover that treats learned models as first-class mathematical objects with a single, precise semantics shared by execution and verification. TorchLean unifies (1) a PyTorch-style verified API with eager and compiled modes that lower to a shared op-tagged SSA/DAG computation-graph IR, (2) explicit Float32 semantics via an executable IEEE-754 binary32 kernel and proof-relevant rounding models, and (3) verification via IBP and CROWN/LiRPA-style bound propagation with certificate checking. We validate TorchLean end-to-end on certified robustness, physics-informed residual bounds for PINNs, and Lyapunov-style neural controller verification, alongside mechanized theoretical results including a universal approximation theorem. These results demonstrate a semantics-first infrastructure for fully formal, end-to-end verification of learning-enabled systems.

TorchLean: Formalizing Neural Networks in Lean

Abstract

Neural networks are increasingly deployed in safety- and mission-critical pipelines, yet many verification and analysis results are produced outside the programming environment that defines and runs the model. This separation creates a semantic gap between the executed network and the analyzed artifact, so guarantees can hinge on implicit conventions such as operator semantics, tensor layouts, preprocessing, and floating-point corner cases. We introduce TorchLean, a framework in the Lean 4 theorem prover that treats learned models as first-class mathematical objects with a single, precise semantics shared by execution and verification. TorchLean unifies (1) a PyTorch-style verified API with eager and compiled modes that lower to a shared op-tagged SSA/DAG computation-graph IR, (2) explicit Float32 semantics via an executable IEEE-754 binary32 kernel and proof-relevant rounding models, and (3) verification via IBP and CROWN/LiRPA-style bound propagation with certificate checking. We validate TorchLean end-to-end on certified robustness, physics-informed residual bounds for PINNs, and Lyapunov-style neural controller verification, alongside mechanized theoretical results including a universal approximation theorem. These results demonstrate a semantics-first infrastructure for fully formal, end-to-end verification of learning-enabled systems.
Paper Structure (51 sections, 22 theorems, 29 equations, 10 figures, 7 tables)

This paper contains 51 sections, 22 theorems, 29 equations, 10 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Computation graphs and our IR (definitions).
  4. NN.Spec
  5. NN.Runtime
  6. Floating Point Semantics
  7. NN.Verification
  8. Results
  9. VNN-COMP-style suites (ONNX + VNNLIB).
  10. Numerical semantics stress tests.
  11. Related work
  12. Conclusion
  13. Acknowledgments
  14. Appendix Overview
  15. TorchLean: a PyTorch-style front end with a single semantic target
  16. ...and 36 more sections

Key Result

Theorem 2.1

Let $G$ be a well-typed SSA/DAG computation graph over $\mathbb{R}$ with input context $\Gamma$. Assume GraphFDeriv Correct$G$. Define $\mathrm{bp}_{G}(x,seed) := \mathrm{backpropVec}(G,x,seed)$. Then for any input $x$ and cotangent $seed$, For non-smooth primitives (e.g., ReLU at $0$), we use the pointwise variant GraphFDeriv CorrectAt under explicit side conditions (Appendix app:torchlean); we

Figures (10)

  • Figure 1: Semantic gap: the standard pipeline (top) chains PyTorch $\to$ Export $\to$ ONNX $\to$ Verifier, with potential drift at each conversion step; TorchLean (bottom) keeps a single shared IR as the semantic target for both execution and verification (Verifier artifacts are checked against this same meaning).
  • Figure 2: Comprehensive architecture of TorchLean. The three modules (NN.Spec, NN.Runtime, NN.Verification) share a single op-tagged SSA/DAG computation-graph IR as the semantic target. Models are defined once (spec), executed/trained via TorchLean (runtime), and verified via IBP/CROWN-style bounds and certificate checking (verification), all against the same graph semantics. Scalar polymorphism ($\alpha$) instantiates the same definitions over reals (proofs), floats/Float32 (execution), and intervals (bounds).
  • Figure 3: PyTorch vs. TorchLean. We highlight comparable blocks: model/setup (yellow), data (blue), loop structure (purple), and per-step updates (green).
  • Figure 4: Execution microbenchmark scaling (warm-up excluded). TorchLeanEager uses the eager tape backend; TorchLeanCompiled uses the SSA/DAG compiled backend (see Section \ref{['sec:method-runtime']}). TorchLeanEagerFast enables a small set of runtime-only fast kernels for hot ops in eager mode.
  • Figure 5: Theory + execution: UAT and IEEE32Exec soundness.
  • ...and 5 more figures

Theorems & Definitions (22)

  • Theorem 2.1: Reverse-mode AD correctness
  • Theorem 2.1: Forward lowering correctness
  • Theorem 2.2: Eager--graph correspondence
  • Theorem 2.3: Graph adjointness (algebraic)
  • Theorem 2.4: Tensor--vector inner-product bridge
  • Theorem 2.5: Backpropagation computes the adjoint derivative
  • Theorem 2.6: Global node correctness (smooth primitives)
  • Theorem 2.7: Pointwise node correctness (non-smooth primitives)
  • Theorem 2.8: Example: affine node rule
  • Theorem 2.9: Example: elementwise lifting (informal)
  • ...and 12 more