Learning-Rate Schedules

Pure learning-rate schedules used by TorchLean training loops and examples. A Config describes a schedule, and lrAt evaluates it at an optimizer step.

References

• PyTorch schedulers: torch.optim.lr_scheduler.* (https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate)

inductive TorchLean.Trainer.Scheduler.Config : Type

Small learning-rate scheduler surface for higher-level training code.

This file keeps the interface compact: a Config is just a description of a schedule, and lrAt cfg t computes the learning rate at step/epoch index t.

PyTorch mapping

Config.step and Config.exponential correspond to the schedule math of:

• torch.optim.lr_scheduler.StepLR
• torch.optim.lr_scheduler.ExponentialLR

Config.warmupCosine is the schedule commonly used for Transformer pretraining: a short linear warm-up followed by cosine decay to a nonzero floor.

• constant (lr : Float) : Config
• step (base : Float) (stepSize : Nat) (gamma : Float := 0.1) : Config
• exponential (base gamma : Float) : Config
• warmupCosine (peak min : Float) (warmupSteps totalSteps : Nat) : Config

@[implicit_reducible]

instance TorchLean.Trainer.Scheduler.instReprConfig : Repr Config

def TorchLean.Trainer.Scheduler.instReprConfig.repr : Config → Nat → Std.Format

def TorchLean.Trainer.Scheduler.constant (lr : Float) : Config

Constant learning-rate schedule.

def TorchLean.Trainer.Scheduler.step (base : Float) (stepSize : Nat) (gamma : Float := 0.1) : Config

Step decay learning-rate schedule.

def TorchLean.Trainer.Scheduler.exponential (base gamma : Float) : Config

Exponential learning-rate schedule.

def TorchLean.Trainer.Scheduler.warmupCosine (peak min : Float) (warmupSteps totalSteps : Nat) : Config

Linearly warm up to peak, then follow a cosine curve down to min.

warmupSteps counts optimizer updates. The first update uses peak / warmupSteps, and the last warm-up update reaches peak. Once totalSteps updates have been scheduled, the learning rate remains at min. A warm-up longer than the run is clamped to totalSteps. When totalSteps = 0, no update belongs to the schedule and lrAt returns min.

def TorchLean.Trainer.Scheduler.lrAt : Config → Nat → Float

Learning rate at a given step or epoch index.