TorchLean: Formalizing Neural Networks in Lean

Robert Joseph George; Jennifer Cruden; Will Adkisson; Xiangru Zhong; Huan Zhang; Anima Anandkumar

cs.MSFeb 26, 2026

TorchLean: Formalizing Neural Networks in Lean

Authors: Robert Joseph George, Jennifer Cruden, Will Adkisson, Xiangru Zhong, Huan Zhang, Anima Anandkumar

Organizations: California Institute of Technology · Washington University in St. Louis · University of Illinois Urbana-Champaign

Abstract

Neural networks are increasingly deployed in scientific, safety critical, and mission critical pipelines, yet verification and analysis are often performed outside the programming environment that defines and runs the model. This creates a semantic gap between the executed network and the analyzed artifact: guarantees can depend on implicit conventions about operator semantics, tensor layouts, preprocessing, floating-point behavior, graph transformations, accelerated kernels, and external certificates. We present TorchLean, a unified framework for formalizing, executing, and verifying neural networks in Lean 4. TorchLean treats learned models as executable programs and mathematical objects with a shared semantics for computation, verification, and theorem proving. The framework provides a PyTorch style API for typed tensors, layers, objectives, optimizers, automatic differentiation, and graph programs, with eager and compiled execution paths that lower to a common computation-graph representation. TorchLean supports exact and finite-precision tensor semantics, verified reverse-mode differentiation, interval and affine bound propagation, CROWN/LiRPA style certificate checking, import/export workflows, and CUDA-backed execution through explicit FFI boundaries. It also includes semantic layers for attention and FlashAttention, state-space sequence models, diffusion and sampling processes, probability kernels, reinforcement-learning objectives and Markov decision processes, and self-supervised objectives such as masked autoencoding, JEPA-style predictive views, and variance/correlation-based anti-collapse losses. Together, these components provide a semantic foundation for verified machine learning, where executable neural network artifacts, verification procedures, runtime boundaries, and mathematical claims can be stated and related inside one theorem-proving environment.

Figures & tables

Figure 1: TorchLean provides a shared semantic layer for verified machine learning. Model code, imported graphs, modern ML workflows, scientific ML artifacts, and certificates are funneled into a common typed tensor and operator-tagged SSA/DAG graph semantics in Lean. Execution, automatic differentiation, bound propagation, and certificate checking are all performed against this shared representation, enabling checked claims in robustness, control, scientific machine learning, and bug detection while keeping external runtimes and solvers as explicit boundaries.

Panel of Figure 2 — Figure 2: 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 3: PyTorch vs. TorchLean . We highlight comparable blocks: model/setup (yellow), data (blue), loop structure (purple), and per-step updates (green). In TorchLean , the same model can be executed, lowered to the shared IR, differentiated, and passed to verification.

Mode	Purpose	Semantics and trust
$\mathbb{R}$	Proofs, reference Autograd	Exact real arithmetic; reference semantics for theorem statements.
Interval / affine	IBP, CROWN/LiRPA bounds	Sound enclosure domains with proved transfer rules where available.
FP32 / NF	Rounding-aware theorems	Round-on- $\mathbb{R}$ proof model with compositional error envelopes.
IEEE32Exec	Executable Float32 model	Lean-defined binary32 semantics for core ops; hardware matching is separate.
Arb/FLINT oracle	Rigorous transcendentals	External validated ball/interval enclosures; explicit non-kernel oracle boundary.
Lean Float32	Fast execution	Runtime implementation opaque to the kernel; treated as a deployment assumption or validated backend.

Table 1: Trust and scalar semantics in TorchLean , including explicit native runtime and validated-numerics boundaries.

Panel of Figure 4 — Figure 4: Comprehensive architecture of TorchLean . The specification, runtime, and verification layers share a single operator-tagged SSA/DAG computation-graph IR as the semantic target. Models are defined once, executed or trained by the runtime, and verified through CROWN bounds and certificate checking against the same graph semantics.

Figure 5: Representative model families exercised in TorchLean . The examples cover vision, transformer-style language models, Mamba/state-space sequence models, diffusion/sampling, Fourier-neural-operator/scientific ML workflows, and reinforcement learning. The plotted curves illustrate executable training and evaluation workflows across these model families. The formal role of the examples is that each can be connected to the typed tensor, runtime, graph-lowering, finite-precision, and checking infrastructure described in the methodology.

Path	Method	Safe/ $n$	Time
Lean	IBP	0/30	6.05s
Lean	CROWN-Obj	6/30	22.60s
Lean	CROWN+ $\alpha$	6/30	14.21s
CPU producer	$\alpha$ -CROWN	13/30	8.39s
CUDA producer	$\alpha$ -CROWN	13/30	7.67s

Table 2: VNN-COMP style MNIST-FC slice. Checked rows are Lean-verified; CUDA is an untrusted GPU producer.

Figure 6: Neural-controller workflow in three execution settings. Certification claims are checked against the shared operator-tagged SSA/DAG IR semantics in Lean; external optimizers and runtime systems act as untrusted producers.

Figure 7: Representative semantic boundary bugs. Left: typed shape semantics distinguishes intended row-wise bias addition from an unintended broadcast. Middle: temporal and masking contracts expose causal-mask or KV-cache leaks in sequence models. Right: explicit finite-precision semantics separates a stable log-sum-exp/log-softmax computation from a naive Float32 implementation that can overflow or produce NaNs.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Component	What the paper claims (scope)	Status	Where
Shared IR semantics	Operator tagged SSA/DAG IR with a precise denotation reused by execution and verification.	Kernel-checked definitions	Sections 2 , B
Reverse-mode AD	Backpropagation equals the adjoint Fréchet derivative for well-typed SSA/DAG graphs over $\mathbb{R}$ , assuming local per-op derivative correctness (with a pointwise variant for non-smooth ops).	Proved	Theorem 2.2 , Appendix B
TorchLean compilation	“Compiled mode” produces the verifier IR; full compiler correctness for arbitrary TorchLean programs is not yet proved.	Proved (fragment) / Future (full)	Theorems B.3 – B.3
IBP / CROWN operators	Sound transfer rules and relaxations for a curated set of operators; additional ops use conservative fallbacks or are treated as unsupported in the strongest theorems.	Mixed (proved subset)	Appendix E
Certificate checking	Small Lean checkers validate certificate structure and the final enclosure constraints needed to discharge a theorem about $\mathopen{[\![}G\mathclose{]\!]}$ .	Checked (plus proved core)	Section 3 , Appendix E
External solvers (e.g. $\alpha/\beta$ -CROWN)	Used only as untrusted certificate producers; full optimizer internals are outside the checked certificate scope in this work.	Assumed/External	Section 3 , Appendix E

Appendix

Table 3: Verified scope in TorchLean (summary). “Proved” means a Lean theorem over the stated semantics; “Checked” means a small executable validator that re-computes or validates constraints; “Assumed/External” marks trust boundaries or future work. Appendix E also tabulates operator level verification coverage.

Panel of Figure 8 — Figure 8: Example SSA/DAG computation graph: $y=\text{ReLU}(Wx+b)$ . Each intermediate value ( $v_{1}$ , $v_{2}$ ) is assigned once. Forward pass (solid arrows) evaluates nodes in topological order. Backward pass (dashed arrows) propagates gradients in reverse topological order, accumulating at each node.

Panel of Figure 9 — Figure 9: Two execution paths to a unified semantic target. The same TorchLean program can be interpreted eagerly (building a tape dynamically) or compiled (building a static graph upfront). Both paths produce computations that share the same well-typed SSA/DAG denotation, enabling unified reasoning about execution and verification.

OpKind / primitive	Global HasFDerivAt	Pointwise (side conditions)	Not covered
Linear / shape ops	add , sub , mul_elem , linear , matmul , conv2d , broadcastTo , reshape , flatten , permute , reduce_sum , reduce_mean , sum , concat (curated axes), swap_first_two , transpose3d_last_two	–	–
Smooth nonlinearities / losses	tanh , sigmoid , exp , square , sinh , cosh , softplus , gelu (tanh approximation), silu , softmax / log_softmax (curated axis), layernorm (curated axis), mse_loss , cross-entropy/NLL/BCE/KL pointwise losses, smooth surrogates safe_log and smooth_abs (with $\varepsilon>0$ )	–	–
Non-smooth / domain-sensitive	–	relu , abs , log , inv , sqrt , max_elem / min_elem (no ties)	–
Structured graph families	fixed-mask dropout, scaled dot-product attention, unmasked MHA, residual MHA, post-norm transformer sublayers, BatchNorm-like graphs, one-step Elman RNN, gather-row / embedding lookup adjointness	–	full GPT/ViT stacks, masked causal decoder blocks, full recurrent/state-space sequence theorems, FNO spectral-conv training paths
Nondifferentiable / workflow ops	–	–	detach , rand_uniform , stochastic masks beyond fixed-mask dropout, pooling ( max_pool2d , avg_pool2d ), executable CUDA/cuBLAS/cuFFT kernels

Appendix

Table 4: Autograd operator coverage (Fréchet-derivative level). “Global” means we have a NodeFDerivCorrect proof (a HasFDerivAt theorem without side conditions). “Pointwise” means we prove NodeFDerivCorrectAt under explicit hypotheses that rule out non-differentiable points (e.g., no zeros/ties). Operators listed under “Not covered” may still be executable (and may have a deterministic backprop convention), but are not claimed correct at the level of classical derivatives.

Panel of Figure 10 — Figure 10: Two-stage autograd proof architecture with primitive coverage. Stage 1 proves adjointness between $\mathrm{JVP}_{G}$ and $\mathrm{VJP}_{G}$ over abstract scalars; Stage 2 bridges tensors to Euclidean spaces and upgrades to the Fréchet-derivative statement $\mathrm{VJP}_{G}(x;\bar{y})=(D\,\mathsf{EvalVec}_{G}(x))^{\top}\bar{y}$ . Bottom row summarizes operator classes covered by local derivative lemmas; see Table 4 for details.

Aspect	PyTorch	TorchLean
Shapes	Dynamic runtime shapes; many errors are runtime exceptions.	Shapes are part of types ( Tensor $\alpha$ s ); many mismatches are untypeable.
Dtypes	Many numeric and integer dtypes; mixed-dtype graphs are common.	Scalar-polymorphic single-dtype graphs (one $\alpha$ per run); integer indices handled via a separate non-differentiable channel ( NatRef / NatVecRef ) in sessions.
Execution modes	Eager by default, with compilation/export toolchains (TorchScript, ONNX, AOT).	Eager tape backend and proof-linked compiled SSA/DAG backend share one semantics; the compiled graph is the verifier target by construction.
Autograd status	Widely tested and trusted, but not formally proved correct.	Reverse-mode correctness theorem for well-typed SSA/DAG graphs; eager runs linked to proved graphs.
Indexing	Tensor-valued indexing/slicing/gather/scatter across dtypes (e.g. LongTensor ).	Typed indexing primitives and session-level Nat channels for labels/indices; not yet PyTorch-complete for tensor-valued integer indices.
Breadth/performance	Very broad op surface and ecosystem; highly optimized CPU/GPU kernels.	Curated op surface focused on verification-relevant primitives; extensible via a “new op” workflow; Lean execution prioritizes clarity/verification over performance.

Appendix

Table 5: High-level comparison between PyTorch Paszke et al. [2019] and TorchLean .

Explore similar work

May 8, 2026

VNN-LIB 2.0: Rigorous Foundations for Neural Network Verification

Neural network verification is an active and rapidly maturing research area, with a growing ecosystem of solvers and tools. The VNN-LIB standard was introduced to support interoperability in this ecosystem, but Version1.0 has several serious short-comings as a formal foundation: it lacks a precise syntax, semantics, and type system, offers limited expressivity, and relies on externally defined ONNX models whose semantics are informal and constantly evolving. The latter distinguishes VNN-LIB from established standards such as SMT-LIB, where queries are self-contained and have fixed semantics. In this paper we address these challenges by developing the theoretical foundations of VNN-LIB2.0. Our key contribution is the introduction of the notion of a network theory, which abstractly characterises the minimal semantic interface required from a neural network model format. This abstraction enables VNN-LIB to be defined independently of any specific ONNX version while remaining compatible with evolving model representations. Building on this foundation, we present a formal syntax for a more expressive query language, a type system for it over the numeric domains provided by the network theory, and finally a formal semantics. To ensure internal consistency, the standard is mechanised in the Agda theorem prover. VNN-LIB~2.0 therefore provides robust and rigorous foundations for trustworthy neural network verification.

Ann Roy, Allen Antony, Andrea Gimelli +1

University of Western Australia, Perth, Australia · University of Genoa, Genoa, Italy

Neural Network Verification Formal Verification

Ask PDF

Jul 31, 2026

Learning Lookahead Lemmas for Neural Network Verification

State-of-the-art neural network verifiers use the branch-and-bound procedure as their core solving mechanism. We introduce an inprocessing framework for neural network verification driven by the lookahead procedure. Under this framework, lookahead derives new lemmas over the phases of unstable ReLUs, which are collected into an implication graph that is used to prune the search space and vivify boolean cuts. We instantiate the framework in two state-of-the-art verifiers, Marabou and

α

β

-CROWN, and demonstrate that it improves performance in both, proving up to 34% more instances unsatisfiable.

Liam Davis, Haoze Wu

Amherst College

Neural Network Verification ReLU Neural Networks Branch and Bound Tree Search

Ask PDF

Mar 25, 2026

The Luna Bound Propagator for Formal Analysis of Neural Networks

The parameterized CROWN analysis, a.k.a., alpha-CROWN has emerged as a practically successful abstract interpretation method for neural network verification. However, existing implementations of alpha-CROWN are limited to Python, which complicates integration into existing DNN verifiers and long-term production-level systems. We introduce Luna, a new abstract-interpretation-based bound propagator implemented in C++. Luna supports Interval Bound Propagation, the DeepPoly/CROWN analysis, and the alpha-CROWN analysis over a general computational graph. We describe the architecture of Luna and show that it outperforms the state-of-the-art alpha-CROWN implementation in terms of both bound tightness and computational efficiency on supported benchmarks from VNN-COMP 2025. Luna is publicly available at https://github.com/ai-ar-research/luna.

Henry LeCates, Haoze Wu

Amherst College, USA

Neural Network Verification Neural Network Robustness

Ask PDF