PoC: the MNIST-MLP train step, proof-tied to the certified SGD step #
The MLP analogue of LinearFold. MainMnistMlpVerified trains on
verified_mlir/mlp_train_step.mlir; this file makes the whole module
pretty(provenGraph) — forward (denseF/reluF), the loss cotangent, the
backward chain (dotOut/selectPos), and the six parameter SGD updates as the
weightSgd/biasSgd SHlo ops added in LinearFold's core extension —
and proves each output's den equals the certified fderiv-derived loss-descent
step, reusing mlp_render_{W,b}*_certified (the per-layer bridges) and
mlpCotOut{0,1}_denote (the explicit chain cotangents).
No new core SHlo ops are needed: the backward chain uses the existing
dotOut/selectPos, and the param updates reuse weightSgd/biasSgd.
Residual (as for linear): per-op pretty lexing; B=1 (the emitted module
batch-contracts; den is per-example); the ReLU smooth-point hypotheses are
inherited from the bridges; ℝ→Float32.
The backward-chain cotangent subgraphs denote the proven mlpCotOut* #
The emitted selectPos/dotOut chain — what the faithful renderer prints for the
per-layer pre-activation cotangents — denotes exactly mlpCotOut{1,0}.denote g.
Layer-1 cotangent subgraph selectPos p₁ (dotOut W₂ dy) denotes mlpCotOut1.denote g.
Layer-0 cotangent subgraph selectPos p₀ (dotOut W₁ cot1) denotes mlpCotOut0.denote g.
The six emitted parameter ops denote the certified SGD step #
Each weightSgd/biasSgd op, fed the right activation (x field) and the
cotangent the chain delivers (the .operand value), denotes θ − lr·(certified per-layer Jacobian · cotangent) — via the op den = emitWeightGrad/emitBiasGrad
(outer / reduce) and the mlp_render_*_certified bridges.
Output-layer weight op weightSgd a1 W₂ (cot = dy) = certified W₂ step.
Hidden-layer weight op weightSgd a0 W₁ (cot = mlpCotOut1) = certified W₁ step.
Input-layer weight op weightSgd x W₀ (cot = mlpCotOut0) = certified W₀ step.
Output-layer bias op = certified b₂ step.
Hidden-layer bias op = certified b₁ step.
Input-layer bias op = certified b₀ step.
Fully tied — the top loss cotangent is the composed softmax-CE of the forward #
The six *_den_certified above hold ∀ g (a symbolic top loss cotangent). The renderer feeds
the cotangent the EMITTED loss graph sub(softmaxDiv(expe(logits)), onehot) produces, where
logits is the real forward output. The lemma below pins that graph's den to the composed
softmax-CE gradient of the forward (∂CE/∂logits, the mlp analogue of lossCotGraph_isCEgrad);
instantiating the six certified theorems at it ties the train step forward→loss→backward, with
the output layer folding to the WHOLE-loss gradient ∂CE/∂W₂ (mlp_output_total_loss_grad).
The emitted loss-cotangent graph denotes the composed softmax-CE gradient of the forward
(= softmax(logits) − onehot = ∂CE/∂logits at the real forward logits).
Output-layer weight op, fully tied to the WHOLE softmax-CE loss. Fed the cotangent the
emitted loss graph denotes (mlpLossCot_den), the weightSgd for W₂ denotes
W₂ − lr·∂(crossEntropy ∘ forward)/∂W₂ — the certified TOTAL-loss-descent step, composed
forward→softmax-CE→gradient, no symbolic cotangent.
Whole mlp train step, tied. With the top loss cotangent g pinned to the composed
softmax-CE gradient of the forward (mlpLossCot_den), all six emitted parameter ops denote
the certified loss-descent step — the output weight W₂ folded to the WHOLE-loss gradient
∂CE/∂W₂ (mlp_W2_tied_totalloss), the other five to θ − lr·(certified ∂layer/∂θ · the backward-chain cotangent the real loss drives) (the *_den_certified at the composed g).
No symbolic cotangent remains; the forward is shared (single render, correctly-threaded SSAs).