TorchLean

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:

  1. the eager engine records operations into a dynamic tape;

  2. each operation has a backward rule, either built in or registered by extension code;

  3. the engine traverses the tape in reverse and accumulates gradients;

  4. 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, and ContinuousLinearMap.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:

  1. compiling g and running Tape.backwardDenseFrom succeeds with the graph's full reverse context;

  2. 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.

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:

  • affineMat for an affine layer;

  • mlpVecMat for a two layer MLP with a hidden nonlinearity;

  • mse and mseGrad for the scalar loss;

  • gradient lemmas for W2, b2, b1, x, and W1.

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:

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_at for residual MHA followed by LayerNorm;

  • seqFfnPostNorm_backpropVec_eq_adjoint_fderiv_at for residual feed forward followed by LayerNorm;

  • postNorm_backpropVec_eq_adjoint_fderiv_at for the common post-norm boundary;

  • twoSublayerPostNormBlock_hasFDerivAt for 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:

  1. The tape algebra soundness API contains Graph.backprop_correct.

  2. The Fréchet derivative tape API contains Graph.backpropVec_eq_adjoint_fderiv.

  3. Local derivative files include the softmax derivative API and log-softmax derivative API.

  4. Model-block proofs include the Transformer post-norm API.

  5. 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.