Documentation

LeanMlir.Proofs.Nets.Small.LinearTrainStep

The linear train step descends the certified softmax-CE gradient #

StableHLO.Basic 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 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 a named, axiom-audited certified quantity — no residual den-of-graph, no trusted optimizer step. This is the denotation half for linear (what the emitted train step computes), for one example x.

The chain-rule fold is here too: by pdiv_comp, the two-factor sum is the single gradient ∂/∂θ (crossEntropy ∘ mnistLinear) (lossWeightGrad_eq_sum, in SmoothedLossCot with the other loss-cotangent lemmas), so the weight update is literally one step of gradient descent on the loss (sgdW_descends_loss_gradient, using crossEntropy_differentiable). The rendering half — that the emitted text is pretty of these graphs — is linTrainStepText and LinearFold.

source
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.

source
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)

Weight update. 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.

source
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)

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

source
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.

source
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

Weight update, 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.

source
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
Instances For
    source
    noncomputable def Proofs.StableHLO.linBiasDen {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (lr : ℝ) (label : Fin n) :
    Vec n
    Equations
    Instances For
      source
      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 (sgdW_descends_loss_gradient).

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