Lean 4 → MLIR → GPU

The book: Verified Deep Learning with Lean 4 (PDF) — the interactive proof blueprint is the book: every theorem clickable, from the pdiv primitives to the whole-network backward passes.

Lean 4 as a specification language for neural networks. Declare the architecture in Lean, render one StableHLO graph — forward, loss, backward and optimizer fused — from proofs that the backward is the Jacobian-transpose of the forward, hand the graph to a trusted lowerer, train end to end. No Python at run time and no autograd library: the gradients are derived at codegen time, in Lean, and machine-checked over the reals.

Companion code for Verified Deep Learning with Lean 4, forthcoming from Apress (Springer Nature) as the follow-up to Convolutional Neural Networks with Swift for TensorFlow (2021). Current version v0.7.0 (release history) · DOI 10.5281/zenodo.20402133

The tour

Four commands, one per scale, in the order the book meets the nets, then the demos. The numbers are the book's, from the verified XLA path on one RTX 4060 Ti unless the row says otherwise; the two Imagenette side quests the book does not quote (ResNet-50, MobileNetV4) are medians of five seeds. Setup is the book's Getting started, one track per tier; the short form is

curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | sh   # Lean 4
python3 -m venv .venv && . .venv/bin/activate && pip install jax-cuda12-pjrt          # the XLA plugin only (jax-rocm7-pjrt on AMD)
gcc -fPIC -O2 -shared ffi/pjrt_ffi.c -ldl -o ffi/libpjrt_ffi.so                        # the shim; needs nothing but libc
lake exe cache get && ./download_mnist.sh && lake run mnist                            # Mathlib oleans, MNIST, tier 1
tiercommandtrainsthe numberchapter
1lake run mnist linear, MLP and CNN on MNIST, 12 epochs each (~1 min) 92.10 · 97.81 · 98.77 % 1 · 2 · 3
2lake run cifar the wide 8-conv net on CIFAR-10: SGD / momentum / AdamW × no-BN / BN, 40 epochs at a constant lr (~19 min) 76.3 % — BN + momentum, median of five 4
3lake run imagenette seven nets on Imagenette at 224², 80 epochs AdamW, book order (~9 h) R34 89.50 · R50 89.71 · MNv2 89.25 · MNv4-Conv-M 86.24 · B0 89.96 · ConvNeXt-T 85.07 · ViT-Tiny 68.74 % 5 · 6 · 7 · 8 · 9
4lake run imagenet the same nets on ImageNet-1k, 4× 4060 Ti, weeks of wall-clock; bare it prints the plan and every row's estimate, start runs it R34 74.16 · R50 (RSB-A3) 78.26 · MNv2 71.90 · MNv4-Conv-M 75.48 · B0 77.15 · ConvNeXt-T 81.53 · ViT-Tiny 72.31 % Track 4

The demos ride on the chapter nets; demos/README.md has the command, the figure and the reasoning for each.

democommandthe number
segmentationlake exe unet-brats-r34, then brats-predict BraTS mIoU 0.742 — a ResNet-34 encoder under a UNet
detectionlake exe yolov1-visdrone-fpn VisDrone [email protected] 0.2363 — ResNet-34 + FPN at 448
diffusionlake exe mnist-ddpm-train, then mnist-ddpm-sample the sample grid
languagelake exe tinygpt-shakespeare (also bigram-shakespeare, tinystories) 2.28 bits/char held-out

Everything else in the repository is the lab — apps/baselines/, the ablation and robustness exes, the tests, the Bestiary — the evidence behind these numbers, one level down.

The proofs

Every layer's backward is proven to be the Jacobian-transpose of its forward over the exact reals (Mathlib's fderiv), composed up to whole-network VJPs for ResNet-34, MobileNetV2, EfficientNet-B0, ConvNeXt-T and ViT-Tiny, with zero project axioms. For every chapter net the committed train-step render in verified_mlir/ is tied to those proofs at the denotational level: each emitted parameter-update node denotes the certified descent step, and the tiers train on exactly those bytes. What stays trusted is the ℝ→Float32 numerics, the per-op text printing, and the lowerer with its runtime. The book's On Verification appendix is the full argument, gap by gap; the dependency graph is the same suite as a clickable DAG, and LeanMlir/Proofs/README.md is the file-level map.

Check them without a GPU:

lake exe cache get           # Mathlib oleans, ~30 s
lake build ProofsMinimal     # the smallest end-to-end tie, ~1 min
lake build Certs             # every certificate CI checks (the long one)

tests/comparator/run.sh re-runs Lean's kernel typechecker over the headline theorems independently, and tests/comparator/Challenge.lean imports Mathlib and nothing else, so those can be read and checked without reading a line of this project.

Two lowerers, one graph

Training runs through XLA/PJRT, and every number above comes from it. IREE is the second trusted lowerer: it is what the differential oracle lowers the Lean side through, so that agreement with the JAX reference is evidence from two independent compilers, and each lake run tier has an -iree twin.