💬 NLP

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

This paper demonstrates that using a Warmup-Stable-Only (WOS) learning rate schedule without decay during large-scale pre-training enhances downstream supervised fine-tuning performance by preserving flatter loss minima, even though decay-based schedulers may yield lower pre-training loss.

Original authors: Kazuki Yano, Shun Kiyono, Sosuke Kobayashi, Sho Takase, Jun Suzuki

Published 2026-03-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kazuki Yano, Shun Kiyono, Sosuke Kobayashi, Sho Takase, Jun Suzuki

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a brilliant student to become a world-class expert. You have two main phases:

The University Years (Pre-training): The student reads millions of books, learns grammar, facts, and how the world works.
The Internship (Fine-tuning): The student learns to apply that knowledge to specific jobs, like writing code, giving medical advice, or following strict company rules.

For years, the standard advice for the "University Years" was: "Start with high energy, then slowly wind down to a complete stop."

In the world of AI, this "energy" is the Learning Rate. The standard method (called Cosine Decay or WSD) tells the AI: "Learn fast at first, but as you get closer to finishing your degree, slow your learning pace down to almost zero." The logic was that slowing down helps the AI memorize the textbooks perfectly, getting the lowest possible error score on its final exams.

This paper says: "Stop slowing down!"

The researchers discovered that if you keep the AI's learning pace steady and constant all the way to the finish line (a method they call WSO or Warmup-Stable-Only), the AI actually becomes a better intern.

Here is the breakdown of their findings using simple analogies:

1. The "Perfect Student" vs. The "Adaptable Student"

The Old Way (Decay): Imagine a student who studies so hard and slows down so much at the end that they memorize the textbook word-for-word. They get a perfect score on the final exam (Pre-training Loss). But when they start their internship, they are rigid. If the boss asks a slightly different question, the student freezes because they are stuck in their "perfect" memorization.
The New Way (WSO): Imagine a student who keeps a steady, energetic pace until the very last day. They might not have memorized the textbook quite as perfectly as the first student (their final exam score is slightly lower). However, because they kept their brain "active" and flexible until the end, they are much better at handling new, real-world problems during their internship.

2. The "Valley" Analogy (Loss Landscape)

To understand why this happens, imagine the learning process is like a hiker trying to find the deepest point in a valley (the best solution).

The Decay Method: When the hiker slows down at the end, they tend to get stuck in a narrow, deep canyon. It's the lowest point, but it's very tight. If you try to move even a little bit (like during an internship), you hit the walls immediately. This is called a "sharp minimum."
The WSO Method: By keeping the pace steady, the hiker finds a wide, flat meadow at the bottom of the valley. It's not the absolute deepest point, but it's spacious. If you move around a bit, you don't hit a wall; you just stay in the meadow. This is called a "flat minimum."

The Magic: When the AI moves from "University" to "Internship," it needs to move around. The AI trained with the "flat meadow" (WSO) can adapt easily without crashing. The AI trained in the "narrow canyon" (Decay) breaks easily when asked to do something new.

3. The Results

The researchers tested this on AI models of different sizes (1 billion and 8 billion parameters).

Before Fine-Tuning: The "Decay" models looked better. They had lower error rates on standard tests.
After Fine-Tuning (The Real Test): The "WSO" models crushed it. They were better at following instructions, telling the truth, and solving complex tasks.

The Big Takeaway

For a long time, AI researchers thought, "If the model learns the basics perfectly, it will be good at everything." This paper proves that being too perfect at the basics can actually hurt your ability to adapt later.

The Recommendation:
If you want to build an AI that is truly useful in the real world, don't let it "coast" to a stop at the end of its training. Keep the learning rate steady. It might not look like the "smartest" student on paper, but it will be the most adaptable and helpful employee once hired.

In short: Don't let your AI slow down too early. Keep it running at full speed so it stays flexible enough to learn new tricks later!

1. Problem Statement

Large Language Model (LLM) pre-training typically employs decay-based learning rate (LR) schedulers (e.g., Cosine decay, Linear decay, or Warmup-Stable-Decay [WSD]) to minimize pre-training loss. The prevailing assumption is that optimizing the pre-training objective directly leads to better downstream performance.

However, the authors identify a critical gap:

Suboptimality for Downstream Tasks: Recent studies suggest that a model with superior pre-training metrics does not necessarily yield better performance after Supervised Fine-Tuning (SFT).
Inflexibility: Decay-based schedulers are often rigid, requiring heuristic tuning when models undergo continual pre-training or mid-training.
The Core Question: Does the standard practice of decaying the learning rate to zero (or near-zero) during pre-training actually harm the model's adaptability for downstream SFT tasks, even if it improves pre-training loss?

2. Methodology

The authors propose and evaluate Warmup-Stable-Only (WSO), a scheduler that removes the decay phase entirely, maintaining a constant learning rate after the warmup period until the end of training.

Experimental Setup

Models: Experiments were conducted on 1B and 8B parameter models following the Llama 3 architecture.
Schedulers Compared:
- WSO: Warmup + Stable (Constant LR, $\alpha_{pre} = 1.0$ ).
- WSD: Warmup + Stable + Decay (Linear decay to $\alpha_{pre} \cdot \eta_{max}$ ).
- Cosine/Linear: Standard decay schedules.
- Variations: The minimum LR factor ( $\alpha_{pre}$ ) was varied ($0.0, 0.1, 1.0$).
Training Regimes:
1. Two-Stage: Standard Pre-training $\to$ SFT.
2. Three-Stage: Pre-training $\to$ Mid-training $\to$ SFT (incorporating modern pipelines like OLMo 2).
3. Over-training: Training on significantly larger token counts (up to 2 Trillion tokens, ~100x Chinchilla optimal) to test scalability.
Evaluation:
- Pre-training/Mid-training: Validation loss and zero-shot benchmarks (ARC, HellaSwag, etc.).
- Post-training (SFT): Performance on instruction-following (AlpacaEval), truthfulness (TruthfulQA), and multitask understanding (MMLU).
Theoretical Analysis: The authors analyzed the loss landscape geometry by measuring Sharpness (Trace of the Hessian matrix) using Hutchinson's estimator to understand why certain schedulers perform better.

3. Key Contributions

Systematic Demonstration of WSO Superiority: The paper provides empirical evidence that WSO consistently outperforms decay-based schedulers on downstream SFT tasks across different model sizes (1B, 8B) and training regimes (standard, mid-training, over-training).
Inversion of Performance Metrics: The study reveals a "performance inversion": models trained with decay-based schedulers often achieve lower validation loss and better zero-shot scores during pre-training, yet they underperform WSO models after SFT.
Loss Landscape Analysis: The authors link the performance gap to the geometry of the loss landscape. They demonstrate that decay-based schedulers drive models into sharper minima, whereas WSO preserves flatter minima, which correlates with higher adaptability during fine-tuning.
Generalizability: The findings hold true even in extreme over-training scenarios (2T tokens) and complex multi-stage pipelines involving mid-training.

4. Key Results

A. Performance Inversion (Pre-training vs. SFT)

Pre-training Metrics: Decay-based schedulers (specifically Linear and WSD with $\alpha_{pre}=0.0$ ) generally achieved the lowest validation loss and competitive zero-shot benchmark scores.
SFT Metrics: WSO ( $\alpha_{pre}=1.0$ $α_{p r e} = 1.0$ ) consistently achieved the highest SFT scores (AlpacaEval, MMLU, TruthfulQA).
- Example (8B Model): WSO achieved a +1.1 relative improvement in SFT Task Average compared to the best decay-based scheduler, despite having a higher pre-training validation loss (+0.127).
- Example (1B Model): WSO showed a +0.3 relative improvement in SFT performance.

B. Multi-Stage and Over-Training

Mid-Training: In three-stage pipelines, applying decay at any stage (pre-training or mid-training) reduced final SFT performance. The best results were achieved when both pre-training and mid-training used WSO (constant LR).
Over-Training: Even when trained on 2 Trillion tokens (far beyond the Chinchilla optimal regime), WSO maintained its advantage, yielding superior SFT performance compared to decay-based counterparts.

C. Loss Landscape Geometry

Sharpness Correlation: There is a strong negative correlation (Pearson $r \approx -0.71$ ) between the sharpness of the loss landscape at the end of pre-training and downstream SFT performance.
Mechanism:
- Decay Schedulers: As the LR decays, the model converges into sharper minima (high sharpness values, ~2–3 $\times$ higher than WSO). This makes the model sensitive to parameter updates, leading to instability or performance drops during SFT.
- WSO: Maintains a flatter loss landscape (low sharpness). Models in flat minima are more robust to parameter perturbations, allowing them to adapt more effectively to new data distributions during SFT without catastrophic forgetting or instability.

5. Significance and Implications

Rethinking Pre-training Objectives: The paper challenges the dogma that minimizing pre-training loss is the ultimate goal. It suggests that optimizing for pre-training loss may compromise downstream adaptability.
Practical Training Strategy: For practitioners aiming to release models for fine-tuning, WSO is a superior strategy. It simplifies the training pipeline (removing the need to tune decay schedules) and yields models that are "more portable" and easier to adapt.
Model Release Recommendations: The authors recommend releasing LLMs trained with WSO to maximize their utility for downstream developers.
Theoretical Insight: The work bridges the gap between optimization theory (flat minima generalization) and large-scale LLM training, providing a concrete mechanism (loss landscape sharpness) for why "over-training" or "decay" can sometimes be detrimental to fine-tuning.

Conclusion

The paper concludes that Warmup-Stable-Only (WSO) is a more effective learning rate schedule for the entire LLM development pipeline than traditional decay-based methods. By avoiding the decay phase, WSO preserves flatter minima in the loss landscape, which significantly enhances the model's ability to adapt during Supervised Fine-Tuning, ultimately leading to better real-world performance despite potentially higher pre-training loss.