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

Titel: Pre-Training van LLM's zonder Learning Rate Decay verbetert Supervised Fine-Tuning

1. Het Probleem

In de huidige praktijk voor het voorvertrouwen (pre-training) van Large Language Models (LLM's) is het gebruik van learning rate (LR) schedulers met decay (zoals Cosine decay, Linear decay naar nul, of Warmup-Stable-Decay) de standaard. Deze methoden worden gebruikt om de loss tijdens het voorvertrouwen te minimaliseren.

Het paper identificeert echter een fundamenteel probleem:

• Suboptimale downstream prestaties: Hoewel decay-based schedulers vaak leiden tot de beste prestaties op de pre-training metrics (zoals validatie loss), garandeert dit niet dat het model goed presteert na Supervised Fine-Tuning (SFT).
• Gebrek aan onderzoek: Er is weinig onderzoek gedaan naar hoe de keuze van de LR-scheduler tijdens het voorvertrouwen de aanpasbaarheid (adaptability) van het model beïnvloedt voor downstream taken.
• Vaste aanname: De gemeenschap gaat er vaak ten onrechte vanuit dat een model dat beter presteert tijdens het voorvertrouwen, ook beter zal zijn na fijnafstemming. Recent onderzoek suggereert echter dat dit niet altijd het geval is.

2. Methodologie

De auteurs onderzoeken een alternatieve aanpak genaamd Warmup-Stable-Only (WSO). In plaats van de learning rate aan het einde van het trainingstraject te laten dalen (decay), houdt WSO de learning rate constant na de warmup-fase tot het einde van het voorvertrouwen.

Experimenteel Opzet:

• Modellen: Experimenten zijn uitgevoerd op twee schalen: 1B en 8B parameter modellen (gebaseerd op de Llama 3 architectuur).
• Schedulers: Er werden vier schedulers vergeleken:
  1. WSO: Warmup gevolgd door een constante learning rate (geen decay).
  2. WSD (Warmup-Stable-Decay): Constante LR gevolgd door een korte decay-fase aan het einde.
  3. Cosine en Linear: Traditionele decay-methoden.
• Variatie in $\alpha_{pre}$: De auteurs varieerden de minimale LR-factor ($\alpha_{pre}$) van 0.0 (decay naar nul) tot 1.0 (geen decay, dus WSO).
• Training Regimes: De studie omvatte drie scenario's:
  1. Twee-staps training (Pre-training -> SFT).
  2. Drie-staps training (Pre-training -> Mid-training -> SFT), waarbij Mid-training een tussenstap is voor domeinuitbreiding.
  3. Over-training (Over-training): Training op 2 Triljoen tokens (ongeveer 100x de Chinchilla-optimaliteit) om de schaalbaarheid te testen.
• Analyse van het Loss Landscape: Om de onderliggende oorzaken te begrijpen, analyseerden de auteurs de sharpness (scherpte) van de minima in het loss landscape. Sharpness wordt gemeten via de trace van de Hessian-matrix (geschat met Hutchinson's estimator). Een "vlakkere" minima (lagere sharpness) wordt geassocieerd met betere generalisatie en aanpasbaarheid.

3. Belangrijkste Bijdragen

1. Empirisch Bewijs voor WSO: Het paper biedt systematisch bewijs dat WSO consistent beter presteert op downstream taken na SFT dan decay-based schedulers, ondanks dat decay-based schedulers vaak betere pre-training metrics laten zien.
2. Generalisatie over Training Regimes: De bevindingen zijn robuust en gelden voor standaard pre-training, mid-training, en zelfs in regimes van over-training (triljoenen tokens).
3. Theoretisch Inzicht: Door analyse van het loss landscape wordt aangetoond dat decay-based schedulers modellen leiden naar "scherpere" minima (hoge sharpness), terwijl WSO "vlakkere" minima behoudt. Vlakkere minima bevorderen de aanpasbaarheid tijdens SFT.

4. Resultaten

De resultaten tonen een duidelijke omkering (inversion) in prestaties tussen de pre-training fase en de SFT-fase:

• Pre-training Metrics: Decay-based schedulers (met name Linear en WSD met $\alpha_{pre} = 0.0$) behalen de laagste validatie loss en vaak de beste zero-shot scores op pre-training benchmarks. WSO presteert hier vaak iets slechter.
• SFT Prestaties (Downstream): Na Supervised Fine-Tuning keert de rangorde om. Modellen getraind met WSO behalen significant betere scores op taken zoals instructie-opvolging (AlpacaEval), waarheidsgetrouwheid (TruthfulQA) en multi-task begrip (MMLU).
  • Bij de 1B en 8B modellen presteerde WSO consistent beter dan alle decay-varianten.
  • Zelfs bij over-training (2T tokens) behield WSO zijn superioriteit na SFT.
• Sharpness Analyse:
  • Modellen getraind met decay convergeren naar minima met een 2-3x hogere sharpness dan WSO-modellen.
  • Er is een sterke negatieve correlatie ($r = -0.709$) tussen de sharpness van het pre-trained model en de SFT-prestaties: hoe scherper het minimum, hoe slechter de aanpasbaarheid.
  • WSO behoudt een vlakkere loss landscape, wat betekent dat het model minder gevoelig is voor parameterupdates tijdens SFT en zijn pre-trained kennis beter kan behouden terwijl het nieuwe taken leert.

5. Betekenis en Conclusie

Dit paper daagt de gevestigde orde van LLM-training uit. Het toont aan dat het optimaliseren van pre-training metrics (zoals het minimaliseren van loss via decay) niet noodzakelijk leidt tot het beste eindresultaat voor praktische toepassingen.

• Praktische Implicatie: Ontwikkelaars en onderzoekers zouden moeten overwegen om WSO (geen learning rate decay) te gebruiken tijdens het voorvertrouwen. Dit is een simpele wijziging die leidt tot modellen met superieure aanpasbaarheid voor downstream taken zonder extra rekenkosten.
• Strategisch Advies: Het paper adviseert om modellen die met WSO zijn getraind te publiceren, omdat deze beter geschikt zijn voor verdere fijnafstemming en adaptatie in diverse scenario's.
• Theoretische Bijdrage: Het versterkt het inzicht dat de geometrie van het loss landscape (vlakkere minima) cruciaal is voor de transfer learning capaciteit van LLM's, en dat decay-based schedulers deze geometrie ongunstig beïnvloeden.

Kortom: Stop met learning rate decay tijdens het voorvertrouwen als je doel is een model te krijgen dat uitstekend presteert na fijnafstemming.