Retrieved article excerpt
Open article · Retrieved 2026-09-15T22:21:57.899249+00:00
1. Lean Verified Transformers
# Lean Verified Transformers
This post explores writing formally verified ML code in Lean.
Since the cost of proofs is declining rapidly and the amount of code generated
is skyrocketing, the value of verified code seems likely to climb.
While understanding proofs remains challenging, collaborating with AI to get
proofs of easy-to-understand properties seems like a natural middle ground.
The goal of this post is to verify foundational properties of Transformers.
These are critical properties that are used for
parallelization and optimization, including tensor parallelism, data parallelism,
batch invariance, permutation invariance, correctness of tiling, and locality
of sparse attention models. Code is available at [srush/lean-transformer](https://github.com/srush/lean-transformer).
The text, comments, and structure of the blog are all human-written;
the proofs are all written by AI. Hopefully it can also serve as an advanced
intro to Lean.
This project is inspired by [TorchLean](https://arxiv.org/abs/2602.22631),
[Verified Deep Learning with Lean 4](https://lean.brettkoonce.com/blueprint/),
and the [Dex Programming Language](https://github.com/google-research/dex-lang).
- [srush](https://rush-nlp.com/)
# Invariance and Equivariance
Different neural network architectures retain different properties of their input.
We generally classify these properties in terms of equivariance and invariance.
These allow researchers to reason about what they can learn, and
implementers to optimize computation while maintaining equivalence.
Our goal will be to prove equivariances and invariances for specific architectures.
Equivariance transforms the output along with the input; invariance leaves the output unchanged.
Notationally, ML definitions often assume the same functions can work on
different input shapes, e.g. batch sizes. For this reason, our Lean definition
will be a bit complex to allow for functions that are polymorphic over the shape.
`def Equivariant
-- Arguments with { } are implicit
{Shape : Type u}
{Input : Shape → Type v} {Output : Shape → Type w}
-- Arguments with ( ) are explicit
(f : {shape : Shape} → Input shape → Output shape)
{source target : Shape}
(T : Input source → Input target)
(S : Output source → Output target)
-- : gives the return type. Here it is a property.
: Prop :=
∀ x, f (T x) = S (f x)``def Invariant
{Shape : Type u}
{Input : Shape → Type v} {Output : Type w}
(f : {shape : Shape} → Input shape → Output)
{source target : Shape}
(T : Input source → Input target) : Prop :=
∀ x, f (T x) = f x`
# Vectors, Matrices, and Neural Networks
We begin by building a simple neural network library in Lean.
The `relu` function takes in a number and returns its non-negative part.
Along with the definition, we prove it does what we claim.
`def relu (z : Rat) : Rat := max z 0``theorem relu_non_negative
-- For all z
(z : Rat) :
-- relu is ≥ 0
relu z ≥ 0
:= byz:Rat⊢ relu z ≥ 0
-- Do a short (grind) search
grind [relu]All goals completed! 🐙`
Following the style of
[JAX](https://docs.jax.dev/en/latest/_autosummary/jax.vmap.html),
we lift scalar functions to operate on vectors.
Vectors (and tensors) are represented as higher-order functions mapping
indices to rational numbers. This makes our proofs easier since we do not
have to care about storage or efficiency.
`-- Vector type. Maps a finite set of {0,...,n-1} to a rational.
abbrev Vector (n : Nat) := Fin n → Rat``-- Examples
-- [10, 10, 10, 10, 10]
def vector_of_tens_example: Vector 5 := fun _ => 10``-- [0, 1, 2, 3]
def arange (n: Nat) : Vector n := fun i => i``-- Greek letters are types.
variable {α : Type u} {β : Type v} {δ : Type w}``-- vmap on 1-arg functions.
def vmap (fn: α -> β) {n : Nat} :
((Fin n -> α) -> (Fin n -> β)) :=
fun a => fun i => fn (a i)``-- Example: vector vmap.
def vector_relu (z: Vector n) : Vector n :=
(vmap relu) z``-- vmap on 2-arg functions
def vmap2 (fn: α -> β -> δ) : ((Fin n -> α) -> (Fin n -> β) -> (Fin n -> δ)) :=
fun a b => vmap (fun i => fn (a i) (b i)) id``-- Add two vectors as + overload
instance : Add (Vector n) where
add := vmap2 (fun a b => a + b)``-- Mul two vectors with * overload
instance : Mul (Vector n) where
mul := vmap2 (fun a b => a * b)`
For aggregations, we define a vector scan. Since we are using rationals for
simplicity, we do not have an exponential, so we define a "softmax-like"
nonlinear normalization instead.
`-- Fold over vectors.
abbrev fori {α : Type u} {n : Nat} (f : Fin n → α) : List α := List.ofFn f``def scan (step : σ → α → σ) (xs : Fin n → α) (initial : σ) : σ :=
Fin.foldl n (fun state i => step state (xs i)) initial``-- Sum is a fold
def Vector.sum (a : Vector n) : Rat :=
-- Alternative: scan (fun a b => a + b) a 0
(fori (fun i => a i)).sum``def softmax_like (z : Vector n) : Vector n :=
let weights : Vector n := vmap (fun x => 1 + relu x) z
let total := weights.sum
vmap (fun w => w / total) weights``def Vector.dot_product (a b : Vector n) : Rat :=
(a * b).sum`
As an exercise, let's look at a simple vector theorem. Click the square □
next to each line of the proof and it will show you the current proof state.
The proof state divides the context from the goal ⊢. Each step will transform
these terms until we can construct the goal.
`-- Theorem: Multiplication distributes.
theorem Vector.mul_add
-- Given vectors a, b, c, of length n
(a b c : Vector n) :
-- then
a * (b + c) = a * b + a * c
:= byn:Nata:Vector nb:Vector nc:Vector n⊢ a * (b + c) = a * b + a * c
-- Strategy: show equiv for all indices i of the output vector
funext in:Nata:Vector nb:Vector nc:Vector ni:Fin n⊢ (a * (b + c)) i = (a * b + a * c) i
-- Apply the rational property to the numbers at position i.
exact Rat.mul_add (a i) (b i) (c i)All goals completed! 🐙`
Matrices are defined similarly. We are basically just stacking
`vmap` calls to get our core operations. Note the implementation of `matmul`
in particular, which will be the target of future proofs.
`abbrev Matrix (n m : Nat) := Fin n → Vector m``instance : Add (Matrix n m) where
add := vmap2 (fun a b => a + b)``instance : Mul (Matrix n m) where
mul := vmap2 (fun a b => a * b)``def Matrix.transpose (a : Matrix n m) : Matrix m n := fun i j => a j i``def Matrix.matvec (a : Matrix n m) (x : Vector m) : Vector n :=
vmap (fun row => row.dot_product x) a``def Matrix.matmul (a : Matrix n m) (b : Matrix m p) : Matrix n p :=
-- Functions can be called directly or with .transpose when the type is clear.
Matrix.transpose (vmap a.matvec b.transpose)`
We now have the full machinery of deep learning.
A neural network is just stacking layers and applying a simple loss function.
`def forward (layer: Matrix hidden hidden) {batch : Nat}
(input: Matrix batch hidden) :
Matrix batch hidden :=
(vmap (vmap relu)) (input.matmul layer)``abbrev Layer {Shape : Type v} (State : Shape → Type u) :=
{shape : Shape} → State shape → State shape``def neural_network {Shape : Type v} {State : Shape → Type u}
(layers : List (Layer State)) : Layer State :=
fun input => layers.foldl (fun state layer => layer state) input``def loss (point_loss : Fin batch → Vector hidden → Rat)
(matrix : Matrix batch hidden) : Rat :=
Vector.sum (vmap2 point_loss id matrix)`
# Properties of Neural Networks
Now let us return to our goal of proving network equivariances.
Our strategy will be to first show that in general equivariances compose,
and then show that they propagate through a neural network.
`-- Equivariances compose
theorem Equivariant.comp
-- Boilerplate
{Shape : Type u}
{A : Shape → Type v} {B : Shape → Type w} {C : Shape → Type z}
{first : {shape : Shape} → A shape → B shape}
{next : {shape : Shape} → B shape → C shape}
{source target : Shape}
{T : A source → A target} {S : B source → B target}
{U : C source → C target}
-- If f(T x) = S f(x)
(hfirst : Equivariant (Input := A) (Output := B) first T S)
-- and g(S x) = U g(x)
(hnext : Equivariant (Input := B) (Output := C) next S U) :
-- then g(f(T x )) = U (g (f (x)))
Equivariant (Input := A) (Output := C)
(fun input => next (first input)) T U := byShape:Type uA:Shape → Type vB:Shape → Type wC:Shape → Type zfirst:{shape : Shape} → A shape → B shapenext:{shape : Shape} → B shape → C shapesource:Shapetarget:ShapeT:A source → A targetS:B source → B targetU:C source → C targethfirst:Equivariant (fun {shape} => first) T Shnext:Equivariant (fun {shape} => next) S U⊢ Equivariant (fun {shape} input => next (first input)) T U
intro inputShape:Type uA:Shape → Type vB:Shape → Type wC:Shape → Type zfirst:{shape : Shape} → A shape → B shapenext:{shape : Shape} → B shape → C shapesource:Shapetarget:ShapeT:A source → A targetS:B source → B targetU:C source → C targethfirst:Equivariant (fun {shape} => first) T Shnext:Equivariant (fun {shape} => next) S Uinput:A source⊢ (fun {shape} input => next (first input)) (T input) = U ((fun {shape} input => next (first input)) input)
exact (congrArg next (hfirst input)).trans (hnext (first input))All goals completed! 🐙``-- Equivariances flow through tuples
theorem Equivariant.prod
-- Boilerplate
{Shape : Type u}
{Input₁ : Shape → Type u₁} {Input₂ : Shape → Type u₂}
{Output₁ : Shape → Type v₁} {Output₂ : Shape → Type v₂}
{f : {shape : Shape} → Input₁ shape → Output₁ shape}
{g : {shape : Shape} → Input₂ shape → Output₂ shape}
{source target : Shape}
{T₁ : Input₁ source → Input₁ target} {S₁ : Output₁ source → Output₁ target}
{T₂ : Input₂ source → Input₂ target} {S₂ : Output₂ source → Output₂ target}
-- If f(T1 x) = S1 f(x)
(hf : Equivariant (Input := Input₁) (Output := Output₁) f T₁ S₁)
-- and g(T2 x) = S2 g(x)
(hg : Equivariant (Input := Input₂) (Output := Output₂) g T₂ S₂) :
-- Then <f,g> <T1 x, T2 y> = <S1 f( x), S2 g( y)>
Equivariant (Input := fun shape => Input₁ shape × Input₂ shape)
(Output := fun shape => Output₁ shape × Output₂ shape)
(fun input => Prod.map f g input) (Prod.map T₁ T₂) (Prod.map S₁ S₂) := byShape:Type uInput₁:Shape → Type u₁Input₂:Shape → Type u₂Output₁:Shape → Type v₁Output₂:Shape → Type v₂f:{shape : Shape} → Input₁ shape → Output₁ shapeg:{shape : Shape} → Input₂ shape → Output₂ shapesource:Shapetarget:ShapeT₁:Input₁ source → Input₁ targetS₁:Output₁ source → Output₁ targetT₂:Input₂ source → Input₂ targetS₂:Output₂ source → Output₂ targethf:Equivariant (fun {shape} => f) T₁ S₁hg:Equivariant (fun {shape} => g) T₂ S₂⊢ Equivariant (fun {shape} input => Prod.map f g input) (Prod.map T₁ T₂) (Prod.map S₁ S₂)
intro xShape:Type uInput₁:Shape → Type u₁Input₂:Shape → Type u₂Output₁:Shape → Type v₁Output₂:Shape → Type v₂f:{shape : Shape} → Input₁ shape → Output₁ shapeg:{shape : Shape} → Input₂ shape → Output₂ shapesource:Shapetarget:ShapeT₁:Input₁ source → Input₁ targetS₁:Output₁ source → Output₁ targetT₂:Input₂ source → Input₂ targetS₂:Output₂ source → Output₂ targethf:Equivariant (fun {shape} => f) T₁ S₁hg:Equivariant (fun {shape} => g) T₂ S₂x:Input₁ source × Input₂ source⊢ (fun {shape} input => Prod.map f g input) (Prod.map T₁ T₂ x) =
Prod.map S₁ S₂ ((fun {shape} input => Prod.map f g input) x)
exact Prod.ext (hf x.1) (hg x.2)All goals completed! 🐙``-- Equivariances flow through neural networks.
theorem neural_network_equivariant
{Shape : Type v} {State : Shape → Type u} {source target : Shape}
(layers : List (Layer State))
(transform : State source → State target)
-- If all layers preserve equivariance
(equivariant : ∀ layer ∈ layers,
Equivariant (Input := State) (Output := State)
layer transform transform) :
-- Then the neural network itself preserves it.
Equivariant (Input := State) (Output := State)
(neural_network layers) transform transform := byShape:Type vState:Shape → Type usource:Shapetarget:Shapelayers:List (Layer State)transform:State source → State targetequivariant:∀ (layer : Layer State), layer ∈ layers → Equivariant (fun {shape} => layer) transform transform⊢ Equivariant (fun {shape} => neural_network layers) transform transform
-- Proof is by induction over layers.
induction layers with
| nil =>nilShape:Type vState:Shape → Type usourc