I want verification to become useful in the engineering of the next generation of AI. An agent proposes a faster kernel or a cheaper training method, checks the properties it needs to preserve, and measures whether the change is worth using. The next agent can build on both the implementation and the argument behind it. That’s the research loop I’d like to help build.
We don’t have a satisfying theory of why deep learning works yet, but not having one doesn’t mean we lack useful engineering knowledge. I want TorchLean to help codify that knowledge and make implementations more dependable. We don’t need to wait for a complete theory of learning to make that worthwhile.
By early 2024, while leading our lab’s AI4Math effort at Caltech through LeanDojo, I expected we would soon have more AI-generated mathematics and Lean code than people could read carefully. Much of my PhD work now concerns what happens when we bring that approach to ML. We rely on approximations, experiments, and arguments that are often much less precise than the code we run.
I’ve always thought mathematics would be one of the first things AI became good at. The abstraction and creativity are difficult, of course. But once someone has found an argument, we can often make the statement precise enough to check it. ML gives us a harder mixture of exact computations, approximations, and experimental claims. We can still start with the code we have and the properties we can state precisely.
That’s why I’ve been building foundations for verifying training code, its arithmetic, and the optimizations used to run it. Some of the checks we can do today are mundane: whether a data pipeline keeps the right tokens, or whether a calculation matches its stated definition. Connecting those checks through a compiler to the code running on a GPU is a much larger ambition. I care about both, and I want to be clear about the distance between them.
Verification in mathematics
If I ask one AI to check another AI’s argument, both might make the same mistake. A formal proof gives me another way to check it. The model supplies an argument, and a checker establishes whether it follows from the definitions and assumptions. The checker doesn’t have to share the model’s confidence.
This is what interests me about AlphaProof. It learns to search for proofs through reinforcement learning in Lean, with formal checking as part of the feedback. Finding a proof and checking it are different jobs. We can try approaches that fail frequently and still retain the checked successes. I’d like to use that separation when agents improve ML software too.
We still have to choose the right statement. A missing assumption or a different definition can change the claim even when the proof checks. For ML, that means deciding what an optimization promises: the same outputs, the same gradient over real numbers, a numerical error bound, or a different algorithm with better measured performance. Those are different claims, and each needs its own evidence.
Large-scale training
Suppose an agent finds a cheaper way to train a transformer. It changes the precision, recomputes activations, and rearranges communication across GPUs. The benchmark improves. Before using it for the next long run, I’d want to ask: are we still training the model we meant to train?
Even the loss gives us something to check. Suppose one worker processes 32 examples and another processes 64, all with equal intended weight. Averaging their local mean gradients gives each worker half the weight; the global example mean needs weights of one third and two thirds. PyTorch’s distributed data-parallel implementation averages gradients across processes. The training code has to account for how the local losses were normalized. The program can run and its loss can decrease while it optimizes a differently weighted objective.
For one coordinate, suppose those workers return mean gradients of 1 and 3. Averaging the workers gives 2. Averaging the examples gives \(7/3\). In the first calculation, each example on the smaller worker has twice the influence of an example on the larger one. Neither result needs to be NaN, crash, or look suspicious in a loss plot. We have to say which objective we intended before we can decide whether the implementation is right.
I’d like that requirement to survive the next optimization. If we change the number of workers, split a batch into smaller pieces, or move the averaging into a fused operation, the same statement about example weights should still apply. A test on equally sized batches won’t distinguish the two formulas. A general proof can force us to account for the unequal case, provided we have stated the intended weighting correctly.
Recomputation has its own requirements. When we regenerate an activation for backward, we need the value that the derivative calculation was supposed to use. Randomness matters here: PyTorch’s checkpointing documentation discusses preserving random-generator state when replaying dropout. A schedule that saves memory but replays a different dropout mask can change the gradient. I’d want an agent searching over schedules to carry that requirement into its search.
Mixed precision can interact with gradient accumulation just as directly. PyTorch’s guidance requires a constant loss scale while accumulating an effective batch. If we change that scale partway through, the buffer combines gradients scaled by different factors, and dividing by the final scale won’t recover the intended sum. This is a useful property to establish before comparing the speed of two accumulation implementations.
These are ordinary engineering requirements. Our own training example had a dialogue-packing error: arbitrary windows could begin inside an assistant answer or cross into another conversation. We repaired the pipeline and proved properties about the input, target, and mask contents produced by the pure packing functions. That doesn’t verify the whole training run, but it checks a requirement we had already got wrong. I want more work like this, where a proof answers a question that matters to the implementation.
The value depends on how often we change this code and what a mistake costs. A property proved for a general reduction or packing routine could cover many uses of it. We would still need tests, profiling, and experiments with the trained model. I’d like proofs to take some recurring correctness questions out of that workload.
Full-stack verification
The longer-term goal is to connect these results from the mathematical training rule down to machine instructions. Start with a loss and an update rule. Establish that the backward program computes the required derivative under stated assumptions. Account for the errors introduced by finite precision. Then justify the transformations that distribute, fuse, and compile that calculation for the hardware.
A theorem at one level doesn’t cover the others automatically. Proving a derivative over the real numbers doesn’t establish the accuracy of a floating-point kernel. A proof about a tensor expression doesn’t establish that an importer translated the original Python correctly, or that the generated GPU instructions implement that expression. TorchLean doesn’t yet give us an end-to-end proof of this stack.
I’d find the local results much more useful if we could connect them. If an optimizer theorem permits bounded gradient error, a numerical analysis of the kernel could establish that premise. A compiler theorem could then carry the kernel’s guarantee down to assembly under a specified instruction set and memory model. The error bounds have to compose too: later calculations may amplify an earlier error.
Take one SGD step as a small example. We start from parameters \(\theta\), use a gradient \(g\), and choose a nonnegative learning rate \(\eta\). An optimized implementation supplies an approximate gradient \(\widehat g\). Over real arithmetic, the two updates are
\[ \theta^+ = \theta - \eta g, \qquad \widehat\theta^+ = \theta - \eta\widehat g. \]
If we establish that the gradient error has norm at most \(\delta\), then the updated parameters differ by at most \(\eta\delta\). If evaluating the update in floating point introduces an additional error of norm at most \(\rho\), the resulting one-step bound is \(\eta\delta + \rho\). This calculation assumes the same starting parameters and learning rate; it says nothing yet about later steps or the final model’s accuracy.
But it gives the pieces of the implementation something precise to contribute. The backward calculation needs an argument about the gradient it computes. The numerical analysis needs to bound its error. The parameter update needs to account for its own rounding. Instead of treating “the optimizer is correct” as one enormous claim, we can ask whether these particular statements connect. For many steps, we would also need to understand how changing the parameters changes the next gradient. An optimizer with momentum brings additional state into that argument.
This is where I hope the libraries earn their cost. An engineer trying a different precision shouldn’t have to reprove the elementary update argument. They should be able to concentrate on the new error bound and whether it is small enough for their purpose. Establishing that bound for a useful kernel is substantial work; writing the inequality alone doesn’t do it.
I don’t expect us to check every multiplication in a huge run by reducing it inside a proof assistant. I’d like general results about operations and transformations, with a way to establish that a particular program meets their conditions. Tensor size still affects the cost, but a proof about gradient reduction should apply to more than one tensor length. A useful compiler would reuse those results as it transforms the program.
AI researching AI
AI is already helping build AI. Google DeepMind’s AlphaEvolve proposes and evaluates program changes, including optimizations to kernels used in model training. I’d like to extend that kind of search so that an agent can propose a change, prove the properties we require, and benchmark the result.
By recursive self-improvement, I mean a loop in which AI helps improve the systems that produce its own capabilities. Faster code can support more experiments; better experiments can produce models that are more useful at improving the code. I don’t know how far that loop will go. But the engineering problem is already interesting before we get to the most ambitious version of it.
Suppose we ask an agent to reduce the memory used by a training step. It proposes a recomputation schedule and writes the implementation. Using a fixed reference calculation, it tries to prove that the schedule supplies the same values to backward and preserves the random state needed for replay. It might reuse a theorem for most of the schedule and need a new argument for one operation. If it finds a counterexample, that becomes feedback for its next proposal. Failing to find a proof would leave the claim unresolved.
For a candidate that meets the requirement, we measure peak memory and step time on the workloads we care about. A proof of equivalence doesn’t tell us whether the extra recomputation makes the code too slow. The agent uses both the correctness result and the measurements to choose its next attempt. We keep the implementation, proof, and benchmark setup together so a later agent can extend the work.
I don’t want every new algorithm to prove that it behaves like the old one. I’d like us to say what may change, what must remain true, and what evidence we need for the new idea.
A new architecture needs new experiments and may call for different properties. An approximate kernel might earn its place through a useful error bound and better measured training performance. The agent should be able to pursue that possibility, while keeping the claim about the approximation explicit.
Alignment and safety
If I give a research agent access to code, data, and compute, I want to be clear about what it’s allowed to change. We can formalize some of those limits and prove that the software around the model enforces them. We still need to choose the goals and limits ourselves.
For example, an agent could propose a code change, but a separate component would require approval before deployment. I’d rather enforce that rule in software than rely on a prompt telling the model to ask first. We could then prove that the component rejects deployment requests without the required approval, under the assumptions we make about its environment.
In proof-carrying code, whoever supplies a program also supplies evidence that it meets the recipient’s requirements. The recipient checks that evidence. I’d like AI-generated software to come with that kind of argument. Guaranteed Safe AI explores a broader approach using models of the world, safety requirements, and verifiers. Writing requirements that cover the situations we care about is still difficult.
In the research loop above, the agent shouldn’t be able to improve its score by changing the reference calculation or removing a difficult benchmark. Changes to the requirements and evaluator need their own review. Proving that a new evaluator implements its definition doesn’t settle whether we should adopt that definition. Similarly, proving a property of the training code doesn’t establish that the trained model will behave safely.
The code that checks a proposal matters too. Lean’s guidance describes sandboxed builds, checking which axioms a proof uses, and validating proofs independently. I’d want the agent to be unable to change its checker or relax the requirements on its own. If a required check hasn’t passed, the proposed change should wait.
Shared foundations
Much of the work is in the libraries. Before proving a property of a model, we need definitions of its operations and the arithmetic they use. If those are missing, even a small claim can turn into a large project.
For ML, there is a lot we should be able to share. Different models repeatedly use tensor contractions, reductions, normalization, differentiation, and state updates. I’d rather prove a general result about one of these and reuse it than build a separate formalization for every architecture. An agent improving a model should be able to use the numerical analysis and probability it needs without rebuilding those subjects inside the model project.
FloatLib is my contribution on the arithmetic side: executable floating-point operations, specifications, proofs connecting certified implementations to those specifications, and rounding and error results. With TorchLean, we’re bringing tensor computations, models, differentiation, and verification tools into Lean. We’re also developing ways to import supported computations from existing ML code. We still need to justify the connections from the source program to an imported graph, and from the formal arithmetic to an external kernel used for execution.
For a concrete starting point, our MNIST walkthrough keeps the training code in JAX. We train a small two-layer classifier on 60,000 handwritten digits, with float32 master parameters and bfloat16 rounding in its forward tensors. After three epochs, that checkpoint classified 9,351 of the separate 10,000 test images correctly. We also evaluate the same weights with float16 and float32 forward tensors, then capture the computation for one test image and replay it through Lean’s reference arithmetic.
On that image, both JAX and the reference predicted 7. The float16 and bfloat16 output words agreed exactly. The float32 outputs differed by at most \(2^{-21}\) per logit. That comparison doesn’t tell us that an implementation is broken: different choices about accumulating a matrix product can produce different rounded results. It gives us a concrete numerical difference to investigate. The training accuracy is an experimental result, and the replay checks this input; neither establishes correctness of the whole training procedure.
A small example shows why those connections matter. Suppose a reference calculation has a unique highest-scoring token, with margin \(m\) over every competitor. If an implementation changes each score by at most \(\varepsilon\), then
\[ m > 2\varepsilon \]
is sufficient to preserve the winner: its score can fall by at most \(\varepsilon\), while a competitor’s can rise by at most \(\varepsilon\). Our batch-invariance example works through this connection between arithmetic error and a model decision. Establishing the error bound for the computation remains the hard part.
For an engineer considering a cheaper numerical implementation, this is a useful question to ask. We may not need every output bit to match, but we might need the selected class or token to remain unchanged. The MNIST comparison supplies the actual difference for one pair of output tensors. To cover a whole input region, we would need an error bound valid throughout that region, together with the required margins. The general decision theorem can be reused; the evidence about the particular model and implementation still has to be supplied.
Looking ahead
Suppose an AI gives me a training method that works better than anything I could have designed. I’d want to use it. For a claim that admits a proof, I’d want a separately checkable argument; for an empirical claim, the measurements, methods, and uncertainty. I’d like the statements, definitions, and tools to be available alongside the code so someone else can question or extend the result.
My ambition is to work toward formalizing all of ML: the mathematics, learning algorithms, architectures, numerical methods, and software that runs them. Proof search is expensive, bounds can be too loose, and good definitions take work. I don’t want a complete proof to be a prerequisite for trying an idea. The near-term test is whether verification helps us make a useful engineering decision: catch a silent objective change, justify a memory-saving schedule, or accept a faster kernel with an error bound we can use.
By the end of next year, 2027, I believe we’ll see in AI research what we’re already seeing in mathematics: AI doing more of the research, with verification helping us check what it produces. At least on the software side; hardware may take a little longer :)
I’d judge this work by whether developing and maintaining the proofs costs less than the debugging, repeated review, and wasted runs they save. I want an ML engineer to be able to change an implementation more confidently and spend more time on the next experiment.
I use AI assistance with some of the writing and Lean code. TorchLean Verified Examples
Cite this essay
Robert Joseph George. “Verification in the age of RSI.” 2026.
@misc{george2026verification,
author = {Robert Joseph George},
title = {Verification in the age of {RSI}},
year = {2026},
url = {https://www.robertj1.com/ai4science/vision/}
}