(How) Learning Rates Regulate Catastrophic Overtraining

Probleemstelling: Catastrophic Overtraining

Het paper adresseert een kritiek probleem in de post-training van Large Language Models (LLMs): catastrophic overtraining. Waar "catastrophic forgetting" (het vergeten van eerder geleerde vaardigheden tijdens finetuning) al bekend is, hebben Springer et al. (2025) ontdekt dat langere pretraining juist leidt tot meer vergeten tijdens de Supervised Finetuning (SFT). Dit is tegenintuïtief: een beter getraind base-model (door meer tokens) resulteert in een slechter finetuned model voor standaard-taken.

De auteurs stellen dat de bestaande hypothese (dat dit te wijten is aan progressieve gevoeligheid voor ruis of capaciteitsbeperkingen) onvolledig is. Ze hypotheseren dat optimatiedynamiek, en specifiek de rol van de leerfrequentie (Learning Rate - LR), de drijvende kracht is achter dit fenomeen.

Methodologie

De auteurs analyseren het probleem vanuit een optimalisatieperspectief, gebruikmakend van een combinatie van empirische experimenten en theoretische inzichten:

1. Experimenteel Opzet:

   • Ze gebruiken diverse modellen (OLMo-1/2, Hubble, Gemma-3, SmolLM) met 1B tot 3B parameters.
   • Ze finetunen deze modellen op instructie-datasets (Anthropic-HH, Tülu 3) met verschillende constante learning rates (van $10^{-6}$ tot $10^{-4}$).
   • Metrieken:
     • Train Loss: Hoe goed het model de SFT-data past.
     • OOD Score: Prestaties op out-of-distribution taken (algemene vaardigheden) om vergeten te meten.
     • Mean Principal Angle (MPA): Een maat voor de "feature drift". Dit meet de hoek tussen de subruimtes van de residual stream-representaties van het base-model en het finetuned model. Een grotere hoek betekent meer afwijking van de oorspronkelijke features.
     • Sharpness: Geschat via de gevoeligheid van de output voor Gaussische ruis in de parameters (proxy voor de eigenwaarden van de Hessian).

2. Theoretische Analyse:

   • Ze analyseren het verlieslandschap (loss landscape) en gebruiken een vereenvoudigd diagonaal netwerk model (sparse regression) om de interactie tussen initialisatie, learning rate en feature learning te modelleren.
   • Ze koppelen hun bevindingen aan het concept van Progressive Sharpening en de Edge of Stability (EoS).

Belangrijkste Bijdragen en Resultaten

1. Learning Rate als Implicit Regularisatie

De auteurs tonen aan dat voor een vaste trainingsloss (d.w.z. wanneer het model even goed de SFT-taak heeft geleerd), de keuze van de learning rate leidt tot kwalitatief verschillende modellen:

• Hoge Learning Rates: Leiden tot een grotere feature drift (hogere MPA). Het model verlaat de regio van het parameter-ruimte waar de pre-trained features behouden blijven, wat resulteert in ernstigere catastrophic forgetting.
• Lage Learning Rates: Opereren in een "lazy regime". Ze behouden de pre-trained features beter (kleinere MPA) en leiden tot minder vergeten, zelfs bij dezelfde trainingsloss.
• Conclusie: Een lage LR werkt als een implicit regularisatie die de stabiliteit van de base-model features bevordert.

2. De Rol van Sharpness (Scherpte)

Het paper legt een causaal verband tussen de sharpness van het base-model en de mate van vergeten:

• Sharpness en Feature Drift: Modellen die in "scherpere" regio's van het verlieslandschap worden geïnitieerd, vertonen bij finetuning een snellere toename in feature drift (MPA) per stap.
• Diagonaal Netwerk Bewijs: In het vereenvoudigde model wordt getoond dat grote learning rates feature learning stimuleren (rich regime), terwijl kleine learning rates features behouden (lazy regime). Modellen met hogere initiële sharpness gedragen zich alsof ze een effectief hogere learning rate hebben, wat leidt tot meer vergeten.

3. Oorzakelijk Verband: Pretraining LR Decay en Overtraining

Dit is de kernbijdrage voor het fenomeen "catastrophic overtraining":

• Tijdens pretraining wordt de learning rate vaak geleidelijk verlaagd (decay). Volgens de Edge of Stability theorie zorgt een dalende LR ervoor dat de sharpness van het model toeneemt (de drempel $2/\eta$ stijgt, en het model oscilleert rond een scherpere top).
• Het Mechanisme:
  1. Lange pretraining met LR-decay $\rightarrow$ Het base-model wordt scherper.
  2. Een scherper base-model is gevoeliger voor de finetuning-stappen.
  3. Zelfs bij een vaste finetuning-LR, veroorzaakt een scherper base-model meer feature drift en dus meer vergeten.
• Empirisch Bewijs: De auteurs tonen aan dat checkpoints genomen later in de pretraining (waar de LR is verlaagd en de sharpness is gestegen) leiden tot een drastischere daling in OOD-prestaties na finetuning. Dit patroon is vooral zichtbaar bij modellen die een Warmup-Stable-Decay (WSD) schedule gebruiken; het vergeten neemt plotseling toe precies op het moment dat de LR-decay begint.

Significantie en Implicaties

De bevindingen bieden een nieuw inzicht in de "stability-plasticity trade-off" bij LLMs en hebben directe praktische implicaties:

1. Herinterpretatie van Scaling: Het probleem van overtraining is niet alleen een kwestie van "te veel tokens", maar een gevolg van de optimalisatiedynamiek (LR-decay) die de modelsharpness verhoogt.
2. Praktische Aanbevelingen:
   • Finetuning: Gebruik de kleinste mogelijke learning rate die nog acceptabele in-distribution prestaties oplevert om catastrophic forgetting te minimaliseren.
   • Pretraining: Vermijd het verlagen van de learning rate tijdens de pretraining als het doel is om een robuust base-model te krijgen voor SFT. In plaats daarvan zou de LR-annealing (verlaging) beter kunnen plaatsvinden tijdens de post-training fase, of men moet rekening houden met de toegenomen sharpness bij het kiezen van de finetuning-LR.
3. Theoretische Inzicht: Het paper verbindt losse concepten zoals implicit regularization, loss landscape sharpness, en catastrophic forgetting tot een coherent verhaal over hoe optimatiedynamiek de behouding van kennis in neurale netwerken bepaalt.

Samenvattend stelt het paper dat catastrophic overtraining wordt veroorzaakt door de toename in sharpness van het base-model als gevolg van learning rate decay tijdens pretraining, en dat dit effect kan worden tegengegaan door het gebruik van lage learning rates tijdens de supervized finetuning.