M2 — the MLP train step: per-layer parameter-gradient assembly #
The MLP (dense → relu → dense → relu → dense) train step updates six parameters
W₀,b₀,W₁,b₁,W₂,b₂. Each layer's gradient is the assembly of backprop:
∂L/∂Wᵢ = (layer-i input) ⊗ (cotangent the backward chain delivers at layer i),
∂L/∂bᵢ = (that cotangent).
The generic per-layer bridges already certify this for any backward subgraph
e : Back (weight_grad_bridge/bias_grad_bridge, IR.lean) — emit*Grad equals
the certified Jacobian of that dense layer contracted with e.denote dy. So the
assembly is choosing the right cotangent subgraph per layer:
- layer 2 (logits): the loss cotangent
gitself —Back.cotangent(weight_grad_bridge … Back.cotangent, withBack.cotangent.denote g = g). - layer 1 (
p₁):relu'(p₁) ⊙ (W₂ · g)—mlpCotOut1(already built;mlp_layer1_weight_grad_bridge). - layer 0 (
p₀):relu'(p₀) ⊙ (W₁ · mlpCotOut1)—mlpCotOut0, below.
This file supplies the only missing piece, mlpCotOut0, and its weight/bias
bridges — completing the three-layer assembly. This is Crux A of
planning/archive/verified_train_step.md: the multi-layer param-grad assembly, the step
linear couldn't show (one layer, no chain). The SGD wrapping θ − lr·∇ on top is
identical to the linear case (StableHLO.sgdW).
Layer-0 cotangent subgraph — the cotangent the backward chain delivers at
the layer-0 dense output p₀: relu'(p₀) ⊙ (W₁ · mlpCotOut1). Prepends one more
relu-back ∘ dense-back to mlpCotOut1, exactly as mlpCotOut1 extends
Back.cotangent.
Equations
- Proofs.IR.mlpCotOut0 W₁ W₂ p₀ p₁ = (Proofs.IR.emitReluBack p₀).subst ((Proofs.IR.emitDenseBack W₁).subst (Proofs.IR.mlpCotOut1 W₂ p₁))
Instances For
MLP layer-0 weight-gradient bridge. The emitted layer-0 weight gradient
(x₀ ⊗ mlpCotOut0) equals the certified Jacobian of the layer-0 dense wrt W₀,
contracted with the cotangent the backward chain delivers there — the deepest
chain (relu'(p₀) ⊙ W₁ · relu'(p₁) ⊙ W₂ · dy). The layer-0 peer of
mlp_layer1_weight_grad_bridge.
MLP layer-0 bias-gradient bridge. Likewise the layer-0 bias gradient is the certified ∂/∂b₀ Jacobian contracted with the same deepest cotangent.
MLP layer-1 bias-gradient bridge (the bias peer of the existing
mlp_layer1_weight_grad_bridge).
MLP layer-2 (output) weight-gradient bridge — the output layer's cotangent
is the loss cotangent dy itself (Back.cotangent, denote dy = dy); the
bridge is the generic one specialized there.
Output-layer total-loss gradient. For the top dense layer on activation a₁
(in the MLP, a₁ = relu(dense W₁ b₁ (relu(dense W₀ b₀ x)))), the single gradient of
the whole softmax-CE loss wrt W₂ equals the certified ∂logits/∂W₂ contracted with
the softmax-CE residual softmax − onehot. Unconditional — a direct instance of the
linear fold lossWeightGrad_eq_sum.
Input-layer total-loss fold (conditional, deepest). The same fold for the
first layer W₀, whose chain runs back through both ReLUs — so it carries both
smoothness hypotheses (the same pair as mlp_has_vjp_at). The total loss gradient
wrt W₀ = certified ∂p₀/∂W₀ contracted with the loss gradient at p₀ (the
deepest cotangent the backward chain delivers, cf. mlpCotOut0_denote).
Whole-network MLP weight-gradient capstone. The three weight layers' total-loss gradients, jointly, under the two smoothness hypotheses.
Rendered %W1n output = W₁ − lr·(certified ∂p₁/∂W₁ · the chain cotangent).
Rendered %W0n output = W₀ − lr·(certified ∂p₀/∂W₀ · the deepest chain cotangent).