Optimizer integration for Runtime.Autograd

This module is the training-loop side of autograd: it takes a gradient map produced by Runtime.Autograd and applies parameter updates.

PyTorch analogy:

• ParameterTable is like an ordered list of parameters, but we key everything by a stable Nat id (closer to state_dict keys than pointer identity).
• ParameterGroup and OptimizerState mirror torch.optim.Optimizer parameter groups and state.
• LearningRateScheduler is a small wrapper around our scheduler implementations, similar to torch.optim.lr_scheduler.*.

All updates are shape checked and implemented using the pure Spec tensor operators, so they can be used in eager execution or lowered into a typed graph.

Formula ownership:

• this file owns the heterogeneous parameter-table handling, parameter groups, lazily created per-parameter buffers, scheduler stepping, and PyTorch-style coupled weight decay at the training-loop boundary;
• NN.Runtime.Optim.Optimizers owns the canonical per-tensor optimizer equations.

The important rule is: this file must not define a second public optimizer-formula surface. Each Optimizer.Internal.<algorithm> function below constructs the canonical optimizer state for one parameter, typed by that parameter's shape, and calls NN.Runtime.Optim.Optimizers directly. The only local algebra left here is training-loop glue that is not represented by the canonical pure states, such as coupled weight-decay preprocessing and PyTorch-style momentum dampening/Nesterov handling.

Per-parameter buffers are stored shape-erased (ParameterState) because the parameter table is heterogeneous; ParameterState.cast is the single place where the stored shape is checked against the live parameter shape.

Parameter table

The declarations below provide the parameter registry used by the training loop.

Unlike PyTorch (where parameters are objects with identity), we use an explicit Nat id so that:

• gradients can be stored in a HashMap Nat _,
• optimizer state buffers can be stored in a HashMap Nat _,
• serialization can be done by a pure state_dict record.

structure Runtime.Autograd.Train.Parameter (α : Type) [TorchLean.Storage α] : Type

A single trainable parameter entry.

This is the Runtime.Autograd equivalent of a "parameter tensor" in PyTorch, except we make the identifier explicit (id : Nat) so we can key gradients and optimizer state in pure maps.

• id : ℕ

  Stable identifier used to key gradients and optimizer state.

• name : Option String

  Optional label, such as a module path; used only for reporting and debugging.

• value : Spec.SomeTensor α

  The shape-erased parameter value.

@[reducible, inline]

abbrev Runtime.Autograd.Train.ParameterTable (α : Type) [TorchLean.Storage α] : Type

A flat runtime array of parameters used by the training loop.

Constructors

def Runtime.Autograd.Train.Parameter.create {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (id : ℕ) (t : TorchLean.Tensor α s) (name : Option String := none) : Parameter α

Create a ParamEntry from a typed tensor.

This is mostly a convenience for assembling a ParamTable from known-shaped tensors.

def Runtime.Autograd.Train.ParameterTable.ids {α : Type} [TorchLean.Storage α] (parameters : ParameterTable α) : Array ℕ

Array of ids for membership checks.

def Runtime.Autograd.Train.ParameterTable.find? {α : Type} [TorchLean.Storage α] (parameters : ParameterTable α) (id : ℕ) : Option (Parameter α)

Find a parameter entry by id.

def Runtime.Autograd.Train.ParameterTable.get {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (tag : String) (parameters : ParameterTable α) (id : ℕ) : Result (TorchLean.Tensor α s)

Get a typed tensor from the table, with shape checking.

def Runtime.Autograd.Train.ParameterTable.set {α : Type} [TorchLean.Storage α] (parameters : ParameterTable α) (id : ℕ) (value : Spec.SomeTensor α) : ParameterTable α

Replace a parameter entry value by id.

def Runtime.Autograd.Train.ParameterTable.Internal.checkedIdSet {α : Type} [TorchLean.Storage α] (parameters : ParameterTable α) : Result (Std.HashMap ℕ Unit)

Build the set of parameter identifiers, rejecting duplicate ids.

Optimizer buffers and gradients are keyed by id, so accepting two table entries with the same id would make both entries consume one gradient and mutate one shared optimizer state.

Scheduler wrapper

inductive Runtime.Autograd.Train.LearningRateScheduler (α : Type) : Type

Learning-rate scheduler wrapper used by the training loop.

PyTorch analogy: this plays the role of torch.optim.lr_scheduler.* objects, except we keep the state as an inductive value and expose a pure current/advance API.

The torch* constructors wrap the schedules from Optim.Scheduler.PyTorch, whose phase and step-count conventions follow PyTorch exactly where the native schedules deliberately differ.

• constant {α : Type} : Optim.Scheduler.Constant α → LearningRateScheduler α
• exponential {α : Type} : Optim.Scheduler.ExponentialDecay α → LearningRateScheduler α
• step {α : Type} : Optim.Scheduler.StepDecay α → LearningRateScheduler α
• cosine {α : Type} : Optim.Scheduler.CosineAnnealing α → LearningRateScheduler α
• linearWarmup {α : Type} : Optim.Scheduler.LinearWarmup α → LearningRateScheduler α
• warmupCosine {α : Type} : Optim.Scheduler.WarmupCosine α → LearningRateScheduler α
• cyclic {α : Type} : Optim.Scheduler.Cyclic α → LearningRateScheduler α
• triangular {α : Type} : Optim.Scheduler.TriangularCycle α → LearningRateScheduler α
• oneCycle {α : Type} : Optim.Scheduler.OneCycle α → LearningRateScheduler α
• rangeTest {α : Type} : Optim.Scheduler.RangeTest α → LearningRateScheduler α
• torchCosineAnnealing {α : Type} : Optim.Scheduler.PyTorch.CosineAnnealing α → LearningRateScheduler α

  PyTorch CosineAnnealingLR: the cosine continues past T_max with period 2 * T_max instead of clamping at the minimum learning rate as the native cosine schedule does.

• torchOneCycle {α : Type} : Optim.Scheduler.PyTorch.OneCycle α → LearningRateScheduler α

  PyTorch OneCycleLR (learning rate only): fractional phase endpoints pct_start * total - 1, min_lr = initial_lr / final_div_factor, cosine or linear annealing, and the optional three-phase variant. The native oneCycle uses linear ramps with a max_lr / final_div_factor floor.

• custom {α : Type} : (ℕ → α) → ℕ → LearningRateScheduler α

  Custom schedule with an explicit step counter.

  custom f k means "use learning rate f k for this step, and increment to k+1 on advance".

def Runtime.Autograd.Train.LearningRateScheduler.current {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] : LearningRateScheduler α → α

Read current learning rate from the scheduler state.

def Runtime.Autograd.Train.LearningRateScheduler.advance {α : Type} : LearningRateScheduler α → LearningRateScheduler α

Advance scheduler state by one step.

Optimizer configuration

inductive Runtime.Autograd.Train.OptimizerAlgorithm : Type

Which optimizer update rule to apply.

PyTorch analogy: these correspond approximately to torch.optim.SGD, Adam, AdamW, etc.

• sgd : OptimizerAlgorithm
• momentum : OptimizerAlgorithm
• adagrad : OptimizerAlgorithm
• rmsprop : OptimizerAlgorithm
• adam : OptimizerAlgorithm
• adamw : OptimizerAlgorithm
• adadelta : OptimizerAlgorithm

def Runtime.Autograd.Train.instReprOptimizerAlgorithm.repr : OptimizerAlgorithm → ℕ → Std.Format

@[instance_reducible]

instance Runtime.Autograd.Train.instReprOptimizerAlgorithm : Repr OptimizerAlgorithm

@[instance_reducible]

instance Runtime.Autograd.Train.instDecidableEqOptimizerAlgorithm : DecidableEq OptimizerAlgorithm

structure Runtime.Autograd.Train.ParameterGroup (α : Type) [Context α] : Type

Optimizer hyperparameters for a subset of parameters.

PyTorch analogy: this is a single entry in the optimizer's param-group list (optimizer.param_groups).

• parameterIds : Array ℕ

  Parameter ids that belong to this group.

• learningRate : α

  Base learning rate (possibly overridden by scheduler on each step).

• weightDecay : α

  $\ell_2$ regularization coefficient (behavior depends on the optimizer kind; see AdamW).

• momentum : α

  Momentum factor (SGD with momentum).

• dampening : α

  Dampening for momentum updates.

• nesterov : Bool

  Use Nesterov variant for momentum updates.

• beta1 : α

  Adam beta1 parameter (exponential decay for the first moment).

• beta2 : α

  Adam beta2 parameter (exponential decay for the second moment).

• epsilon : α

  Numerical stability term used by adaptive optimizers.

• rho : α

  "Rho" decay parameter for RMSProp/AdaDelta style optimizers.

• scheduler : Option (LearningRateScheduler α)

  Optional learning-rate scheduler for this group.

Per-parameter optimizer buffers

inductive Runtime.Autograd.Train.ParameterBuffers (α : Type) [TorchLean.Storage α] (s : Spec.Shape) : Type

Optimizer buffers for one parameter of shape s.

Each constructor belongs to one algorithm family, so the buffers are typed by the parameter shape and there is no separate map per buffer kind. Buffers stored by a different family (for example after switching algorithm on a restored state) are treated as absent and re-initialised lazily, which is what the previous per-kind maps did implicitly.

• momentum {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (buffer : TorchLean.Tensor α s) : ParameterBuffers α s

  SGD with momentum or Nesterov: the velocity buffer.

• squaredGradient {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (accumulator : TorchLean.Tensor α s) : ParameterBuffers α s

  AdaGrad and RMSProp: the squared-gradient accumulator.

• adam {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (stepCount : ℕ) (firstMoment secondMoment : TorchLean.Tensor α s) : ParameterBuffers α s

  Adam and AdamW: the parameter-local step count and both moment estimates.

• adadelta {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (squaredGradientAverage squaredUpdateAverage : TorchLean.Tensor α s) : ParameterBuffers α s

  Adadelta: squared-gradient and squared-update running averages.

structure Runtime.Autograd.Train.ParameterState (α : Type) [TorchLean.Storage α] : Type

Shape-erased optimizer buffers for one parameter.

The parameter table is heterogeneous, so buffers are stored with their runtime shape and cast back to the live parameter shape by ParameterState.cast.

• shape : Spec.Shape

  Shape of the parameter the buffers belong to.

• buffers : ParameterBuffers α self.shape

  The typed buffers.

def Runtime.Autograd.Train.ParameterState.ofBuffers {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (buffers : ParameterBuffers α s) : ParameterState α

Erase the shape of typed buffers.

def Runtime.Autograd.Train.ParameterState.cast {α : Type} [TorchLean.Storage α] (id : ℕ) (state : ParameterState α) (s : Spec.Shape) : Result (ParameterBuffers α s)

Recover typed buffers for a parameter of shape s.

This is the only shape cast on optimizer state. It fails when a checkpoint is reloaded into a model whose parameter id now has a different shape.

structure Runtime.Autograd.Train.OptimizerState (α : Type) [TorchLean.Storage α] [Context α] : Type

Full optimizer state used by the training loop.

This mirrors PyTorch's optimizer state:

• a global optimizer-call counter,
• hyperparameter groups, and
• per-parameter buffers keyed by parameter id (Nat), including the parameter-local Adam step.

• algorithm : OptimizerAlgorithm

  Which update rule to apply on step.

• parameterGroups : Array (ParameterGroup α)

  Parameter groups (hyperparameters + membership).

• stepCount : ℕ

  Global optimizer-call counter (increments once per step, including sparse steps).

• parameterStates : Std.HashMap ℕ (ParameterState α)

  Per-parameter buffers keyed by parameter id.

  Adam bias correction is parameter-local: a parameter whose gradient is absent does not advance its moment step. The step lives inside ParameterBuffers.adam, so it cannot drift from the moments it corrects.

structure Runtime.Autograd.Train.OptimizerSnapshot (α : Type) [TorchLean.Storage α] [Context α] : Type

A pure state snapshot for saving/restoring optimizer state.

PyTorch analogy: this is the data carried by optimizer.state_dict() (modulo naming/layout). We use association lists instead of HashMap so the result is deterministic and easy to serialize.

• algorithm : OptimizerAlgorithm

  Optimizer algorithm used to interpret the stored buffers.

• stepCount : ℕ

  Global optimizer step at the time the snapshot was taken.

• parameterGroups : Array (ParameterGroup α)

  Parameter groups, including scheduler state and hyperparameters.

• parameterStates : Array (ℕ × ParameterState α)

  Per-parameter buffers keyed by parameter id.

def Runtime.Autograd.Train.OptimizerState.snapshot {α : Type} [TorchLean.Storage α] [Context α] (optimizerState : OptimizerState α) : OptimizerSnapshot α

Serialize optimizer state to a pure record.

PyTorch analogy: this is the "export" step for state_dict().

def Runtime.Autograd.Train.OptimizerState.restore {α : Type} [TorchLean.Storage α] [Context α] (snapshot : OptimizerSnapshot α) : OptimizerState α

Restore optimizer state from a state dict.

PyTorch analogy: this is the "import" step for load_state_dict(...).

def Runtime.Autograd.Train.OptimizerState.parameterStepCount? {α : Type} [TorchLean.Storage α] [Context α] (optimizerState : OptimizerState α) (id : ℕ) : Option ℕ

Number of Adam/AdamW updates applied to parameter id, if it has Adam-family buffers.

Optimizer step

def Runtime.Autograd.Train.Optimizer.Internal.addWeightDecay {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (parameters gradients : TorchLean.Tensor α s) (weightDecay : α) : TorchLean.Tensor α s

Add an $\ell_2$ regularization term to the gradient: $g+\operatorname{weightDecay}\,\operatorname{param}$.

Note: this is the coupled weight decay used by classic SGD-style updates. For AdamW the integration step delegates to the canonical optimizer's decoupled update.

def Runtime.Autograd.Train.Optimizer.Internal.zeroBuffer {α : Type} [TorchLean.Storage α] [Context α] (s : Spec.Shape) : TorchLean.Tensor α s

Zero buffer used when a parameter has no stored state yet (PyTorch initialises lazily).

def Runtime.Autograd.Train.Optimizer.Internal.sgd {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (group : ParameterGroup α) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s

Plain SGD: p - lr * (g + wd * p). Stateless.

def Runtime.Autograd.Train.Optimizer.Internal.momentum {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (group : ParameterGroup α) (buffers : Option (ParameterBuffers α s)) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s × ParameterBuffers α s

SGD with momentum, PyTorch conventions.

The momentum buffer is initialised from the first raw (decayed) gradient; dampening is applied only once a buffer already exists. Nesterov uses g + momentum * buffer as the update direction.

def Runtime.Autograd.Train.Optimizer.Internal.squaredGradientBuffer {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (buffers : Option (ParameterBuffers α s)) : TorchLean.Tensor α s

Squared-gradient accumulator shared by AdaGrad and RMSProp, or zeros when absent.

def Runtime.Autograd.Train.Optimizer.Internal.adagrad {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (group : ParameterGroup α) (buffers : Option (ParameterBuffers α s)) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s × ParameterBuffers α s

AdaGrad with coupled weight decay.

def Runtime.Autograd.Train.Optimizer.Internal.rmsprop {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (group : ParameterGroup α) (buffers : Option (ParameterBuffers α s)) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s × ParameterBuffers α s

RMSProp with coupled weight decay.

def Runtime.Autograd.Train.Optimizer.Internal.adamBuffers {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (buffers : Option (ParameterBuffers α s)) : ℕ × TorchLean.Tensor α s × TorchLean.Tensor α s

Adam-family buffers (stepCount, firstMoment, secondMoment), zero-initialised when absent.

def Runtime.Autograd.Train.Optimizer.Internal.adam {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (group : ParameterGroup α) (buffers : Option (ParameterBuffers α s)) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s × ParameterBuffers α s

Adam with coupled weight decay and parameter-local bias correction.

def Runtime.Autograd.Train.Optimizer.Internal.adamw {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (group : ParameterGroup α) (buffers : Option (ParameterBuffers α s)) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s × ParameterBuffers α s

AdamW: decoupled weight decay handled by the canonical update, no gradient preprocessing.

def Runtime.Autograd.Train.Optimizer.Internal.adadelta {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (group : ParameterGroup α) (buffers : Option (ParameterBuffers α s)) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s × ParameterBuffers α s

Adadelta with coupled weight decay.

def Runtime.Autograd.Train.Optimizer.Internal.updateParameter {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (algorithm : OptimizerAlgorithm) (group : ParameterGroup α) (buffers : Option (ParameterBuffers α s)) (parameters gradient : TorchLean.Tensor α s) : TorchLean.Tensor α s × Option (ParameterBuffers α s)

Apply the configured update rule to one parameter.

Returns the new parameter value and the buffers to store for it (none for stateless SGD, which leaves any stored buffers untouched).

def Runtime.Autograd.Train.Optimizer.Internal.castGradient {α : Type} [TorchLean.Storage α] (id : ℕ) (gradient : Spec.SomeTensor α) (s : Spec.Shape) : Result (TorchLean.Tensor α s)

Shape-check a gradient delivered as a shape-erased tensor against its parameter.

def Runtime.Autograd.Train.Optimizer.Internal.advanceSchedulers {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (parameterGroups : Array (ParameterGroup α)) : Array (ParameterGroup α)

Update each group's learning rate from its scheduler (if present) and advance the scheduler state.

This matches the common training-loop pattern: "read LR, then call scheduler.step()".

def Runtime.Autograd.Train.Optimizer.Internal.parameterGroupMap {α : Type} [Context α] (parameterGroups : Array (ParameterGroup α)) : Result (Std.HashMap ℕ (ParameterGroup α))

Build a map from parameter id to its ParameterGroup.

Fails if an id appears in multiple groups (PyTorch also disallows overlapping param groups).

structure Runtime.Autograd.Train.Optimizer.Step (α : Type) [TorchLean.Storage α] [Context α] : Type

Updated dynamic optimizer state and parameter table produced by one step.

• optimizerState : OptimizerState α

  Optimizer state after applying the gradients.

• parameters : ParameterTable α

  Updated parameter tensors.

def Runtime.Autograd.Train.Optimizer.step {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (optimizerState : OptimizerState α) (parameters : ParameterTable α) (gradients : Std.HashMap ℕ (Spec.SomeTensor α)) : Result (Step α)

Apply one optimizer step to a parameter table.

Inputs:

• optimizerState is the current optimizer state (including per-parameter buffers),
• parameters is the current parameter table,
• gradients maps parameter ids to gradients (as produced by autograd).

Behavior:

• rejects duplicate parameter ids, overlapping groups, and group ids missing from the table,
• applies learning-rate schedulers (if configured) per group,
• shape-checks gradients and stored buffers against each parameter,
• updates per-parameter buffers (momentum, Adam moments, and accumulators),
• returns the updated optimizer state and an updated parameter table.

Parameters without a gradient are left untouched, including their buffers.