Weight Decay Improves Language Model Plasticity

Technisches Resümee: Weight Decay verbessert die Plastizität von Sprachmodellen

Problemstellung
Die gängige Praxis beim Training von Large Language Models (LLMs) sieht typischerweise eine Entkopplung von Pretraining und Post-Training (Feinabstimmung, Alignment) vor. Hyperparameter-Optimierung und Scaling Laws während des Pretrainings werden vorwiegend basierend auf dem Validierungsverlust (Cross-Entropy) des Basismodells bewertet, unter der Annahme, dass ein niedrigerer Pretraining-Loss mit einem leistungsfähigeren Downstream-Modell korreliert. Dieser Ansatz vernachlässigt jedoch die Modell-Plastizität – die Fähigkeit eines trainierten Modells, sich bei weiterem Training ohne Reinitialisierung erfolgreich an neue Daten und Aufgaben anzupassen. Die Arbeit untersucht, ob die Optimierung der Pretraining-Hyperparameter ausschließlich auf die Pretraining-Performance hin die Anpassungsfähigkeit des Modells während der anschließenden Feinabstimmung unbeabsichtigt beeinträchtigt. Konkret konzentrieren sich die Autoren auf Weight Decay (Gewichtszerfall), einen Standard-Regularisierungshyperparameter, um zu bestimmen, ob sein optimaler Wert zur Minimierung des Pretraining-Losses vom optimalen Wert zur Maximierung der Downstream-Adaptivität abweicht.

Methodik
Die Autoren führen eine systematische empirische Studie über zwei Modellfamilien (Llama-2 und OLMo-2) und mehrere Modellgrößen (bis zu 4 Mrd. Parameter) durch. Das experimentelle Design variiert den Weight-Decay-Hyperparameter ($\lambda$) während des Pretrainings, während andere Faktoren kontrolliert werden.

1. Trainingsregime: Modelle werden in zwei unterschiedlichen Regimen trainiert:
   • Rechenoptimal (compute-optimal): 20 Tokens pro Parameter (TPP).
   • Übertrainiert (overtrained): 140 TPP.
2. Feinabstimmungs-Evaluierung: Die vortrainierten Modelle werden auf einer vielfältigen Suite von Downstream-Aufgaben evaluiert, darunter:
   • Sechs Chain-of-Thought (CoT) Reasoning-Datensätze (z. B. MetaMathQA, MMLUProCoT).
   • Fünf Benchmarks für Sprachverständnis und Commonsense-Reasoning (z. B. HellaSwag, WinoGrande).
   • Eine Safety-Alignment-Aufgabe.
3. Metriken: Die Downstream-Performance wird mit mehreren Metriken bewertet, einschließlich Greedy (Pass@1), Pass@16, Majority Vote (Maj@16), Reward Model Scoring (RM@16) und Correct Ratio.
4. Mechanistische Analyse: Um zu verstehen, warum Weight Decay die Plastizität beeinflusst, analysieren die Autoren:
   • Linear Probing: Messung der linearen Separabilität interner Repräsentationen (Sentiment- und Topic-Klassifikation).
   • Attention Rank: Analyse des Pseudo-Rangs von Query-Key- ($W_{QK}$) und Value-Projection-Matrizen ($W_{VP}$).
   • Overfitting: Berechnung der Differenz zwischen Trainings- und Validierungsverlust (Train-Val Gap).

Zentrale Beiträge

1. Weight Decay als Treiber der Plastizität: Die Arbeit zeigt, dass Weight Decay ein entscheidender Faktor bei der Gestaltung der Modell-Plastizität ist. Im Gegensatz zu Standard-Defaults führen größere Weight-Decay-Werte während des Pretrainings konsistent zu einer besseren Downstream-Performance nach der Feinabstimmung.
2. Entkopplung der Zielsetzungen: Die Studie liefert Belege dafür, dass der Hyperparameterwert, der den Pretraining-Validierungsverlust minimiert, nicht zwangsläufig der Wert ist, der die Downstream-Performance maximiert. In mehreren Regimen erzielen Modelle mit höherem Pretraining-Loss (aufgrund von größerem Weight Decay) bessere Ergebnisse nach der Feinabstimmung als Modelle mit niedrigerem Pretraining-Loss.
3. Mechanistische Einblicke: Die Autoren identifizieren drei spezifische Mechanismen, durch die Weight Decay die Plastizität erhöht:
   • Es fördert linear separierbare Repräsentationen, wodurch interne Merkmale für Downstream-Aufgaben zugänglicher werden.
   • Es regularisiert Attention-Matrizen, indem es insbesondere den Rang des Query-Key-Produkts ($W_{QK}$) reduziert, was zu kompakteren und robusteren Strukturen führt.
   • Es reduziert Overfitting auf Pretraining-Daten, indem es effektiv das „Vergessen“ spezifischer Trainings-Rauschanteile fördert und so die Kapazität des Modells zur Aufnahme neuer Informationen bewahrt.

Zentrale Ergebnisse

• Verschiebung des optimalen Weight Decay: Für das rechenoptimale Training (20 TPP) ist der optimale Weight Decay für die Downstream-Performance signifikant größer als der Standardwert von 0,1. Beispielsweise wurden optimale Werte bei 0,5–1,0 für Llama-2- und OLMo-2-Modelle in diesem Regime gefunden. Im übertrainierten Regime (140 TPP) liegt der optimale Wert für die Downstream-Performance (0,3) ebenfalls höher als der Wert, der den Pretraining-Loss minimiert (0,1).
• Kontraintuitive Trade-offs: Die Arbeit hebt einen Trade-off hervor, bei dem „schlechtere“ Basismodelle (höherer Pretraining-Loss) nach der Feinabstimmung zu „besseren“ Endmodellen werden können. Beispielsweise erzielte ein OLMo-2-1B-140x, das mit $\lambda=0,3$ oder $\lambda=1,0$ vortrainiert wurde, einen höheren Pretraining-Loss als das $\lambda=0,1$ Baseline-Modell, erreichte aber eine deutlich bessere Downstream-Genauigkeit.
• Robustheit: Diese Erkenntnisse halten über verschiedene Modellfamilien, Größen und auch bei Variation anderer Hyperparameter (Learning Rate, Batch Size) während des Pretrainings und der Feinabstimmung stand.
• Mechanistische Korrelationen:
  • Höheres Weight Decay korreliert mit höherer Linear-Probing-Genauigkeit, was auf strukturiertere Repräsentationen hindeutet.
  • Höheres Weight Decay korreliert mit einem niedrigeren Pseudo-Rang in den Attention-Matrizen (insbesondere $W_{QK}$), was auf eine Bewegung hin zu Low-Rank-, robusten Hypothesen hindeutet.
  • Höheres Weight Decay korreliert mit einem kleineren Train-Val Gap, was auf reduziertes Overfitting hindeutet.

Bedeutung und Behauptungen
Die Arbeit argumentet, dass das Feld die exklusive Abhängigkeit vom Pretraining-Validierungsverlust als einzigem Metrik für die Hyperparameter-Optimierung überdenken muss. Die Autoren behaupten, dass die Modell-Plastizität eine eigenständige und entscheidende Eigenschaft ist, die oft geopfert wird, wenn rein auf den Pretraining-Loss optimiert wird.

Die Bedeutung dieser Arbeit liegt in:

1. Neubewertung von Defaults: Sie legt nahe, dass Standard-Hyperparameter-Wahlen (z. B. $\lambda=0,1$) suboptimal für den gesamten Lebenszyklus der LLM-Entwicklung sein können, insbesondere wenn die Downstream-Adaptivität das ultimative Ziel ist.
2. Limitierungen von Cross-Entropy: Sie verdeutlicht, dass der Cross-Entropy-Loss ein unvollkommener Stellvertreter für die Downstream-Fähigkeiten ist, da Modelle mit ähnlichen Verlusten eine völlig unterschiedliche Plastizität aufweisen können.
3. Multifunktionale Rolle von Hyperparametern: Sie unterstreicht, dass ein einzelner Hyperparameter wie Weight Decay eine komplexe, vielschichtige Rolle bei der Gestaltung des Modellverhaltens spielt, die nicht nur die Generalisierung, sondern auch die strukturellen Eigenschaften (Linearität, Rang) beeinflusst, die zukünftiges Lernen erleichtern.

Die Autoren bleiben bescheiden und merken an, dass ihre Ergebnisse zwar über die getesteten Skalen hinweg robust sind, die spezifischen optimalen Werte jedoch in extrem groß angelegten oder stark übertrainierten Szenarien jenseits ihres aktuellen experimentellen Rahmens variieren können. Sie fordern künftige Arbeiten dazu auf, den Stabilitäts-Plastizitäts-Trade-off sowie die Anwendbarkeit dieser Erkenntnisse auf andere Modalitäten von Foundation-Modellen weiter zu untersuchen.