Search papers, labs, and topics across Lattice.
This paper investigates the impact of learning rate (LR) scheduling during large language model (LLM) pre-training on subsequent supervised fine-tuning (SFT) performance. They introduce and evaluate a "Warmup-Stable-Only" (WSO) LR schedule, which maintains a constant LR after the warmup phase, contrasting it with decay-based schedules. Experiments on 1B and 8B parameter models demonstrate that WSO consistently outperforms decay-based schedules on downstream SFT tasks, despite the latter achieving better pre-training loss, with loss landscape analysis suggesting that WSO leads to flatter, more adaptable minima.
Ditch the learning rate decay: LLMs pre-trained with a stable learning rate after warmup show superior adaptability and downstream performance after supervised fine-tuning compared to those trained with decay-based schedules.
We investigate the role of learning rate scheduling in the large-scale pre-training of large language models, focusing on its influence on downstream performance after supervised fine-tuning (SFT). Decay-based learning rate schedulers are widely used to minimize pre-training loss. However, despite their widespread use, how these schedulers affect performance after SFT remains underexplored. In this paper, we examine Warmup-Stable-Only (WSO), which maintains a constant learning rate after warmup without any decay. Through experiments with 1B and 8B parameter models, we show that WSO consistently outperforms decay-based schedulers in terms of performance after SFT, even though decay-based schedulers may exhibit better performance after pre-training. The result also holds across different regimes with mid-training and over-training. Loss landscape analysis further reveals that decay-based schedulers lead models into sharper minima, whereas WSO preserves flatter minima that support adaptability. These findings indicate that applying LR decay to improve pre-training metrics may compromise downstream adaptability. Our work also provides practical guidance for training and model release strategies, highlighting that pre-training models with WSO enhances their adaptability for downstream tasks.