Lean Verified Transformers

A blog post titled "Lean Verified Transformers" presents a Lean formalization of foundational Transformer properties — including tensor parallelism, data parallelism, batch invariance, permutation invariance, tiling correctness, and sparse-attention locality — with code available at srush/lean-transformer. The post states the text, comments, and structure are human-written while all proofs were written by AI, and argues that as proof costs decline and generated code volume rises, the value of verified code is likely to climb. The work builds a simple neural network library in Lean, defining relu, vector and matrix operations via higher-order functions over Finn → Rat, and a softmax-like normalization, to prove equivariance and invariance properties for specific architectures.

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. 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. 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. 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. defEquivariant-- Arguments with { } are implicit{Shape:Typeu}{Input:Shape→Typev}{Output:Shape→Typew}-- Arguments with are explicit f:{shape:Shape}→Inputshape→Outputshape {sourcetarget:Shape} T:Inputsource→Inputtarget S:Outputsource→Outputtarget -- : gives the return type. Here it is a property.:Prop:=∀x,f Tx =S fx defInvariant{Shape:Typeu}{Input:Shape→Typev}{Output:Typew} f:{shape:Shape}→Inputshape→Output {sourcetarget:Shape} T:Inputsource→Inputtarget :Prop:=∀x,f Tx =fx 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. defrelu z:Rat :Rat:=maxz0theoremrelu non negative-- For all z z:Rat :-- relu is ≥ 0reluz≥0:=z:Rat⊢ reluz≥0-- Do a short grind searchAll goals completed 🐙 Following the style of JAX, 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.abbrevVector n:Nat :=Finn→Rat-- Examples-- 10, 10, 10, 10, 10 defvector of tens example:Vector5:=fun = 10-- 0, 1, 2, 3 defarange n:Nat :Vectorn:=funi= i-- Greek letters are types.variable{α:Typeu}{β:Typev}{δ:Typew}-- vmap on 1-arg functions.defvmap fn:α- β {n:Nat}: Finn- α - Finn- β :=funa= funi= fn ai -- Example: vector vmap.defvector relu z:Vectorn :Vectorn:= vmaprelu z-- vmap on 2-arg functionsdefvmap2 fn:α- β- δ : Finn- α - Finn- β - Finn- δ :=funab= vmap funi= fn ai bi id-- Add two vectors as + overloadinstance:Add Vectorn whereadd:=vmap2 funab= a+b -- Mul two vectors with overloadinstance:Mul Vectorn wheremul:=vmap2 funab= 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.abbrevfori{α:Typeu}{n:Nat} f:Finn→α :Listα:=List.ofFnfdefscan step:σ→α→σ xs:Finn→α initial:σ :σ:=Fin.foldln funstatei= stepstate xsi initial-- Sum is a folddefVector.sum a:Vectorn :Rat:=-- Alternative: scan fun a b = a + b a 0 fori funi= ai .sumdefsoftmax like z:Vectorn :Vectorn:=letweights:Vectorn:=vmap funx= 1+relux zlettotal:=weights.sumvmap funw= w/total weightsdefVector.dot product ab:Vectorn :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.theoremVector.mul add-- Given vectors a, b, c, of length n abc:Vectorn :-- thena b+c =a b+a c:=n:Nata:Vectornb:Vectornc:Vectorn⊢ a b+c =a b+a c-- Strategy: show equiv for all indices i of the output vectorn:Nata:Vectornb:Vectornc:Vectorni:Finn⊢ a b+c i= a b+a c i-- Apply the rational property to the numbers at position 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. 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 composetheoremEquivariant.comp-- Boilerplate{Shape:Typeu}{A:Shape→Typev}{B:Shape→Typew}{C:Shape→Typez}{first:{shape:Shape}→Ashape→Bshape}{next:{shape:Shape}→Bshape→Cshape}{sourcetarget:Shape}{T:Asource→Atarget}{S:Bsource→Btarget}{U:Csource→Ctarget}-- If f T x = S f x hfirst:Equivariant Input:=A Output:=B firstTS -- and g S x = U g x hnext:Equivariant Input:=B Output:=C nextSU :-- then g f T x = U g f x Equivariant Input:=A Output:=C funinput= next firstinput TU:=Shape:Type uA:Shape→Type vB:Shape→Type wC:Shape→Type zfirst:{shape:Shape}→Ashape→Bshapenext:{shape:Shape}→Bshape→Cshapesource:Shapetarget:ShapeT:Asource→AtargetS:Bsource→BtargetU:Csource→Ctargethfirst:Equivariant fun{shape}= first TShnext:Equivariant fun{shape}= next SU⊢ Equivariant fun{shape}input= next firstinput TUShape:Type uA:Shape→Type vB:Shape→Type wC:Shape→Type zfirst:{shape:Shape}→Ashape→Bshapenext:{shape:Shape}→Bshape→Cshapesource:Shapetarget:ShapeT:Asource→AtargetS:Bsource→BtargetU:Csource→Ctargethfirst:Equivariant fun{shape}= first TShnext:Equivariant fun{shape}= next SUinput:Asource⊢ fun{shape}input= next firstinput Tinput =U fun{shape}input= next firstinput input All goals completed 🐙-- Equivariances flow through tuplestheoremEquivariant.prod-- Boilerplate{Shape:Typeu}{Input₁:Shape→Typeu₁}{Input₂:Shape→Typeu₂}{Output₁:Shape→Typev₁}{Output₂:Shape→Typev₂}{f:{shape:Shape}→Input₁shape→Output₁shape}{g:{shape:Shape}→Input₂shape→Output₂shape}{sourcetarget: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₁ fT₁S₁ -- and g T2 x = S2 g x hg:Equivariant Input:=Input₂ Output:=Output₂ gT₂S₂ :-- Then <f,g <T1 x, T2 y = <S1 f x , S2 g y Equivariant Input:=funshape= Input₁shape×Input₂shape Output:=funshape= Output₁shape×Output₂shape funinput= Prod.mapfginput Prod.mapT₁T₂ Prod.mapS₁S₂ :=Shape: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.mapfginput Prod.mapT₁T₂ Prod.mapS₁S₂ Shape: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.mapfginput Prod.mapT₁T₂x =Prod.mapS₁S₂ fun{shape}input= Prod.mapfginput x All goals completed 🐙consShape:Type vState:Shape→Type usource:Shapetarget:Shapetransform:Statesource→Statetargetlayer:LayerStaterest:List LayerState ih: ∀ layer:LayerState ,layer∈rest→Equivariant fun{shape}= layer transformtransform →Equivariant fun{shape}= neural networkrest transformtransformequivariant:∀ layer 1:LayerState ,layer 1∈layer::rest→Equivariant fun{shape}= layer 1 transformtransformcomposed:Equivariant fun{shape}input= neural networkrest layerinput transformtransform⊢ Equivariant fun{shape}= neural network layer::rest transformtransformexactcomposedAll goals completed 🐙 We can use these properties to show that our neural network is selection equivariant, roughly that each individual result should be the same no matter how batches are built or ordered. -- Select m arbitrary elements of a set of n elements.defselect selection:Finm→Finn a:Finn→α :Finm→α:=funi= a selectioni -- Slice out a fixed-size group.defslice startcount:Nat -- Note that this takes a proof that the selection is in-bounds as an arg h:start+count≤n xs:Finn→α :Fincount→α:=select funi= ⟨i.val+start,byα:Type uβ:Type vδ:Type wn:Natstart:Natcount:Nath:start+count≤nxs:Finn→αi:Fincount⊢ ↑i+start<nomegaAll goals completed 🐙⟩ xs-- Equivariance under every selection, including selection across lengths.defSelectionEquivariant op:{n:Nat}→ Finn→α → Finn→β :Prop:=∀{nm} selection:Finm→Finn ,Equivariant Input:=funn= Finn→α Output:=funn= Finn→β op selectselection selectselection -- Under vmap selection doesn't matter.theoremvmap selection equivariant fn:α→β :SelectionEquivariant vmapfn :=byα:Type uβ:Type vfn:α→β⊢ SelectionEquivariantfun{n}= vmapfnintronmselectionaα:Type uβ:Type vfn:α→βn:Natm:Natselection:Finm→Finna:Finn→α⊢ fun{shape}{n}= vmapfn selectselectiona =selectselection fun{shape}{n}= vmapfn a rflAll goals completed 🐙@ simp theoremvmap2 apply fn:α→β→δ a:Finn→α b:Finn→β i:Finn :vmap2fnabi=fn ai bi :=rfl From these individual results, we directly build up to our first main result. A simple neural network does not depend on the order or content of its batch. While these properties so far seem basic, they are essential for designing large-scale LLMs. These properties provide the means for parallelizing and optimizing these systems. They also are properties that are commonly broken when new low-level optimizations are introduced. Let's look at a couple of these in more detail. Batch Invariance Batch invariance ensures that the final loss of the system is independent of the size of the batch used. This property can ensure replicability across systems. See Horace He's beautifully described blog about why batch invariance is useful and how it is often sacrificed under different optimizations. Here we prove that selection equivariance implies a simple form of batch invariance. Basically, you get the same loss independent of the batch. -- The same example has the same scalar loss in a batch or on its own.theoremnn batch invariant{nn:{batch:Nat}→ Finbatch→α → Finbatch→β } equivariant:SelectionEquivariantnn point loss:β→Rat input:Finbatch→α b:Finbatch :point loss nn select fun :Fin1= b input 0 =point loss nninputb :=byα:Type uβ:Type vbatch:Natnn:{batch:Nat}→ Finbatch→α →Finbatch→βequivariant:SelectionEquivariantfun{n}= nnpoint loss:β→Ratinput:Finbatch→αb:Finbatch⊢ point loss nn select funx= b input 0 =point loss nninputb exactcongrArgpoint loss congrFun equivariant fun :Fin1= b input 0 All goals completed 🐙 Tensor Parallel Tensor parallelism is a common optimization for distributed neural networks. It's a fancy way of saying that instead of doing a matrix multiplication on one host, you can instead split it into two or more parts, do those multiplications separately, and then merge them. defMatrix.row split a:Matrix k+k m :Matrixkm×Matrixkm:=-- Note here that i ∈ {0..k-1} but to index row need i ∈ {0..2 k-1}.-- These functions handle that cast.⟨select funi= i.castAddk a,select funi= i.natAddk a⟩-- sum is splittabletheoremVector.sum split values:Vector m+n :Vector.sum select funi:Finm= i.castAddn values +Vector.sum select funi:Finn= i.natAddm values =values.sum:=bym:Natn:Natvalues:Vector m+n ⊢ sum select funi= Fin.castAddni values +sum select funi= Fin.natAddmi values =values.sumsimponly Vector.sum,fori,select,List.ofFn add,List.sum append,Fin.castAdd,Fin.castLE All goals completed 🐙defMatrix.tensor parallel a:Matrixn k+k b:Matrix k+k p :Matrixnp:=let a₁,a₂ :=a.transpose.row splitlet b₁,b₂ :=b.row splitletc₁:=a₁.transpose.matmulb₁letc₂:=a₂.transpose.matmulb₂c₁+c₂theoremMatrix.tensor parallel correct-- For any splittable a, b a:Matrixn k+k b:Matrix k+k p :-- running tensor parallel gives the same result as matmula.tensor parallelb=a.matmulb:=byn:Natk:Natp:Nata:Matrixn k+k b:Matrix k+k p⊢ a.tensor parallelb=a.matmulb-- Show each final i, j ends up the same.funextijn:Natk:Natp:Nata:Matrixn k+k b:Matrix k+k pi:Finnj:Finp⊢ a.tensor parallelbij=a.matmulbijexactVector.sum split funt= ait btj All goals completed 🐙 Data Parallel Data parallelism says that, in training, we can split the data into different groups, run the full neural network and loss on different machines, and then combine. We need to ensure that we get the same result by running things separately as together. We next study a simple bidirectional Transformer with attention. The sequence becomes an additional dimension of our tensor. We first define attention. One of the more surprising properties of the vanilla bidirectional Transformer is that it is a set-to-set model, i.e. it is permutation equivariant in sequence length. Let's define first what that means formally. -- A permutation is a bijective map from {0..n-1} = {0..n-1}structurePositionPermutation n:Nat whereindex:Finn→Finnvalid: foriindex .Perm forifuni:Finn= i -- Apply a permutation.defpermute π:PositionPermutationn a:Finn→α :Finn→α:=selectπ.indexadefpermute both π:PositionPermutationn a:Finn→Finn→α :Finn→Finn→α:=permuteπ vmap permuteπ a defpermute qkv{hidden:Nat} π:PositionPermutationseq :=Prod.map Prod.map permute α:=Vectorhidden π permute α:=Vectorhidden π permute α:=Vectorhidden π -- Main property.defPermuteEquivariant op:α→β -- If permutation is applied to our input, inputAction:PositionPermutationn→α→α:=byexactpermute -- The same permutation applied somehow to output yields the same result. outputAction:PositionPermutationn→β→β:=byexactpermute :Prop:=∀π:PositionPermutationn,Equivariant Input:=fun :Unit= α Output:=fun :Unit= β op source:= target:= inputActionπ outputActionπ Most of the core operations we have defined have the necessary equivariance already. The main additional property we need is for our softmax, which follows directly from addition. Of course, in practice, we add additional information that breaks this property. The simplest way is through the use of positional features. We can show that even simple positional features break equivariance with a direct counterexample.