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 step advances 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.ConstantScheduler (α : Type) : Type

Constant scheduler (no learning rate changes).

• lr : α

  Fixed learning rate.

def Optim.ConstantScheduler.getLr {α : Type} (scheduler : ConstantScheduler α) : ℕ → α

Get the learning rate for a constant schedule.

The step argument is ignored (the LR never changes).

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

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

Advance a constant scheduler by one step.

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

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

def Optim.constantScheduler {α : Type} (lr : α) : ConstantScheduler α

Create a constant learning-rate scheduler.

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

Exponential decay

structure Optim.ExponentialDecayScheduler (α : Type) : Type

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

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

• initialLr : α

  Learning rate at step 0.

• decayRate : α

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

• currentStep : ℕ

  Current step counter (0-indexed).

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

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

Formula: initial_lr * decay_rate ^ current_step.

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

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

Advance the exponential decay scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.exponentialDecayScheduler {α : Type} (initialLr decay_rate : α) : ExponentialDecayScheduler α

Create an exponential decay scheduler starting at step 0.

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

Step decay

structure Optim.StepDecayScheduler (α : Type) : Type

Piecewise-constant decay: every step_size steps, multiply the learning rate by decay_factor.

• initialLr : α

  Learning rate at step 0.

• decayFactor : α

  Multiplicative decay factor applied every step_size steps.

• stepSize : ℕ

  Number of steps between decays.

• currentStep : ℕ

  Current step counter (0-indexed).

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

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

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

PyTorch analogy: torch.optim.lr_scheduler.StepLR.

theorem Optim.StepDecayScheduler.getLr_zero_stepSize {α : Type} [Context α] (initialLr decayFactor : α) (currentStep : ℕ) : { initialLr := initialLr, decayFactor := decayFactor, stepSize := 0, currentStep := currentStep }.getLr = initialLr

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.StepDecayScheduler.step {α : Type} (scheduler : StepDecayScheduler α) : StepDecayScheduler α

Advance the step-decay scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.stepDecayScheduler {α : Type} (initialLr decay_factor : α) (stepSize : ℕ) : StepDecayScheduler α

Create a step-decay scheduler starting at step 0.

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

Cosine annealing

structure Optim.CosineAnnealingScheduler (α : Type) : Type

Cosine annealing down to min_lr over max_steps steps.

PyTorch analogy: CosineAnnealingLR (without restarts).

• initialLr : α

  Learning rate at step 0.

• minLr : α

  Minimum learning rate after annealing completes.

• maxSteps : ℕ

  Number of steps over which to anneal.

• currentStep : ℕ

  Current step counter (0-indexed).

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

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

We anneal from initial_lr to min_lr over max_steps steps (clamping once we pass max_steps).

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

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

Advance the cosine annealing scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.cosineAnnealingScheduler {α : Type} [Context α] (initialLr : α) (maxSteps : ℕ) (minLr : α := 0) : CosineAnnealingScheduler α

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.LinearWarmupScheduler (α : Type) : Type

Linear warmup from start_lr to initial_lr over warmup_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.

• initialLr : α

  Target learning rate after warmup.

• warmupSteps : ℕ

  Number of warmup steps.

• startLr : α

  Starting learning rate during warmup.

• currentStep : ℕ

  Current step counter (0-indexed).

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

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

Before warmup_steps, linearly interpolate from start_lr to initial_lr. Afterwards, keep initial_lr fixed.

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

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

Advance the linear warmup scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.linearWarmupScheduler {α : Type} [Context α] (initialLr : α) (warmupSteps : ℕ) (startLr : α := 0) : LinearWarmupScheduler α

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.WarmupCosineScheduler (α : 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.

• initialLr : α

  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.WarmupCosineScheduler.getLr {α : Type} [Context α] (scheduler : WarmupCosineScheduler α) : α

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

• During warmup, LR increases linearly from 0 to initial_lr.
• After warmup, LR follows a cosine anneal over the remaining steps.
• At and after total_steps, 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.WarmupCosineScheduler.step {α : Type} (scheduler : WarmupCosineScheduler α) : WarmupCosineScheduler α

Advance the warmup+cosine scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.warmupCosineScheduler {α : Type} (initialLr : α) (warmupSteps totalSteps : ℕ) : WarmupCosineScheduler α

Create a warmup+cosine scheduler starting at step 0.

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

Cyclic LR

structure Optim.CyclicScheduler (α : Type) : Type

Cyclic learning rate schedule.

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

We keep mode as a String so this runtime layer can be configured from simple config files or CLI arguments (mirroring how training scripts are usually written).

• baseLr : α

  Minimum learning rate within the cycle.

• maxLr : α

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

• stepSize : ℕ

  Half-cycle size (in steps).

• mode : String

  "triangular", "triangular2", or "exp_range".

• gamma : α

  Decay factor used by "exp_range".

• currentStep : ℕ

  Current step counter (0-indexed).

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

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.CyclicScheduler.step {α : Type} (scheduler : CyclicScheduler α) : CyclicScheduler α

Advance the cyclic scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.cyclicScheduler {α : Type} [Context α] (baseLr maxLr : α) (stepSize : ℕ) (mode : String := "triangular") (gamma : α := 1) : CyclicScheduler α

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.TriangularCycleScheduler (α : Type) : Type

A specialized cyclic schedule with fixed amplitude.

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

• baseLr : α

  Minimum learning rate within the cycle.

• maxLr : α

  Maximum learning rate within the cycle.

• stepSize : ℕ

  Half-cycle size (in steps).

• currentStep : ℕ

  Current step counter (0-indexed).

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

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.TriangularCycleScheduler.step {α : Type} (scheduler : TriangularCycleScheduler α) : TriangularCycleScheduler α

Advance the triangular cycle scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.triangularCycleScheduler {α : Type} (baseLr maxLr : α) (stepSize : ℕ) : TriangularCycleScheduler α

Create a triangular cycle scheduler starting at step 0.

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

1cycle Learning Rate Schedule

structure Optim.OneCycleScheduler (α : Type) : Type

One-cycle learning-rate schedule.

• increase LR from initial_lr to max_lr over the first pct_start fraction of steps,
• then decrease to final_lr over the rest.

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

• maxLr : α

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

• totalSteps : ℕ

  Total number of steps in the schedule.

• initialLr : α

  Learning rate at step 0.

• finalLr : α

  Learning rate after the full schedule finishes.

• divFactor : α

  Divides max_lr to get initial_lr in the factory constructor.

• pctStart : α

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

• currentStep : ℕ

  Current step counter (0-indexed).

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

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

This ramps up to max_lr over the pct_start fraction of the schedule, then anneals down to final_lr.

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

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

Advance the 1cycle scheduler by one step.

PyTorch analogy: scheduler.step().

def Optim.oneCycleScheduler {α : Type} [Context α] (maxLr : α) (totalSteps : ℕ) (divFactor pctStart finalDivFactor : α) : OneCycleScheduler α

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.LRFinder (α : Type) : Type

Learning-rate finder schedule: exponential sweep from initial_lr to final_lr over num_steps.

• initialLr : α

  Learning rate at step 0.

• finalLr : α

  Target learning rate at the end of the sweep.

• numSteps : ℕ

  Number of steps in the sweep.

• currentStep : ℕ

  Current step counter (0-indexed).

def Optim.LRFinder.getLr {α : Type} [Context α] (finder : LRFinder α) : α

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

This increases LR exponentially from initial_lr toward final_lr across num_steps.

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

def Optim.LRFinder.step {α : Type} (finder : LRFinder α) : LRFinder α

Advance the LR finder by one step.

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

def Optim.lrFinder {α : Type} (initialLr finalLr : α) (numSteps : ℕ) : LRFinder α

Create an LR finder schedule starting at step 0.

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