M1 — the linear train step descends the certified softmax-CE gradient

StableHLO.lean proves the Chapter-1 linear train step piecewise: the forward graph (fwdGraph_faithful), the loss cotangent (lossCotGraph_isCEgrad), the per-parameter Jacobians (wGrad/bGrad_isWeightJacobian), and the plain-SGD update (sgdW/sgdB_isCertifiedGradStep). Each of those, however, still mentions the emitted cotangent as den (lossCotGraph …) — a denotation of an emitted graph, not yet a named closed form.

This file bundles them into a single statement per parameter: the emitted SGD update subtracts lr times [the certified ∂logits/∂θ Jacobian] contracted with [the certified closed-form softmax-CE gradient softmax − onehot]. Every factor is now a named, axiom-audited certified quantity — no residual den-of-graph, no trusted optimizer step. This is the denotation half of milestone M1 for linear (what the emitted train step computes).

Two things are deliberately NOT done here, and are tracked in planning/archive/verified_train_step.md:

• The chain-rule fold. The two-factor sum below is, by pdiv_comp, the single gradient ∂/∂θ (crossEntropy ∘ mnistLinear) — i.e. literally one step of gradient descent on the loss. Stating it in folded form needs DifferentiableAt for crossEntropy (no such lemma exists yet) and for the dense-wrt-flattened-weights map; left as the next proof step. The unfolded form here carries the same content with no smoothness obligation.
• The rendering half. den/SHlo is a single-example semantics with no constructors for the batched weight-grad dot_general, bias-grad reduce, or SGD multiply/subtract; that tail of verified_mlir/linear_train_step.mlir is still hand-written string concat. Closing emitted text = render(provenGraph) needs the batched multi-output AST (Stage 1 of the plan).

theorem Proofs.StableHLO.lossCot_eq_softmax_sub_onehot {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (label k : Fin n) : den (lossCotGraph W b x (oneHot n label)) k = softmax n (mnistLinear W b x) k - oneHot n label k

The emitted loss cotangent is the certified closed-form softmax-CE gradient softmax(logits) − onehot, by lossCotGraph_isCEgrad then softmaxCE_grad.

theorem Proofs.StableHLO.sgdW_descends_softmaxCE_grad {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label : Fin n) (i : Fin m) (j : Fin n) : sgdW W b x lr label i j = W i j - lr * ∑ k : Fin n, pdiv (fun (v : Vec (m * n)) => dense (Mat.unflatten v) b x) W.flatten (finProdFinEquiv (i, j)) k * (softmax n (mnistLinear W b x) k - oneHot n label k)

M1 (weight). The emitted linear SGD weight update subtracts lr times the certified ∂logits/∂W Jacobian contracted with the certified closed-form softmax-CE gradient softmax − onehot.

theorem Proofs.StableHLO.sgdB_descends_softmaxCE_grad {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label j : Fin n) : sgdB W b x lr label j = b j - lr * ∑ i : Fin n, pdiv (fun (b' : Vec n) => dense W b' x) b j i * (softmax n (mnistLinear W b x) i - oneHot n label i)

M1 (bias). The emitted linear SGD bias update subtracts lr times the certified ∂logits/∂b Jacobian contracted with the same certified softmax-CE gradient.

theorem Proofs.StableHLO.crossEntropy_differentiable (c : ℕ) (label : Fin c) : Differentiable ℝ fun (z : Vec c) => crossEntropy c z label

crossEntropy is differentiable in the logits. The standalone form of the differentiability infrastructure inside softmaxCE_grad: softmax > 0 lets Real.log (hence crossEntropy = -log(softmax · label)) inherit smoothness.

theorem Proofs.StableHLO.denseWeightMap_differentiable {m n : ℕ} (b : Vec n) (x : Vec m) : Differentiable ℝ fun (v : Vec (m * n)) => dense (Mat.unflatten v) b x

The dense layer is differentiable in its (flattened) weights. The map v ↦ dense (unflatten v) b x is affine — a finite sum of coordinate evaluations scaled by x, plus the constant bias.

theorem Proofs.StableHLO.lossWeightMap_differentiable {m n : ℕ} (b : Vec n) (x : Vec m) (label : Fin n) : Differentiable ℝ fun (v : Vec (m * n)) (x_1 : Fin 1) => crossEntropy n (dense (Mat.unflatten v) b x) label

The softmax-CE loss is differentiable in the (flattened) weights.

theorem Proofs.StableHLO.lossWeightGrad_eq_sum {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (label : Fin n) (i : Fin m) (j : Fin n) : pdiv (fun (v : Vec (m * n)) (x_1 : Fin 1) => crossEntropy n (dense (Mat.unflatten v) b x) label) W.flatten (finProdFinEquiv (i, j)) 0 = ∑ k : Fin n, pdiv (fun (v : Vec (m * n)) => dense (Mat.unflatten v) b x) W.flatten (finProdFinEquiv (i, j)) k * (softmax n (mnistLinear W b x) k - oneHot n label k)

The total-loss gradient wrt a weight entry equals the certified (∂logits/∂W) · (softmax − onehot) contraction (the chain rule, pdiv_comp).

theorem Proofs.StableHLO.sgdW_descends_loss_gradient {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label : Fin n) (i : Fin m) (j : Fin n) : sgdW W b x lr label i j = W i j - lr * pdiv (fun (v : Vec (m * n)) (x_1 : Fin 1) => crossEntropy n (dense (Mat.unflatten v) b x) label) W.flatten (finProdFinEquiv (i, j)) 0

M1 (weight, folded). The emitted linear SGD weight update is literally one step of gradient descent on the certified softmax-CE loss: W − lr·∂(crossEntropy ∘ mnistLinear)/∂W.

structure Proofs.StableHLO.TrainOut : Type

One updated-parameter output of a multi-result module: its MLIR result type, the SSA name it binds, and the lines computing it from the rendered cotangent %dy. Renderable (computable); the ℝ-valued denotation lives separately.

• tyStr : String
• result : String
• emit : String → String

structure Proofs.StableHLO.TrainStepModule (B : ℕ) : Type

A multi-output train-step module: the forward+loss cotangent subgraph (an SHlo, rendered ONCE → shared %dy) plus the updated-parameter outputs. The multi-result generalization of renderModule.

• fname : String
• argSig : String
• cotLen : ℕ
• cot : SHlo self.cotLen
• outs : List TrainOut

def Proofs.StableHLO.renderModuleN {B : ℕ} (M : TrainStepModule B) : String

renderModuleN — render a multi-output module: cotangent once (shared %dy), each output's lines, then a tuple return.

Equations

• One or more equations did not get rendered due to their size.

def Proofs.StableHLO.linWeightOut {m n : ℕ} (B : ℕ) (lrStr : String) : TrainOut

Weight output W0' = W0 − lr·dot_general(x, dy) (batch-contracting outer product), rendered.

Equations

• One or more equations did not get rendered due to their size.

def Proofs.StableHLO.linBiasOut {n : ℕ} (B : ℕ) (lrStr : String) : TrainOut

Bias output b0' = b0 − lr·reduce(dy) (batch-sum), rendered.

Equations

• One or more equations did not get rendered due to their size.

def Proofs.StableHLO.linTrainStepModule {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (B : ℕ) (lrStr : String) : TrainStepModule B

The linear train-step module (renderable; the %onehot value is a runtime input that pretty ignores, so the placeholder cotangent renders identically to the live one). Structural peer of linearTrainStepModuleV.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.StableHLO.linWeightDen {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label : Fin n) : Vec (m * n)

The two outputs' ℝ denotations: the flattened certified weight update and the certified bias update.

Equations

• Proofs.StableHLO.linWeightDen W b x lr label = (Proofs.StableHLO.sgdW W b x lr label).flatten

noncomputable def Proofs.StableHLO.linBiasDen {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label : Fin n) : Vec n

Equations

• Proofs.StableHLO.linBiasDen W b x lr label = Proofs.StableHLO.sgdB W b x lr label

noncomputable def Proofs.StableHLO.linTrainStepDenN {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label : Fin n) : List ((k : ℕ) × Vec k)

denN — the tuple of per-example output denotations the module computes.

Equations

• Proofs.StableHLO.linTrainStepDenN W b x lr label = [⟨m * n, Proofs.StableHLO.linWeightDen W b x lr label⟩, ⟨n, Proofs.StableHLO.linBiasDen W b x lr label⟩]

theorem Proofs.StableHLO.linWeightDen_is_loss_descent {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label : Fin n) (i : Fin m) (j : Fin n) : linWeightDen W b x lr label (finProdFinEquiv (i, j)) = W i j - lr * pdiv (fun (v : Vec (m * n)) (x_1 : Fin 1) => crossEntropy n (dense (Mat.unflatten v) b x) label) W.flatten (finProdFinEquiv (i, j)) 0

Faithfulness, output 0 (weights). The rendered weight output denotes literally W − lr·∂(softmax-CE loss)/∂W (M1 sgdW_descends_loss_gradient).

theorem Proofs.StableHLO.linBiasDen_is_certified {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label j : Fin n) : linBiasDen W b x lr label j = b j - lr * ∑ i : Fin n, pdiv (fun (b' : Vec n) => dense W b' x) b j i * (softmax n (mnistLinear W b x) i - oneHot n label i)

Faithfulness, output 1 (bias). The rendered bias output denotes the certified b − lr·(∂logits/∂b · (softmax − onehot)) (M1).