TorchLean logo

Formalizing Neural Networks in Lean

Robert Joseph George 1, Jennifer Cruden1, Will Adkisson2, Xiangru Zhong3, Huan Zhang3, Anima Anandkumar1
1California Institute of Technology; 2Washington University in St. Louis; 3University of Illinois Urbana-Champaign
PDF Code Project Site

Abstract

Neural networks are increasingly deployed through software stacks whose behavior is defined by tensor layouts, preprocessing choices, compiler boundaries, floating-point arithmetic, GPU kernels, and external model exporters. TorchLean is a Lean 4 codebase for neural network specification, execution, and verification that keeps those interfaces explicit. It gives models a shared formal vocabulary for typed tensors, shapes, layers, graph programs, objectives, gradients, finite-precision behavior, certificate checking, and runtime trust boundaries. The current codebase includes a PyTorch-style API, eager and compiled execution through an op-tagged SSA/DAG graph IR, reverse-mode autograd, optimizers, CUDA/GPU execution with CPU stubs and sanitizer/profiling support, PyTorch import/export workflows, Float32/IEEE-754 semantics, IBP and CROWN/LiRPA-style verification, and a model zoo covering MLPs, CNNs, ResNets, ViTs, GPT-style text models, GPT-2, Mamba/SSM examples, diffusion, MAE-style self-supervised learning, reinforcement learning, PINNs, and 3D-geometry certificate demos.


TorchLean system overview

Key Contributions

  • Typed tensor and layer APIs for writing neural-network programs directly in Lean.
  • Eager runtime, reverse-mode autograd, optimizers, logging, and compiled execution through a shared graph IR.
  • CUDA/GPU backend with native kernels, CPU stubs, Compute Sanitizer checks, and Nsight profiling hooks.
  • Explicit trust boundaries for external FFI, PyTorch import/export, Python/Julia/Arb oracles, α,β-CROWN certificates, and CUDA kernels.
  • Float32 and IEEE-754-style semantics with proof layers for runtime approximation and finite-precision reasoning.
  • Verification workflows for IBP, CROWN/LiRPA, robustness certificates, PINNs, Lyapunov-style controllers, VNN-COMP-style artifacts, and 3D geometry certificates.
  • Runnable modern ML examples including GPT-2-style text, Mamba/SSM sequence models, diffusion, ResNet, ViT, MAE/self-supervised learning, and reinforcement learning.

System Modules

TorchLean is organized around a shared semantics that connects model code, runtime execution, verification, and external tools:


Frontend and Models

  • PyTorch-shaped API for tensors, modules, optimizers, datasets, and training loops.
  • Examples for supervised, vision, sequence, generative, self-supervised, RL, and scientific ML models.
  • Interop paths for PyTorch-style exporters and round-trip artifacts.

Runtime and CUDA

  • Eager and compiled execution share an op-tagged graph representation.
  • Reverse-mode autograd and backend-polymorphic model runners support CPU and CUDA paths.
  • Native CUDA code is isolated behind explicit FFI symbols, CPU stubs, sanitizer checks, and profiling scripts.

Proof and Verification

  • Tensor algebra, autograd correctness, runtime approximation, and IEEE/Float32 bridge theorems.
  • Certificate checkers for robustness, LiRPA/CROWN, PINNs, control, and geometry workflows.
  • Bug-zoo examples make failures such as masking, RoPE/KV-cache, tokenizer, compiler, and projection bugs concrete.