Lean Verified Transformers — Lean Verified Transformers

131 min read Original article ↗

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.

This project is inspired by TorchLean, Verified Deep Learning with Lean 4, and the Dex Programming Language.

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 0theorem relu_non_negative -- For all z (z : Rat) : -- relu is ≥ 0 relu z ≥ 0 := z:Rat⊢ relu z ≥ 0 -- Do a short (grind) search All 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. 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 fdef 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)).sumdef 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) weightsdef 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 := n: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 n: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. 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 minstance : 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 idef Matrix.matvec (a : Matrix n m) (x : Vector m) : Vector n := vmap (fun row => row.dot_product x) adef 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 shapedef neural_network {Shape : Type v} {State : Shape → Type u} (layers : List (Layer State)) : Layer State := fun input => layers.foldl (fun state layer => layer state) inputdef 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 := Shape: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 Shape: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) 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₂) := 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.map f g input) (Prod.map T₁ T₂) (Prod.map S₁ 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.map f g input) (Prod.map T₁ T₂ x) = Prod.map S₁ S₂ ((fun {shape} input => Prod.map f g input) x) All goals completed! 🐙Shape:Type vState:Shape → Type usource:Shapetarget:Shapetransform:State source → State targetlayer:Layer Staterest:List (Layer State)ih:(∀ (layer : Layer State), layer ∈ rest → Equivariant (fun {shape} => layer) transform transform) → Equivariant (fun {shape} => neural_network rest) transform transformequivariant:∀ (layer_1 : Layer State), layer_1 ∈ layer :: rest → Equivariant (fun {shape} => layer_1) transform transformcomposed:Equivariant (fun {shape} input => neural_network rest (layer input)) transform transform⊢ Equivariant (fun {shape} => neural_network (layer :: rest)) transform transform All 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.

Selecting, reordering, and repeating positions commutes with a selection-equivariant function.

-- Select m arbitrary elements of a set of n elements. def select (selection : Fin m → Fin n) (a : Fin n → α) : Fin m → α := fun i => a (selection i)-- Slice out a fixed-size group. def slice (start count : Nat) -- Note that this takes a proof that the selection is in-bounds as an arg! (h : start + count ≤ n) (xs : Fin n → α) : Fin count → α := select (fun i => ⟨i.val + start, α:Type uβ:Type vδ:Type wn:Natstart:Natcount:Nath:start + count ≤ nxs:Fin n → αi:Fin count⊢ ↑i + start < n All goals completed! 🐙⟩) xs-- Equivariance under every selection, including selection across lengths. def SelectionEquivariant (op : {n : Nat} → (Fin n → α) → (Fin n → β)) : Prop := ∀ {n m} (selection : Fin m → Fin n), Equivariant (Input := fun n => Fin n → α) (Output := fun n => Fin n → β) op (select selection) (select selection)-- Under vmap selection doesn't matter. theorem vmap_selection_equivariant (fn : α → β) : SelectionEquivariant (vmap fn) := α:Type uβ:Type vfn:α → β⊢ SelectionEquivariant fun {n} => vmap fn α:Type uβ:Type vfn:α → βn:Natm:Natselection:Fin m → Fin na:Fin n → α⊢ (fun {shape} {n} => vmap fn) (select selection a) = select selection ((fun {shape} {n} => vmap fn) a) All goals completed! 🐙@[simp] theorem vmap2_apply (fn : α → β → δ) (a : Fin n → α) (b : Fin n → β) (i : Fin n) : vmap2 fn a b i = fn (a i) (b i) := 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.

System Optimization

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.

Batch invariance: selecting an example before or after the same network gives the same output.

-- The same example has the same scalar loss in a batch or on its own. theorem nn_batch_invariant {nn : {batch : Nat} → (Fin batch → α) → (Fin batch → β)} (equivariant : SelectionEquivariant nn) (point_loss : β → Rat) (input : Fin batch → α) (b : Fin batch) : point_loss (nn (select (fun _ : Fin 1 => b) input) 0) = point_loss (nn input b) := α:Type uβ:Type vbatch:Natnn:{batch : Nat} → (Fin batch → α) → Fin batch → βequivariant:SelectionEquivariant fun {n} => nnpoint_loss:β → Ratinput:Fin batch → αb:Fin batch⊢ point_loss (nn (select (fun x => b) input) 0) = point_loss (nn input b) 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.

Tensor parallelism: split A by columns and B by rows, multiply each pair independently, and add the results.

def Matrix.row_split (a : Matrix (k + k) m) : Matrix k m × Matrix k m := -- Note here that i ∈ {0..k-1} but to index row need i ∈ {0..2 k-1}. -- These functions handle that cast. ⟨select (fun i => i.castAdd k) a, select (fun i => i.natAdd k) a⟩-- sum is splittable theorem Vector.sum_split (values : Vector (m + n)) : Vector.sum (select (fun i : Fin m => i.castAdd n) values) + Vector.sum (select (fun i : Fin n => i.natAdd m) values) = values.sum := m:Natn:Natvalues:Vector (m + n)⊢ sum (select (fun i => Fin.castAdd n i) values) + sum (select (fun i => Fin.natAdd m i) values) = values.sum All goals completed! 🐙def Matrix.tensor_parallel (a : Matrix n (k + k)) (b : Matrix (k + k) p) : Matrix n p := let (a₁, a₂) := a.transpose.row_split let (b₁, b₂) := b.row_split let c₁ := a₁.transpose.matmul b₁ let c₂ := a₂.transpose.matmul b₂ c₁ + c₂theorem Matrix.tensor_parallel_correct -- For any splittable a, b (a : Matrix n (k + k)) (b : Matrix (k + k) p) : -- running tensor_parallel gives the same result as matmul a.tensor_parallel b = a.matmul b := n:Natk:Natp:Nata:Matrix n (k + k)b:Matrix (k + k) p⊢ a.tensor_parallel b = a.matmul b -- Show each final i, j ends up the same. n:Natk:Natp:Nata:Matrix n (k + k)b:Matrix (k + k) pi:Fin nj:Fin p⊢ a.tensor_parallel b i j = a.matmul b i j 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.

Data parallelism: apply the same network to each batch half, sum point losses using their original batch indices, and add the two losses.

theorem loss_row_split (point_loss : Fin (k + k) → Vector hidden → Rat) (output : Matrix (k + k) hidden) : loss (fun b row => point_loss (b.castAdd k) row) output.row_split.1 + loss (fun b row => point_loss (b.natAdd k) row) output.row_split.2 = loss point_loss output := Vector.sum_split (fun b => point_loss b (output b))-- Split data points in half and run on separate machines. def data_parallel_loss (layers : List (Layer (fun batch => Fin batch → Vector hidden))) (point_loss : Fin (k + k) → Vector hidden → Rat) (input : Matrix (k + k) hidden) : Rat := let (first, second) := input.row_split loss (fun b row => point_loss (b.castAdd k) row) (neural_network layers first) + loss (fun b row => point_loss (b.natAdd k) row) (neural_network layers second)theorem data_parallel_loss_correct (layers : List (Layer (fun batch => Fin batch → Vector hidden))) (equivariant : ∀ layer ∈ layers, SelectionEquivariant layer) (point_loss : Fin (k + k) → Vector hidden → Rat) (input : Matrix (k + k) hidden) : data_parallel_loss layers point_loss input = loss point_loss (neural_network layers input) := hidden:Natk:Natlayers:List (Layer fun batch => Fin batch → Vector hidden)equivariant:∀ (layer : Layer fun batch => Fin batch → Vector hidden), layer ∈ layers → SelectionEquivariant fun {n} => layerpoint_loss:Fin (k + k) → Vector hidden → Ratinput:Matrix (k + k) hidden⊢ data_parallel_loss layers point_loss input = loss point_loss (neural_network layers input) hidden:Natk:Natlayers:List (Layer fun batch => Fin batch → Vector hidden)equivariant:∀ (layer : Layer fun batch => Fin batch → Vector hidden), layer ∈ layers → SelectionEquivariant fun {n} => layerpoint_loss:Fin (k + k) → Vector hidden → Ratinput:Fin (k + k) → Vector hidden⊢ data_parallel_loss layers point_loss input = loss point_loss (neural_network layers input) hidden:Natk:Natlayers:List (Layer fun batch => Fin batch → Vector hidden)equivariant:∀ (layer : Layer fun batch => Fin batch → Vector hidden), layer ∈ layers → SelectionEquivariant fun {n} => layerpoint_loss:Fin (k + k) → Vector hidden → Ratinput:Fin (k + k) → Vector hiddenequiv:∀ (selection : Fin k → Fin (k + k)), Equivariant (fun {shape} {n} => neural_network layers) (select selection) (select selection)⊢ data_parallel_loss layers point_loss input = loss point_loss (neural_network layers input) hidden:Natk:Natlayers:List (Layer fun batch => Fin batch → Vector hidden)equivariant:∀ (layer : Layer fun batch => Fin batch → Vector hidden), layer ∈ layers → SelectionEquivariant fun {n} => layerpoint_loss:Fin (k + k) → Vector hidden → Ratinput:Fin (k + k) → Vector hiddenequiv:∀ (selection : Fin k → Fin (k + k)) (x : (fun n => Fin n → Vector hidden) (k + k)), (fun {shape} {n} => neural_network layers) (select selection x) = select selection ((fun {shape} {n} => neural_network layers) x)⊢ data_parallel_loss layers point_loss input = loss point_loss (neural_network layers input) All goals completed! 🐙

Transformers and attention

We next study a simple bidirectional Transformer with attention. The sequence becomes an additional dimension of our tensor. We first define attention.

-- A mixer takes <<q,k>,v> as an arg and returns the result. abbrev Mixer (seq hidden : Nat) := (Matrix seq hidden × Matrix seq hidden) × Matrix seq hidden → Matrix seq hidden-- The famed softmax(Q K^T) V formula. def base_attention (s : Matrix seq seq → Matrix seq seq) : Mixer seq hidden := fun ((q,k), v) => Matrix.matmul (s (q.matmul k.transpose)) vdef attention_layer : Mixer seq hidden := fun input => base_attention (vmap softmax_like) input

Attention is included in the main network through a parameterized Transformer block.

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} structure PositionPermutation (n : Nat) where index : Fin n → Fin n valid : (fori index).Perm (fori fun i : Fin n => i)-- Apply a permutation. def permute (π : PositionPermutation n) (a : Fin n → α) : Fin n → α := select π.index adef permute_both (π : PositionPermutation n) (a : Fin n → Fin n → α) : Fin n → Fin n → α := permute π (vmap (permute π) a)def permute_qkv {hidden : Nat} (π : PositionPermutation seq) := Prod.map (Prod.map (permute (α := Vector hidden) π) (permute (α := Vector hidden) π)) (permute (α := Vector hidden) π)-- Main property. def PermuteEquivariant (op : α → β) -- If permutation is applied to our input, (inputAction : PositionPermutation n → α → α := by exact permute) -- The same permutation applied somehow to output yields the same result. (outputAction : PositionPermutation n → β → β := by exact permute) : Prop := ∀ π : PositionPermutation n, 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.

Row equivariance: swapping the rows before matmul gives the same result as swapping the output rows.

Now we can show that the vanilla Transformer is permutation equivariant.

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.

def position_features (i : Rat) : Vector 3 := fun d => if d.val = 0 then i else if d.val = 1 then i * i else 1def Matrix.add_positions (input : Matrix seq 3) : Matrix seq 3 := fun s => input s + position_features (s.val : Rat)equivariant:PermuteEquivariant (neural_network [@Matrix.add_positions, fun {shape} input => transformer_block attention_layer { weight := fun i j => if i = j then 1 else 0, wq := fun i j => if i = j then 1 else 0, wk := fun i j => if i = j then 1 else 0, wv := fun i j => if i = j then 1 else 0 } input]) permute permuteinput:Matrix 2 3 := fun s d => if ↑d = 0 then ↑↑s else 0swap:PositionPermutation 2 := { index := fun s => ⟨1 - ↑s, ⋯⟩, valid := ⋯ }⊢ False equivariant:PermuteEquivariant (neural_network [@Matrix.add_positions, fun {shape} input => transformer_block attention_layer { weight := fun i j => if i = j then 1 else 0, wq := fun i j => if i = j then 1 else 0, wk := fun i j => if i = j then 1 else 0, wv := fun i j => if i = j then 1 else 0 } input]) permute permuteinput:Matrix 2 3 := fun s d => if ↑d = 0 then ↑↑s else 0swap:PositionPermutation 2 := { index := fun s => ⟨1 - ↑s, ⋯⟩, valid := ⋯ }same:(fun {shape} => neural_network [@Matrix.add_positions, fun {shape} input => transformer_block attention_layer { weight := fun i j => if i = j then 1 else 0, wq := fun i j => if i = j then 1 else 0, wk := fun i j => if i = j then 1 else 0, wv := fun i j => if i = j then 1 else 0 } input]) (permute swap input) 0 0 = permute swap ((fun {shape} => neural_network [@Matrix.add_positions, fun {shape} input => transformer_block attention_layer { weight := fun i j => if i = j then 1 else 0, wq := fun i j => if i = j then 1 else 0, wk := fun i j => if i = j then 1 else 0, wv := fun i j => if i = j then 1 else 0 } input]) input) 0 0⊢ False equivariant:PermuteEquivariant (neural_network [@Matrix.add_positions, fun {shape} input => transformer_block attention_layer { weight := fun i j => if i = j then 1 else 0, wq := fun i j => if i = j then 1 else 0, wk := fun i j => if i = j then 1 else 0, wv := fun i j => if i = j then 1 else 0 } input]) permute permuteinput:Matrix 2 3 := fun s d => if ↑d = 0 then ↑↑s else 0swap:PositionPermutation 2 := { index := fun s => ⟨1 - ↑s, ⋯⟩, valid := ⋯ }⊢ (fun {shape} => neural_network [@Matrix.add_positions, fun {shape} input => transformer_block attention_layer { weight := fun i j => if i = j then 1 else 0, wq := fun i j => if i = j then 1 else 0, wk := fun i j => if i = j then 1 else 0, wv := fun i j => if i = j then 1 else 0 } input]) (permute swap input) 0 0 = permute swap ((fun {shape} => neural_network [@Matrix.add_positions, fun {shape} input => transformer_block attention_layer { weight := fun i j => if i = j then 1 else 0, wq := fun i j => if i = j then 1 else 0, wk := fun i j => if i = j then 1 else 0, wv := fun i j => if i = j then 1 else 0 } input]) input) 0 0 → False All goals completed! 🐙

Sparse Attention

An alternative to full attention over the sequence is a sparse attention over a local region. Sliding window attention only looks at the surrounding window. We implement this with a windowed mask.

abbrev InWindow (radius : Nat) (s t : Fin seq) : Prop := s.val ≤ t.val + radius ∧ t.val ≤ s.val + radiusdef window_mask (radius : Nat) : Matrix seq seq := fun s t => if InWindow radius s t then 1 else 0def matrix_mask (radius: Nat) (a: Matrix m m) : Matrix m m := let mask: Matrix m m := window_mask radius a * maskdef swa (radius : Nat) : Mixer seq hidden := base_attention (matrix_mask radius)

The sliding window has the property that each position is only impacted by the region around it. We formalize this idea below.

Inputs agreeing within a window give the same output at its center, even when outside values differ.

-- Value is not impacted outside a region of a given radius. def RegionInvariant (radius : Nat) (op : (Fin seq → α) → Fin seq → β) : Prop := ∀ (input other : Fin seq → α) (s : Fin seq), (∀ t, InWindow radius s t → input t = other t ) → op input s = op other sAll goals completed! 🐙 _ = op other i := congrFun (equivariant (fun _ : Fin 1 => i) other) 0-- vmap has radius 0 theorem vmap_region_invariant (fn : α → β) : RegionInvariant (seq := seq) 0 (vmap fn) := SelectionEquivariant.region_invariant (vmap_selection_equivariant fn)

The interesting new aspect of region invariance will be how it composes. Region invariance is additive under composition.

-- Composing sliding windows increases the radius. theorem RegionInvariant.comp {first : (Fin seq → α) → Fin seq → β} {second : (Fin seq → β) → Fin seq → δ} (hfirst : RegionInvariant r first) (hsecond : RegionInvariant t second) : RegionInvariant (r + t) (fun input => second (first input)) := α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t second⊢ RegionInvariant (r + t) fun input => second (first input) α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1⊢ (fun input => second (first input)) input s = (fun input => second (first input)) other s α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1⊢ ∀ (t_1 : Fin seq), InWindow t s t_1 → first input t_1 = first other t_1 α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:InWindow t s u⊢ first input u = first other u α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:InWindow t s u⊢ ∀ (t : Fin seq), InWindow r u t → input t = other t α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:InWindow t s uv:Fin seqhv:InWindow r u v⊢ input v = other v α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:InWindow t s uv:Fin seqhv:InWindow r u v⊢ InWindow (r + t) s v α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:↑s ≤ ↑u + t ∧ ↑u ≤ ↑s + tv:Fin seqhv:↑u ≤ ↑v + r ∧ ↑v ≤ ↑u + r⊢ ↑s ≤ ↑v + (r + t) ∧ ↑v ≤ ↑s + (r + t) α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:↑s ≤ ↑u + t ∧ ↑u ≤ ↑s + tv:Fin seqhv:↑u ≤ ↑v + r ∧ ↑v ≤ ↑u + r⊢ ↑s ≤ ↑v + (r + t)α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:↑s ≤ ↑u + t ∧ ↑u ≤ ↑s + tv:Fin seqhv:↑u ≤ ↑v + r ∧ ↑v ≤ ↑u + r⊢ ↑v ≤ ↑s + (r + t) α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:↑s ≤ ↑u + t ∧ ↑u ≤ ↑s + tv:Fin seqhv:↑u ≤ ↑v + r ∧ ↑v ≤ ↑u + r⊢ ↑s ≤ ↑v + (r + t)α:Type uβ:Type vδ:Type wseq:Natr:Natt:Natfirst:(Fin seq → α) → Fin seq → βsecond:(Fin seq → β) → Fin seq → δhfirst:RegionInvariant r firsthsecond:RegionInvariant t secondinput:Fin seq → αother:Fin seq → αs:Fin seqagree:∀ (t_1 : Fin seq), InWindow (r + t) s t_1 → input t_1 = other t_1u:Fin seqhu:↑s ≤ ↑u + t ∧ ↑u ≤ ↑s + tv:Fin seqhv:↑u ≤ ↑v + r ∧ ↑v ≤ ↑u + r⊢ ↑v ≤ ↑s + (r + t) All goals completed! 🐙α:Type useq:Natradius:Natlayer:Layer fun seq => Fin seq → αrest:List (Layer fun seq => Fin seq → α)ih:(∀ (layer : Layer fun seq => Fin seq → α), layer ∈ rest → RegionInvariant radius layer) → RegionInvariant (rest.length * radius) (neural_network rest)invariant:∀ (layer_1 : Layer fun seq => Fin seq → α), layer_1 ∈ layer :: rest → RegionInvariant radius layer_1composed:RegionInvariant (radius + rest.length * radius) fun input => neural_network rest (layer input)⊢ RegionInvariant ((layer :: rest).length * radius) (neural_network (layer :: rest)) All goals completed! 🐙

Finally, we need to show the radius of SWA. We first show that our implementation of masking leads to a given radius and then apply this to the sparse attention implementation.

attribute [local simp] Rat.add_zero Rat.zero_add Rat.zero_mul Rat.mul_zeroAll goals completed! 🐙 seq:Nathidden:Natradius:Natweights:Matrix seq seqinput:Fin seq → Vector hiddenother:Fin seq → Vector hiddens:Fin seqagree:∀ (t : Fin seq), InWindow radius s t → input t = other tj:Fin hiddent:Fin seqht:¬InWindow radius s t⊢ weights s t * window_mask radius s t * input t j = weights s t * window_mask radius s t * other t j All goals completed! 🐙hidden:Natseq:Natradius:Natwq:Matrix hidden hiddenwk:Matrix hidden hiddenwv:Matrix hidden hiddeninput:Fin seq → Vector hiddenother:Fin seq → Vector hiddens:Fin seqagree:∀ (t : Fin seq), InWindow radius s t → input t = other tcenter:input s = other scontribution:Vector hidden → Vector hidden → Vector hidden := fun query row j => (Vector.dot_product (fun d => query.dot_product (wq.transpose d)) fun d => row.dot_product (wk.transpose d)) * row.dot_product (wv.transpose j)masked:∀ (x : Matrix seq hidden), swa radius (project_qkv x wq wk wv) s = (matrix_mask radius fun x x_1 => 1).matmul (vmap (contribution (x s)) x) s⊢ (fun input => swa radius (project_qkv input wq wk wv)) input s = (fun input => swa radius (project_qkv input wq wk wv)) other s hidden:Natseq:Natradius:Natwq:Matrix hidden hiddenwk:Matrix hidden hiddenwv:Matrix hidden hiddeninput:Fin seq → Vector hiddenother:Fin seq → Vector hiddens:Fin seqagree:∀ (t : Fin seq), InWindow radius s t → input t = other tcenter:input s = other scontribution:Vector hidden → Vector hidden → Vector hidden := fun query row j => (Vector.dot_product (fun d => query.dot_product (wq.transpose d)) fun d => row.dot_product (wk.transpose d)) * row.dot_product (wv.transpose j)masked:∀ (x : Matrix seq hidden), swa radius (project_qkv x wq wk wv) s = (matrix_mask radius fun x x_1 => 1).matmul (vmap (contribution (x s)) x) ssame:(fun values => (matrix_mask radius fun x x_1 => 1).matmul values) (vmap (contribution (input s)) input) s = (fun values => (matrix_mask radius fun x x_1 => 1).matmul values) (vmap (contribution (input s)) other) s⊢ (fun input => swa radius (project_qkv input wq wk wv)) input s = (fun input => swa radius (project_qkv input wq wk wv)) other s exact (masked input).trans (same.trans (hidden:Natseq:Natradius:Natwq:Matrix hidden hiddenwk:Matrix hidden hiddenwv:Matrix hidden hiddeninput:Fin seq → Vector hiddenother:Fin seq → Vector hiddens:Fin seqagree:∀ (t : Fin seq), InWindow radius s t → input t = other tcenter:input s = other scontribution:Vector hidden → Vector hidden → Vector hidden := fun query row j => (Vector.dot_product (fun d => query.dot_product (wq.transpose d)) fun d => row.dot_product (wk.transpose d)) * row.dot_product (wv.transpose j)masked:∀ (x : Matrix seq hidden), swa radius (project_qkv x wq wk wv) s = (matrix_mask radius fun x x_1 => 1).matmul (vmap (contribution (x s)) x) ssame:(fun values => (matrix_mask radius fun x x_1 => 1).matmul values) (vmap (contribution (input s)) input) s = (fun values => (matrix_mask radius fun x x_1 => 1).matmul values) (vmap (contribution (input s)) other) s⊢ (fun values => (matrix_mask radius fun x x_1 => 1).matmul values) (vmap (contribution (input s)) other) s = (fun input => swa radius (project_qkv input wq wk wv)) other s All goals completed! 🐙))

Flash Attention

Once we have standard attention implemented, we can begin to consider optimizations. Flash attention uses a tiling approach to split the full sequence into groups which can be computed separately to reduce memory. We define a typed tiling.

-- Select one of tiles groups, each of size n. def tile (n : Nat) (xs : Fin (tiles * n) → α) (c : Fin tiles) : Fin n → α := -- The proof here shows that it is legal to make this slicing. slice (c.val * n) n (α:Type uβ:Type vδ:Type wtiles:Natn:Natxs:Fin (tiles * n) → αc:Fin tiles⊢ ↑c * n + n ≤ tiles * n α:Type uβ:Type vδ:Type wtiles:Natn:Natxs:Fin (tiles * n) → αc:Fin tilesthis:(↑c).succ * n ≤ tiles * n⊢ ↑c * n + n ≤ tiles * n All goals completed! 🐙) xs-- Apply an accumulator across each of these. def tile_fold (tiles n : Nat) (step : σ → α → σ) (xs : Fin (tiles * n) → α) (initial : σ) : σ := scan (fun state chunk => scan step chunk state) (tile n xs) initial

Once we have this primitive, we can compute and aggregate a value for each tile.

Process equal-size key/value tiles, carry the two accumulators, and normalize once at the end.

structure FlashAcc (hidden : Nat) where scoreSum : Rat weighted : Vector hiddendef flash_step (r : α → Rat) (v : α → Vector hidden) (state : FlashAcc hidden) (t : α) : FlashAcc hidden := let score := r t { scoreSum := state.scoreSum + score weighted := fun j => state.weighted j + score * v t j }def flash_attention (tiles n : Nat) (input : (Matrix (tiles * n) hidden × Matrix (tiles * n) hidden) × Matrix (tiles * n) hidden) : Matrix (tiles * n) hidden := let ⟨⟨q, k⟩, v⟩ := input fun s => let r := fun t => 1 + relu (Vector.sum (fun d => q s d * k t d)) let initial : FlashAcc hidden := ⟨0, fun _ => 0⟩ let state := tile_fold tiles n (flash_step r v) (fun t : Fin (tiles * n) => t) initial fun j => state.weighted j / state.scoreSum

For our equivalence proof, we need to define some basic properties of scans and sums. These are not unique to Flash attention, but would be in a core mathematics library.

theorem scan_slice (step : σ → α → σ) (xs : Fin n → α) (start count : Nat) (h : start + count ≤ n) (initial : σ) : scan step (slice start count h xs) (scan step (slice 0 start (α:Type uβ:Type vδ:Type wσ:Sort ?u.7n:Natstep:σ → α → σxs:Fin n → αstart:Natcount:Nath:start + count ≤ ninitial:σ⊢ 0 + start ≤ n All goals completed! 🐙) xs) initial) = scan step (slice 0 (start + count) (α:Type uβ:Type vδ:Type wσ:Sort ?u.7n:Natstep:σ → α → σxs:Fin n → αstart:Natcount:Nath:start + count ≤ ninitial:σ⊢ 0 + (start + count) ≤ n All goals completed! 🐙) xs) initial := α:Type uσ:Sort u_1n:Natstep:σ → α → σxs:Fin n → αstart:Natcount:Nath:start + count ≤ ninitial:σ⊢ scan step (slice start count h xs) (scan step (slice 0 start ⋯ xs) initial) = scan step (slice 0 (start + count) ⋯ xs) initial All goals completed! 🐙theorem scan_full (step : σ → α → σ) (xs : Fin n → α) (initial : σ) (h : count = n) : scan step (slice 0 count (α:Type uβ:Type vδ:Type wσ:Sort ?u.8n:Natcount:Natstep:σ → α → σxs:Fin n → αinitial:σh:count = n⊢ 0 + count ≤ n All goals completed! 🐙) xs) initial = scan step xs initial := α:Type uσ:Sort u_1n:Natcount:Natstep:σ → α → σxs:Fin n → αinitial:σh:count = n⊢ scan step (slice 0 count ⋯ xs) initial = scan step xs initial α:Type uσ:Sort u_1n:Natstep:σ → α → σxs:Fin n → αinitial:σ⊢ scan step (slice 0 n ⋯ xs) initial = scan step xs initial All goals completed! 🐙theorem Vector.sum_succ (f : Vector (n + 1)) : f.sum = f 0 + Vector.sum (fun i : Fin n => f i.succ) := n:Natf:Vector (n + 1)⊢ f.sum = f 0 + sum fun i => f i.succ All goals completed! 🐙All goals completed! 🐙theorem Vector.sum_last (f : Vector (n + 1)) : f.sum = Vector.sum (fun i : Fin n => f i.castSucc) + f (Fin.last n) := n:Natf:Vector (n + 1)⊢ f.sum = (sum fun i => f i.castSucc) + f (Fin.last n) All goals completed! 🐙n:Natih:∀ (f : Vector n), (sum fun t => f t.rev) = f.sumf:Vector (n + 1)⊢ (f (Fin.rev 0) + sum fun i => f i.succ.rev) = (sum fun i => f i.castSucc) + f (Fin.last n) n:Natih:∀ (f : Vector n), (sum fun t => f t.rev) = f.sumf:Vector (n + 1)h:(sum fun t => f t.rev.castSucc) = sum fun i => f i.castSucc⊢ (f (Fin.rev 0) + sum fun i => f i.succ.rev) = (sum fun i => f i.castSucc) + f (Fin.last n) All goals completed! 🐙α:Type uσ:Sort u_1n:Natstep:σ → α → σtiles:Natih:∀ (xs : Fin (tiles * n) → α) (initial : σ), tile_fold tiles n step xs initial = scan step xs initialxs:Fin ((tiles + 1) * n) → αinitial:σshape:(tiles + 1) * n = tiles * n + n⊢ scan step (slice 0 (tiles * n + n) ⋯ xs) initial = scan step xs initial All goals completed! 🐙

The main proof is that Flash attention is equivalent to our original attention. This is done by showing a lemma over the internal fold that we are accumulating the correct values.

α:Type uhidden:Natr:α → Ratv:α → Vector hiddenn:Natih:∀ (xs : Fin n → α) (state : FlashAcc hidden), scan (flash_step r v) xs state = { scoreSum := state.scoreSum + Vector.sum fun i => r (xs i), weighted := fun j => state.weighted j + Vector.sum fun i => r (xs i) * v (xs i) j }xs:Fin (n + 1) → αstate:FlashAcc hidden⊢ { scoreSum := (flash_step r v state (xs 0)).scoreSum + Vector.sum fun i => r (xs i.succ), weighted := fun j => (flash_step r v state (xs 0)).weighted j + Vector.sum fun i => r (xs i.succ) * v (xs i.succ) j } = { scoreSum := state.scoreSum + Vector.sum fun i => r (xs i), weighted := fun j => state.weighted j + Vector.sum fun i => r (xs i) * v (xs i) j } All goals completed! 🐙theorem flash_attention_eq (tiles n : Nat) (input : (Matrix (tiles * n) hidden × Matrix (tiles * n) hidden) × Matrix (tiles * n) hidden) : flash_attention tiles n input = attention_layer input := hidden:Nattiles:Natn:Natinput:(Matrix (tiles * n) hidden × Matrix (tiles * n) hidden) × Matrix (tiles * n) hidden⊢ flash_attention tiles n input = attention_layer input hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hidden⊢ flash_attention tiles n ((q, k), v) = attention_layer ((q, k), v) hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hidden⊢ flash_attention tiles n ((q, k), v) s j = attention_layer ((q, k), v) s j hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hidden⊢ ((Vector.sum fun i => (1 + relu (Vector.sum fun d => q s d * k i d)) * v i j) / Vector.sum fun t => 1 + relu (Vector.sum fun d => q s d * k t d)) = attention_layer ((q, k), v) s j hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hiddenr:Fin (tiles * n) → Rat := fun t => 1 + relu (Vector.sum fun d => q s d * k t d)⊢ ((Vector.sum fun i => (1 + relu (Vector.sum fun d => q s d * k i d)) * v i j) / Vector.sum fun t => 1 + relu (Vector.sum fun d => q s d * k t d)) = attention_layer ((q, k), v) s j hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hiddenr:Fin (tiles * n) → Rat := fun t => 1 + relu (Vector.sum fun d => q s d * k t d)⊢ (Vector.sum fun t => r t * v t j) / Vector.sum r = Vector.sum fun t => r t / Vector.sum r * v t j hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hiddenr:Fin (tiles * n) → Rat := fun t => 1 + relu (Vector.sum fun d => q s d * k t d)⊢ (Vector.sum fun i => r i * v i j * (Vector.sum r)⁻¹) = Vector.sum fun t => r t * (Vector.sum r)⁻¹ * v t j hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hiddenr:Fin (tiles * n) → Rat := fun t => 1 + relu (Vector.sum fun d => q s d * k t d)⊢ (fun i => r i * v i j * (Vector.sum r)⁻¹) = fun t => r t * (Vector.sum r)⁻¹ * v t j hidden:Nattiles:Natn:Natv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hiddenr:Fin (tiles * n) → Rat := fun t => 1 + relu (Vector.sum fun d => q s d * k t d)t:Fin (tiles * n)⊢ r t * v t j * (Vector.sum r)⁻¹ = r t * (Vector.sum r)⁻¹ * v t j All goals completed! 🐙

State Space Models

Another alternative to standard attention is to use a state space model or linear attention approach. These can be defined by the following recurrence.

The recurrent, masked linear-attention, and chunkwise forms compute the same decayed SSM.

def ssm_scan (a : Rat) (values : Vector n) (initial : Rat := 0) : Rat := scan (fun state x => a * state + x) values initialdef ssm_state (a : Rat) (updates : Fin n → Matrix hidden hidden) (incoming : Matrix hidden hidden := fun _ _ => 0) : Matrix hidden hidden := fun d j => ssm_scan a (fun t => updates t d j) (incoming d j)def ssm_layer (a : Rat) : Mixer seq hidden := fun input => let ⟨⟨q, k⟩, v⟩ := input let updates := fun t => (fun d j => k t d * v t j : Matrix hidden hidden) fun s => (ssm_state a (slice 0 (s.val + 1) (α:Type uβ:Type vδ:Type wseq:Nathidden:Nata:Ratinput:(Matrix seq hidden × Matrix seq hidden) × Matrix seq hiddenq:Matrix seq hiddenk:Matrix seq hiddenv:Matrix seq hiddenupdates:Fin seq → Fin hidden → Fin hidden → Rat := fun t d j => k t d * v t js:Fin seq⊢ 0 + (↑s + 1) ≤ seq All goals completed! 🐙) updates)).transpose.matvec (q s)a:Ratn:Natih:∀ (values : Vector n) (z : Rat), ssm_scan a values z = a ^ n * z + Vector.sum fun t => a ^ (n - 1 - ↑t) * values tvalues:Vector (n + 1)z:Rat⊢ (a ^ n * (a * z + values 0) + Vector.sum fun t => a ^ (n - 1 - ↑t) * values t.succ) = a ^ (n + 1) * z + (a ^ (n + 1 - 1 - ↑0) * values 0 + Vector.sum fun i => a ^ (n + 1 - 1 - ↑i.succ) * values i.succ) All goals completed! 🐙

Unlike vanilla attention, state space models clearly induce an ordering on the sequence in the multiplicative term a. However, in the special case where that term is 1 and we use a bidirectional SSM, we can show that the permutation equivariance remains.

def bidirectional_ssm_scan (a : Rat) (values : Vector n) : Vector n := fun s => let suffix := slice s.val (n - s.val) (α:Type uβ:Type vδ:Type wn:Nata:Ratvalues:Vector ns:Fin n⊢ ↑s + (n - ↑s) ≤ n All goals completed! 🐙) values ssm_scan a (slice 0 (s.val + 1) (α:Type uβ:Type vδ:Type wn:Nata:Ratvalues:Vector ns:Fin nsuffix:Fin (n - ↑s) → Rat := slice (↑s) (n - ↑s) ⋯ values⊢ 0 + (↑s + 1) ≤ n All goals completed! 🐙) values) + ssm_scan a (fun t => suffix t.rev)def bidirectional_ssm_layer (α : Rat) : Mixer seq hidden := fun input => let ⟨⟨q, k⟩, v⟩ := input fun s => let state : Matrix hidden hidden := fun d j => bidirectional_ssm_scan α (fun t => k t d * v t j) s state.transpose.matvec (q s)n:Natf:Fin n → Rats:Fin none_pow:∀ (m : Nat), 1 ^ m = 1split:Vector.sum (select (fun i => Fin.castAdd (n - ↑s) i) fun t => f ⟨↑t, ⋯⟩) + Vector.sum (select (fun i => Fin.natAdd (↑s) i) fun t => f ⟨↑t, ⋯⟩) = Vector.sum ftotal:(Vector.sum fun t => f ⟨↑t, ⋯⟩) = Vector.sum fsplit':Vector.sum (slice 0 ↑s ⋯ f) + Vector.sum (slice (↑s) (n - ↑s) ⋯ f) = Vector.sum f⊢ Vector.sum (slice 0 ↑s ⋯ f) + f s + Vector.sum (slice (↑s) (n - ↑s) ⋯ f) = Vector.sum f + f s All goals completed! 🐙seq:Nathidden:Natπ:PositionPermutation seqv:Matrix seq hiddenq:Matrix seq hiddenk:Matrix seq hiddens:Fin seqj:Fin hiddend:Fin hidden⊢ ((Vector.sum fun t => k (π.index t) d * v (π.index t) j) + k (π.index s) d * v (π.index s) j) * q (π.index s) d = ((Vector.sum fun t => k t d * v t j) + k (π.index s) d * v (π.index s) j) * q (π.index s) d All goals completed! 🐙

These models also have the property that we can chunk them into groups which can be computed separately. Here we consider a simplified version of chunking in order to distribute across machines.

Equal-size chunks pass a carry between boundaries; each position combines the decayed incoming state with its local weighted updates.

def ssm_chunk (a : Rat) (values : Vector n) (incoming : Rat) : Rat := a ^ n * incoming + Vector.sum (fun t => a ^ (n - 1 - t.val) * values t)def ssm_chunk_carry (tiles n : Nat) (a : Rat) (values : Vector (tiles * n)) : Rat := scan (fun state values => ssm_chunk a values state) (tile n values) 0def chunkwise_ssm_layer (tiles n : Nat) (a : Rat) : Mixer (tiles * n) hidden := fun input => let ⟨⟨q, k⟩, v⟩ := input fun s => let start := s.val / n * n let state : Matrix hidden hidden := fun d j => let updates : Vector (tiles * n) := fun t => k t d * v t j let incoming := ssm_chunk_carry (s.val / n) n a (slice 0 start (α:Type uβ:Type vδ:Type whidden:Nattiles:Natn:Nata:Ratinput:(Matrix (tiles * n) hidden × Matrix (tiles * n) hidden) × Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddenv:Matrix (tiles * n) hiddens:Fin (tiles * n)start:Nat := ↑s / n * nd:Fin hiddenj:Fin hiddenupdates:Vector (tiles * n) := fun t => k t d * v t j⊢ 0 + start ≤ tiles * n α:Type uβ:Type vδ:Type whidden:Nattiles:Natn:Nata:Ratinput:(Matrix (tiles * n) hidden × Matrix (tiles * n) hidden) × Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddenv:Matrix (tiles * n) hiddens:Fin (tiles * n)start:Nat := ↑s / n * nd:Fin hiddenj:Fin hiddenupdates:Vector (tiles * n) := fun t => k t d * v t jthis:↑s / n * n + ↑s % n = ↑s⊢ 0 + start ≤ tiles * n All goals completed! 🐙) updates) ssm_chunk a (slice start (s.val % n + 1) (α:Type uβ:Type vδ:Type whidden:Nattiles:Natn:Nata:Ratinput:(Matrix (tiles * n) hidden × Matrix (tiles * n) hidden) × Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddenv:Matrix (tiles * n) hiddens:Fin (tiles * n)start:Nat := ↑s / n * nd:Fin hiddenj:Fin hiddenupdates:Vector (tiles * n) := fun t => k t d * v t jincoming:Rat := ssm_chunk_carry (↑s / n) n a (slice 0 start ⋯ updates)⊢ start + (↑s % n + 1) ≤ tiles * n α:Type uβ:Type vδ:Type whidden:Nattiles:Natn:Nata:Ratinput:(Matrix (tiles * n) hidden × Matrix (tiles * n) hidden) × Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddenv:Matrix (tiles * n) hiddens:Fin (tiles * n)start:Nat := ↑s / n * nd:Fin hiddenj:Fin hiddenupdates:Vector (tiles * n) := fun t => k t d * v t jincoming:Rat := ssm_chunk_carry (↑s / n) n a (slice 0 start ⋯ updates)this:↑s / n * n + ↑s % n = ↑s⊢ start + (↑s % n + 1) ≤ tiles * n All goals completed! 🐙) updates) incoming state.transpose.matvec (q s)

We can prove that this yields the same result as our simple implementation.

theorem ssm_chunk_carry_eq (tiles n : Nat) (a : Rat) (values : Vector (tiles * n)) : ssm_chunk_carry tiles n a values = ssm_scan a values := tiles:Natn:Nata:Ratvalues:Vector (tiles * n)⊢ ssm_chunk_carry tiles n a values = ssm_scan a values tiles:Natn:Nata:Ratvalues:Vector (tiles * n)⊢ scan (fun state values => scan (fun state x => a * state + x) values state) (tile n values) 0 = scan (fun state x => a * state + x) values 0 All goals completed! 🐙hidden:Nattiles:Natn:Nata:Ratv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hiddend:Fin hiddenupdates:Vector (tiles * n) := fun t => k t d * v t jbound:↑s / n * n + (↑s % n + 1) ≤ tiles * noffset:↑s / n * n + (↑s % n + 1) = ↑s + 1⊢ scan (fun state x => a * state + x) (slice 0 (↑s / n * n + (↑s % n + 1)) ⋯ updates) 0 = ssm_scan a (slice 0 (↑s + 1) ⋯ updates) exact scan_full _ (slice 0 (s.val + 1) (hidden:Nattiles:Natn:Nata:Ratv:Matrix (tiles * n) hiddenq:Matrix (tiles * n) hiddenk:Matrix (tiles * n) hiddens:Fin (tiles * n)j:Fin hiddend:Fin hiddenupdates:Vector (tiles * n) := fun t => k t d * v t jbound:↑s / n * n + (↑s % n + 1) ≤ tiles * noffset:↑s / n * n + (↑s % n + 1) = ↑s + 1⊢ 0 + (↑s + 1) ≤ tiles * n All goals completed! 🐙) updates) 0 offset

Conclusion

This blog considers the use of Lean as a method to verify elementary properties about Transformers and related models. There are many additional things one might consider here, including bounding errors introduced from quantization, aggregation, backpropagation, and training-inference mismatch. Additionally, there are likely many ways to simplify these proofs or develop libraries to make them more minimal.

At a high level, the main recent change is not the machinery for proving these properties; it is the ease with which a person can specify "what" they want to be proven and receive a certificate that a property is true. Understanding how that interface should work and how it can be used will be an extremely interesting challenge over the next year.