Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Pre-training LLM without Learning Rate Decay Enhances Supervised Fine-Tuning

Kazuki Yano, Shun Kiyono, Sosuke Kobayashi, Sho Takase, Jun Suzuki

Abstract

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.

Pre-training LLM without Learning Rate Decay Enhances Supervised Fine-Tuning

Abstract

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.
Paper Structure (36 sections, 6 equations, 4 figures, 4 tables)

This paper contains 36 sections, 6 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Pre-training.
  4. Post-training.
  5. Task Definition
  6. Further considerations
  7. Mid-training.
  8. Over-training.
  9. Current LR Scheduling Practices
  10. Formalization of Learning Rate Schedulers
  11. LR Scheduling in Mid-training.
  12. Experiment 1: Two-stage (Pre- and Post-training) Setting
  13. Experimental Setup
  14. Model Architectures.
  15. Pre-training Configuration.
  16. ...and 21 more sections

Figures (4)

  • Figure 1: Learning rate schedulers used in pre-training and their impact on performance after supervised fine-tuning (SFT). Warmup-Stable-Only (WSO), which removes the decay phase, achieves the highest performance after SFT.
  • Figure 2: Mid-training LR schedulers with different $\alpha_{\text{pre}}$ and $\alpha_{\text{mid}}$ values.
  • Figure 3: $\text{Sharpness}(\theta_t)$ during pre-training of the 1B model. Vertical line at step $T_{\text{stable}}$ indicating where WSD decays LR. Decay-based schedulers ($\alpha_{\text{pre}} = 0$ or $\alpha_{\text{pre}} = 0.1$) lead to sharper minima, while WSO ($\alpha_{\text{pre}} = 1.0$) maintains flatter landscapes.
  • Figure 4: Pre-training sharpness negatively correlates with downstream SFT performance.

Theorems & Definitions (1)

  • Definition 6.1: Sharpness