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.

Only schedules whose semantics differ from the native ones live here. StepLR is not duplicated: the native Scheduler.StepDecay already computes base_lr * gamma ^ (step / step_size) with the same zero-indexed counter, so it is the PyTorch-compatible step schedule as well.

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 currentStep : Nat := 0 counter. Think: currentStep 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.

CosineAnnealingLR

structure Optim.Scheduler.PyTorch.CosineAnnealing (α : Type) : Type

PyTorch-compatible CosineAnnealingLR.

Key behavior difference from TorchLean's native Scheduler.CosineAnnealing:

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

• baseLearningRate : α

  Base learning rate (base_lrs[i]).

• halfCycleSteps : ℕ

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

• minimumLearningRate : α

  Minimum learning rate (eta_min).

• currentStep : ℕ

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

def Optim.Scheduler.PyTorch.CosineAnnealing.current {α : Type} [Context α] (scheduler : CosineAnnealing α) : α

Current learning rate for PyTorch-compatible cosine annealing.

def Optim.Scheduler.PyTorch.CosineAnnealing.advance {α : Type} (scheduler : CosineAnnealing α) : CosineAnnealing α

Advance PyTorch-compatible cosine annealing by one step.

def Optim.Scheduler.PyTorch.CosineAnnealing.create {α : Type} [Context α] (baseLearningRate : α) (halfCycleSteps : ℕ) (minimumLearningRate : α := 0) : CosineAnnealing α

Create PyTorch-compatible cosine annealing at step zero.

OneCycleLR (LR-only)

inductive Optim.Scheduler.PyTorch.AnnealingStrategy : Type

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

• cosine : AnnealingStrategy
• linear : AnnealingStrategy

def Optim.Scheduler.PyTorch.instReprAnnealingStrategy.repr : AnnealingStrategy → ℕ → Std.Format

@[instance_reducible]

instance Optim.Scheduler.PyTorch.instReprAnnealingStrategy : Repr AnnealingStrategy

@[instance_reducible]

instance Optim.Scheduler.PyTorch.instDecidableEqAnnealingStrategy : DecidableEq AnnealingStrategy

structure Optim.Scheduler.PyTorch.OneCycle (α : 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.

• maximumLearningRate : α

  Peak learning rate (max_lr).

• totalSteps : ℕ

  Total number of steps (total_steps).

• increasingFraction : α

  Fraction of steps spent increasing LR (pct_start).

• divisionFactor : α

  div_factor used to derive initial_lr = max_lr / div_factor.

• finalDivisionFactor : α

  final_div_factor used to derive min_lr = initial_lr / final_div_factor.

• annealingStrategy : AnnealingStrategy

  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.Scheduler.PyTorch.OneCycle.initialLearningRate {α : Type} [Context α] (scheduler : OneCycle α) : α

Derived initial LR (max_lr / div_factor).

def Optim.Scheduler.PyTorch.OneCycle.minimumLearningRate {α : Type} [Context α] (scheduler : OneCycle α) : α

Derived minimum LR (initial_lr / final_div_factor).

def Optim.Scheduler.PyTorch.OneCycle.anneal {α : Type} [Context α] (scheduler : OneCycle α) (startingLearningRate endingLearningRate fraction : α) : α

PyTorch-compatible anneal helper (no clamping).

def Optim.Scheduler.PyTorch.OneCycle.current {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (scheduler : OneCycle α) : α

Current learning rate for PyTorch-compatible one-cycle scheduling (LR-only).

def Optim.Scheduler.PyTorch.OneCycle.advance {α : Type} (scheduler : OneCycle α) : OneCycle α

Advance PyTorch-compatible one-cycle scheduling by one step.

def Optim.Scheduler.PyTorch.OneCycle.create {α : Type} (maximumLearningRate : α) (totalSteps : ℕ) (increasingFraction divisionFactor finalDivisionFactor : α) (annealingStrategy : AnnealingStrategy := AnnealingStrategy.cosine) (threePhase : Bool := false) : OneCycle α

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