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)=ab+ac:=n:Nata:Vectornb:Vectornc:Vectorn⊢ a*(b+c)=ab+ac-- Strategy: show equiv for all indices i of the output vectorn:Nata:Vectornb:Vectornc:Vectorni:Finn⊢ (a*(b+c))i=(ab+ac)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.