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 demos 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.

Where this file sits in the stack:

• 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.

That separation is deliberate: the formula appears once, while runtime adapters and public 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.
• More importantly, 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, and Muon/GaLore carry backend functions.
• 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.

So this file intentionally favors small canonical state records over 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.

@[simp]

theorem Optim.scalarPowNat_zero {α : Type} [One α] [Mul α] (x : α) : scalarPowNat x 0 = 1

Scalar exponentiation starts at 1.

@[simp]

theorem Optim.scalarPowNat_succ {α : Type} [One α] [Mul α] (x : α) (n : ℕ) : scalarPowNat x (n + 1) = scalarPowNat x n * x

Successor case for scalar exponentiation.

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 μ * 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 lr / (sqrt(denom) + ε).

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.

@[simp]

theorem Optim.SGD.init_lr {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr : α) (params : Spec.Tensor α s) : (init lr params).lr = lr

SGD initialization records exactly the requested learning rate.

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 μ. Update rule:

• buf ← μ buf + g
• p ← p - lr * buf

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

• lr : α

  Learning rate.

• momentum : α

  Momentum coefficient μ.

• 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.

@[simp]

theorem Optim.MomentumSGD.init_buf {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr momentum : α) (params : Spec.Tensor α s) : (init lr momentum params).buf = Spec.fill 0 s

Momentum-SGD starts with a zero momentum 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) + ε).

• lr : α

  Base learning rate.

• epsilon : α

  Numerical stability constant ε.

• 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.

@[simp]

theorem Optim.AdaGrad.init_accumulator {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr epsilon : α) (params : Spec.Tensor α s) : (init lr epsilon params).accumulator = Spec.fill 0 s

AdaGrad starts with a zero squared-gradient 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² (often called alpha).

• epsilon : α

  Numerical stability constant ε.

• 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.

@[simp]

theorem Optim.RMSProp.init_accumulator {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr decay epsilon : α) (params : Spec.Tensor α s) : (init lr decay epsilon params).accumulator = Spec.fill 0 s

RMSProp starts with a zero running average of squared gradients.

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 β₁.

• beta2 : α

  Second moment decay β₂.

• epsilon : α

  Numerical stability constant ε.

• 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.

@[simp]

theorem Optim.Adam.init_t {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr beta1 beta2 epsilon : α) (params : Spec.Tensor α s) : (init lr beta1 beta2 epsilon params).t = 0

Adam starts at step 0.

@[simp]

theorem Optim.Adam.init_m {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr beta1 beta2 epsilon : α) (params : Spec.Tensor α s) : (init lr beta1 beta2 epsilon params).m = Spec.fill 0 s

Adam starts with a zero first-moment buffer.

@[simp]

theorem Optim.Adam.init_v {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr beta1 beta2 epsilon : α) (params : Spec.Tensor α s) : (init lr beta1 beta2 epsilon params).v = Spec.fill 0 s

Adam starts with a zero second-moment buffer.

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 ← β₁ m + (1-β₁) g
• v ← β₂ v + (1-β₂) g²
• m̂ ← m / (1-β₁ᵗ)
• v̂ ← v / (1-β₂ᵗ)
• p ← p - lr * m̂ / (sqrt(v̂) + ε)

The ε placement matches Kingma and Ba: it is added after sqrt(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 β₁.

• beta2 : α

  Second moment decay β₂.

• epsilon : α

  Numerical stability constant ε.

• 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).

@[simp]

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

AdamW initialization records the requested decoupled weight-decay coefficient.

@[simp]

theorem Optim.AdamW.init_t {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr weightDecay beta1 beta2 epsilon : α) (params : Spec.Tensor α s) : (init lr weightDecay beta1 beta2 epsilon params).t = 0

AdamW starts at step 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 ρ.

• epsilon : α

  Numerical stability constant ε.

• 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).

@[simp]

theorem Optim.Adadelta.init_v {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr rho epsilon : α) (params : Spec.Tensor α s) : (init lr rho epsilon params).v = Spec.fill 0 s

Adadelta starts with a zero squared-gradient EMA.

@[simp]

theorem Optim.Adadelta.init_u {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr rho epsilon : α) (params : Spec.Tensor α s) : (init lr rho epsilon params).u = Spec.fill 0 s

Adadelta starts with a zero squared-update EMA.

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 ← ρ v + (1-ρ) g²
• Δp ← - lr * (sqrt(u + ε) / sqrt(v + ε)) ⊙ g
• p ← p + Δp
• u ← ρ u + (1-ρ) (Δp)²

The ε placement is inside the RMS terms, matching Zeiler's Adadelta update.

Projected / low-rank optimizers

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 deliberately only 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, useful for tests and for the theorem that projected SGD reduces to SGD.

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)).

theorem Optim.GaLore.projectedSGDUpdate_identity_eq_sgd {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr : α) (params grads : Spec.Tensor α s) : projectedSGDUpdate { lr := lr, projector := identityProjector } params grads = SGD.update { lr := lr } params grads

With the identity projector, projected SGD is exactly ordinary SGD.

This is the main sanity check for the GaLore extension point: adding a projection backend cannot silently change the base optimizer when the backend is the identity.

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; with this backend Muon reduces to momentum SGD.

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.

theorem Optim.Muon.update_identity_param_eq_momentumSGD {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (lr momentum : α) (buf params grads : Spec.Tensor α s) : (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := identityOrthogonalizer } params grads).2 = (MomentumSGD.update { lr := lr, momentum := momentum, buf := buf } params grads).2

With the identity orthogonalizer, Muon's parameter update is exactly momentum SGD's parameter update.

The state records are different because Muon carries an orthogonalizer backend, but the parameter direction is the same when that backend is id.