6.3. Autograd Proofs
Most ML users learn gradients operationally: run the forward pass, call
loss.backward(), inspect a few numbers, and move on. That workflow is powerful, and PyTorch and
JAX made it practical at enormous scale. TorchLean keeps the same programmer intuition and puts an
explicit theorem underneath it.
The difference is not that TorchLean "has autograd" while PyTorch and JAX do not. They do, and they do it very well. The difference is that ordinary framework use usually relies on an implementation of reverse mode AD, its derivative registrations, compiler rewrites, and native kernels. TorchLean's proof tree asks a more explicit question:
For the supported graph fragment, does the reverse pass compute the adjoint of the derivative of
the forward denotation?
The proof architecture lives in NN/Proofs/Autograd. We built it in layers so the auditor can stop at the boundary they care about: algebraic tape soundness, Frechet derivative statements, proofs for particular operators, theorem entry points for model blocks, and algebra for training steps.
The proof pipeline is:
local op derivative -> local JVP/VJP adjointness -> graph backprop correctness -> Frechet derivative theorem -> model block theorem -> training step algebra
That ordering is the proof structure. A theorem about a Transformer sublayer is assembled from local derivative rules, graph composition, and the analytic bridge to the derivative of the denotation.
6.3.1. From Runtime Confidence To Theorem Obligations
In PyTorch, the default runtime model is approximately:
-
the eager engine records operations into a dynamic tape;
-
each operation has a backward rule, either built in or registered by extension code;
-
the engine traverses the tape in reverse and accumulates gradients;
-
tests, numerical checks, and framework maintenance give confidence that the result is right.
JAX moves the same idea into a functional transformation pipeline: grad, vjp, jvp, jit, and
lowering passes cooperate to produce differentiated and compiled programs. See the PyTorch autograd
overview at https://pytorch.org/docs/stable/autograd.html and JAX's autodiff guide at
https://jax.readthedocs.io/en/latest/automatic-differentiation.html for the user facing version of
that workflow.
TorchLean keeps the same operational picture: local derivative rules plus reverse accumulation. What changes is that runtime confidence is split into named theorem obligations.
-
A registered backward rule becomes a local JVP/VJP or Fréchet derivative lemma for the op.
-
A tape reversal becomes global tape soundness by induction over the graph.
-
A scalar loss training step becomes an explicit theorem about the loss seed and parameter update algebra.
-
A model block such as attention or an RNN cell becomes a packaged theorem with stated hypotheses.
The resulting style is slower to author, but easier to audit. We can say exactly which part is a Lean theorem and which part is a runtime, compiler, or finite precision agreement statement.
6.3.2. Tape Soundness: The Algebraic Core
The central algebraic file is NN.Proofs.Autograd.Tape.Algebra.Soundness API. It defines the small tape language used by the rest of the autograd proofs, together with the local proof data carried by each node.
The objects to track are:
-
TensorPack: typed tensor payloads indexed by a list of shapes; internally, the autograd algebra still has a small context datatype with the same shape-indexed structure. -
Idx: pointer into a context, carrying the needed proof data. -
NodeData: forward, JVP, and VJP data for one local operation. -
Node: a node plus the local inner product soundness law. -
GraphData: executable snoc list graph data. -
Graph: graph whose nodes can share earlier values and carry local proof obligations. -
Graph.backprop_correct: global dot product soundness theorem.
The theorem Graph.backprop_correct is the first big hinge. Informally, it says that reverse
accumulation is the adjoint of forward sensitivity:
\langle \operatorname{jvp}_G(x, dx), seed\rangle = \langle dx, \operatorname{backprop}_G(x, seed)\rangle
This dot product identity is the algebraic essence of reverse mode AD. A forward sensitivity dx
pushed through the graph and then paired with an output cotangent gives the same scalar as pairing
dx with the cotangent produced by reverse accumulation.
The dot product statement is the common language between implementation and analysis. Runtime
engineers recognize it as the VJP law. Mathematicians recognize it as the adjoint property. Lean can
prove it compositionally: each node gives the local adjoint law, and graph
soundness follows by induction over Graph.
The comparison to PyTorch is most direct here. PyTorch's engine performs a reverse walk over a dynamic graph and accumulates cotangents into inputs. TorchLean's algebraic graph does the same conceptual work, but the graph object carries enough structure for Lean to prove that the accumulation is sound for every input context in the supported fragment.
6.3.3. From Dot Products To Fréchet Derivatives
The algebraic theorem is not the last word. A dot product VJP law still has to be connected to the function being differentiated. The Fréchet derivative bridge provides that link.
That file vectorizes shaped tensor contexts into Euclidean spaces and connects three views:
-
shaped tensors:
Tensor Real s,TensorPack Real Gamma; -
flat Euclidean vectors:
CtxVec Gamma,flattenCtx,unflattenCtx; -
analytic derivatives:
HasFDerivAt,fderiv, andContinuousLinearMap.adjoint.
The main theorem is Graph.backpropVec_eq_adjoint_fderiv. In plain English:
If every node in the graph has the stated Fréchet derivative, then graph backprop equals the
adjoint of the Fréchet derivative of graph evaluation.
There is also a pointwise version, Graph.backpropVec_eq_adjoint_fderiv_at, for hypotheses that
only hold at a particular input. That distinction matters for neural networks. ReLU, normalization,
division, logarithms, and square roots all have domain or nondifferentiability issues. TorchLean
states those conditions explicitly instead of using a blanket "autograd works" slogan. The theorem
can demand exactly the local smoothness or nonzero hypotheses needed by the graph being
differentiated.
6.3.4. Connecting The Compiled Tape To The Derivative
The derivative theorem above is stated for the real analytic graph. The compiler correctness theorem was originally stated for a more general algebraic graph: its scalar type is abstract, and each node may read a non-differentiable environment. Those are useful abstractions, but leaving the two results side by side would not prove that the compiled tape computes the Fréchet derivative.
NN.Proofs.Autograd.Runtime.Link.FDeriv
closes that gap. At scalar type Real and environment Unit, the algebraic and analytic node types
convert in both directions. Both node and graph conversions round-trip, and evaluation, JVP, and
reverse accumulation commute with the conversion.
The algebraic reverse pass returns cotangents for the inputs and every intermediate value. The
analytic theorem needs only the input cotangent. TList.takeLeft selects that input prefix, and
takeLeft_backpropAllCtx proves that it is exactly the inputs-only reverse pass used by the
derivative theorem.
The two public endpoints can be inspected directly:
import NN.Proofs.Autograd.Runtime.Link.FDeriv #check Proofs.Autograd.Algebra.Graph.backwardDenseFrom_compileAux_adjoint_fderiv #check Proofs.Autograd.Algebra.Graph.backwardDenseFrom_compileAux_adjoint_fderiv_at
Suppose g is an algebraic graph over Real, x is its typed input context, d is its fixed
environment, and seed is an output cotangent. The first theorem returns a conjunction:
-
compiling
gand runningTape.backwardDenseFromsucceeds with the graph's full reverse context; -
the input prefix of that context, after flattening, is
(\operatorname{fderiv}\,\operatorname{eval}(x))^\dagger seed.
The _at theorem asks for differentiability only at x. It is the useful form for graphs
containing piecewise-smooth operators, provided the execution point avoids their non-differentiable
or invalid cases.
These are theorems about the exact tape instantiated over Real. A native Float or CUDA run needs
an additional numerical-refinement argument; the rounded-runtime chapter develops that separate
layer rather than folding it into the exact derivative claim.
For the running example, take
forward(x) = softmax(Wx + b). The scalar loss supplies an output cotangent
seed = dL/dforward, and the reverse pass returns the input and parameter cotangents
dL/dx, dL/dW, and dL/db.
In a framework, the registered kernels and the engine are expected to compose into the right answer. In this proof layer, the statement is that the backpropagated cotangent is the adjoint action of the derivative of the forward denotation. That theorem is the bridge from an executed reverse pass to a mathematical gradient object.
6.3.5. Operator Specs: Softmax And LogSoftmax
Softmax and log-softmax are good examples because they look familiar but carry real analytic content. The two derivative APIs are:
The softmax API defines softmaxVec, softmaxDerivCLM, and softmaxJvp. The theorem
softmaxJvp_eq_deriv identifies the implemented JVP formula with the derivative formula, and
hasFDerivAt_softmaxVec states the Frechet derivative of the vector softmax. The theorem
inner_softmaxJvp_comm packages the self adjoint structure of the softmax Jacobian, which is the
reason the VJP can reuse the same formula shape.
The log-softmax API follows the same discipline with logSoftmaxVec, logSoftmaxJvp, and
logSoftmaxVjp. The theorem logSoftmaxJvp_eq_deriv gives the derivative formula, while
inner_logSoftmaxJvp_vjp states the adjoint relationship between the JVP and the VJP:
For softmax, if
s_i=\frac{e^{x_i}}{\sum_j e^{x_j}},
then the directional derivative has the familiar form
D\,\operatorname{softmax}(x)[dx]_i
=
s_i\left(dx_i-\sum_j s_j dx_j\right).
For log-softmax, the formula is even cleaner:
D\,\log\operatorname{softmax}(x)[dx]_i
=
dx_i-\sum_j \operatorname{softmax}(x)_j\,dx_j.
The theorem proves that the formula used by the reverse rule is the adjoint of the derivative of the forward function. Framework bugs and mistakes in custom ops often live at exactly that boundary: the program might execute successfully while returning a subtly wrong gradient.
Their comments cite the PyTorch API reference for naming alignment, not as proof sources:
Documentation tells us what users expect the op to mean; Lean proves the derivative law for TorchLean's mathematical definition.
6.3.5.1. A Small Operator Audit Pattern
When adding a new differentiable operator, the proof obligation is intentionally mechanical. The operator should expose the same three objects that the softmax files expose:
-- Mathematical forward map. def forward : Vec n -> Vec m := ... -- Directional derivative used by the forward-mode view. def jvp : Vec n -> Vec n -> Vec m := ... -- Reverse rule used by backprop. def vjp : Vec n -> Vec m -> Vec n := ...
The local theorem should then say two things:
-- The JVP is the derivative of the forward denotation. #check NN.Proofs.Autograd.FDeriv.Softmax.hasFDerivAt_softmaxVec -- The VJP is adjoint to the JVP. #check NN.Proofs.Autograd.FDeriv.LogSoftmax.inner_logSoftmaxJvp_vjp
The exact names differ by operator, but the contract should not. A runtime rule is not considered an autograd theorem merely because the code returns a tensor of the right shape. It needs a mathematical forward function, a derivative statement, and an adjointness statement that lets the graph theorem compose the local rule with surrounding nodes.
This also explains why nondifferentiable points are not swept under the carpet. ReLU can be used in graphs, but a theorem phrased as a Fréchet derivative at a point must either avoid coordinates where the pre-activation is zero or state a subgradient convention in a different theorem. The current real-analysis statements take the first route: the hypothesis says where the derivative exists.
6.3.6. Runtime Autograd Link
The theorem stack above is mathematical. The runtime autograd engine is another object. TorchLean therefore keeps a link layer between executed reverse graphs and proof graphs:
Those files are about representation agreement. They do not prove a new derivative formula. They say that the runtime graph, its saved forward values, and its accumulation discipline can be read as the proof-level graph when the required invariants hold. That is the right granularity for auditing a new runtime node:
-
local derivative theorem for the mathematical operator;
-
runtime link theorem that the recorded node carries the same forward/VJP structure;
-
finite precision or backend theorem if the runtime scalar path is not the ideal real path.
Write an autograd claim by naming each layer of evidence:
ideal VJP theorem + runtime graph/link invariant + scalar/runtime approximation bridge = claim about executed gradients
Without the middle bridge, we only have a theorem about the proof graph. Without the last bridge, we only have a theorem about ideal arithmetic.
6.3.7. MLP And MSE Gradients: A Concrete Scalar Loss Story
The MLP/MSE derivative API turns the abstract autograd theorem into a familiar training example: a small MLP followed by mean squared error.
The definitions are close by design to what a reader would write on a whiteboard:
-
affineMatfor an affine layer; -
mlpVecMatfor a two layer MLP with a hidden nonlinearity; -
mseandmseGradfor the scalar loss; -
gradient lemmas for
W2,b2,b1,x, andW1.
The theorem pattern is:
\operatorname{backpropGradient}
=
\left(D\,\operatorname{loss}\right)^{\!*}(1)
Equivalently, for parameters \theta and a scalar loss L(\theta)=\ell(f_\theta(x),t), the gradient
statement has the shape
\nabla_\theta L(\theta)
=
\left(D_\theta f_\theta(x)\right)^{\!*}\nabla_y \ell(y,t).
For the last layer, the statement is clean. For the hidden ReLU layer, the theorem carries hypotheses such as "the value before activation is nonzero" at the coordinates being differentiated. Runtime systems usually leave that condition implicit. PyTorch chooses a subgradient convention at zero; TorchLean's real analysis statement names the differentiability condition instead.
Read the claim at that level. The theorem proves the ideal real-valued MLP/MSE gradient. A float32 training run connects to that statement through a finite-precision bridge, which belongs to runtime approximation and to the runtime approximation proof API.
6.3.8. Model Coverage: Attention, Transformers, And Recurrent Cells
The autograd APIs for model blocks give theorem entry points for selected fragments rather than a claim that every modern model is fully verified end to end. We built them to show how the algebra scales to the shapes users care about while keeping boundaries explicit.
Representative theorem entry points:
-
NN.Proofs.Autograd.Tape.Ops.Attention.MultiHeadSelfAttention API
-
NN.Proofs.Autograd.Tape.Ops.Transformer.ResidualAttention API
Loss-node proofs are maintained separately for binary cross entropy with logits, cross entropy,
KL divergence, MSE, and NLL under
Tape/Nodes/Losses.
Coverage is theorem-by-theorem: the existence of an operator proof does not imply that every
runtime registration, backend kernel, mask convention, or whole-model unroll has been connected to
it.
For attention, the named theorem
backprop_eq_adjoint_fderiv_scaledDotProduct states the desired attention theorem directly: the
graph reverse pass for scaled dot product attention agrees with the adjoint derivative of the
forward attention map. The masked theorem uses the hard-mask semantics named by its hypotheses;
backprop_eq_adjoint_fderiv_maskedScaledDotProduct must not be transferred to a finite additive-bias
mask without a separate equivalence result. Multi-head attention and residual attention then
package that structure at a wider interface.
For Transformer post-norm blocks, the post-norm API contains several theorem layers:
-
mhaPostNorm_backpropVec_eq_adjoint_fderiv_atfor residual MHA followed by LayerNorm; -
seqFfnPostNorm_backpropVec_eq_adjoint_fderiv_atfor residual feed forward followed by LayerNorm; -
postNorm_backpropVec_eq_adjoint_fderiv_atfor the common post-norm boundary; -
twoSublayerPostNormBlock_hasFDerivAtfor the analytic composition of two post-norm sublayers; -
named interfaces for residual attention and residual feed-forward post-norm variants.
The recurrent file is kept modest. elmanCell_backpropVec_eq_adjoint_fderiv proves the
reverse mode theorem for one Elman RNN cell, and elmanTwoStep_hasFDerivAt shows the
shape of a short unrolled composition. Full BPTT over an arbitrary sequence length is the next
induction over the unroll. That boundary is stated explicitly so the current theorem scope is clear.
6.3.9. Training Step Algebra
Backprop correctness is about gradients. Training correctness also needs a clean account of how those gradients are seeded and consumed. That role belongs to NN.Proofs.Autograd.Training.StepAlgebra API.
There are two pieces to keep separate.
First, Graph.scalarLoss_grad_correct specializes the global graph theorem to scalar losses. The
seed is the scalar cotangent 1, represented by seedScalarLoss. Formally, this is
loss.backward() seeding d loss / d loss = 1.
Second, step defines the algebra behind a simple optimizer update:
\theta_{t+1} = \theta_t - \eta\,\nabla_\theta L(\theta_t)
The theorem step_cons says the head tensor of the parameter list is updated by exactly that
formula, and step_nil handles the empty parameter list. These are compact theorems, but they keep
the training loop from becoming an opaque execution artifact. A realistic optimizer will add momentum,
Adam statistics, clipping, or weight decay; this file gives the simple algebraic core that those
extensions can refine.
The optimizer theory page extends this idea for larger update records. The autograd page should not be read as proving optimizer convergence. It proves the gradient side of the handoff:
\text{reverse pass returns }(D_\theta L(\theta))^{\ast} 1.
Optimization theory then decides what an update using that gradient means under step-size, smoothness, convexity, or backend-certification hypotheses.
6.3.10. Autograd Proof Dependencies
An autograd claim is easiest to audit from the graph theorem down to the local derivative rules and then back up to training. The proof stack has five levels:
-
The tape algebra soundness API contains
Graph.backprop_correct. -
The Fréchet derivative tape API contains
Graph.backpropVec_eq_adjoint_fderiv. -
Local derivative files include the softmax derivative API and log-softmax derivative API.
-
Model-block proofs include the Transformer post-norm API.
-
The training step algebra API connects differentiation to parameter updates.
The pattern should feel familiar if you know PyTorch internals or JAX transformations: local rules, graph traversal, cotangent accumulation, and scalar loss seeding. The proof contribution is that TorchLean turns those engineering moves into named Lean statements instead of leaving them as a large opaque execution layer.
The next question is numerical rather than differential. We know how the ideal reverse pass is assembled from local rules; the runtime-approximation chapter asks how far an executable reverse pass can drift when those rules run with rounded arithmetic.