Learning-Rate Schedules

Pure learning-rate schedules used by TorchLean training loops and examples. A Config describes a schedule, and learningRateAt 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 learningRateAt config stepIndex computes the learning rate at that optimizer step or epoch.

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 (learningRate : Float) : Config
• step (baseLearningRate : Float) (stepSize : Nat) (decayFactor : Float := 0.1) : Config
• exponential (baseLearningRate decayFactor : Float) : Config
• warmupCosine (peakLearningRate minimumLearningRate : Float) (warmupSteps totalSteps : Nat) : Config

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

@[instance_reducible]

instance TorchLean.Trainer.Scheduler.instReprConfig : Repr Config

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

Constant learning-rate schedule.

def TorchLean.Trainer.Scheduler.step (baseLearningRate : Float) (stepSize : Nat) (decayFactor : Float := 0.1) : Config

Step decay learning-rate schedule.

def TorchLean.Trainer.Scheduler.exponential (baseLearningRate decayFactor : Float) : Config

Exponential learning-rate schedule.

def TorchLean.Trainer.Scheduler.warmupCosine (peakLearningRate minimumLearningRate : Float) (warmupSteps totalSteps : Nat) : Config

Linearly warm up to peakLearningRate, then follow a cosine curve down to minimumLearningRate.

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

def TorchLean.Trainer.Scheduler.Internal.requireRate (name : String) (value : Float) : Except String Unit

Reject a learning rate that is not a finite number at or above zero.

def TorchLean.Trainer.Scheduler.Internal.requireDecay (name : String) (value : Float) : Except String Unit

Reject a decay factor outside [0, 1]. A factor above one grows the learning rate instead of decaying it, which in practice is a typo rather than an intention.

def TorchLean.Trainer.Scheduler.validate : Config → Except String Unit

Validate a learning-rate schedule before it is attached to optimizer state.

The checks keep every scheduled rate finite and nonnegative. Step and exponential schedules use decay factors in [0, 1]; a zero step size is rejected instead of silently changing the schedule to a constant rate.

def TorchLean.Trainer.Scheduler.validateFloat32 (config : Config) : Except String Unit

Also check that the largest scheduled rate remains finite in binary32 training.

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

Learning rate at a given step or epoch index.