Executed Backward Pass versus Proved Backward Pass #
The eager trainer executes Tape.backwardDenseAll, which runs Tape.backwardDense: an
Option-valued reverse sweep that runs a node's VJP only when the node has received a cotangent,
then zero-fills the unreached slots. The link theorems in BackwardGraph and FDeriv are about
Tape.backwardDenseFrom, which starts from a total gradient array and runs every VJP.
This file closes that gap at the tape level. The hypothesis is ZeroPreserving: every VJP on the
tape sends the node's zero cotangent to zero contributions of the parents' shapes. Under it,
backwardDenseAll_eq_backwardDenseFrom : Tape.backwardDenseAll t outId seed = Tape.backwardDenseFrom t (oneHotGrads t outId seed)
holds for any tape, any valid output id and any seed of the output's shape, including the error
cases. The proof is a step-by-step simulation: totalizeGrads maps the optional accumulator of
backwardDense to the total accumulator of backwardDenseFrom, and each reverse step commutes
with it. A skipped (unreached) node on the executed side corresponds to a VJP call on the zero
cotangent on the proved side, which by ZeroPreserving adds only zeros.
BackwardDenseGraph proves that every tape produced by lowerGraphToTape is ZeroPreserving
and derives the corollaries for the executed backward pass.
Zero cotangents and totalization #
The zero cotangent of a runtime node: an all-zero tensor of the node's value shape.
Instances For
A zero cotangent has the same shape as the node's value, by construction.
Totalize the optional gradient array of backwardDense: reached nodes keep their gradient, every
other node gets its zero cotangent. This is exactly the post-processing backwardDenseAll does.
Instances For
The initial total gradient array of backwardDenseFrom: seed at outId, zeros elsewhere.
Instances For
backwardDenseAll is backwardDense followed by totalizeGrads.
Totalization produces exactly one entry per tape node.
So does the one-hot seed array, which is why both can be indexed by node id without bounds checks in the invariants below.
Entry id of the totalized array, expressed through getNode?.
Entry id of the one-hot array, expressed through getNode?.
A node that the backward pass reached keeps its computed gradient.
A node the backward pass never reached gets its zero cotangent.
Taken together, these two lemmas say totalization loses no information: none in the sparse array
means the node genuinely received no contribution, so filling in zero is not an approximation.
Totalization commutes with writing a present gradient.
Invariants #
Well-formedness of the optional accumulator used by backwardDense: one slot per tape node, and
every present gradient has its node's value shape.
One slot per tape node.
- shape (id : ℕ) (g : Spec.SomeTensor α) : grads[id]? = some (some g) → ∃ (node : Runtime.Autograd.Node α), t.getNode? id = some node ∧ g.shape = node.value.shape
Present gradients have the shape of their node's value.
Instances For
Well-formedness of the total accumulator used by backwardDenseFrom: one slot per tape node, and
every slot has its node's value shape.
One slot per tape node.
- shape (id : ℕ) (node : Runtime.Autograd.Node α) : t.getNode? id = some node → ∃ (g : Spec.SomeTensor α), grads[id]? = some g ∧ g.shape = node.value.shape
Every slot has the shape of its node's value.
Instances For
Zero preservation of a tape: every node's VJP sends the node's zero cotangent to an array of contributions each of which is the zero cotangent of an existing parent node.
This is the exact condition under which skipping unreached nodes (backwardDense) and running
every node (backwardDenseFrom) produce the same gradients. It holds for linear VJPs of correct
parent shapes, in particular for every tape produced by lowerGraphToTape.
- backward_zero (id : ℕ) (node : Runtime.Autograd.Node α) : t.getNode? id = some node → node.requiresGrad = true → ∃ (contribs : Array (ℕ × Spec.SomeTensor α)), node.backward (zeroCotangent node) = Except.ok contribs ∧ ∀ (pid : ℕ) (pg : Spec.SomeTensor α), (pid, pg) ∈ contribs → ∃ (pnode : Runtime.Autograd.Node α), t.getNode? pid = some pnode ∧ pg = zeroCotangent pnode
The VJP at the zero cotangent succeeds and emits only parent zero cotangents.
Instances For
A successful getNode? witnesses that the id is in range.
Conversely, any in-range id resolves to a node. The pair lets the invariants below be phrased in
terms of getNode? alone, without carrying array bounds around.
Totalizing a well-formed optional accumulator yields a well-formed total accumulator.
Writing a correctly shaped gradient preserves OptGradsOk.
Adding a zero contribution is a no-op #
Adding the all-zero tensor on the left is the identity.
Adding the all-zero tensor on the right is the identity.
Accumulating onto a node's zero cotangent returns the contribution unchanged.
Accumulating a node's zero cotangent onto an existing gradient leaves it unchanged.
addGradAll with a parent's zero cotangent is the identity on a well-formed accumulator.
Folding zero contributions through addGradAll is the identity on a well-formed accumulator.
One contribution: addGradDense simulates addGradAll #
addGradDense preserves the optional accumulator invariant.
One addGradDense step, mapped through totalizeGrads, is the matching addGradAll step.
Folding contributions #
Except.map commutes with bind when it commutes with the first action and, on success, with
the continuation. This is the single lemma behind every "executed step simulates proved step"
composition below.
Folding addGradDense preserves the optional accumulator invariant.
Folding addGradDense, mapped through totalizeGrads, is folding addGradAll.
One node: the executed step simulates the proved step #
The per-node step of Tape.backwardDense, written out as a function.
backwardDense folds this step over node ids in reverse order (backwardDense_eq_foldlM).
Unreached nodes (acc[id] = some none) are skipped without running their VJP.
Instances For
The executed step preserves the optional accumulator invariant.
The executed step at a node, mapped through totalizeGrads, is the proved step
backwardDenseFromStep at the same node, provided the tape is zero preserving.
When the node was reached both sides run the same VJP on the same cotangent. When it was not, the executed side does nothing and the proved side runs the VJP on the zero cotangent, whose contributions are all zeros and therefore accumulate to nothing.
The reverse loop #
Peel the first (largest) id off a reverse-range fold.
The executed reverse loop over the first n node ids, mapped through totalizeGrads, is the
proved loop backwardDenseFromLoop over the same ids.
Main theorem #
The seeded optional accumulator backwardDense starts from.
Totalizing the seeded optional accumulator gives the one-hot total accumulator.
backwardDense is the reverse fold of backwardDenseStep from the seeded accumulator.
Executed backward pass = proved backward pass. On a zero-preserving tape, the totalized
executed sweep backwardDenseAll (which skips unreached nodes) returns exactly what the proved
sweep backwardDenseFrom returns when started from the one-hot seed array, including agreement
of the error cases. The hypotheses on outId and seed are the checks backwardDense performs
before traversing; without them backwardDense fails while backwardDenseFrom may not.