WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training

2026-07-12

AI Summary

This work proposes WSqD, a novel learning rate scheduling strategy that overcomes the inflexibility of existing methods which rely on predetermined training durations. WSqD introduces, for the first time in large model training, a time-independent inverse square root base schedule, integrated into a three-phase design comprising warmup, the proposed base schedule, and terminal linear decay. By requiring only the specification of the decay onset, WSqD decouples learning rate adaptation from total training length. Grounded in stochastic convex optimization theory, WSqD achieves competitive or superior performance compared to strong baselines such as WSD across varying training durations on the SlimPajama corpus, using a single peak learning rate. This approach substantially enhances scheduling flexibility and training robustness.

Abstract

Standard learning rate schedules such as cosine annealing are tied to a fixed training horizon, limiting their ability to accommodate post hoc horizon extension. Warmup-stable-decay (WSD) partially addresses this issue by maintaining a long constant-rate phase before a short linear cooldown, allowing training to resume from a pre-decay checkpoint. However, its peak learning rate is still tuned based on the original training horizon and can become suboptimal when training is extended. Motivated by stochastic convex optimization, we propose WSqD (Warmup with Square-root base and linear Decay), a learning rate schedule that replaces WSD's constant stable phase with a shifted inverse-square-root base while retaining the final linear cooldown. In the stochastic convex setting, WSqD provably attains the minimax-optimal $O(1/\sqrt{T})$ last-iterate convergence rate. Importantly, its base learning rate schedule is horizon-independent, and the training horizon is needed only to determine when to begin the final cooldown. Empirically, on language-model pretraining using the SlimPajama corpus, WSqD matches or outperforms carefully tuned WSD and other baselines across multiple training horizons while reusing a single peak learning rate.

Problem

• learning rate schedule
• training horizon
• horizon-free
• large model training
• stochastic optimization

Innovation

• horizon-free
• learning rate schedule
• square-root decay
• stochastic convex optimization
• last-iterate convergence