The MNIST-CNN train step, proof-tied to the certified SGD step #

The CNN analogue of LinearFold / MlpFold. MainMnistCnnVerified trains on verified_mlir/cnn_train_step.mlir; this file states what the parameter updates of that module denote: each emitted SGD op denotes θ − lr·(certified per-layer Jacobian · the cotangent the rendered backward chain feeds it), and the output weight W₅ is folded to the whole-loss gradient (cnn_W5_tied_totalloss).

The CNN has two kinds of parameter: the dense classifier head (W₃,W₄,W₅ + biases — structurally a 3-layer MLP over the flattened pool output) and the convolution kernels/biases (W₁,W₂ + biases). The dense head reuses the weightSgd/biasSgd SHlo ops added in LinearFold (its dens certified via weight_grad_bridge/bias_grad_bridge at the mlpCotOut-style chain cotangents — the head is a 3-layer MLP, so the IR mlpCotOut0/1 apply verbatim). The conv layers use the core ops convWeightSgd/convBiasSgd (StableHLO/Basic.lean): their den is flatten(W − lr·conv2dWeightGrad…) / b − lr·conv2dBiasGrad…, proven = certified by the chain-pinned conv bridges cnn_render_conv{W,b}{1,2}_chain_certified (CnnChainClose.lean) at the cotangents the CNN backward chain delivers (cnnChainCotW1/cnnChainCotW2).

(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.)

What is proved here (kernel, [propext, Classical.choice, Quot.sound]) #

• cW1_den/cb1_den/cW2_den/cb2_den — the four emitted conv param ops (convWeightSgd/convBiasSgd), fed the chain cotangent c, denote θ − lr·(certified ∂conv/∂θ · c).
• dW5_den — the output-layer dense-head op (weightSgd) denotes W₅ − lr·(certified ∂dense/∂W₅ · dy); the other five head ops are SgdNode.denseW_den / SgdNode.denseB_den at their layer.
• cnn_W5_tied_totalloss — at the emitted loss cotangent, the W₅ op denotes W₅ − lr·∂(crossEntropy ∘ forward)/∂W₅.
• cnn_train_step_tied_certified — all ten parameter ops (the six dense-head ops and the four conv ops) at the real forward activations and the chain cotangents driven by the emitted loss cotangent.

Scope #

• Chain cotangent vs loss gradient. Below the output layer the cotangents are the rendered chain (mlpCotOut1/mlpCotOut0 in the head; cnnChainCotW1, cnnChainCotW2: relu masks, select-and-scatter pool-back, conv-back). The loss-gradient statement is cnn_net_lossGrad (CnnParamGrad), at the un-fused *Grad nodes these ops step by (θ − lr·node, SmallParamGrad.convWeightSgd_eq_grad) and with the pool's cotangent routed to one maximal cell of each window, as the rendered select_and_scatter does. cnnChainCotW2 reads the pool backward as maxPoolBackDenote, which routes to the first maximal cell (maxPool2Argmax), the op's own choice, so it is the capstone's chain at that selection, ties included (cnnChainCotW2_eq_sel).
• Cotangent subgraph ⇄ rendered SHlo. The chain cotangents (cnnChainCotW1/2, mlpCotOut0/1) are proven = the rendered backward form in CnnChainClose (cnnChainCotW1_eq, cnnChainCotW2_eq) and MlpTrainStep (mlpCotOut0_denote, mlpCotOut1_denote); they are not pinned to the emitted selectPos/dotOut/convBack/maxPoolBack SHlo subgraph (as MlpFold.cot0_den/MlpFold.cot1_den do for the MLP).
• Per-op pretty lexing (shared with the whole suite) + ℝ → Float32.

Convolution layers — the new convWeightSgd/convBiasSgd ops denote certified #

den (convWeightSgd … (.operand _ c)) is by construction flatten W − lr·conv2dWeightGrad(b,x)·c (and likewise for the bias); pinning c to the cotangent the chain delivers and applying the chain-certified conv bridge gives θ − lr·(certified ∂conv/∂θ · the-chain-cotangent). (The den reduction is definitional — rfl — exactly as LinFold.poc_weightSgd_den_eq.)

Dense classifier head — reuse weightSgd/biasSgd (the head is a 3-layer MLP) #

The pool-output pool : Vec (c·h·w) flows through W₃→relu→W₄→relu→W₅; the per-layer cotangents are exactly the IR mlpCotOut0/1 (with (W₅,W₄,W₃) playing the MLP's (W₂,W₁,W₀)). Every head op's den = certified is SgdNode.denseW_den / SgdNode.denseB_den at that layer, and cnn_train_step_tied_certified states all six at the real forward; the output-layer weight op is stated here on its own because the whole-loss fold cnn_W5_tied_totalloss rewrites with it.

Tie (dense head) — the top loss cotangent is the composed softmax-CE of the CONV forward #

The *_den theorems above hold for a free top cotangent dy and a free pool output. The renderer feeds the cotangent the emitted loss graph sub(softmaxDiv(expe(logits)), onehot) produces, with logits the REAL conv-forward output mnistCnnNoBnForward … x. The lemma below pins that graph to the composed softmax-CE gradient of the conv forward (the cnn analogue of mlpLossCot_den), and the headline folds the dense output weight W₅ to the whole-loss gradient ∂CE/∂W₅ — so the output layer is tied forward(conv+dense)→softmax-CE→gradient. Every parameter op, at the chain cotangent this g drives, is stated in the next section.

The whole step — every parameter op at the real forward #

The *_den theorems above hold for FREE activations (ac1/ac2/hc2, the pool output) and a free cotangent. The capstone below instantiates them at the real forward (ac1/hc1/hc2/ac2 = the actual conv₁/relu/conv₂/relu outputs, pool their max-pool, h3/h4 the dense pre-activations) and the chain cotangents driven by the composed top cotangent g = softmax(mnistCnnNoBnForward x) − onehot (cnnLossCot_den). Each of the ten parameter ops denotes θ − lr·(certified ∂layer/∂θ · c) with c the rendered backward-chain cotangent: g for W₅/b₅, mlpCotOut1 for W₄/b₄, mlpCotOut0 for W₃/b₃, cnnChainCotW2 for conv₂ and cnnChainCotW1 W₂ hc1 cotW2 for conv₁ (it crosses one more conv-back). The loss-gradient form below the output is cnn_net_lossGrad (see the module's Scope); at the output layer cnn_W5_tied_totalloss folds W₅ to ∂CE/∂W₅. The correspondence between the hand-written conv-backward SSA values and cnnChainCotW1/cnnChainCotW2 is the per-op trust the whole suite carries.