Documentation

LeanMlir.Proofs.Training.SgdDescentLinear

Lipschitz constants for the linear softmax-CE loss #

The missing hypothesis of sgd_descends, discharged for the Chapter-1 net: the gradient of v ↦ crossEntropy(dense(unflatten v, b, x), label) is segment-Lipschitz with the explicit constant 2a²/(1 − 2aD) (pixels bounded by a, step ℓ1-radius D, small-step condition 2aD < 1).

No Hessian appears. The route is the same elementary ratio argument as FloatBridge.lean's: the loss gradient is xᵢ·(softmax(z)ⱼ − onehotⱼ) (linear_loss_gradAt, assembled from the suite's certified lossWeightGrad_eq_sum + pdiv_dense_W), the logits move linearly in the parameters (dense_unflatten_drift), and FloatModel.softmax_perturb + the γ-form FloatModel.exp_sub_one_le turn the logit drift into a gradient drift that is linear in t along the segment — exactly the shape descent_segment consumes.

linear_sgd_descends is the capstone: an η-accurate gradient oracle (e.g. the float budgets), the small-step condition, and the two dominance conditions ⇒ one inexact SGD step on the MNIST-linear classifier provably decreases the cross-entropy loss by ≥ lr·‖∇L‖₂²/2. Every hypothesis is checkable arithmetic; smoothness is proven, not assumed.

theorem Proofs.gradAt_eq_pdiv {p : } (f : Vec p) (v : Vec p) (hf : DifferentiableAt f v) (idx : Fin p) :
gradAt f v idx = pdiv (fun (w : Vec p) (x : Fin 1) => f w) v idx 0

gradAt agrees with the suite's Vec 1-codomain pdiv convention.

theorem Proofs.linear_loss_gradAt {m n : } (b : Vec n) (x : Vec m) (label : Fin n) (v : Vec (m * n)) (i : Fin m) (j : Fin n) :
gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) v (finProdFinEquiv (i, j)) = x i * (softmax n (dense (Mat.unflatten v) b x) j - oneHot n label j)

Closed form of the linear softmax-CE loss gradient at any parameter point: ∂L/∂W_{ij} = xᵢ·(softmax(z)ⱼ − onehotⱼ) — the suite's certified contraction (lossWeightGrad_eq_sum + pdiv_dense_W), re-expressed through gradAt.

theorem Proofs.dense_unflatten_drift {m n : } (b : Vec n) (x : Vec m) {a : } (ha : 0 a) (hx : ∀ (i : Fin m), |x i| a) (v d : Vec (m * n)) (k : Fin n) :
|dense (Mat.unflatten (v + d)) b x k - dense (Mat.unflatten v) b x k| a * idx : Fin (m * n), |d idx|

The logits move linearly in the parameters: a parameter perturbation of ℓ1 mass ‖d‖₁ moves every logit by at most a·‖d‖₁.

theorem Proofs.linear_loss_grad_lipschitz {m n : } (b : Vec n) (x : Vec m) (label : Fin n) {a D : } (ha : 0 a) (hx : ∀ (i : Fin m), |x i| a) (v d : Vec (m * n)) (hd : idx : Fin (m * n), |d idx| D) (hsmall : 2 * (a * D) < 1) (t : ) (ht : t Set.Icc 0 1) (idx : Fin (m * n)) :
|gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) (v + t d) idx - gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) v idx| 2 * a ^ 2 / (1 - 2 * (a * D)) * (t * D)

Segment-Lipschitz gradient for the linear softmax-CE loss, explicit constant. Under the small-step condition 2aD < 1, the gradient entries drift by at most (2a²/(1−2aD))·(t·D) along [v, v+d] — the exact shape descent_segment consumes. The exponential softmax perturbation is linearized by the γ-form, not the mean value theorem.

theorem Proofs.linear_sgd_descends {m n : } (W : Mat m n) (b : Vec n) (x : Vec m) (label : Fin n) (gh : Vec (m * n)) {lr η a : } (ha : 0 a) (hx : ∀ (i : Fin m), |x i| a) (hlr : 0 lr) ( : 0 η) (hgh : ∀ (idx : Fin (m * n)), |gh idx - gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| η) (hsmall : 2 * (a * (lr * (idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| + ↑(m * n) * η))) < 1) (h1 : lr * η * idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| (lr * idx : Fin (m * n), gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx ^ 2) / 4) (h2 : 2 * a ^ 2 / (1 - 2 * (a * (lr * (idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| + ↑(m * n) * η)))) * (lr * (idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| + ↑(m * n) * η)) ^ 2 (lr * idx : Fin (m * n), gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx ^ 2) / 4) :
crossEntropy n (dense (Mat.unflatten (W.flatten - lr gh)) b x) label crossEntropy n (dense (Mat.unflatten W.flatten) b x) label - (lr * idx : Fin (m * n), gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx ^ 2) / 2

One inexact SGD step on the MNIST-linear classifier provably decreases the cross-entropy loss. All of sgd_descends' hypotheses discharged for the Chapter-1 net: differentiability is lossWeightMap_differentiable, the segment-Lipschitz constant is the explicit C = 2a²/(1−2aD) at step radius D = lr·(‖∇L‖₁ + mn·η). Remaining hypotheses are checkable arithmetic: the oracle accuracy η (supplied by the float budgets), the small-step condition, and the two dominance conditions. Conclusion: the loss drops by ≥ lr·‖∇L‖₂²/2.

noncomputable def Proofs.FloatModel.linearFloatGrad (M : FloatModel) {m n : } (W : Mat m n) (b : Vec n) (x : Vec m) (fexp : ) (label : Fin n) :
Vec (m * n)

The binary32 gradient of the MNIST-linear loss, exactly as the rendered trainer computes it: float forward logits z̃ = M.dense W b x, the rounded softmax−onehot cotangent head, and one final rounded multiply by the (exact) input xᵢ to form the outer-product weight gradient ∂L/∂Wᵢⱼ = xᵢ·(softmax(z)ⱼ − onehotⱼ). Flattened to the Vec (m*n) parameter layout that gradAt/linear_sgd_descends use.

Equations
Instances For
    @[simp]
    theorem Proofs.linearFloatGrad_apply (M : FloatModel) {m n : } (W : Mat m n) (b : Vec n) (x : Vec m) (fexp : ) (label : Fin n) (i : Fin m) (j : Fin n) :
    M.linearFloatGrad W b x fexp label (finProdFinEquiv (i, j)) = M.mul (x i) (M.softmaxCECotF fexp (M.dense W b x) label j)
    theorem Proofs.linear_grad_close {m n : } (M : FloatModel) (W : Mat m n) (b : Vec n) (x : Vec m) (label : Fin n) (fexp : ) {eexp δ a : } (hx : ∀ (i : Fin m), |x i| a) (heexp0 : 0 eexp) (heexp1 : eexp 1) (hfexp : ∀ (t : ), |fexp t - Real.exp t| eexp * Real.exp t) (hρ1 : FloatModel.smRho M.u eexp n < 1) ( : ∀ (k' : Fin n), |M.dense W b x k' - dense W b x k'| δ) (i : Fin m) (j : Fin n) :
    |M.linearFloatGrad W b x fexp label (finProdFinEquiv (i, j)) - gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten (finProdFinEquiv (i, j))| FloatModel.mulErr M.u a 1 0 (FloatModel.cotErr M.u eexp δ n)

    The binary32 gradient is within mulErr u a 1 0 (cotErr …) of the certified real gradient, per entry. The head accuracy is the existing softmax_ce_cot_close (cotErr); the final input-multiply is one mul_close with an exact left operand (ea = 0) bounded by a, and a right operand softmax−onehot ∈ [−1,1] (C = 1). This is the bridge that discharges linear_sgd_descends' abstract η.

    theorem Proofs.linear_float_sgd_descends {m n : } (M : FloatModel) (W : Mat m n) (b : Vec n) (x : Vec m) (label : Fin n) (fexp : ) {lr a eexp δ : } (ha : 0 a) (hx : ∀ (i : Fin m), |x i| a) (hlr : 0 lr) (heexp0 : 0 eexp) (heexp1 : eexp 1) (hδ0 : 0 δ) (hfexp : ∀ (t : ), |fexp t - Real.exp t| eexp * Real.exp t) (hρ1 : FloatModel.smRho M.u eexp n < 1) ( : ∀ (k' : Fin n), |M.dense W b x k' - dense W b x k'| δ) (hsmall : 2 * (a * (lr * (idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| + ↑(m * n) * FloatModel.mulErr M.u a 1 0 (FloatModel.cotErr M.u eexp δ n)))) < 1) (h1 : lr * FloatModel.mulErr M.u a 1 0 (FloatModel.cotErr M.u eexp δ n) * idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| (lr * idx : Fin (m * n), gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx ^ 2) / 4) (h2 : 2 * a ^ 2 / (1 - 2 * (a * (lr * (idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| + ↑(m * n) * FloatModel.mulErr M.u a 1 0 (FloatModel.cotErr M.u eexp δ n))))) * (lr * (idx : Fin (m * n), |gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx| + ↑(m * n) * FloatModel.mulErr M.u a 1 0 (FloatModel.cotErr M.u eexp δ n))) ^ 2 (lr * idx : Fin (m * n), gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx ^ 2) / 4) :
    crossEntropy n (dense (Mat.unflatten (W.flatten - lr M.linearFloatGrad W b x fexp label)) b x) label crossEntropy n (dense (Mat.unflatten W.flatten) b x) label - (lr * idx : Fin (m * n), gradAt (fun (w : Vec (m * n)) => crossEntropy n (dense (Mat.unflatten w) b x) label) W.flatten idx ^ 2) / 2

    One binary32 SGD step on the MNIST-linear classifier provably decreases the cross-entropy loss — with NO abstract gradient-accuracy parameter. This is Item D / G1, the η-composition: the descent side (linear_sgd_descends) and the rounding side (FloatBridge's cotErr/mulErr head budget) are fused into one statement. The gradient gh is the actual float-computed gradient (M.linearFloatGrad), and its accuracy η = mulErr u a 1 0 (cotErr …) is proven by linear_grad_close, not assumed.

    What remains as hypotheses is exactly the honest residue: the input bound a, 0 ≤ lr, the GPU exp accuracy eexp and the a-posteriori logit drift δ (the documented FloatModel → kernel trust boundary, softmax_ce_cot_close), and the checkable-arithmetic small-step + two dominance conditions. Depth-1 means there is no per-layer η-threading — the clean pilot for the chain binary32 → proximity → smoothness → descent, closed end-to-end for one net.