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 after the backward pass, so they are not recorded on the autograd tape.
• 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

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.

• setLr : α → self.State → self.State

  The same state with every parameter's learning rate set to the given value; schedules use it between steps.

• native? : self.State → Option (Torch.NativeOptimizer α)

  The update this state names among the trainer's native optimizers, if any. A trainer whose backend keeps parameters, gradients and optimizer moments on the device (eager CUDA) runs it in place of the generic backward + step path, and the state is then kept as is, since the moments live on the device. none always takes the generic path.

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

Run optimizer's native update on a nonempty batch when both the optimizer and the trainer's backend have one; one sample is a batch of one. The outer option reports whether the native path ran, and 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.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.mapStateList {State : (α : Type) → [TorchLean.Storage α] → Spec.Shape → Type} {α : Type} [TorchLean.Storage α] (f : {s : Spec.Shape} → State α s → State α s) {ss : List Spec.Shape} : StateList State α ss → StateList State α ss

Apply f to every per-parameter optimizer state.

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 α] (lr : α) {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 α] (lr 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 α] (lr 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 α] (lr 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 α] (lr 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 α] (lr 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 α] (lr 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 α] (lr : α) (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 α] (lr 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.