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) [TorchLean.Storage α] [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} : TorchLean.Tensor α s → self.State s

  Initialize optimizer state from the current parameter tensor.

• update {s : Spec.Shape} : self.State s → TorchLean.Tensor α s → TorchLean.Tensor α s → Step α s (self.State s)

  One update from state, parameters, and gradients.

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

Package plain SGD as a TensorOptimizer.

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

Package momentum SGD as a TensorOptimizer.

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

Package AdaGrad as a TensorOptimizer.

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

Package RMSProp as a TensorOptimizer.

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

Package Adam as a TensorOptimizer.

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

Package AdamW as a TensorOptimizer.

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

Package Adadelta as a TensorOptimizer.

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

Package Muon-style orthogonalized momentum as a TensorOptimizer.

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

Run one optimizer step on its current state and parameters.

def Optim.TensorOptimizer.runSteps {α : Type} [TorchLean.Storage α] [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : Step α s (opt.State s)) (gradients : Array (TorchLean.Tensor α s)) : Step α s (opt.State s)

Run a finite stream of gradients through an optimizer.

theorem Optim.TensorOptimizer.runSteps_append {α : Type} [TorchLean.Storage α] [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : Step α s (opt.State s)) (left right : Array (TorchLean.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} [TorchLean.Storage α] [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : Step α s (opt.State s)) (gradients : Array (TorchLean.Tensor α s)) : opt.State s

Optimizer state after a finite gradient stream.

def Optim.TensorOptimizer.parametersAfter {α : Type} [TorchLean.Storage α] [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : Step α s (opt.State s)) (gradients : Array (TorchLean.Tensor α s)) : TorchLean.Tensor α s

Optimizer parameters after a finite gradient stream.

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

State projection of runSteps_append.

theorem Optim.TensorOptimizer.parametersAfter_append {α : Type} [TorchLean.Storage α] [Context α] (opt : TensorOptimizer α) {s : Spec.Shape} (current : Step α s (opt.State s)) (left right : Array (TorchLean.Tensor α s)) : opt.parametersAfter current (left ++ right) = opt.parametersAfter (opt.runSteps current left) right

Parameter projection of runSteps_append.

Generic step specifications

structure Optim.StepSpec {α : Type} [TorchLean.Storage α] [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 → TorchLean.Tensor α s → TorchLean.Tensor α s → opt.State s

  Spec equation for the next optimizer state.

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

  Spec equation for the next parameter tensor.

• update_eq {s : Spec.Shape} (state : opt.State s) (parameters gradients : TorchLean.Tensor α s) : opt.update state parameters gradients = { optimizerState := self.nextState state parameters gradients, parameters := self.nextParameters state parameters gradients }

  The executable optimizer update agrees with the stated step equations.

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

Run one step through the proof layer equations.

def Optim.StepSpec.runSteps {α : Type} [TorchLean.Storage α] [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} (current : Step α s (opt.State s)) (gradients : Array (TorchLean.Tensor α s)) : Step α s (opt.State s)

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

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

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

theorem Optim.StepSpec.runSteps_eq_optimizer_runSteps {α : Type} [TorchLean.Storage α] [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} (current : Step α s (opt.State s)) (gradients : Array (TorchLean.Tensor α s)) : law.runSteps current gradients = opt.runSteps current gradients

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} [TorchLean.Storage α] [Context α] {opt : TensorOptimizer α} (law : StepSpec opt) {s : Spec.Shape} (current : Step α s (opt.State s)) (left right : Array (TorchLean.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} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) (happly : state.orthogonalizer.apply (updateMomentumBuffer state.momentumBuffer state.momentum gradients) = updateMomentumBuffer state.momentumBuffer state.momentum gradients) : (update state parameters gradients).parameters = (MomentumSGD.update { learningRate := state.learningRate, momentum := state.momentum, momentumBuffer := state.momentumBuffer } parameters gradients).parameters

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} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate momentum : α) (orthogonalizer : Orthogonalizer α s) (parameters gradients : TorchLean.Tensor α s) (happly : orthogonalizer.apply (updateMomentumBuffer (TorchLean.Tensor.full s 0) momentum gradients) = updateMomentumBuffer (TorchLean.Tensor.full s 0) momentum gradients) : (update (init learningRate momentum orthogonalizer parameters) parameters gradients).parameters = (MomentumSGD.update (MomentumSGD.init learningRate momentum parameters) parameters gradients).parameters

Initialized version of update_params_eq_momentumSGD_of_apply_eq.