[CVPR/CADE 2024] TORCHLEAN: Closing the Semantic Gap in AI Verification via Lean 4
TorchLean: Formalizing Neural Networks in Lean
The paper introduces TORCHLEAN, a formal framework in the Lean 4 theorem prover that unifies neural network definition, execution, and verification. It provides a PyTorch-style API that compiles models into a shared, op-tagged SSA/DAG Intermediate Representation (IR) with precise mathematical semantics, enabling SOTA-level certificate checking for transition-critical AI systems.
Executive Summary
TL;DR: Researchers have developed TORCHLEAN, the first framework that brings PyTorch's "eager execution" workflow into the Lean 4 formal theorem prover. By compiling models into a shared, mathematically precise Intermediate Representation (IR), it ensures that the model you train is exactly the same as the model you verify, effectively eliminating the "semantic drift" that plagues current safety-critical AI.
Background Positioning: This work is a foundational infrastructure project. Rather than just proposing a faster verifier, it builds the "missing link" between deep learning engineering and formal methods, positioning itself as the SOTA substrate for verified machine learning.
The Problem: The "Silent Killers" of Formal Guarantees
In a typical mission-critical pipeline (e.g., a neural controller for a drone), a model is trained in PyTorch and then exported to ONNX or TorchScript for verification. This process introduces two fatal flaws:
- Semantic Drift: The ONNX interpreter might handle a specific operator (like a
LayerNormat the edge of its domain) slightly differently than the training runtime. - Floating-Point Mirages: Most verifiers reason about Real Numbers (), but hardware runs on FP32. Rounding errors, NaNs, and signed zeros can be exploited to bypass "proven" robustness, as shown in previous adversarial attacks on verifiers.
Figure 1: The standard pipeline (top) vs. TORCHLEAN's unified approach (bottom), where execution and verification share one ground-truth IR.
Methodology: Neural Networks as Mathematical Proofs
TORCHLEAN’s core innovation is its Three-Pillar Architecture:
1. The Shared IR (SSADAG)
Every model in TORCHLEAN—whether an MLP, CNN, or Transformer—lowers to a Static Single Assignment (SSA) graph. Because it is a Directed Acyclic Graph (DAG), Lean can perform induction over the nodes. The authors proved a Reverse-mode AD Correctness Theorem, mathematically guaranteeing that the gradients used for training are the true adjoint Fréchet derivatives of the network's denotation.
2. Verified Numerics (IEEE32Exec)
Instead of treating floats as "opaque" values, TORCHLEAN includes IEEE32Exec, a bit-level implementation of the IEEE-754 binary32 standard.
- Proofs use : For high-level logic.
- Execution uses IEEE32Exec: For hardware-matching behavior.
- Refinement: The authors proved that when values are finite (no overflow), the IEEE results align with the mathematical rounding models.
3. Native & Certificate Verification
The framework implements IBP (Interval Bound Propagation) and CROWN/LiRPA (linear relaxations) natively in Lean. Crucially, it uses a Producer-Checker model: a fast, untrustworthy external solver (like α-CROWN) generates a "certificate," which a tiny, trusted Lean checker then validates against the IR semantics.
Figure 2: The system splits into NN.Spec (definition), NN.Runtime (execution), and NN.Verification (proofs), all anchored by the shared IR.
Experiments: Real-World Safety
The authors didn't just test toys; they applied TORCHLEAN to three complex scenarios:
- Certified Robustness: Certified an margin for digits classification, checking results in just 0.032ms.
- PINNs (Physics-Informed Neural Networks): Verified that a network solving the Burgers equation stays within a specific residual bound ().
- Neural Controllers: Proved the stability of a feedback controller using Lyapunov inequalities—proving and systematically within Lean.
Table 1: Benchmarking TORCHLEAN on VNN-COMP suites. While pure Lean IBP is conservative, the certificate checker successfully validates complex CROWN bounds.
Critical Insights & Future Outlook
The Takeaway: TORCHLEAN is a major step toward Certified AI. It proves that we can have the ergonomic flexibility of PyTorch without sacrificing the rigor of formal verification.
Limitations:
- Performance: Lean execution is currently CPU-bound and slower than CUDA-optimized kernels. However, this is acceptable for verification tasks where "correctness-first" is the priority.
- Operator Coverage: While it supports common layers (Linear, Conv2D, ReLU, Tanh), complex custom ops still require manual formalization in Lean.
Future Work: The authors envision a future where CSLib (Lean's emerging computer science library) provides the foundation for fully verified autonomous systems—from the neural perception engine down to the hardware-level control logic.
