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

1. Problemstellung

Die Vorverarbeitung (Pre-Training) von Large Language Models (LLMs) nutzt üblicherweise Lernraten-Scheduler mit einem Decay-Mechanismus (z. B. Cosine-Decay, Linear-Decay oder Warmup-Stable-Decay), um den Pre-Training-Verlust (Loss) zu minimieren. Die gängige Praxis besteht darin, die Lernrate gegen Ende des Trainings schrittweise auf einen kleinen Bruchteil (oft 0–10 %) der maximalen Lernrate zu senken.

Das zentrale Problem, das dieses Paper adressiert, ist die Diskrepanz zwischen der Optimierung des Pre-Training-Verlusts und der tatsächlichen Leistung des Modells nach dem anschließenden überwachtem Fine-Tuning (Supervised Fine-Tuning, SFT). Bisherige Studien zeigen, dass ein besseres Pre-Training-Ergebnis nicht zwangsläufig zu einer besseren Anpassungsfähigkeit an Downstream-Aufgaben führt. Die Autoren hinterfragen, ob das Herabsetzen der Lernrate (Decay), das auf Pre-Training-Metriken optimiert ist, die Adaptivität des Modells für nachfolgende Trainingsphasen wie SFT beeinträchtigt.

2. Methodik

Die Autoren untersuchen systematisch den Einfluss verschiedener Lernraten-Scheduler auf die Endleistung nach SFT. Sie vergleichen etablierte Decay-basierte Scheduler mit einem neuartigen Ansatz, den sie Warmup-Stable-Only (WSO) nennen.

• Warmup-Stable-Only (WSO): Dieser Scheduler entfernt die Decay-Phase vollständig. Nach der Warmup-Phase bleibt die Lernrate bis zum Ende des Trainings konstant auf dem Maximum ($\alpha_{pre} = 1.0$).
• Vergleichs-Scheduler: Als Baseline dienen Warmup-Stable-Decay (WSD), Cosine-Decay und Linear-Decay, wobei die Lernrate am Ende auf 0% ($\alpha_{pre}=0.0$) oder 10% ($\alpha_{pre}=0.1$) der Maximalrate gesenkt wird.
• Experimentelle Setup:
  • Modelle: Es wurden Modelle mit 1 Mrd. und 8 Mrd. Parametern (basierend auf der Llama-3-Architektur) trainiert.
  • Trainingsphasen: Die Experimente umfassen drei Szenarien:
    1. Zwei-Stufen-Pipeline (Pre-Training + SFT).
    2. Drei-Stufen-Pipeline (Pre-Training + Mid-Training + SFT), wobei Mid-Training als Zwischenschritt für Domänenanpassung und Kontextverlängerung dient.
    3. Over-Training-Szenario (Training auf 2 Billionen Tokens, weit über dem Chinchilla-Optimum).
  • Daten: FineWeb-Edu für Pre-Training, diverse Mid-Training-Datasets (z. B. dolmino-mix) und Tulu-3 SFT-Mixture für das Fine-Tuning.
  • Analyse der Loss-Landschaft: Um die Ursachen der Leistungsunterschiede zu verstehen, analysieren die Autoren die geometrische Struktur der Loss-Landschaft mittels der Sharpness (Schärfenmessung), definiert als Spur der Hesse-Matrix (Hessian Trace).

3. Wichtige Beiträge

1. Systematischer Nachweis der Überlegenheit von WSO: Die Studie liefert umfassende empirische Beweise, dass WSO in allen getesteten Szenarien (1B/8B Modelle, verschiedene Trainingslängen) konsistent bessere Ergebnisse nach SFT erzielt als Decay-basierte Scheduler, obwohl diese oft bessere Pre-Training-Metriken aufweisen.
2. Generalisierung über Trainingsparadigmen: Die Vorteile von WSO gelten nicht nur für das Standard-Pre-Training, sondern auch für moderne Pipelines mit Mid-Training und extremem Over-Training.
3. Theoretische Erklärung durch Loss-Landschaft: Die Autoren zeigen durch Analyse der Sharpness, dass Decay-basierte Scheduler Modelle in schärfere Minima (sharper minima) treiben, während WSO flachere Minima (flatter minima) erhält. Flache Minima korrelieren mit einer höheren Adaptivität und Stabilität während des Fine-Tunings.

4. Ergebnisse

• Leistungsumkehr (Inversion): Es zeigt sich ein deutlicher Trendumbruch: Decay-basierte Scheduler (insbesondere Linear-Decay zu 0) erzielen die besten Validierungs-Verluste und Task-Scores während des Pre-Trainings. Sobald jedoch SFT durchgeführt wird, schneiden WSO-Modelle signifikant besser ab.
  • Beispiel 8B-Modell: WSO erreicht nach SFT einen um +1,1 % höheren Durchschnittswert über die Aufgaben hinweg im Vergleich zum besten Decay-Scheduler, obwohl es im Pre-Training einen schlechteren Loss hatte.
• Robustheit bei Over-Training: Selbst bei Training auf 2 Billionen Tokens (100x Chinchilla-Optimum) behält WSO seinen Vorteil. Decay-basierte Modelle werden hier sogar noch schwerer zu fine-tunen.
• Korrelation Sharpness vs. SFT-Leistung: Eine negative Korrelation (Pearson $r = -0.709$) wurde zwischen der Sharpness des Pre-Training-Modells und der SFT-Leistung festgestellt. Modelle mit WSO weisen eine geringere Sharpness (flachere Landschaft) auf, was ihre Fähigkeit, sich an neue Daten anzupassen, ohne das vorherige Wissen zu verlieren, erklärt.
• Mid-Training: Auch in der Mid-Training-Phase führt die Vermeidung von Decay ($\alpha_{mid} = 1.0$) zu besseren SFT-Ergebnissen, selbst wenn die Mid-Training-Metriken selbst leicht darunter leiden.

5. Bedeutung und Implikationen

Diese Arbeit stellt eine fundamentale Infragestellung der aktuellen Best Practices im Training von LLMs dar.

• Paradigmenwechsel: Sie zeigt, dass die Optimierung von Pre-Training-Metriken (wie Validierungs-Loss) nicht das primäre Ziel sein sollte, wenn das Modell für Downstream-Aufgaben vorgesehen ist. Stattdessen sollte die Adaptivität (Anpassungsfähigkeit) priorisiert werden.
• Praktische Empfehlung: Die Autoren empfehlen, Modelle mit WSO vorzu- und zu veröffentlichen. Dies ermöglicht Praktikern, Modelle zu erhalten, die sich effizienter und effektiver für spezifische Anwendungen (via SFT, DPO etc.) anpassen lassen.
• Theoretischer Beitrag: Die Verbindung zwischen der Geometrie der Loss-Landschaft (Flachheit) und der Transfer-Learning-Fähigkeit wird hier empirisch für LLMs untermauert. Es wird gezeigt, dass ein konstantes hohes Lernraten-Niveau am Ende des Trainings hilft, das Modell in einem Bereich der Parameterraum zu halten, der für zukünftige Gradienten-Updates (Fine-Tuning) robuster ist.

Zusammenfassend demonstriert das Paper, dass der Verzicht auf Lernraten-Decay (WSO) eine einfache, aber hochwirksame Strategie ist, um die Endleistung von LLMs nach dem Fine-Tuning zu maximieren, indem die zugrundeliegende Adaptivität des Modells durch flachere Loss-Minima bewahrt wird.