{"slug": "lean-verified-transformers", "title": "Lean Verified Transformers", "summary": "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.", "body_md": "This post explores writing formally verified ML code in Lean.\nSince the cost of proofs is declining rapidly and the amount of code generated\nis skyrocketing, the value of verified code seems likely to climb.\nWhile understanding proofs remains challenging, collaborating with AI to get\nproofs of easy-to-understand properties seems like a natural middle ground.\n\nThe goal of this post is to verify foundational properties of Transformers.\nThese are critical properties that are used for\nparallelization and optimization, including tensor parallelism, data parallelism,\nbatch invariance, permutation invariance, correctness of tiling, and locality\nof sparse attention models. Code is available at srush/lean-transformer.\nThe text, comments, and structure of the blog are all human-written;\nthe proofs are all written by AI. Hopefully it can also serve as an advanced\nintro to Lean.\n\nDifferent neural network architectures retain different properties of their input.\nWe generally classify these properties in terms of equivariance and invariance.\nThese allow researchers to reason about what they can learn, and\nimplementers to optimize computation while maintaining equivalence.\nOur goal will be to prove equivariances and invariances for specific architectures.\n\nNotationally, ML definitions often assume the same functions can work on\ndifferent input shapes, e.g. batch sizes. For this reason, our Lean definition\nwill be a bit complex to allow for functions that are polymorphic over the shape.\n\ndefEquivariant-- 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\n\nVectors, Matrices, and Neural Networks\n\nWe begin by building a simple neural network library in Lean.\n\nThe relu function takes in a number and returns its non-negative part.\nAlong with the definition, we prove it does what we claim.\n\ndefrelu(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! 🐙\n\nFollowing the style of\nJAX,\nwe lift scalar functions to operate on vectors.\nVectors (and tensors) are represented as higher-order functions mapping\nindices to rational numbers. This makes our proofs easier since we do not\nhave to care about storage or efficiency.\n\n-- 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)\n\nFor aggregations, we define a vector scan. Since we are using rationals for\nsimplicity, we do not have an exponential, so we define a \"softmax-like\"\nnonlinear normalization instead.\n\n-- 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\n\nAs an exercise, let's look at a simple vector theorem. Click the square □\nnext to each line of the proof and it will show you the current proof state.\nThe proof state divides the context from the goal ⊢. Each step will transform\nthese terms until we can construct the goal.\n\n-- 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! 🐙\n\nMatrices are defined similarly. We are basically just stacking\nvmap calls to get our core operations. Note the implementation of matmul\nin particular, which will be the target of future proofs.\n\nNow let us return to our goal of proving network equivariances.\nOur strategy will be to first show that in general equivariances compose,\nand then show that they propagate through a neural network.\n\n-- 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! 🐙\n\nWe can use these properties to show that our neural network\nis selection equivariant, roughly that each individual result should be the\nsame no matter how batches are built or ordered.\n\n-- 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\n\nFrom these individual results, we directly build up to our first main result.\nA simple neural network does not depend on the order or content of its batch.\n\nWhile these properties so far seem basic, they are essential for\ndesigning large-scale LLMs. These properties provide the means for parallelizing\nand optimizing these systems. They also are properties that are commonly broken when\nnew low-level optimizations are introduced. Let's look at a couple of these\nin more detail.\n\nBatch Invariance\n\nBatch invariance ensures that the final loss of the system is independent\nof the size of the batch used. This property can ensure replicability across\nsystems. See\nHorace He's\nbeautifully described blog about why batch invariance is useful and how it\nis often sacrificed under different optimizations.\n\nHere we prove that selection equivariance implies a simple form of batch\ninvariance. Basically, you get the same loss independent of the batch.\n\n-- 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! 🐙\n\nTensor Parallel\n\nTensor parallelism is a common optimization for distributed neural networks.\nIt's a fancy way of saying that instead of doing a matrix multiplication on one\nhost, you can instead split it into two or more parts, do those multiplications\nseparately, and then merge them.\n\ndefMatrix.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! 🐙\n\nData Parallel\n\nData parallelism says that, in training, we can split the data into\ndifferent groups, run the full neural network and loss on different machines,\nand then combine. We need to ensure that we get the same result\nby running things separately as together.\n\nWe next study a simple bidirectional Transformer with attention.\nThe sequence becomes an additional dimension of our tensor.\nWe first define attention.\n\nOne of the more surprising properties of the vanilla (bidirectional) Transformer\nis that it is a set-to-set model, i.e. it is permutation equivariant in\nsequence length. Let's define first what that means formally.\n\n-- 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π)\n\nMost of the core operations we have defined have the necessary equivariance already.\nThe main additional property we need is for our softmax, which follows directly from\naddition.\n\nOf course, in practice, we add additional information that breaks this property.\nThe simplest way is through the use of positional features. We can show that even\nsimple positional features break equivariance with a direct counterexample.", "url": "https://wpnews.pro/news/lean-verified-transformers", "canonical_source": "https://srush.github.io/lean-transformer/", "published_at": "2026-09-22 03:55:18+00:00", "updated_at": "2026-09-22 04:24:23.541233+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "neural-networks", "large-language-models", "developer-tools"], "entities": ["Lean", "srush/lean-transformer", "JAX", "Transformer", "relu"], "alternates": {"html": "https://wpnews.pro/news/lean-verified-transformers", "markdown": "https://wpnews.pro/news/lean-verified-transformers.md", "text": "https://wpnews.pro/news/lean-verified-transformers.txt", "jsonld": "https://wpnews.pro/news/lean-verified-transformers.jsonld"}}