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; in that module every line that feeds a returned
parameter is pretty of a node of the proven graph — 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. The module also carries
a hand-written, report-only %loss block that reads the logits and %onehot and feeds no
parameter. This file proves that the output weight's den is W₂ − lr·∂(crossEntropy ∘ forward)/∂W₂ at the
emitted loss cotangent (mlp_W2_tied_totalloss), and that each of the other five is
θ − lr·(certified per-layer Jacobian · the rendered chain cotangent) — SgdNode.denseW_den /
SgdNode.denseB_den at the chain cotangents mlpCotOut1/mlpCotOut0, whose emitted
selectPos/dotOut subgraphs cot1_den/cot0_den pin.
No new core SHlo ops are needed: the backward chain uses the existing
dotOut/selectPos, and the param updates reuse weightSgd/biasSgd.
Scope: one example (the emitted module batch-contracts; den is per-example); per-op pretty
lexing; ℝ→Float32. At the same chain cotangents mlpCotOut1.denote g / mlpCotOut0.denote g,
mlp_net_lossGrad (MlpParamGrad) states each layer's gradient node as the loss gradient in its
parameter, at smooth points.
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 SgdNode.denseW_den / SgdNode.denseB_den.
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₂ (StableHLO.lossWeightGrad_eq_sum).
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), the output-weight op denotes
W₂ − lr·∂CE/∂W₂, the WHOLE-loss gradient (mlp_W2_tied_totalloss), and the other five
denote θ − lr·(certified ∂layer/∂θ · the backward-chain cotangent the real loss drives) (the
*_den_certified at the composed g); mlp_net_lossGrad states the nodes at these chain
cotangents as the loss gradients.
No symbolic cotangent remains; the forward is shared (single render, correctly-threaded SSAs).