PyTorch-Compatible Learning-Rate Schedulers

Schedulers whose phase boundaries and step counters follow the corresponding torch.optim.lr_scheduler behavior. They remain pure Lean state machines, so a training run can store, inspect, and reason about the exact scheduler state without calling PyTorch.

Schedulers.Core documents the zero-indexed counter convention, shared scalar operations, and literature. The Native module provides simpler total schedules when compatibility is not the contract.

PyTorch-compatible scheduler variants

The schedulers below use formulas and step-count conventions chosen to match PyTorch's torch.optim.lr_scheduler.* semantics more directly.

Important convention note (PyTorch last_epoch):

• In modern PyTorch, schedulers effectively start at last_epoch = 0 right after construction (fresh run with last_epoch = -1 in the constructor triggers an initial internal step).
• We model that behavior by using a current_step : Nat := 0 counter. Think: current_step corresponds to PyTorch's last_epoch after construction.

These schedulers are LR-only (they do not mutate optimizer momentum/betas). If you need the full PyTorch OneCycle momentum behavior, consider adding a separate momentum schedule and stepping both in lockstep.

StepLR

structure Optim.StepLR (α : Type) : Type

PyTorch-compatible StepLR.

Semantics:

• current_step = 0 yields base_lr.
• Every step_size steps, multiply LR by gamma.
• When step_size = 0, this degenerates to a constant schedule (total, no exceptions).

PyTorch reference: torch.optim.lr_scheduler.StepLR.

• baseLr : α

  Base learning rate (what PyTorch calls base_lrs[i]).

• stepSize : ℕ

  Step interval (step_size).

• gamma : α

  Multiplicative decay factor (gamma).

• currentStep : ℕ

  Step counter matching PyTorch last_epoch after construction (0-indexed).

def Optim.StepLR.getLr {α : Type} [Context α] (scheduler : StepLR α) : α

Current learning rate for StepLR at current_step.

theorem Optim.StepLR.getLr_zero_stepSize {α : Type} [Context α] (baseLr gamma : α) (currentStep : ℕ) : { baseLr := baseLr, stepSize := 0, gamma := gamma, currentStep := currentStep }.getLr = baseLr

The PyTorch-compatible StepLR is also totalized to a constant when step_size = 0.

def Optim.StepLR.step {α : Type} (scheduler : StepLR α) : StepLR α

Advance StepLR by one step.

def Optim.stepLR {α : Type} (baseLr : α) (stepSize : ℕ) (gamma : α) : StepLR α

Constructor for StepLR starting at current_step = 0.

CosineAnnealingLR

structure Optim.CosineAnnealingLR (α : Type) : Type

PyTorch-compatible CosineAnnealingLR.

Key behavior difference from TorchLean's CosineAnnealingScheduler above:

• PyTorch's CosineAnnealingLR continues the cosine curve past T_max (it is periodic with period 2*T_max), rather than clamping to eta_min.

PyTorch reference: torch.optim.lr_scheduler.CosineAnnealingLR.

• baseLr : α

  Base learning rate (base_lrs[i]).

• tMax : ℕ

  Maximum number of steps in a half-cycle (T_max).

• etaMin : α

  Minimum learning rate (eta_min).

• currentStep : ℕ

  Step counter matching PyTorch last_epoch after construction (0-indexed).

def Optim.CosineAnnealingLR.getLr {α : Type} [Context α] (scheduler : CosineAnnealingLR α) : α

Current learning rate for CosineAnnealingLR at current_step.

def Optim.CosineAnnealingLR.step {α : Type} (scheduler : CosineAnnealingLR α) : CosineAnnealingLR α

Advance CosineAnnealingLR by one step.

def Optim.cosineAnnealingLR {α : Type} [Context α] (baseLr : α) (tMax : ℕ) (etaMin : α := Numbers.zero) : CosineAnnealingLR α

Constructor for CosineAnnealingLR starting at current_step = 0.

OneCycleLR (LR-only)

inductive Optim.OneCycleAnnealStrategy : Type

Anneal strategy used by OneCycleLR (matches PyTorch "cos" or "linear").

• cos : OneCycleAnnealStrategy
• linear : OneCycleAnnealStrategy

@[implicit_reducible]

instance Optim.instReprOneCycleAnnealStrategy : Repr OneCycleAnnealStrategy

def Optim.instReprOneCycleAnnealStrategy.repr : OneCycleAnnealStrategy → ℕ → Std.Format

@[implicit_reducible]

instance Optim.instDecidableEqOneCycleAnnealStrategy : DecidableEq OneCycleAnnealStrategy

structure Optim.OneCycleLR (α : Type) : Type

PyTorch-compatible OneCycleLR (LR-only).

Notes:

• This mirrors PyTorch's OneCycleLR learning-rate schedule only. PyTorch can also cycle momentum (or Adam's beta1); TorchLean keeps this scheduler pure and LR-only.
• PyTorch defines:
  • initial_lr = max_lr / div_factor
  • min_lr = initial_lr / final_div_factor (note: min_lr is derived from initial_lr, not directly from max_lr).
• PyTorch uses "phase end steps" that are floats:
  • phase 1 ends at pct_start * total_steps - 1
  • phase 2 ends at total_steps - 1 (and three_phase inserts a middle phase). This means the boundary can be fractional; the schedule uses interpolation ratios (pct) computed from these float endpoints. We match that behavior using α arithmetic.

PyTorch reference: torch.optim.lr_scheduler.OneCycleLR.

• maxLr : α

  Peak learning rate (max_lr).

• totalSteps : ℕ

  Total number of steps (total_steps).

• pctStart : α

  Fraction of steps spent increasing LR (pct_start).

• divFactor : α

  div_factor used to derive initial_lr = max_lr / div_factor.

• finalDivFactor : α

  final_div_factor used to derive min_lr = initial_lr / final_div_factor.

• annealStrategy : OneCycleAnnealStrategy

  Anneal strategy (cos or linear).

• threePhase : Bool

  Use PyTorch's three_phase variant when true.

• currentStep : ℕ

  Step counter matching PyTorch last_epoch after construction (0-indexed).

def Optim.OneCycleLR.initialLr {α : Type} [Context α] (s : OneCycleLR α) : α

Derived initial LR (max_lr / div_factor).

def Optim.OneCycleLR.minLr {α : Type} [Context α] (s : OneCycleLR α) : α

Derived minimum LR (initial_lr / final_div_factor).

def Optim.OneCycleLR.anneal {α : Type} [Context α] (s : OneCycleLR α) (startLR endLR pct : α) : α

PyTorch-compatible anneal helper (no clamping).

def Optim.OneCycleLR.getLr {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (s : OneCycleLR α) : α

Current learning rate for OneCycleLR at current_step (LR-only).

def Optim.OneCycleLR.step {α : Type} (s : OneCycleLR α) : OneCycleLR α

Advance OneCycleLR by one step.

def Optim.oneCycleLR {α : Type} (maxLr : α) (totalSteps : ℕ) (pctStart divFactor finalDivFactor : α) (annealStrategy : OneCycleAnnealStrategy := OneCycleAnnealStrategy.cos) (threePhase : Bool := false) : OneCycleLR α

Constructor for OneCycleLR starting at current_step = 0 (LR-only).

This mirrors the PyTorch parameterization:

• initial_lr = max_lr / div_factor
• min_lr = initial_lr / final_div_factor
• phase endpoints computed as pct_start * total_steps - 1 and total_steps - 1 (with the optional three_phase middle phase).