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

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.

structure Optim.AdamPowers {α : Type} [One α] [Mul α] (beta1 beta2 : α) (stepCount : ℕ) : Type

Adam bias-correction powers with their coefficients and step count in the type.

The certificates preserve the left-associated multiplication order of scalarPowNat, including for rounded scalars without associative multiplication. Changing a coefficient or counter requires a cache for the new indices; scalar Boolean equality is never used to validate a cache.

• first : α

  First moment bias-correction power.

• second : α

  Second moment bias-correction power.

• first_eq : self.first = scalarPowNat beta1 stepCount

  The first power follows the canonical scalar recurrence.

• second_eq : self.second = scalarPowNat beta2 stepCount

  The second power follows the canonical scalar recurrence.

def Optim.AdamPowers.compute {α : Type} [One α] [Mul α] (beta1 beta2 : α) (stepCount : ℕ) : AdamPowers beta1 beta2 stepCount

Reconstruct bias-correction powers, for initialization or a restored counter.

def Optim.AdamPowers.advance {α : Type} [One α] [Mul α] {beta1 beta2 : α} {stepCount : ℕ} (powers : AdamPowers beta1 beta2 stepCount) : AdamPowers beta1 beta2 (stepCount + 1)

Advance both powers with exactly one multiplication per coefficient.

Shared equations

def Optim.Internal.weightedSquare {α : Type} [TorchLean.Storage α] [Zero α] [Sub α] [Mul α] [BEq α] {s : Spec.Shape} (values : TorchLean.Tensor α s) (weight : α) : TorchLean.Tensor α s

Weight squared entries before an overflowing square loses their finite contribution. Finite squares keep their usual multiplication order.

def Optim.Internal.correctMoment {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (moment : TorchLean.Tensor α s) (denominator : α) : TorchLean.Tensor α s

Correct a moment without requiring the reciprocal correction to fit the scalar format.

def Optim.Internal.correctedRms {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (moment : TorchLean.Tensor α s) (denominator : α) : TorchLean.Tensor α s

Take the corrected RMS without requiring the corrected second moment to be representable. The fallback requires a positive correction denominator and preserves the clamped-root convention.

def Optim.Internal.normalizedUpdate {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (lr epsilon : α) (moment standardDeviation : TorchLean.Tensor α s) : TorchLean.Tensor α s

Normalize before applying the learning rate, replaying detected range loss when supported.

Replay uses the original finite coefficient trees and the scalar's checked rounding. It can recover an update lost by either association of the product and quotient, including its dual coefficients. Nonfinite inputs, zero divisors and unsupported carriers retain ordinary arithmetic; rounding before the supplied inputs were formed is not reversed.

def Optim.updateMomentumBuffer {α : Type} [TorchLean.Storage α] [Add α] [Mul α] {s : Spec.Shape} (momentumBuffer : TorchLean.Tensor α s) (momentum : α) (gradients : TorchLean.Tensor α s) : TorchLean.Tensor α s

Next momentum buffer $b\mu+g$, preserving the scalar order used by tensor scaling.

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

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

This factor can underflow before a later gradient multiplication. Optimizer updates instead use Internal.normalizedUpdate, which applies the learning rate after normalization.

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

Single-pass form of adaptiveLr. The reference definition builds two constant tensors and three intermediates per call; this one maps the scalar formula once.

@[csimp]

theorem Optim.adaptiveLr_eq_fused : @adaptiveLr = @adaptiveLrFused

Compiled code runs the single-pass form. Proofs keep unfolding the reference definition.

SGD

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

SGD state (per parameter tensor).

We only store the learning rate here.

• lr : α

  Learning rate.

def Optim.SGD.init {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (lr : α) : 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 α] [Sub α] [Mul α] {s : Spec.Shape} (state : State α s) (parameters gradients : TorchLean.Tensor α s) : Step α s (State α s)

One SGD step: p ← p - g * lr, with the learning rate on the right of each gradient element.

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 b\mu+g$,
• $p\gets p-b\,\mathtt{lr}$.

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

• lr : α

  Learning rate.

• momentum : α

  Momentum coefficient $\mu$.

• momentumBuffer : TorchLean.Tensor α s

  Momentum buffer.

def Optim.MomentumSGD.init {α : Type} [TorchLean.Storage α] [Zero α] {s : Spec.Shape} (lr momentum : α) : TorchLean.Tensor α s → State α s

Initialize momentum SGD with a zero buffer.

def Optim.MomentumSGD.update {α : Type} [TorchLean.Storage α] [Add α] [Sub α] [Mul α] {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)$.

• lr : α

  Base learning rate.

• epsilon : α

  Numerical stability constant $\varepsilon$.

• squaredGradientSum : TorchLean.Tensor α s

  Accumulated squared gradients.

def Optim.AdaGrad.init {α : Type} [TorchLean.Storage α] [Zero α] {s : Spec.Shape} (lr 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.

• lr : α

  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 α] [Zero α] {s : Spec.Shape} (lr 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 α] [One α] [Mul α] (s : Spec.Shape) : Type

Adam state (per parameter tensor).

We store first and second moment averages, a step counter, and certified bias-correction powers. The default powers reconstruct the original recurrence when creating a state at a nonzero step. When updating either beta or the counter in an existing record, also set powers to AdamPowers.compute at the new values. Updates to other fields preserve the cache.

• lr : α

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

• powers : AdamPowers self.beta1 self.beta2 self.stepCount

  Cached powers for these coefficients and this counter.

def Optim.Adam.init {α : Type} [TorchLean.Storage α] [Zero α] [One α] [Mul α] {s : Spec.Shape} (lr 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}$. Weighted squares, moment corrections and the corrected square root avoid unnecessary overflow. The learning rate is applied after normalization; supported scalars replay finite coefficient trees if that product or quotient loses range. Stored moments and final updates can still exceed the chosen scalar format.

@[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 α] [One α] [Mul α] (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. Its certified powers follow Adam.State: replacing either beta or the counter also requires powers := AdamPowers.compute at the new values.

• lr : α

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

• powers : AdamPowers self.beta1 self.beta2 self.stepCount

  Cached powers for these coefficients and this counter.

def Optim.AdamW.init {α : Type} [TorchLean.Storage α] [Zero α] [One α] [Mul α] {s : Spec.Shape} (lr 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:

• apply weight decay directly to the parameters (p ← p - lr * wd * p),
• apply the canonical Adam update using the raw gradient g,
• retain the AdamW decay coefficient alongside the updated moments and powers.

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.

• lr : α

  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 α] [Zero α] {s : Spec.Shape} (lr 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.

• lr : α

  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 α] [Sub α] [Mul α] {full low : Spec.Shape} (state : SGDState α full low) (parameters gradients : TorchLean.Tensor α full) : Step α full (SGDState α full low)

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

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 update rule can be used with different orthogonalization algorithms.

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

• lr : α

  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 α] [Zero α] {s : Spec.Shape} (lr 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 α] [Add α] [Sub α] [Mul α] {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.