Optimizer Law Interface

This module gives TorchLean optimizers a small proof layer interface.

Runtime optimizers live in NN.Runtime.Optim.Optimizers as executable tensor equations. The definitions below package those equations as shape-polymorphic optimizers and provide a common interface for independent update specifications.

The pattern for adding an optimizer is:

1. define a pure per-tensor init and update equation;
2. package it as a TensorOptimizer;
3. state an independent StepSpec when a proof-facing recurrence is needed;
4. prove optimizer-specific algebraic facts as consequences of that generic interface.

TorchLean does not register a second, definitionally identical copy of every runtime update. Such a copy would add a theorem name without adding an independent claim. Higher-level trainer proofs can instead quantify over any TensorOptimizer, reason about whole gradient streams via runSteps, and introduce a StepSpec only when its equations come from a separate mathematical description.

structure Optim.TensorOptimizer (α : Type) [Context α] : Type 1

A shape-polymorphic per-tensor optimizer.

• State : Spec.Shape → Type

  Per-parameter optimizer state for a tensor of shape s.

• init {s : Spec.Shape} : Spec.Tensor α s → self.State s

  Initialize optimizer state from the current parameter tensor.

• update {s : Spec.Shape} : self.State s → Spec.Tensor α s → Spec.Tensor α s → self.State s × Spec.Tensor α s

  One update from state, parameters, and gradients.

def Optim.TensorOptimizer.sgd {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr : α) : TensorOptimizer α

Package plain SGD as a TensorOptimizer.

def Optim.TensorOptimizer.momentumSGD {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr momentum : α) : TensorOptimizer α

Package momentum SGD as a TensorOptimizer.

def Optim.TensorOptimizer.adagrad {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr epsilon : α) : TensorOptimizer α

Package AdaGrad as a TensorOptimizer.

def Optim.TensorOptimizer.rmsprop {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr decay epsilon : α) : TensorOptimizer α

Package RMSProp as a TensorOptimizer.

def Optim.TensorOptimizer.adam {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr beta1 beta2 epsilon : α) : TensorOptimizer α

Package Adam as a TensorOptimizer.

def Optim.TensorOptimizer.adamw {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr weightDecay beta1 beta2 epsilon : α) : TensorOptimizer α

Package AdamW as a TensorOptimizer.

def Optim.TensorOptimizer.adadelta {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr rho epsilon : α) : TensorOptimizer α

Package Adadelta as a TensorOptimizer.

def Optim.TensorOptimizer.muon {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (lr momentum : α) (orthogonalizer : {s : Spec.Shape} → Muon.Orthogonalizer α s := fun {s : Spec.Shape} => Muon.identityOrthogonalizer) : TensorOptimizer α

Package Muon-style orthogonalized momentum as a TensorOptimizer.

@[reducible, inline]

abbrev Optim.TensorOptimizer.Step {α : Type} [Context α] (opt : TensorOptimizer α) (s : Spec.Shape) : Type

State/parameter pair threaded by an optimizer for one fixed tensor shape.

def Optim.TensorOptimizer.step {α : Type} [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : opt.Step s) (grads : Spec.Tensor α s) : opt.Step s

Run one optimizer step on a state/parameter pair.

def Optim.TensorOptimizer.runSteps {α : Type} [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} : opt.Step s → List (Spec.Tensor α s) → opt.Step s

Run a finite stream of gradients through an optimizer.

theorem Optim.TensorOptimizer.runSteps_append {α : Type} [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : opt.Step s) (left right : List (Spec.Tensor α s)) : opt.runSteps current (left ++ right) = opt.runSteps (opt.runSteps current left) right

Splitting a gradient stream and running the two pieces sequentially gives the same state and parameters as running the concatenated stream.

def Optim.TensorOptimizer.stateAfter {α : Type} [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : opt.Step s) (grads : List (Spec.Tensor α s)) : opt.State s

Optimizer state after a finite gradient stream.

def Optim.TensorOptimizer.paramsAfter {α : Type} [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : opt.Step s) (grads : List (Spec.Tensor α s)) : Spec.Tensor α s

Optimizer parameters after a finite gradient stream.

theorem Optim.TensorOptimizer.stateAfter_append {α : Type} [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : opt.Step s) (left right : List (Spec.Tensor α s)) : opt.stateAfter current (left ++ right) = opt.stateAfter (opt.runSteps current left) right

State projection of runSteps_append.

theorem Optim.TensorOptimizer.paramsAfter_append {α : Type} [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : opt.Step s) (left right : List (Spec.Tensor α s)) : opt.paramsAfter current (left ++ right) = opt.paramsAfter (opt.runSteps current left) right

Parameter projection of runSteps_append.

Generic step specifications

structure Optim.StepSpec {α : Type} [Context α] (opt : TensorOptimizer α) : Type

Proof-facing specification of one optimizer step.

An optimizer-specific file only has to identify the next-state and next-parameter equations once. The generic theorems below then lift that one-step fact to whole finite gradient streams.

• nextState {s : Spec.Shape} : opt.State s → Spec.Tensor α s → Spec.Tensor α s → opt.State s

  Spec equation for the next optimizer state.

• nextParams {s : Spec.Shape} : opt.State s → Spec.Tensor α s → Spec.Tensor α s → Spec.Tensor α s

  Spec equation for the next parameter tensor.

• update_eq {s : Spec.Shape} (state : opt.State s) (params grads : Spec.Tensor α s) : opt.update state params grads = (self.nextState state params grads, self.nextParams state params grads)

  The executable optimizer update agrees with the stated step equations.

def Optim.StepSpec.step {α : Type} [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} (current : opt.Step s) (grads : Spec.Tensor α s) : opt.Step s

Run one step through the proof layer equations.

def Optim.StepSpec.runSteps {α : Type} [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} : opt.Step s → List (Spec.Tensor α s) → opt.Step s

Run a finite stream of gradients through the proof layer equations.

theorem Optim.StepSpec.step_eq_optimizer_step {α : Type} [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} (current : opt.Step s) (grads : Spec.Tensor α s) : law.step current grads = opt.step current grads

A registered step spec agrees with the executable optimizer for one step.

theorem Optim.StepSpec.runSteps_eq_optimizer_runSteps {α : Type} [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} (current : opt.Step s) (grads : List (Spec.Tensor α s)) : law.runSteps current grads = opt.runSteps current grads

A registered one-step optimizer spec agrees with the executable optimizer over any finite gradient stream. This is the general theorem optimizer-specific registrations feed into.

theorem Optim.StepSpec.runSteps_append {α : Type} [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} (current : opt.Step s) (left right : List (Spec.Tensor α s)) : law.runSteps current (left ++ right) = law.runSteps (law.runSteps current left) right

The proof layer equations compose over concatenated gradient streams just like the executable optimizer.

Muon comparison laws

theorem Optim.Muon.update_params_eq_momentumSGD_of_apply_eq {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) (happly : state.orthogonalizer.apply (OptimizerUtils.updateMomentumBuf state.buf state.momentum grads) = OptimizerUtils.updateMomentumBuf state.buf state.momentum grads) : (update state params grads).2 = (MomentumSGD.update { lr := state.lr, momentum := state.momentum, buf := state.buf } params grads).2

If a Muon backend returns the fresh momentum buffer unchanged on this step, then the parameter update agrees with momentum SGD for this step.

theorem Optim.Muon.init_update_params_eq_momentumSGD_of_apply_eq {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr momentum : α) (orthogonalizer : Orthogonalizer α s) (params grads : Spec.Tensor α s) (happly : orthogonalizer.apply (OptimizerUtils.updateMomentumBuf (Spec.fill 0 s) momentum grads) = OptimizerUtils.updateMomentumBuf (Spec.fill 0 s) momentum grads) : (update (init lr momentum orthogonalizer params) params grads).2 = (MomentumSGD.update (MomentumSGD.init lr momentum params) params grads).2

Initialized version of update_params_eq_momentumSGD_of_apply_eq.