Optim

TorchLean optimizer wrappers.

This connects the pure tensor optimizers in NN/Runtime/Optim/Optimizers.lean to the runtime training structures (Torch.ParamList + gradient TorchLean.TensorPack).

Design notes:

• Optimizer state is stored in a shape-indexed list aligned with the parameter shapes.
• Trainable storage aliases share one history and one update from their summed gradients.
• Updates run on plain tensors (not via the autograd tape), so they work the same for eager and typed graph training loops.
• Parameters marked requiresGrad := false are left unchanged (state is preserved).

PyTorch references

• torch.optim overview: https://pytorch.org/docs/stable/optim.html
• torch.optim.SGD: https://pytorch.org/docs/stable/generated/torch.optim.SGD.html
• torch.optim.Adam: https://pytorch.org/docs/stable/generated/torch.optim.Adam.html
• torch.optim.AdamW: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html

For the core math and algorithm-level citations (Adam, AdamW, RMSProp, etc.), see NN/Runtime/Optim/Optimizers.lean.

Generic optimizer interface

structure Runtime.Autograd.Model.Optim.StepWithLoss (α : Type) [TorchLean.Storage α] (OptimizerState : Type) : Type

Optimizer state paired with the loss produced by the same training step.

• optimizerState : OptimizerState

  Optimizer state to use for the next training step.

• loss : TorchLean.Tensor α (Spec.Shape.ofList [])

  Scalar loss whose backward pass produced this update.

inductive Runtime.Autograd.Model.Optim.StateList (State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type) (α : Type) [TorchLean.Storage α] : List Spec.Shape → Type

A shape-indexed list of optimizer state values.

This mirrors the parameter-shape list used by Torch.ParamList. Trainable storage aliases carry copies of one immutable optimizer state. Initialize through the optimizer and retain the complete returned list between steps; alias entries must not be edited independently.

• nil {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} {α : Type} [TorchLean.Storage α] : StateList State α []
• cons {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} {α : Type} [TorchLean.Storage α] {s : Spec.Shape} {ss : List Spec.Shape} : State α s → StateList State α ss → StateList State α (s :: ss)

structure Runtime.Autograd.Model.Optim.Optimizer (α : Type) [TorchLean.Storage α] [Context α] (paramShapes : List Spec.Shape) : Type 1

Runtime-facing optimizer interface.

This is the analogue of a PyTorch torch.optim.Optimizer, but made explicit about:

• which parameter shapes it manages (paramShapes), and
• how it stores internal state (State) aligned with those shapes.

• State : Type

  The optimizer's own state type: momentum buffers, second moments, step counts. Leaving it abstract here is what lets SGD and Adam share one interface despite storing different things.

• init : Torch.ParamList α paramShapes → IO self.State

  Allocate initial state for the given parameters.

• step : self.State → Torch.ParamList α paramShapes → TorchLean.TensorPack α paramShapes → IO self.State

  One update: state, parameters and gradients in, new state out. Parameters are updated in place through ParamList, which is why only the state is returned.

• trainerStep? {β : Type} [TorchLean.Storage β] {inputShapes dataInputShapes : List Spec.Shape} : Torch.ScalarTrainer α β paramShapes inputShapes dataInputShapes → self.State → TorchLean.TensorPack α inputShapes → TorchLean.TensorPack β dataInputShapes → IO (Option self.State)

  Optional trainer-native step.

  Most optimizers are implemented by first materializing a gradient TorchLean.TensorPack and then updating host parameter tensors. Some trainers can do better. In eager CUDA mode, for example, the trainer can keep gradients and optimizer moments on device for SGD/Adam. When this hook returns some st', callers should treat the step as complete and use st' as the next optimizer state. Returning none asks the caller to fall back to the generic backward + step path.

• trainerStepWithLoss? {β : Type} [TorchLean.Storage β] {inputShapes dataInputShapes : List Spec.Shape} : Torch.ScalarTrainer α β paramShapes inputShapes dataInputShapes → self.State → TorchLean.TensorPack α inputShapes → TorchLean.TensorPack β dataInputShapes → IO (Option (StepWithLoss α self.State))

  Optional trainer-native step that returns the loss used for the update.

  This is the logging/inspection counterpart of trainerStep?. The returned scalar was evaluated on the same tape that produced the gradients; none requests the generic lossAndBackward + step path.

• trainerBatchStep? {β : Type} [TorchLean.Storage β] {inputShapes dataInputShapes : List Spec.Shape} : Torch.ScalarTrainer α β paramShapes inputShapes dataInputShapes → self.State → Array (TorchLean.TensorPack α inputShapes × TorchLean.TensorPack β dataInputShapes) → (readLoss : Bool) → IO (Option (self.State × Option (TorchLean.Tensor α [])))

  Optional native mean-gradient update. The outer option reports whether the backend handled the update; the inner loss is present exactly when readLoss was requested.

structure Runtime.Autograd.Model.Optim.Internal.SomeState (State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type) (α : Type) [TorchLean.Storage α] : Type

An optimizer state with its shape, used to copy a canonical state into later alias slots.

• shape : Spec.Shape
• state : State α self.shape

def Runtime.Autograd.Model.Optim.Internal.canonicalState {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} {α : Type} [TorchLean.Storage α] (shape : Spec.Shape) (states : Array (SomeState State α)) (index : ℕ) : IO (State α shape)

Retrieve an earlier canonical state without comparing state or tensor contents.

def Runtime.Autograd.Model.Optim.Internal.nativeBatchStep {α β State : Type} [TorchLean.Storage α] [TorchLean.Storage β] {paramShapes inputShapes dataInputShapes : List Spec.Shape} (optimizer : Torch.NativeOptimizer α) (trainer : Torch.ScalarTrainer α β paramShapes inputShapes dataInputShapes) (state : State) (batch : Array (TorchLean.TensorPack α inputShapes × TorchLean.TensorPack β dataInputShapes)) (readLoss : Bool) : IO (Option (State × Option (TorchLean.Tensor α [])))

Invoke a backend batch update while retaining the wrapper's scheduled optimizer state.

def Runtime.Autograd.Model.Optim.Internal.firstStateValue {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} {α β : Type} [TorchLean.Storage α] (fallback : β) (get : {s : Spec.Shape} → State α s → β) {ss : List Spec.Shape} : StateList State α ss → β

Read a shape-independent field from the first optimizer state, or use fallback when there are no parameters.

def Runtime.Autograd.Model.Optim.Internal.initStateList {α : Type} [TorchLean.Storage α] [Context α] {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} {ss : List Spec.Shape} (initOne : {s : Spec.Shape} → TorchLean.Tensor α s → State α s) (parameters : Torch.ParamList α ss) : IO (StateList State α ss)

Initialize once per trainable storage and copy that immutable state into its alias slots.

Frozen slots retain independent placeholder states so the public shape-indexed list is unchanged. All alias decisions use the same canonical-slot map as SGD and checkpoints.

def Runtime.Autograd.Model.Optim.Internal.initStateList.buildStates {α : Type} [TorchLean.Storage α] {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} (initOne : {s : Spec.Shape} → TorchLean.Tensor α s → State α s) (slots : Array (Option ℕ)) {shapes : List Spec.Shape} : Torch.ParamList α shapes → ℕ → Array (SomeState State α) → IO (StateList State α shapes)

def Runtime.Autograd.Model.Optim.Internal.stepStateList {α : Type} [TorchLean.Storage α] [Context α] {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} {ss : List Spec.Shape} (updateOne : {s : Spec.Shape} → State α s → TorchLean.Tensor α s → TorchLean.Tensor α s → Optim.Step α s (State α s)) (parameters : Torch.ParamList α ss) (states : StateList State α ss) (gradients : TorchLean.TensorPack α ss) : IO (StateList State α ss)

Run one optimizer update per trainable storage using its summed occurrence gradients.

updateOne is called only at canonical slots. Its one immutable returned state is copied into all aliases with a shape check; frozen states are preserved.

Alias states must remain coherent, as produced by initialization or previous steps for the same storage layout. Generic State has no equality operation: divergent external histories are outside this contract and are not detected or merged. Checkpoint restoration must preserve this invariant.

def Runtime.Autograd.Model.Optim.Internal.stepStateList.update {α : Type} [TorchLean.Storage α] {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} (updateOne : {s : Spec.Shape} → State α s → TorchLean.Tensor α s → TorchLean.Tensor α s → Optim.Step α s (State α s)) (slots : Array (Option ℕ)) (sums : Array (Option (Spec.SomeTensor α))) {shapes : List Spec.Shape} : Torch.ParamList α shapes → StateList State α shapes → ℕ → Array (SomeState State α) → IO (StateList State α shapes)

Concrete optimizers

def Runtime.Autograd.Model.Optim.sgd {α : Type} [TorchLean.Storage α] [Context α] (learningRate : α) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

Stochastic gradient descent.

PyTorch analogy: torch.optim.SGD(lr=lr) without momentum.

def Runtime.Autograd.Model.Optim.momentumSGD {α : Type} [TorchLean.Storage α] [Context α] (learningRate momentum : α) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

SGD with classical momentum.

PyTorch analogy: torch.optim.SGD(lr=lr, momentum=momentum).

def Runtime.Autograd.Model.Optim.adagrad {α : Type} [TorchLean.Storage α] [Context α] (learningRate epsilon : α) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

AdaGrad (per-parameter learning rate scaling by accumulated squared gradients).

PyTorch analogy: torch.optim.Adagrad(lr=lr, eps=epsilon).

def Runtime.Autograd.Model.Optim.rmsprop {α : Type} [TorchLean.Storage α] [Context α] (learningRate decay epsilon : α) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

RMSProp (exponentially-decayed second moment / running average of squared gradients).

PyTorch analogy: torch.optim.RMSprop(lr=lr, alpha=decay, eps=epsilon) (we use the common naming decay for alpha).

def Runtime.Autograd.Model.Optim.adam {α : Type} [TorchLean.Storage α] [Context α] (learningRate beta1 beta2 epsilon : α) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

Adam (first/second moment estimates).

PyTorch analogy: torch.optim.Adam(lr=lr, betas=(beta1,beta2), eps=epsilon).

def Runtime.Autograd.Model.Optim.adamw {α : Type} [TorchLean.Storage α] [Context α] (learningRate weightDecay beta1 beta2 epsilon : α) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

AdamW (Adam with decoupled weight decay).

PyTorch analogy: torch.optim.AdamW(lr=lr, weight_decay=weightDecay, betas=(beta1,beta2), eps=epsilon).

def Runtime.Autograd.Model.Optim.adadelta {α : Type} [TorchLean.Storage α] [Context α] (learningRate rho epsilon : α) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

AdaDelta (adaptive learning rate method similar to RMSProp but with a running RMS of updates).

PyTorch analogy: torch.optim.Adadelta(lr=lr, rho=rho, eps=epsilon).

Optimizer extension points

def Runtime.Autograd.Model.Optim.projectedSGD {α : Type} [TorchLean.Storage α] [Context α] (learningRate : α) (projector : {s : Spec.Shape} → Optim.GaLore.Projector α s s := fun {s : Spec.Shape} => Optim.GaLore.identityProjector) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

Projected SGD.

This is the runtime-safe part of a GaLore-style optimizer: every parameter gets a same-shape projector/lift pair, and the update applies p ← p - lr * lift(project(g)).

Full GaLore also needs a rank-changing projector and a refresh schedule. Those pieces require matrix-specific state and SVD/randomized-SVD infrastructure, so they are not hidden inside this generic constructor.

def Runtime.Autograd.Model.Optim.muon {α : Type} [TorchLean.Storage α] [Context α] (learningRate momentum : α) (orthogonalizer : {s : Spec.Shape} → Optim.Muon.Orthogonalizer α s := fun {s : Spec.Shape} => Optim.Muon.identityOrthogonalizer) {paramShapes : List Spec.Shape} : Optimizer α paramShapes

Muon-style momentum with a caller-supplied same-shape orthogonalization backend.

Using the identity backend gives ordinary momentum-SGD behavior. A production Muon backend should provide a matrix-specific Newton-Schulz orthogonalizer and optional CUDA kernels.