The MNIST-linear train step, proof-tied to the certified SGD step #
MainMnistLinearVerified
trains on verified_mlir/linear_train_step.mlir, which is written by
Proofs.StableHLO.linTrainStepText (the #eval writer at the end of
StableHLO.Basic). This file certifies that renderer for a single example x: the emitted
weight output denotes W − lr·∂(crossEntropy ∘ dense)/∂W (Mathlib-fderiv-derived) and the bias
output denotes b − lr·(certified ∂dense/∂b · (softmax − onehot)). The committed module
batch-contracts over B examples, which this file does not state.
(Namespace/name lengths are kept short on purpose: tests/AuditAxioms.lean's
three-axiom closure check greps #print axioms output per line, which Lean wraps
past ~120 cols — long qualified names would split the benign triple across lines
and false-fail the check. Keep future per-chapter capstone names short.)
What is closed here (kernel, [propext, Classical.choice, Quot.sound]) #
fwdGraph_faithful+poc_fwd_is_render— the forward-eval module isrenderModuleof a graph whosedenismnistLinear(text =render(graph)∧den(graph) = math). Forward: end-to-end tied.poc_train_step_tail_certified— fully tied. The two emitted SGD ops consume the provenlossCotGraphnode directly (no SSA-name pin), so the weight output'sdenis proven =W − lr·∂CE/∂Wand the bias output's = the certified bias Jacobian contracted withsoftmax − onehot, with the forward = the provenfwdGraph(nested inlossCotGraph)._fwdand_train_stepshare the same forward graph.
The committed-bytes tie (verified_mlir/linear_train_step.mlir == linTrainStepText(…)) is enforced in CI (regenerate + git diff, the
"Verified-render drift guard" step in proofs.yml), not here.
Scope (the boundary shared with the forward SHlo den) #
- Per-op
den⇄ MLIR text forweightSgd/biasSgd: that the printeddot_general/reduce/multiply/subtractcompute theirdenis the same trusted op-level modelling the forwarddenrelies on (the weight-grad piece iswGrad, tied toIR.emitWeightGradbywGrad_faithful). - Single example (B = 1):
wGrad x dy = x ⊗ dy; the emitted module batch-contracts. The mean-loss cotangent makes the batch sum the mean gradient. - ℝ → Float32: not stated here;
FloatBridge.leanandMlpFloatBridge.leancover the linear/MLP rounding budget separately.
Forward eval — end-to-end tied #
The committed linear_fwd.mlir generator linearFwdModuleV is renderModule
applied to the proven fwdGraph — the emitted bytes are literally the print of
the graph fwdGraph_faithful is about.
The tail fold (closed) — the emitted tail ops are pretty(provenNode) #
StableHLO.linTrainStepText (what generates verified_mlir/linear_train_step.mlir)
renders the whole module as pretty of SHlo nodes, fully tied: each of
SHlo.weightSgd / SHlo.biasSgd consumes the proven lossCotGraph node DIRECTLY
(not a .operand %dy <placeholder> name-pin), so the forward = the proven fwdGraph
(nested inside lossCotGraph) and den(output) = certified is one composed theorem
below — no trusted SSA-name linkage between the cotangent and the SGD ops. The shared
cotangent is rendered once per output (2×); iree CSEs the duplicate. (lr/W,b,x
values are skel-erased, so the render is value-independent — the placeholders
linTrainStepText passes print the same text as the live graph here.)
The emitted weightSgd op — consuming the proven lossCotGraph node DIRECTLY (the
fully-tied render) — denotes linWeightDen (the certified sgdW step).
The emitted biasSgd op — consuming the proven lossCotGraph node DIRECTLY — denotes
linBiasDen (the certified bias step).
Tail fold. For one example x, the two emitted tail ops weightSgd/biasSgd — the
actual SHlo nodes linTrainStepText prints — denote, respectively,
W − lr·∂(crossEntropy ∘ dense)/∂W and b − lr·(certified ∂dense/∂b · (softmax − onehot)).