PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2603.16127v1cs.CLcs.LG

TL;DR

The paper asks whether learning-rate decay chosen to optimize pre-training metrics also benefits downstream SFT, where the evidence remains limited. It evaluates WSO, which keeps the learning rate constant after warmup, across model sizes and training regimes, and finds consistently better post-SFT performance than decay-based schedules. The analysis attributes this advantage to flatter minima, while the paper recommends WSO for more adaptable and portable models.

  • Problem

    Whether decay-based learning-rate schedules that improve pre-training metrics also maximize downstream performance after SFT remains underexplored.

  • Method

    The paper empirically evaluates WSO, which omits decay and maintains a constant learning rate, across 1B and 8B models and standard, mid-training, and over-training regimes.

  • Results

    WSO consistently outperforms decay-based schedulers on downstream tasks after SFT, despite decay-based schedules sometimes achieving better pre-training or intermediate metrics.

  • Takeaways & Limitations

    Pre-training with WSO preserves flatter minima and offers a promising strategy for constructing more adaptable, portable models for downstream use.

  • Takeaways & Limitations

    The study focuses on SFT as the canonical post-training stage and leaves complex learning-rate combinations across multiple post-training stages for future work.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper questions selecting decay-based learning rates for pre-training when downstream SFT performance is the practical objective. It investigates WSO and finds that removing decay consistently improves post-SFT performance while preserving flatter loss minima associated with adaptability.

  • Motivation: Decay-based scheduling can improve pre-training metrics without guaranteeing stronger performance after supervised fine-tuning.This motivates evaluating learning-rate schedules by downstream SFT performance rather than pre-training performance alone.
  • Approach: WSO removes the decay phase and maintains a constant learning rate through the end of training.WSO is presented as a simplified variant of WSD.
  • Findings: WSO consistently outperforms decay-based schedulers after SFT across 1B and 8B models, mid-training, and over-training regimes.The reported conclusion spans standard pre-training experiments and modern multi-stage training settings.
  • Findings: Figure 1 identifies WSO as the scheduler with the highest performance after SFT.The figure directly contrasts scheduler choices used during pre-training with their downstream SFT outcomes.
  • Mechanism: WSO-trained models occupy flatter loss-landscape regions than models trained with decay-based schedulers, supporting better post-training adaptability.The paper links the flatter minima observed under WSO to improved adaptability for downstream tasks.

2 PRELIMINARIES

The paper frames LLM development as a staged pipeline in which pre-training and optional mid-training precede post-training, while questioning whether intermediate-stage model selection optimizes final performance. It formalizes learning-rate schedulers and examines non-decaying schedules across these training stages.

  • Training stages: Pre-training learns general language representations, while post-training adapts models to target tasks through methods including supervised fine-tuning.SFT is treated as the canonical post-training stage evaluated in this paper.
  • Scope: The paper leaves complex combinations of learning-rate scheduling across multiple post-training stages for future work.
  • Task definition: Selecting the best-performing model at each intermediate stage does not guarantee the best final model after subsequent training.The paper therefore formulates final-model selection as a search over model configurations across training stages.
  • Task definition: The study empirically evaluates learning-rate schedulers during large-scale training stages that precede post-training.
  • Training stages: Mid-training occupies an intermediate position between large-scale pre-training and task-specific post-training, supporting objectives such as domain expansion and long-context extension.
  • Training regimes: Over-training extends training far beyond the compute-optimal regime, trading additional training compute for improved inference efficiency at deployment.
  • Learning-rate schedules: WSO omits the decay phase of WSD by setting αpre = 1.0, while WSD, Cosine, and Linear provide comparison schedules.The scheduler notation uses ηScheduler(t, αpre), with αpre controlling the minimum pre-training learning-rate factor.
  • Learning-rate schedules: In mid-training, αmid = 0.0 applies Linear decay, whereas αmid = 1.0 keeps the learning rate constant throughout the stage.Combining αpre = 1.0 and αmid = 1.0 extends WSO across pre-training and mid-training.

3 EXPERIMENT 1: TWO-STAGE (PRE- AND POST-TRAINING) SETTING

This experiment tests whether decaying learning rates during pre-training improve downstream SFT performance across 1B and 8B models. Although decay-based schedulers perform best on pre-training metrics, WSO performs best after SFT.

  • Experimental Setup: The experiment compares WSO, WSD, Cosine, and Linear learning-rate schedulers across 1B and 8B models.Models are evaluated after pre-training and after supervised fine-tuning using zero-shot and fine-tuned-task benchmarks.
  • Results: Pre-training metrics favor decay-based schedulers, with Linear and WSD at αpre = 0 achieving the best PT Task Avg for 1B and 8B models, respectively.The study reports an inversion between intermediate pre-training performance and downstream SFT performance.
  • Results: WSO achieves the best SFT performance for both model sizes despite underperforming decay-based schedulers on pre-training metrics.Relative results are reported for validation loss, zero-shot task averages, and averages across AlpacaEval, TruthfulQA, and MMLU.
  • Results: The results suggest that selecting a pre-training checkpoint by pre-training metrics alone can miss the model that performs best after SFT.The paper frames downstream post-training performance as the more relevant criterion for the overall training pipeline.

4 EXPERIMENT 2: THREE-STAGE (PRE-, MID-, AND POST-TRAINING) SETTING

This experiment evaluates learning-rate scheduling across pre-training, mid-training, and SFT. Across the multi-stage pipeline, WSO yields the best downstream SFT performance even when decay-based schedules produce stronger mid-training metrics.

  • Experimental Setup: The three-stage setup systematically varies pre-training and mid-training schedulers before evaluating downstream SFT performance.Mid-training follows an OLMo 2-style pipeline using either further Linear decay or a constant learning rate.
  • Experimental Setup: Mid-training uses αmid = 0.0 for further Linear decay and αmid = 1.0 for a constant learning rate.The study compares scheduler choices at both pre-training and mid-training stages.
  • Results: Decay-based scheduling achieves the best mid-training metrics, but WSO achieves the best downstream SFT performance.This reproduces the inversion observed between intermediate training metrics and post-training performance.
  • Results: Introducing decay at any stage reduces SFT performance, while avoiding mid-training decay improves both mid-training and SFT results for models pre-trained with decay.The latter comparison uses αpre = 0.1 and contrasts αmid = 1.0 with applying decay during mid-training.
  • Results: WSO throughout pre-training and mid-training gives the strongest overall training-pipeline performance, including mid-training and SFT.The result extends the scheduler comparison beyond the two-stage setting.

5 EXPERIMENT 3: THREE-STAGE SETTING IN THE OVER-TRAINING

The over-training experiment tests whether WSO remains advantageous with much larger data budgets. Using 2T tokens, the inversion persists: decay-based schedulers perform better on intermediate metrics, whereas WSO performs better after SFT.

  • Experimental Setup: The over-training regime evaluates 1B models on 2T pre-training tokens, approximately 100× the Chinchilla-optimal data amount for this model size.The study also evaluates a 500B-token mid-training condition using the same scheduler comparisons.
  • Experimental Setup: All schedulers are evaluated both after pre-training or mid-training and after SFT.This preserves the comparison between intermediate metrics and downstream performance used in earlier experiments.
  • Results: Across single-stage and multi-stage over-training, WSO produces worse intermediate metrics but superior SFT performance than decay-based schedulers.The same inversion holds after 2T tokens and after the 2T + 500B-token pipeline.
  • Results: The WSO advantage persists across different data scales and training configurations.This includes both single-stage over-training and over-training combined with mid-training.

6 UNDERSTANDING ADAPTABILITY THROUGH LOSS LANDSCAPE GEOMETRY

The paper analyzes loss-landscape geometry to explain WSO’s downstream advantage. It measures sharpness through Hessian trace and finds that WSO preserves flatter regions associated with better SFT adaptability.

  • Sharpness Analysis: The analysis uses sharpness to characterize loss-landscape curvature around converged model parameters.Sharpness is defined as the trace of the Hessian with respect to model parameters.
  • Sharpness Analysis: Lower sharpness corresponds to flatter regions where parameter updates during SFT produce more stable loss behavior and better downstream adaptability.The paper connects flatter landscapes with reduced fluctuation during fine-tuning.
  • Sharpness Analysis: Because the full Hessian trace is impractical for billion-parameter models, the study estimates it with Hutchinson’s unbiased estimator using Hessian-vector products.Hessian-vector products are computed efficiently through automatic differentiation.
  • Results: During pre-training, Cosine, Linear, and WSD show increasing or rising sharpness as learning rates decay, whereas WSO maintains lower sharpness.Figure 3 reports this pattern for the 1B model on validation sets from both pre-training and SFT datasets.
  • Results: Pre-training sharpness negatively correlates with downstream SFT performance.This relationship supports flatter regions as a geometric explanation for WSO’s stronger post-training adaptability.

7 RELATED WORK

Prior work largely favors learning-rate decay for minimizing pre-training loss, while studying its implications for downstream adaptability. Research also links flatter loss regions with better generalization and interprets WSD through valley-like loss landscapes.

  • Cosine decay remains a standard LLM pre-training scheduler, while Linear decay to zero can achieve lower pre-training loss in compute-optimal settings.
  • WSD delays learning-rate decay until the final training phase rather than maintaining a constant rate throughout.
  • Continual pre-training studies examine learning-rate scheduling for domain adaptation and model-version updates.
  • Flatter loss regions have been associated with better generalization, motivating sharpness-aware methods and stochastic weight averaging.
  • A river-valley view describes stable training as exploration along the valley floor and decay as convergence toward its center.

8 CONCLUSION

The study evaluates learning-rate schedulers by their downstream performance after SFT rather than only pre-training outcomes. WSO consistently improves post-SFT performance across standard, mid-training, and over-training regimes, while preserving flatter minima associated with adaptability.

  • WSO consistently outperforms decay-based schedulers on downstream tasks after SFT across standard pre-training, mid-training, and over-training regimes.
  • Loss-landscape analysis finds that WSO preserves flatter minima, explaining its superior adaptability.
  • WSO is presented as a simple alternative for constructing more portable models and is recommended for releasing adaptable LLMs.

ETHICS STATEMENT

This material describes the study’s responsible-use statement and selected training-scheduler specifications. It includes safety guidance, reproducibility choices, scheduler formulations, and experimental settings for pre-training, mid-training, over-training, and SFT.

  • ETHICS STATEMENT: The authors encourage responsible deployment with appropriate safety measures during post-training.
  • ETHICS STATEMENT: The models were trained entirely from scratch using publicly available datasets for pre-training, SFT, and evaluation.
  • Learning-rate scheduler formulations: WSD keeps the learning rate stable after warmup until Tstable, then decays it linearly to αpre · ηmax at step T.
  • Learning-rate scheduler formulations: WSO is obtained by setting αpre = 1 in WSD, so the learning rate remains constant after warmup.
  • Learning-rate scheduler formulations: All schedulers share the same warmup phase, with decay controlled by the minimum learning-rate factor αpre ∈[0.0, 1.0].
  • Learning-rate scheduler formulations: Mid-training extends the pre-training learning-rate schedulers across the intermediate training stage.
  • Experimental settings: The WSD stable ratio ρ = 0.75 keeps the learning rate stable for 75% of training after warmup when decay is enabled.
  • Experimental settings: SFT uses a sweep over specified learning rates and selects the best value based on AlpacaEval performance.

E EVALUATION DETAILS

Evaluation covers zero-shot pre-trained models, benchmark-based mid-trained models, and SFT models evaluated with task-specific configurations. In over-training, Cosine shows slightly better zero-shot performance and lower validation loss than WSO, despite the paper’s broader post-SFT findings.

  • Evaluation configurations: Pre-trained models are evaluated on all benchmarks in a zero-shot setting.
  • Evaluation configurations: Mid-trained models are assessed on reasoning, reading-comprehension, and mathematical-reasoning benchmarks using the stated prompting configurations.
  • Evaluation configurations: AlpacaEval uses a same-architecture WSO reference model fine-tuned with the lowest swept learning rate for stable comparisons between schedules.
  • Evaluation results: The appendix reports absolute per-task results alongside aggregated metrics and relative comparisons for pre-trained and fine-tuned models.
  • Over-training evaluation: Table 10 reports that, with 2T tokens, Cosine with decay achieves slightly better zero-shot task performance and lower validation loss than WSO.

F.2 SFT EVALUATION RESULTS

SFT evaluation prioritizes downstream instruction-following performance rather than pre-training metrics. Across standard, over-training, and mid-training settings, WSO generally matches or outperforms decay-based schedules after SFT, despite weaker or similar validation metrics before SFT.

  • Evaluation setup: SFT learning rates are selected using AlpacaEval, and every pre-trained model receives the same learning-rate sweep.This design targets instruction-following performance and avoids giving any scheduler a selective hyperparameter advantage.
  • Pre-training versus SFT: More decay generally lowers validation loss, but does not consistently improve zero-shot task performance or downstream SFT performance.The reported comparison separates pre-training metrics from the performance objective evaluated after supervised fine-tuning.
  • Standard and over-trained models: WSO or moderate decay (αpre = 0.1) often matches or exceeds aggressive decay (αpre = 0.0) after SFT, despite worse pre-training metrics.This pattern is reported for the standard SFT evaluation and its corresponding 1B-model tables.
  • Mid-training and over-training: After over-training with 2T tokens, WSO achieves superior SFT performance compared with decay-based schedulers, including when models are subsequently mid-trained.The same overall result is reported for mid-training and over-training combinations, including nearly identical validation-loss settings.
  • Loss-landscape analysis: Sharpness is measured as Hessian trace with Hutchinson’s estimator at regular intervals on held-out pre-training and SFT data.The estimator uses random Rademacher vectors and automatic differentiation for Hessian-vector products.
Loading 2603.16127v1…