Checking a small transformer in Lean

Robert Joseph George · TorchLean Verified Examples

Neel Somani's verifiable transformer has a task I really like: match a quote or close a bracket. There are only 256 possible prompts, so we can examine every one and check which delimiter the model prefers. Neel designs the calculations as circuits for a satisfiability modulo theories (SMT) solver, a tool that reasons about logical and arithmetic constraints. I wanted to bring the checkpoint and its evidence into Lean.

I imported the checkpoint and its exported scores, then added a forward pass using those weights. Now we can compare the saved decisions with the ones we recompute in Lean. We also check that the export contains the right prompts: repeating an easy case in place of a difficult one should not count as checking the whole task.

There are two checks to follow here. Lean's kernel verifies that the exported integer-score records cover the task and give every intended candidate a positive margin. Separately, a compiled Lean program recomputes the floating-point forward pass from the imported weights. The upstream solver handles its circuit queries; our trace proofs don't certify that solver or prove the rounded forward pass correct. We'll keep those results separate as we go.

The model has two transformer layers, one attention head, hidden width 16, and a 32-token vocabulary. Its input is six tokens long. The claim concerns a choice between two candidates for each prompt: the intended quote or bracket must score higher than the other choice. That pairwise comparison is the property checked throughout the experiment; it does not rank the intended token against the other 30 vocabulary entries.

The small input domain is what makes this example useful. For a large language model, even short prompts give us far too many possibilities to enumerate. Here we fix the length, content vocabulary, and delimiter positions, then check every combination. Our result is about that defined task; adding longer prompts or nested brackets would give us a new task to check.

We'll follow one prompt through the task definition, the exported record, and the forward calculation. Along the way, I'll explain why the model uses different attention and normalization operations from GPT-2, and why checking a saved score is different from recomputing it from the weights.

A task with 256 inputs

A quote prompt has the form [BOS, quote-task, x₁, quote, x₂, x₃], where BOS is the beginning-of-sequence token and quote-task selects the task. Each content token xᵢ is one of IDs 5, 6, 7, or 8. The quote is ID 9 or 10, and the model should prefer that same quote at the prediction position. A bracket prompt has the same layout with a bracket-task token and an opening bracket. Opening IDs 11 and 12 should lead to closing IDs 13 and 14 respectively.

Three content positions with four choices and one delimiter with two choices give 128 prompts per task:

Take the quote prompt [1, 2, 5, 9, 5, 5]. Token 1 starts the sequence; token 2 selects the quote task; the three 5s occupy content positions. The delimiter is 9, so the intended prediction is 9 and the alternative is 10. Changing only the delimiter to 10 exchanges those roles. For the bracket task, [1, 3, 5, 11, 5, 5] instead asks for closing token 13, compared with token 14. The IDs are defined by this synthetic vocabulary, not by a natural language tokenizer.

The content positions vary independently. We cannot replace their 64 combinations with four examples in which all three content tokens happen to agree. Nor can we check one delimiter and assume the other behaves symmetrically: the learned embeddings and output weights are allowed to treat their IDs differently. We enumerate every combination, then take two coordinates of each prompt's 32 raw scores, or logits, and ask for a strict inequality:

We do not need softmax on those final scores to compare the candidates: softmax preserves their order over real numbers. We also do not need the two probabilities to sum to one when viewed apart from the other 30 tokens. If a third token has a higher score than either candidate, the pairwise property can still pass. An unrestricted greedy-generation claim would therefore need a stronger check than the one used here.

The saved upstream evaluation makes this distinction concrete. Its bracket task has candidate accuracy 1.0 but full-vocabulary accuracy 0.0. In other words, 100% on the two-candidate task and 0% when choosing from the full vocabulary. The correct closing bracket wins against the other closing bracket, while an unrestricted choice among all 32 tokens selects something else. The upstream evaluation code computes both decisions separately. A decoder restricted to the two task candidates uses the property we're checking; a general text generator needs the stronger vocabulary-wide property.

Circuits, sparse attention, and normalization

Neel describes the design as one whose:

forward pass is built from SMT-representable primitives

That choice reaches into the model itself. We have to follow the sparse attention, normalization, and activation he used, rather than substitute the operators from a standard GPT implementation. Let's look at those choices before interpreting a solver result.

A circuit here is a selected computation inside the model: the retained dependencies and operations used for a task decision. It isn't a physical logic chip. Neel's formal definitions separate agreement with the reference program from questions about whether retained connections are necessary or irrelevant content can change the decision.

For our pairwise task, a counterexample query would ask for an allowed prompt whose intended candidate fails to beat its alternative:

If a solver proves that query unsatisfiable, it excludes such a prompt under the encoded semantics. That still leaves the connection from the encoding to the actual checkpoint execution to establish. Our Lean trace check instead verifies the exported records, and the replay recomputes the small checkpoint's forward pass. This article follows that small-model artifact; it doesn't import all the circuit properties or later GPT-2-scale results now described upstream.

Neel's model changes three familiar transformer operations to suit an SMT encoding. Sparsemax replaces softmax, SignedL1BandNorm replaces LayerNorm, and LeakyReLU replaces GELU. The upstream project supplies the model family, task, circuit extraction, and solver workflow.

An SMT solver checks formulas combining logical conditions with theories such as arithmetic. For this model, the upstream encoder describes the operators using explicit comparisons and arithmetic branches. A solver can then reason about the encoded task. The operator choices determine those formulas, so we'll examine them before checking the exported evidence.

Sparsemax and zero attention weights

Ordinary softmax exponentiates every allowed score and normalizes the results. For finite real scores, each resulting weight is positive. Sparsemax, introduced by Martins and Astudillo, instead subtracts a threshold and clips negative results to zero. The threshold is chosen so that the remaining weights sum to one:

For illustrative scores [1, 0.8, 0], a threshold of 0.4 gives [0.6, 0.4, 0]. The last position contributes nothing to the weighted value sum. With [2, 1, 0], a threshold of 1 leaves [1, 0, 0]. These examples use exact arithmetic; the replay also has to specify how it computes the threshold in floating point.

To find the threshold, the implementation sorts scores from largest to smallest and scans prefix sums. For a proposed group of k active entries, subtracting the same threshold from all of them gives total mass one when τ = (sum of the top k scores − 1)/k. The scan checks which entries remain above that threshold. Once the active entries are fixed, their weights are affine expressions in the scores: weighted sums plus constants. Comparisons determine the branches, which is useful for the upstream solver encoding.

SignedL1BandNorm, including the zero vector

SignedL1BandNorm first subtracts the vector's mean. It splits the centered coordinates into positive magnitudes and negative magnitudes, then treats those two nonnegative vectors separately. Their L1 masses are the sums of their coordinates. If a mass is too high, a projection reduces it; if it is too low, the implementation distributes the deficit over the positive coordinates, called its active support. An empty support needs an explicit fallback. The replay then recombines the signed parts, centers again, and applies the learned scale and shift.

We can see what the fallback does by giving the layer a constant vector. Subtracting its mean leaves 16 zeros, so both signed halves have empty support. The exported lower mass target is approximately 4.4 for each half. The positive fallback activates the eight even-numbered coordinates; the negative fallback activates the eight odd-numbered ones. Each active coordinate receives approximately 4.4/8 = 0.55. Before the learned scale and shift, recombination therefore produces alternating +0.55 and −0.55, rather than a zero vector. This is a deliberate part of the upstream normalization, and the Lean replay follows it. These decimal calculations describe the branch; the executable still rounds them in floating point.

Those branches matter even for a hidden vector of length 16. Replacing the operation with LayerNorm would introduce a different denominator and different outputs. Omitting the empty-support fallback would also change the computation on inputs for which an entire half is zero. The replay spells out those cases instead of leaving normalization as a name.

LeakyReLU is easier to describe. Positive inputs pass through unchanged; negative inputs are multiplied by a fixed slope. It contributes another explicit branch without GELU's smooth nonlinear calculation. “Solver-friendly” here means these operations have useful branchwise encodings, not that every possible transformer built from them becomes easy to verify.

Those details belong in the claim. Substituting softmax or LayerNorm would change the model being checked, even if the tensor dimensions stayed the same. The Lean export summary therefore checks the operator tags as well as the dimensions, tensor inventory, and parameter count. The small checkpoint contains 7,712 scalar parameters across 37 tensors.

What the exported evidence proves

The export produces several useful objects. A summary describes the architecture and tensor layout. A finite trace records each prompt's target and alternative scores. Circuit summaries record the results of the upstream task-local extraction. The full parameter payload lets me recompute the forward pass.

Checking only the row count would miss duplicates. A file could repeat one prompt and omit another while retaining the right length. In FiniteEval.lean, the checker derives the expected prompt from the row index. The first 128 rows enumerate quotes; the next 128 enumerate brackets. At each index, the task, all six input IDs, target, and alternative must match.

The scores are integers in micro-logit units. The certificate must declare a scale of one million, and every target score must strictly exceed its alternative score. The generated summary and trace have Lean theorems showing that their checkers accept them.

The first quote row is the example above. The second changes the quote from 9 to 10; the third returns to quote 9 and advances the final content position from 5 to 6. The enumeration continues through every content combination, then begins the bracket task at global index 128. Dividing the index by two removes the delimiter choice; the remaining base-four digits select the three content tokens.

That ordering gives the checker an independently computed expectation at every row. The following excerpt is the row-identity check in FiniteEval.lean:

/-- Row `idx` must be exactly the corresponding prompt in the finite domain. -/
def CandidateEval.matchesExpectedIndex (idx : Nat) (row : CandidateEval) : Bool :=
  row.task == expectedTask idx &&
  row.input == expectedInput idx &&
  row.target == expectedTarget idx &&
  row.alternate == expectedAlternate idx

The double equals signs return Boolean comparisons, and && requires all four to pass. expectedInput computes the canonical prompt; it does not read the claimed prompt back from the evidence. Length 256 and successful identity checks at all indices therefore exclude a trace that quietly repeats an easy case in place of a difficult one.

We can inspect the two prompts that begin the quote and bracket tasks. The commands below also show the required target for each. Hover over matchesExpectedIndex to see the row type it checks.

import VerifiableTransformers.Certificate.FiniteEval

open VerifiableTransformers.Certificate.FiniteEval

-- The first quote prompt and the first bracket prompt in the finite task.
#eval expectedInput 0
#eval expectedTarget 0
#eval expectedInput 128
#eval expectedTarget 128
#check CandidateEval.matchesExpectedIndex

[1, 2, 5, 9, 5, 5]

9

[1, 3, 5, 11, 5, 5]

13

VerifiableTransformers.Certificate.FiniteEval.CandidateEval.matchesExpectedIndex (idx : ℕ) (row : CandidateEval) : Bool

The score check comes afterward. Integer micro-logits let Lean compare saved margins without executing a floating-point subtraction. A score of one million denotes one unit under the declared scale. Acceptance proves a positive difference between the recorded integer scores. It does not prove that exporting and rounding the scores preserved a particular native margin; the independent floating-point replay is useful for that reason.

The fingerprint deserves the same care. The export scripts compute a SHA-256 hash. Lean compares the identifier recorded in the trace with the one recorded in the summary. It does not hash the checkpoint bytes in this development. Those comparisons help keep artifacts from different exports apart, provided the exporter attached the right identifiers. To authenticate arbitrary newly supplied files, we would also need to connect their actual bytes to those identifiers.

A different prompt order, wrong target, missing case, changed recorded fingerprint, or nonpositive candidate margin fails the corresponding check. The circuit-summary checks are small checks of the saved summaries; the upstream SMT solver remains responsible for producing its solver evidence.

These checks let us inspect the inputs behind an exported accuracy number. We know the canonical prompt and token IDs occur at each index, and we know the recorded margins are positive. We still need the separate connection from those scores to the checkpoint.

If we rewrite the Python exporter, the new evidence still has to pass the same Lean checks for the domain, targets, and candidate margins. Adding another prompt requires changing the expected enumeration too; changing a label in the export cannot enlarge the theorem. These checks establish what the saved records say about the task. To see whether the weights produce those decisions, let's run the model itself.

Recomputing the scores from the weights

The forward replay starts from the imported parameter arrays. It adds token and position embeddings, runs both residual attention and feed-forward blocks, applies the final normalization, and computes the two candidate logits from the last hidden state.

The replay specifies its arithmetic directly. Dot products use a left fold in coordinate order; the custom normalization and sparsemax threshold scans are plain Lean Float functions. We can inspect each numerical operation without a Python or GPU call in the replay itself. The payload's vector and matrix shapes are checked separately in Lean.

For one prompt, token IDs select six embedding vectors of width 16. Position embeddings are added before the first block. At each position, the block normalizes the hidden vector and computes its query, key, and value projections. Position five may attend to positions zero through five; position zero has only itself available. Sparsemax is applied to each position's allowed score vector, and its weights combine the corresponding values.

The attention output is projected and added back to the hidden sequence. The feed-forward part applies another normalization, expands each width-16 vector to width 64, applies LeakyReLU, and projects back to width 16 before another residual addition. The second block repeats this calculation with its own weights. Finally, a normalization and language-model head turn the last position's hidden vector into candidate scores. Earlier positions affect that decision through attention, even though we inspect the final position's logits.

The stored matrices are row-major. One affine output coordinate takes a weight row, forms its coordinate-order dot product with the input vector, and adds the bias. The exported inventory distinguishes the attention projections, feed-forward projections, embeddings, normalization parameters, and language-model head. Counting 7,712 numbers is a useful consistency check, but shapes and names are necessary too: the right number of values could still be arranged in the wrong matrix.

The recorded replay passes all 128 quote prompts and all 128 bracket prompts. Its smallest target-minus-alternative margin is 0.019695. A nonpositive margin makes the executable fail. These are results of running the compiled Lean Float program, alongside the kernel-checked shape and integer-trace results. There is no theorem here identifying that floating-point run with the rational SMT computation or every Python and CUDA execution.

The integer-trace theorem establishes a fact about the saved integer trace. The replay executable instead runs floating-point operations on the imported values and reports its decisions. That second route reduces our dependence on the supplied score columns, but it still uses the compiled program and Lean's native Float operations. It does not turn a test of 256 compiled forward passes into a kernel proof of those floating-point computations. Btw, “kernel” here means Lean's proof checker, not a CUDA kernel.

There is also a rational specification for the solver-facing operators in UpstreamSmallGPT.lean. Exact rational arithmetic avoids rounding, whereas the replay rounds its operations. To transfer a rational decision theorem to the replay, we would need an error argument strong enough to preserve each candidate margin, or a direct proof about the rounded computation. The minimum recorded margin tells us what such an argument would have to protect; it does not supply the bound.

We also retained an archived CUDA trace from our conventional causal transformer with softmax, LayerNorm, and GELU. It uses the same finite task and dimensions, and its certificate passes the same domain checker. It has its own weights and fingerprint, so it is checked as a separate model.

The two models make a useful comparison for this task, but their agreement on the delimiter choice would not establish equality of their attention weights or hidden vectors. Sparsemax and softmax can produce different mixtures while both lead to the right candidate. Keeping their evidence separate lets us ask the same task question without pretending they are the same computation.

If we extend the task, the input definition has to change before the evidence does. A seventh token, another content ID, or nested brackets creates cases outside these 256 prompts. We would need a new enumeration and new scores, and the old accepted trace would no longer cover the claim. For this experiment, the complete task definition is small enough to keep beside the checkpoint checks and the replay source.

Conclusion

The next proof I want is the connection between the imported weights, their rounded forward calculation, and these decisions. The trace theorem checks the saved integer scores; the replay runs the floating-point model. FloatLib gives us tools to bound the rounding error between calculations and compare that bound with the winning margins. Here we have all 256 prompts and the checkpoint available, so we can attempt that connection on the entire defined task. I'd be happy to compare approaches with others working on verified transformers.

References

• Neel Somani. Verifiable Transformers. The checkpoint, synthetic task, custom normalization, and SMT encoding used here.
• André F. T. Martins and Ramón Fernandez Astudillo. From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification, 2016. Sparsemax and its threshold construction.
• Ashish Vaswani and colleagues. Attention Is All You Need, 2017. The attention calculation underlying the transformer blocks; Somani's model changes several of its operators.

@misc{george2026verifiable,
  author = {Robert Joseph George},
  title = {{Checking a small transformer in Lean}},
  year = {2026},
  howpublished = {TorchLean Verified Examples},
  url = {https://robertboy18.github.io/TorchLean-Verified-Examples/week-02-verifiable-transformer-checkpoint/}
}