StepAlgebra

Pure training-step algebra for proved autograd graphs.

This file stays on the mathematical side of training:

• Graph.scalarLoss_grad_correct specializes the tape/DAG backprop theorem to the scalar-loss convention used by training loops.
• SGD.step lifts the single-tensor SGD equation to a heterogeneous TorchLean.TensorPack of parameters.

It is not another runtime optimizer implementation. Runtime optimizers update parameter references or device buffers; the pure context equation here states their mathematical target, not a proof that those mutations implement it.

PyTorch correspondence / citations

• “Scalar loss” convention: backprop starts with upstream gradient 1. https://pytorch.org/docs/stable/autograd.html
• SGD update rule corresponds to torch.optim.SGD at the level of pure parameter math. https://pytorch.org/docs/stable/generated/torch.optim.SGD.html

def Proofs.Autograd.Algebra.Graph.seedScalarLoss {α : Type} [CommSemiring α] {Γ ss : List Spec.Shape} : TorchLean.TensorPack α (Γ ++ (ss ++ [Spec.Shape.scalar]))

Seed cotangent for a scalar-loss graph.

The full context has shape list Γ ++ ss ++ [scalar]: original inputs/parameters, intermediate nodes, and the final scalar loss. Reverse-mode training seeds all non-output cotangents with zero and the scalar-loss output cotangent with 1, matching loss.backward() in PyTorch. The seed depends only on the shape lists Γ and ss, not on the graph or its inputs.

theorem Proofs.Autograd.Algebra.Graph.scalarLoss_grad_correct {α : Type} [CommSemiring α] {Δ : Type} {Γ ss : List Spec.Shape} (g : Graph Δ Γ (ss ++ [Spec.Shape.scalar])) (x dx : TorchLean.TensorPack α Γ) (d : Δ) : have seed := seedScalarLoss; TensorPack.dotList (g.jvpCtx x dx d) seed = TensorPack.dotList dx (g.backpropCtx x d seed)

Scalar-loss specialization of Graph.backprop_correct.

The left side differentiates the graph by a forward-mode perturbation dx and pairs the result with the scalar-loss seed. The right side pairs the same perturbation with the context cotangent computed by reverse mode. This is the exact algebraic statement behind “the gradient returned by backprop is the cotangent for the scalar loss”.

def Proofs.Autograd.Algebra.SGD.step {α : Type} [CommRing α] {Γ : List Spec.Shape} (params grads : TorchLean.TensorPack α Γ) (lr : α) : TorchLean.TensorPack α Γ

One pure SGD step over a typed parameter context:

params := params - lr * grads

This is the context-level version of the ordinary tensor update p := p - lr * g. Runtime training code may implement the same equation by mutating parameter references, materializing tensors eagerly, or launching CUDA kernels; this definition is the side-effect-free algebraic target those implementations are meant to realize.

@[simp]

theorem Proofs.Autograd.Algebra.SGD.step_nil {α : Type} [CommRing α] (lr : α) : step TorchLean.TensorPack.nil TorchLean.TensorPack.nil lr = TorchLean.TensorPack.nil

The empty parameter context is unchanged by an SGD step.

@[simp]

theorem Proofs.Autograd.Algebra.SGD.step_cons {α : Type} [CommRing α] {s : Spec.Shape} {Γ : List Spec.Shape} (p g : TorchLean.Tensor α s) (ps gs : TorchLean.TensorPack α Γ) (lr : α) : step (TorchLean.TensorPack.cons p ps) (TorchLean.TensorPack.cons g gs) lr = TorchLean.TensorPack.cons (p.subSpec (g.scaleSpec lr)) (step ps gs lr)

An SGD step updates the head tensor by subSpec p (scaleSpec g lr) and applies the same learning rate recursively to the remaining parameters.