Optimizers

Optimizers for TorchLean runtime training.

This file implements the core math of common gradient-based optimizers as pure functions on typed tensors Tensor α s.

Why “pure functions”?

In PyTorch, optimizers mutate parameters in-place and keep state in Python objects. In TorchLean, we want the update rule itself to be explicit and easy to reuse:

• eager examples can call the update directly,
• the runtime training engine can store state in maps keyed by parameter ids,
• and proofs can refer to the same update equations.

The intent is to mimic the standard textbook formulas closely. We do not try to reproduce every implementation detail of torch.optim.* (e.g. foreach kernels, fused updates, or every optional flag); those live at a different layer than the math we specify here.

How this file fits with the runtime and API:

• this file owns the scalar-polymorphic, per-tensor update equations;
• NN.Runtime.Autograd.Model.Optim lifts those equations to runtime parameter lists; and
• NN.API.Runtime exposes ergonomic optim.sgd, optim.adam, and related configuration helpers.

With this separation, the formula appears once while runtime adapters and API configuration can evolve independently around it.

Why each optimizer has its own State structure:

• Lean structures do not inherit from one another the way Python classes do.
• Optimizer state is not uniform: SGD stores only lr, momentum SGD stores a buffer, Adam/AdamW store two moment buffers and a step counter, Adadelta stores gradient/update EMAs, Muon carries an orthogonalization backend, and GaLore-style projected updates carry a projection backend.
• Keeping these as separate typed states makes impossible states unrepresentable. For example, an SGD state cannot accidentally contain a stale Adam v buffer, and AdamW cannot forget its decoupled weightDecay coefficient.

The generic abstraction lives one layer up:

• Runtime.Autograd.Model.Optim.Optimizer packages init/step for shape-indexed parameter lists, like a typed analogue of a PyTorch optimizer object.
• Runtime.Autograd.Train.OptimizerState handles dynamic parameter groups and checkpoint-style maps for the training-loop API.

The result is a collection of canonical state records rather than an inheritance hierarchy.

References (original algorithms / common variants):

• AdaGrad (Duchi–Hazan–Singer, 2011): https://jmlr.org/papers/v12/duchi11a.html
• RMSProp (Hinton lecture notes; widely used variant): https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
• Adam (Kingma–Ba, 2015): https://arxiv.org/abs/1412.6980
• AdamW / decoupled weight decay (Loshchilov–Hutter, 2019): https://arxiv.org/abs/1711.05101
• Adadelta (Zeiler, 2012): https://arxiv.org/abs/1212.5701
• SGD + momentum in deep learning practice (Sutskever et al., 2013): https://arxiv.org/abs/1301.4083
• GaLore / low-rank gradient projection (Zhao et al., 2024): https://arxiv.org/abs/2403.03507
• Muon-style momentum with orthogonalized matrix updates (Jordan et al., 2024): https://kellerjordan.github.io/posts/muon/

PyTorch references (for API/parameter naming):

• 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.Adagrad: https://pytorch.org/docs/stable/generated/torch.optim.Adagrad.html
• torch.optim.RMSprop: https://pytorch.org/docs/stable/generated/torch.optim.RMSprop.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
• torch.optim.Adadelta: https://pytorch.org/docs/stable/generated/torch.optim.Adadelta.html

structure Optim.Step (α : Type) [TorchLean.Storage α] (shape : Spec.Shape) (OptimizerState : Type) : Type

Optimizer state and parameters produced by one tensor update.

• optimizerState : OptimizerState

  Optimizer state to use for the next update.

• parameters : TorchLean.Tensor α shape

  Updated parameter tensor.

def Optim.scalarPowNat {α : Type} [One α] [Mul α] (x : α) : ℕ → α

Integer exponentiation for scalar optimizer coefficients.

We use an explicit Nat → α recursion instead of x ^ (n : Nat) because Context α provides Pow α α (for runtime scalar exponentiation), but not Pow α Nat.

Shared equations

def Optim.updateMomentumBuffer {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (momentumBuffer : TorchLean.Tensor α s) (momentum : α) (gradients : TorchLean.Tensor α s) : TorchLean.Tensor α s

Next momentum buffer $\mu b+g$.

def Optim.adaptiveLearningRate {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate epsilon : α) (denominator : TorchLean.Tensor α s) : TorchLean.Tensor α s

Elementwise adaptive learning-rate tensor $\mathtt{learningRate}/(\sqrt{\mathtt{denominator}}+\varepsilon)$.

This is shared by AdaGrad/RMSProp/Adam-style optimizers.

SGD

structure Optim.SGD.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

SGD state (per parameter tensor).

We only store the learning rate here.

• learningRate : α

  Learning rate.

def Optim.SGD.init {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate : α) : TorchLean.Tensor α s → State α s

Initialize SGD state.

The parameter tensor is unused; we keep it in the signature so optimizers share the same “init from parameters” calling convention.

def Optim.SGD.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One SGD step: p ← p - lr * g.

PyTorch analogy: the core of torch.optim.SGD without momentum/weight-decay extras.

Momentum SGD

structure Optim.MomentumSGD.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

Momentum SGD state (per parameter tensor).

We store a momentum buffer and a momentum coefficient $\mu$. Update rule:

• $b\gets\mu b+g$,
• $p\gets p-\mathtt{learningRate}\,b$.

This matches PyTorch's SGD momentum behavior when dampening = 0 and nesterov = false.

• learningRate : α

  Learning rate.

• momentum : α

  Momentum coefficient $\mu$.

• momentumBuffer : TorchLean.Tensor α s

  Momentum buffer.

def Optim.MomentumSGD.init {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate momentum : α) : TorchLean.Tensor α s → State α s

Initialize momentum SGD with a zero buffer.

def Optim.MomentumSGD.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One momentum-SGD step.

AdaGrad

structure Optim.AdaGrad.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

AdaGrad state (per parameter tensor).

We store an accumulator $G$ of squared gradients (same shape as the parameters). The effective step size is scaled by $1/(\sqrt G+\varepsilon)$.

• learningRate : α

  Base learning rate.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• squaredGradientSum : TorchLean.Tensor α s

  Accumulated squared gradients.

def Optim.AdaGrad.init {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate epsilon : α) : TorchLean.Tensor α s → State α s

Initialize AdaGrad with zero accumulator.

def Optim.AdaGrad.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One AdaGrad step.

RMSProp

structure Optim.RMSProp.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

RMSProp state (per parameter tensor).

We store an exponential moving average of squared gradients.

• learningRate : α

  Learning rate.

• decay : α

  Decay coefficient for the EMA of $g^2$ (often called alpha).

• epsilon : α

  Numerical stability constant $\varepsilon$.

• squaredGradientAverage : TorchLean.Tensor α s

  EMA of squared gradients.

def Optim.RMSProp.init {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate decay epsilon : α) : TorchLean.Tensor α s → State α s

Initialize RMSProp with zero accumulator.

def Optim.RMSProp.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One RMSProp step.

Adam

structure Optim.Adam.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

Adam state (per parameter tensor).

We store first and second moment averages and a step counter used for bias correction.

• learningRate : α

  Learning rate.

• beta1 : α

  First moment decay $\beta_1$.

• beta2 : α

  Second moment decay $\beta_2$.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• firstMoment : TorchLean.Tensor α s

  First moment EMA.

• secondMoment : TorchLean.Tensor α s

  Second moment EMA.

• stepCount : ℕ

  Step counter (used for bias correction).

def Optim.Adam.init {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate beta1 beta2 epsilon : α) : TorchLean.Tensor α s → State α s

Initialize Adam with zero moments and a zero step count.

def Optim.Adam.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One Adam step.

Equations (elementwise):

• $m\gets\beta_1m+(1-\beta_1)g$,
• $v\gets\beta_2v+(1-\beta_2)g^2$,
• $\widehat m\gets m/(1-\beta_1^t)$,
• $\widehat v\gets v/(1-\beta_2^t)$,
• $p\gets p-\mathtt{lr}\,\widehat m/(\sqrt{\widehat v}+\varepsilon)$.

The $\varepsilon$ placement matches Kingma and Ba: it is added after $\sqrt{\widehat v}$.

@[simp]

theorem Optim.Adam.update_stepCount {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : (update state parameters gradients).optimizerState.stepCount = state.stepCount + 1

Adam increments its step counter by one on every update.

AdamW

structure Optim.AdamW.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

AdamW state (per parameter tensor).

AdamW is “Adam + decoupled weight decay”. Weight decay is applied as a separate parameter decay term rather than being folded into the gradient that feeds the moments.

• learningRate : α

  Learning rate.

• beta1 : α

  First moment decay $\beta_1$.

• beta2 : α

  Second moment decay $\beta_2$.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• weightDecay : α

  Weight decay coefficient wd.

• firstMoment : TorchLean.Tensor α s

  First moment EMA.

• secondMoment : TorchLean.Tensor α s

  Second moment EMA.

• stepCount : ℕ

  Step counter (used for bias correction).

def Optim.AdamW.init {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate weightDecay beta1 beta2 epsilon : α) : TorchLean.Tensor α s → State α s

Initialize AdamW state for a parameter tensor (moments start at 0).

def Optim.AdamW.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One AdamW step.

We implement the decoupled form from the AdamW paper:

• update Adam moments using the raw gradient g,
• apply weight decay directly to the parameters (p ← p - lr * wd * p),
• then apply the Adam update.

This is the same single-step ordering used by torch.optim.AdamW.

@[simp]

theorem Optim.AdamW.update_stepCount {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : (update state parameters gradients).optimizerState.stepCount = state.stepCount + 1

AdamW increments its step counter by one on every update.

Adadelta

structure Optim.Adadelta.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

Adadelta state (per parameter tensor).

We store two EMAs:

• squaredGradientAverage,
• squaredUpdateAverage.

• learningRate : α

  Learning rate (often set to 1 in some presentations; we keep it explicit).

• rho : α

  Decay coefficient $\rho$.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• squaredGradientAverage : TorchLean.Tensor α s

  EMA of squared gradients.

• squaredUpdateAverage : TorchLean.Tensor α s

  EMA of squared updates.

def Optim.Adadelta.init {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (learningRate rho epsilon : α) : TorchLean.Tensor α s → State α s

Initialize Adadelta state for a parameter tensor (EMAs start at 0).

def Optim.Adadelta.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One Adadelta step.

Elementwise equations:

• $v\gets\rho v+(1-\rho)g^2$,
• $\Delta p\gets \dfrac{\sqrt{u+\varepsilon}}{\sqrt{v+\varepsilon}}\odot g$,
• $p\gets p-\mathtt{lr}\,\Delta p$,
• $u\gets\rho u+(1-\rho)(\Delta p)^2$.

The $\varepsilon$ placement is inside the RMS terms, matching Zeiler's Adadelta update.

Projected / low-rank gradient transforms

structure Optim.GaLore.Projector (α : Type) [TorchLean.Storage α] (full low : Spec.Shape) : Type

A shape-safe gradient projector.

GaLore-style training periodically builds a low-rank subspace for a large matrix parameter, projects the gradient into that subspace, runs a base optimizer there, and lifts the update back to the original parameter shape. This record is the algebraic interface; the expensive policy that computes or refreshes the projector belongs to the runtime layer.

• project : TorchLean.Tensor α full → TorchLean.Tensor α low

  Project a full gradient into the low-rank optimizer space.

• lift : TorchLean.Tensor α low → TorchLean.Tensor α full

  Lift a low-rank update back to the full parameter shape.

def Optim.GaLore.identityProjector {α : Type} [TorchLean.Storage α] {s : Spec.Shape} : Projector α s s

Identity projector, used when projected SGD is requested without a projection backend.

structure Optim.GaLore.SGDState (α : Type) [TorchLean.Storage α] (full low : Spec.Shape) : Type

GaLore-style projected SGD state for one tensor.

This is not a full GaLore implementation by itself: it specifies the update once a projector is available. A practical trainer still needs a refresh schedule and a way to build projectors for large matrix parameters.

• learningRate : α

  Learning rate used after the gradient has been projected and lifted.

• projector : Projector α full low

  Current gradient projector.

def Optim.GaLore.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {full low : Spec.Shape} (state : SGDState α full low) (parameters gradients : TorchLean.Tensor α full) : Step α full (SGDState α full low)

One projected-SGD update: p ← p - learningRate * lift(project(g)).

Muon-style orthogonalized momentum

structure Optim.Muon.Orthogonalizer (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

Orthogonalization backend for a matrix-shaped update.

Muon uses a momentum buffer and then replaces the raw momentum direction by an approximately orthogonalized update, commonly via Newton-Schulz iterations. TorchLean keeps this as an explicit backend so the pure update rule is testable before CUDA kernels are introduced.

• apply : TorchLean.Tensor α s → TorchLean.Tensor α s

  Convert a momentum buffer into the direction used for the parameter update.

def Optim.Muon.identityOrthogonalizer {α : Type} [TorchLean.Storage α] {s : Spec.Shape} : Orthogonalizer α s

The identity orthogonalizer, used when Muon is requested without a matrix backend.

structure Optim.Muon.State (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

Per-parameter state for Muon-style momentum with an explicit orthogonalization backend.

• learningRate : α

  Learning rate.

• momentum : α

  Momentum coefficient.

• momentumBuffer : TorchLean.Tensor α s

  Momentum buffer.

• orthogonalizer : Orthogonalizer α s

  Backend that turns the momentum buffer into the update direction.

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

Initialize Muon-style state with a zero momentum buffer.

def Optim.Muon.update {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One Muon-style update:

• update the momentum buffer,
• orthogonalize the buffer,
• subtract the scaled orthogonalized direction.

For actual Muon, use a matrix-shaped s and a Newton-Schulz orthogonalizer. The generic shape here keeps the definition reusable for tests and for future batched matrix layouts.