Native Learning-Rate Schedulers

TorchLean-native schedules with explicit state and total formulas. Zero-length warmup or cycle phases have defined fallback behavior, which makes the schedules convenient for direct execution and theorem statements. currentStep is zero-indexed and advance increments it once.

Schedulers.Core documents the shared arithmetic, state convention, and literature. Use the separate PyTorch module when exact PyTorch phase and step-count behavior is required.

Native Schedulers

structure Optim.Scheduler.Constant (α : Type) : Type

Constant scheduler (no learning rate changes).

• learningRate : α

  Fixed learning rate.

def Optim.Scheduler.Constant.current {α : Type} (scheduler : Constant α) : α

Get the learning rate for a constant schedule.

PyTorch analogy: no scheduler (or a scheduler that keeps LR fixed).

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

Advance a constant scheduler by one step.

This is the identity since there is no state to update.

PyTorch analogy: scheduler.advance() for a scheduler that does nothing.

def Optim.Scheduler.Constant.create {α : Type} (learningRate : α) : Constant α

Create a constant learning-rate scheduler.

PyTorch analogy: constructing training code with a fixed lr and no lr_scheduler.

Exponential decay

structure Optim.Scheduler.ExponentialDecay (α : Type) : Type

Exponential decay scheduler: lr(step) = initial_lr * decayRate^step.

PyTorch analogy: similar spirit to ExponentialLR, but we keep state as a simple counter.

• initialLearningRate : α

  Learning rate at step 0.

• decayRate : α

  Multiplicative decay factor per step (gamma in PyTorch terminology).

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for an exponential decay schedule at the current step.

Formula: initial_lr * decayRate ^ current_step.

PyTorch analogy: torch.optim.lr_scheduler.ExponentialLR (but here kept as a pure counter-based record).

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

Advance the exponential decay scheduler by one step.

PyTorch analogy: scheduler.advance().

def Optim.Scheduler.ExponentialDecay.create {α : Type} (initialLearningRate decayRate : α) : ExponentialDecay α

Create an exponential decay scheduler starting at step 0.

PyTorch analogy: torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=decayRate).

Step decay

structure Optim.Scheduler.StepDecay (α : Type) : Type

Piecewise-constant decay: every stepSize steps, multiply the learning rate by decayFactor.

• initialLearningRate : α

  Learning rate at step 0.

• decayFactor : α

  Multiplicative decay factor applied every stepSize steps.

• stepSize : ℕ

  Number of steps between decays.

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for step decay at the current step.

Every stepSize steps, the LR is multiplied by decayFactor. When step_size = 0, this falls back to a constant LR.

PyTorch analogy: torch.optim.lr_scheduler.StepLR.

theorem Optim.Scheduler.StepDecay.current_zero_stepSize {α : Type} [Context α] (initialLearningRate decayFactor : α) (currentStep : ℕ) : { initialLearningRate := initialLearningRate, decayFactor := decayFactor, stepSize := 0, currentStep := currentStep }.current = initialLearningRate

The totalized step_size = 0 case is constant.

PyTorch would reject this configuration; TorchLean keeps scheduler evaluation total so configs can be validated separately from pure schedule semantics.

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

Advance the step-decay scheduler by one step.

PyTorch analogy: scheduler.advance().

def Optim.Scheduler.StepDecay.create {α : Type} (initialLearningRate decayFactor : α) (stepSize : ℕ) : StepDecay α

Create a step-decay scheduler starting at step 0.

PyTorch analogy: torch.optim.lr_scheduler.StepLR(optimizer, step_size=..., gamma=decayFactor).

Cosine annealing

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

Cosine annealing down to minimumLearningRate over maxSteps steps.

PyTorch analogy: CosineAnnealingLR (without restarts).

• initialLearningRate : α

  Learning rate at step 0.

• minimumLearningRate : α

  Minimum learning rate after annealing completes.

• maxSteps : ℕ

  Number of steps over which to anneal.

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for cosine annealing at the current step.

We anneal from initialLearningRate to minimumLearningRate over maxSteps steps, clamping once the step counter passes maxSteps.

PyTorch analogy: torch.optim.lr_scheduler.CosineAnnealingLR (without restarts).

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

Advance the cosine annealing scheduler by one step.

PyTorch analogy: scheduler.advance().

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

Create a cosine annealing scheduler starting at step 0.

PyTorch analogy: torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_steps, eta_min=min_lr).

Linear warmup

structure Optim.Scheduler.LinearWarmup (α : Type) : Type

Linear warmup from startingLearningRate to initialLearningRate over warmupSteps steps, then constant.

Warmup is a practical trick commonly used when training large models (e.g. Transformers) to avoid instability at the start of training.

• initialLearningRate : α

  Target learning rate after warmup.

• warmupSteps : ℕ

  Number of warmup steps.

• startingLearningRate : α

  Starting learning rate during warmup.

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for linear warmup (then constant).

Before warmupSteps, linearly interpolate from startingLearningRate to initialLearningRate. Afterwards, keep initialLearningRate fixed.

PyTorch analogy: warmup logic commonly implemented in training scripts (and in some scheduler helpers).

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

Advance the linear warmup scheduler by one step.

PyTorch analogy: scheduler.advance().

def Optim.Scheduler.LinearWarmup.create {α : Type} [Context α] (initialLearningRate : α) (warmupSteps : ℕ) (startingLearningRate : α := 0) : LinearWarmup α

Create a linear warmup scheduler starting at step 0.

PyTorch analogy: a warmup wrapper around an optimizer or a base scheduler.

Warmup + cosine

structure Optim.Scheduler.WarmupCosine (α : Type) : Type

Warmup followed by cosine annealing.

This is a common “default” schedule for Transformer-style training: warm up for a few thousand steps, then gradually anneal.

• initialLearningRate : α

  Peak learning rate (reached at the end of warmup).

• warmupSteps : ℕ

  Number of warmup steps.

• totalSteps : ℕ

  Total number of steps for the whole schedule (warmup + anneal).

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for the warmup-then-cosine schedule at the current step.

• During warmup, LR increases linearly from 0 to initialLearningRate.
• After warmup, LR follows a cosine anneal over the remaining steps.
• At and after totalSteps, LR remains at 0 instead of beginning another cosine period.

PyTorch analogy: a common Transformer schedule, often implemented by composing warmup with cosine decay.

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

Advance the warmup+cosine scheduler by one step.

PyTorch analogy: scheduler.advance().

def Optim.Scheduler.WarmupCosine.create {α : Type} (initialLearningRate : α) (warmupSteps totalSteps : ℕ) : WarmupCosine α

Create a warmup+cosine scheduler starting at step 0.

PyTorch analogy: composing a warmup schedule with cosine annealing in a training script.

Cyclic LR

inductive Optim.Scheduler.CyclicMode : Type

Cyclic learning rate schedule.

This corresponds to the “triangular” family of schedules where the LR increases linearly from base_lr to maximumLearningRate and then decreases back, repeating in cycles.

The mode is an enum, so unsupported schedule variants cannot enter the runtime state.

• triangular : CyclicMode

  Fixed-amplitude triangular cycles.

• shrinkingTriangular : CyclicMode

  Triangular cycles whose amplitude halves after each cycle.

• exponentialRange : CyclicMode

  Triangular cycles with an exponential amplitude factor.

@[instance_reducible]

instance Optim.Scheduler.instReprCyclicMode : Repr CyclicMode

def Optim.Scheduler.instReprCyclicMode.repr : CyclicMode → ℕ → Std.Format

@[instance_reducible]

instance Optim.Scheduler.instDecidableEqCyclicMode : DecidableEq CyclicMode

structure Optim.Scheduler.Cyclic (α : Type) : Type

State of a cyclic learning-rate schedule (Smith, "Cyclical Learning Rates for Training Neural Networks", WACV 2017), matching PyTorch CyclicLR.

The step counter lives in the structure rather than being passed in, so advance is a pure state transition and a checkpoint can round-trip a schedule mid-cycle.

• baseLearningRate : α

  Minimum learning rate within the cycle.

• maximumLearningRate : α

  Maximum learning rate within the cycle (before any mode-specific adjustment).

• stepSize : ℕ

  Half-cycle size (in steps).

• mode : CyclicMode

  Cycle amplitude policy.

• decayFactor : α

  Decay factor used by exponentialRange.

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for the cyclic schedule at the current step.

Supports the common "triangular", "triangular2", and "exp_range" variants (matching the flavor of PyTorch's CyclicLR).

PyTorch analogy: torch.optim.lr_scheduler.CyclicLR.

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

Advance the cyclic scheduler by one step.

PyTorch analogy: scheduler.advance().

def Optim.Scheduler.Cyclic.create {α : Type} [Context α] (baseLearningRate maximumLearningRate : α) (stepSize : ℕ) (mode : CyclicMode := CyclicMode.triangular) (decayFactor : α := 1) : Cyclic α

Create a cyclic learning-rate scheduler starting at step 0.

PyTorch analogy: torch.optim.lr_scheduler.CyclicLR(base_lr=..., max_lr=..., step_size_up=...).

Triangular cycle (special case)

structure Optim.Scheduler.TriangularCycle (α : Type) : Type

A specialized cyclic schedule with fixed amplitude.

This is essentially Cyclic in "triangular" mode, but we provide it as a separate type so callers don't have to thread mode strings around.

• baseLearningRate : α

  Minimum learning rate within the cycle.

• maximumLearningRate : α

  Maximum learning rate within the cycle.

• stepSize : ℕ

  Half-cycle size (in steps).

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for the triangular cycle schedule at the current step.

This is the canonical "triangle up then down" schedule with fixed amplitude.

PyTorch analogy: CyclicLR in "triangular" mode.

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

Advance the triangular cycle scheduler by one step.

PyTorch analogy: scheduler.advance().

def Optim.Scheduler.TriangularCycle.create {α : Type} (baseLearningRate maximumLearningRate : α) (stepSize : ℕ) : TriangularCycle α

Create a triangular cycle scheduler starting at step 0.

PyTorch analogy: CyclicLR(base_lr=..., max_lr=..., mode=\"triangular\").

1cycle Learning Rate Schedule

structure Optim.Scheduler.OneCycle (α : Type) : Type

One-cycle learning-rate schedule.

• increase LR from initialLearningRate to maximumLearningRate over the first increasingFraction of the steps,
• then decrease to finalLearningRate over the rest.

In the original 1cycle policy, momentum is also scheduled; we keep this runtime version LR-only.

• maximumLearningRate : α

  Peak learning rate (reached at increasingFraction of the schedule).

• totalSteps : ℕ

  Total number of steps in the schedule.

• initialLearningRate : α

  Learning rate at step 0.

• finalLearningRate : α

  Learning rate after the full schedule finishes.

• divisionFactor : α

  Divides maximumLearningRate to get initialLearningRate in the factory constructor.

• increasingFraction : α

  Fraction of the schedule spent increasing LR (0..1).

• currentStep : ℕ

  Current step counter (0-indexed).

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

Get the learning rate for the one-cycle schedule at the current step.

This ramps up to maximumLearningRate over the increasingFraction part of the schedule, then anneals down to finalLearningRate.

PyTorch analogy: torch.optim.lr_scheduler.OneCycleLR, restricted here to the learning-rate curve.

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

Advance the 1cycle scheduler by one step.

PyTorch analogy: scheduler.advance().

def Optim.Scheduler.OneCycle.create {α : Type} [Context α] (maximumLearningRate : α) (totalSteps : ℕ) (divisionFactor increasingFraction finalDivisionFactor : α) : OneCycle α

Create a simplified 1cycle schedule starting at step 0.

We derive initial_lr := max_lr / div_factor and final_lr := max_lr / final_div_factor.

PyTorch analogy: torch.optim.lr_scheduler.OneCycleLR(max_lr=..., total_steps=...).

LR finder

structure Optim.Scheduler.RangeTest (α : Type) : Type

Learning-rate finder schedule: an exponential sweep from initialLearningRate to finalLearningRate over totalSteps steps.

• initialLearningRate : α

  Learning rate at step 0.

• finalLearningRate : α

  Target learning rate at the end of the sweep.

• totalSteps : ℕ

  Number of steps in the sweep.

• currentStep : ℕ

  Current step counter (0-indexed).

def Optim.Scheduler.RangeTest.current {α : Type} [Context α] (finder : RangeTest α) : α

Get the learning rate for the LR-finder exponential sweep at the current step.

This increases LR exponentially from initialLearningRate toward finalLearningRate across totalSteps steps.

PyTorch analogy: LR finder utilities used by libraries like fastai, often implemented as a custom schedule.

def Optim.Scheduler.RangeTest.advance {α : Type} (finder : RangeTest α) : RangeTest α

Advance the LR finder by one step.

PyTorch analogy: stepping a custom LR schedule inside a training loop.

def Optim.Scheduler.RangeTest.create {α : Type} (initialLearningRate finalLearningRate : α) (totalSteps : ℕ) : RangeTest α

Create an LR finder schedule starting at step 0.

PyTorch analogy: setting up an LR finder run to sweep learning rates.