Warmup-Stable-Decay Schedule

The problem, from first principles

The cosine schedule needs to know the total number of training steps in advance, because it decays all the way to the end. That is inconvenient when you want to continue training, add data or branch off a checkpoint. The warmup-stable-decay (WSD) schedule has three phases. The learning rate warms up linearly to its peak, stays flat for most of the run, and is decayed only in a short final stretch. Any checkpoint taken during the stable phase can be turned into a finished model by running a decay phase from it, and the plateau avoids committing to a total length in advance.

From theory to code

Implement wsd_lr.

Constraints

• wsd_lr(step, peak_lr, warmup_steps, decay_start, total_steps, final_ratio) with 0 <= warmup_steps <= decay_start <= total_steps.
• For step < warmup_steps the rate is peak_lr * (step + 1) / warmup_steps.
• For warmup_steps <= step < decay_start the rate is peak_lr.
• For step >= decay_start the rate falls linearly from peak_lr to peak_lr * final_ratio: with p = (step - decay_start) / (total_steps - decay_start) clipped to [0, 1] the rate is peak_lr * (1 - p * (1 - final_ratio)). If decay_start == total_steps the decay phase never begins.

Hints