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.TorchLean.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 weight_decay coefficient.

The generic abstraction lives one layer up:

• Runtime.Autograd.TorchLean.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

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 utilities

def Optim.OptimizerUtils.updateMomentumBuf {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (buf : Spec.Tensor α s) (momentum : α) (grads : Spec.Tensor α s) : Spec.Tensor α s

Momentum-style buffer update $\mu\,\mathtt{buf}+g$.

def Optim.OptimizerUtils.mkAdaptiveLR {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr epsilon : α) (denom : Spec.Tensor α s) : Spec.Tensor α s

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

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

SGD

structure Optim.SGD.State (α : Type) (s : Spec.Shape) : Type

SGD state (per parameter tensor).

We only store the learning rate here.

• lr : α

  Learning rate.

def Optim.SGD.init {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr : α) : Spec.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} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : Spec.Tensor α 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) (s : Spec.Shape) : Type

Momentum SGD state (per parameter tensor).

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

• $\mathtt{buf}\gets\mu\,\mathtt{buf}+g$,
• $p\gets p-\mathtt{lr}\,\mathtt{buf}$.

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

• lr : α

  Learning rate.

• momentum : α

  Momentum coefficient $\mu$.

• buf : Spec.Tensor α s

  Momentum buffer buf.

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

Initialize momentum SGD with a zero buffer.

def Optim.MomentumSGD.update {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : State α s × Spec.Tensor α s

One momentum-SGD step (returns updated state and parameters).

AdaGrad

structure Optim.AdaGrad.State (α : Type) (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)$.

• lr : α

  Base learning rate.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• accumulator : Spec.Tensor α s

  Accumulated squared gradients.

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

Initialize AdaGrad with zero accumulator.

def Optim.AdaGrad.update {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : State α s × Spec.Tensor α s

One AdaGrad step (returns updated state and parameters).

RMSProp

structure Optim.RMSProp.State (α : Type) (s : Spec.Shape) : Type

RMSProp state (per parameter tensor).

We store an EMA of squared gradients (accumulator), often called square_avg in PyTorch code.

• lr : α

  Learning rate.

• decay : α

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

• epsilon : α

  Numerical stability constant $\varepsilon$.

• accumulator : Spec.Tensor α s

  EMA of squared gradients.

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

Initialize RMSProp with zero accumulator.

def Optim.RMSProp.update {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : State α s × Spec.Tensor α s

One RMSProp step (returns updated state and parameters).

Adam

structure Optim.Adam.State (α : Type) (s : Spec.Shape) : Type

Adam state (per parameter tensor).

We store first/second moment EMAs (m, v) and a step counter t used for bias correction.

• lr : α

  Learning rate.

• beta1 : α

  First moment decay $\beta_1$.

• beta2 : α

  Second moment decay $\beta_2$.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• m : Spec.Tensor α s

  First moment EMA.

• v : Spec.Tensor α s

  Second moment EMA.

• t : ℕ

  Step counter (used for bias correction).

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

Initialize Adam with $m=0$, $v=0$, and $t=0$.

def Optim.Adam.update {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : State α s × Spec.Tensor α s

One Adam step (returns updated state and parameters).

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_t {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : (update state params grads).1.t = state.t + 1

Adam increments its step counter by one on every update.

AdamW

structure Optim.AdamW.State (α : Type) (s : Spec.Shape) : Type

AdamW state (per parameter tensor).

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

• lr : α

  Learning rate.

• beta1 : α

  First moment decay $\beta_1$.

• beta2 : α

  Second moment decay $\beta_2$.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• weight_decay : α

  Weight decay coefficient wd.

• m : Spec.Tensor α s

  First moment EMA.

• v : Spec.Tensor α s

  Second moment EMA.

• t : ℕ

  Step counter (used for bias correction).

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

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

def Optim.AdamW.update {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : State α s × Spec.Tensor α s

One AdamW step (returns updated state and parameters).

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_t {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : (update state params grads).1.t = state.t + 1

AdamW increments its step counter by one on every update.

Adadelta

structure Optim.Adadelta.State (α : Type) (s : Spec.Shape) : Type

Adadelta state (per parameter tensor).

We store two EMAs:

• v: EMA of squared gradients,
• u: EMA of squared updates.

• lr : α

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

• rho : α

  Decay coefficient $\rho$.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• v : Spec.Tensor α s

  EMA of squared gradients.

• u : Spec.Tensor α s

  EMA of squared updates.

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

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

def Optim.Adadelta.update {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : State α s × Spec.Tensor α s

One Adadelta step (returns updated state and parameters).

Elementwise equations:

• $v\gets\rho v+(1-\rho)g^2$,
• $\Delta p\gets-\mathtt{lr}\, \dfrac{\sqrt{u+\varepsilon}}{\sqrt{v+\varepsilon}}\odot g$,
• $p\gets p+\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) (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 : Spec.Tensor α full → Spec.Tensor α low

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

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

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

def Optim.GaLore.identityProjector {α : Type} {s : Spec.Shape} : Projector α s s

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

structure Optim.GaLore.SGDState (α : Type) (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.

• lr : α

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

• projector : Projector α full low

  Current gradient projector.

def Optim.GaLore.projectedSGDUpdate {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {full low : Spec.Shape} (state : SGDState α full low) (params grads : Spec.Tensor α full) : Spec.Tensor α full

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

Muon-style orthogonalized momentum

structure Optim.Muon.Orthogonalizer (α : Type) (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 : Spec.Tensor α s → Spec.Tensor α s

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

def Optim.Muon.identityOrthogonalizer {α : Type} {s : Spec.Shape} : Orthogonalizer α s

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

structure Optim.Muon.State (α : Type) (s : Spec.Shape) : Type

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

• lr : α

  Learning rate.

• momentum : α

  Momentum coefficient.

• buf : Spec.Tensor α s

  Momentum buffer.

• orthogonalizer : Orthogonalizer α s

  Backend that turns the momentum buffer into the update direction.

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

Initialize Muon-style state with a zero momentum buffer.

def Optim.Muon.update {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (state : State α s) (params grads : Spec.Tensor α s) : State α s × Spec.Tensor α 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.