Weight Decay Improves Language Model Plasticity

Technische Samenvatting: Weight Decay Verbetert de Plasticiteit van Taalmodellen

Probleemstelling
De huidige praktijken voor het trainen van Large Language Models (LLMs) koppelen pretraining doorgaans los van post-training (fine-tuning, alignment). Hyperparameteroptimalisatie en schalingwetten tijdens pretraining worden voornamelijk geëvalueerd op basis van de validatieverlies (cross-entropy) van het basemodel, onder de aanname dat een lager pretraining-verlies correleert met een krachtiger downstream-model. Deze aanpak negeert echter modelplasticiteit—het vermogen van een getraind model om zich succesvol aan te passen aan nieuwe data en taken bij verdere training zonder herinitialisatie. Het onderzoek onderzoekt of het optimaliseren van pretraining-hyperparameters uitsluitend voor pretraining-prestaties onbedoeld de adaptiviteit van het model tijdens daaropvolgende fine-tuning in gevaar brengt. Specifiek richten de auteurs zich op weight decay, een standaard regularisatiehyperparameter, om te bepalen of de optimale waarde voor het minimaliseren van het pretraining-verlies verschilt van de optimale waarde voor het maximaliseren van downstream-adaptiviteit.

Methodologie
De auteurs voeren een systematische empirische studie uit over twee modelfamilies (Llama-2 en OLMo-2) en meerdere modelgroottes (tot 4B parameters). Het experimentele ontwerp varieert de weight decay-hyperparameter ($\lambda$) tijdens pretraining, terwijl andere factoren constant worden gehouden.

1. Trainingsregimes: Modellen worden getraind in twee verschillende regimes:
   • Compute-optimal: 20 tokens-per-parameter (TPP).
   • Overtrained: 140 TPP.
2. Fine-tuning Evaluatie: Pretrained modellen worden gefine-tuned op een diverse reeks downstream-taken, waaronder:
   • Zes Chain-of-Thought (CoT) redeneerdatasets (bijv. MetaMathQA, MMLUProCoT).
   • Vijf taalbegrip- en common sense reasoning-benchmarks (bijv. HellaSwag, WinoGrande).
   • Eén veiligheid-alignment taak.
3. Metrieken: Downstream-prestaties worden geëvalueerd met meerdere metrieken, waaronder Greedy (Pass@1), Pass@16, Majority Vote (Maj@16), Reward Model scoring (RM@16) en Correct Ratio.
4. Mechanistische Analyse: Om te begrijpen waarom weight decay de plasticiteit beïnvloedt, analyseren de auteurs:
   • Linear Probing: Het meten van de lineaire scheidbaarheid van interne representaties (sentiment en topic classificatie).
   • Attention Rank: Het analyseren van de pseudo-rang van query-key ($W_{QK}$) en value-projection ($W_{VP}$) matrices.
   • Overfitting: Het berekenen van het verschil tussen de trainingsverlies en de validatieverlies (Train-Val Gap).

Belangrijkste Bijdragen

1. Weight Decay als Driver voor Plasticiteit: Het paper demonstreert dat weight decay een cruciale factor is in het vormgeven van modelplasticiteit. In tegenstelling tot standaard defaults leiden grotere weight decay-waarden tijdens pretraining consistent tot betere downstream-prestaties na fine-tuning.
2. Ontkoppeling van Doelstellingen: De studie levert bewijs dat de hyperparameterwaarde die het pretraining-validatieverlies minimaliseert, niet noodzakelijkerwijs de waarde is die de downstream-prestaties maximaliseert. In verschillende regimes presteren modellen met een hoger pretraining-verlies (door grotere weight decay) beter na fine-tuning dan modellen met een lager pretraining-verlies.
3. Mechanistische Inzichten: De auteurs identificeren drie specifieke mechanismen waardoor weight decay de plasticiteit verbetert:
   • Het stimuleert lineair scheidbare representaties, waardoor interne kenmerken toegankelijker zijn voor downstream-taken.
   • Het reguleert attention-matrices, specifiek door de rang van de query-key product ($W_{QK}$) te verlagen, wat leidt tot meer gecomprimeerde en robuuste structuren.
   • Het vermindert overfitting op pretraining-data, waardoor het effectief specifieke trainingsruis "vergeet" en de capaciteit van het model om nieuwe informatie te leren behoudt.

Belangrijkste Resultaten

• Verschuiving van Optimale Weight Decay: Voor compute-optimal (20 TPP) training is de optimale weight decay voor downstream-prestaties aanzienlijk groter dan de standaard default van 0.1. Bijvoorbeeld, optimale waarden werden gevonden bij 0.5–1.0 voor Llama-2 en OLMo-2 modellen in dit regime. In het overtrained (140 TPP) regime blijft de optimale waarde voor downstream-prestaties (0.3) hoger dan de waarde die het pretraining-verlies minimaliseert (0.1).
• Contraintuitieve Trade-offs: Het paper benadrukt een trade-off waarbij "slechtere" basemodellen (hoger pretraining-verlies) "betere" uiteindelijke modellen kunnen worden na fine-tuning. Zo hadden een OLMo-2-1B-140x model, getraind met $\lambda=0.3$ of $1.0$, een hoger pretraining-verlies dan de $\lambda=0.1$ baseline, maar behaalde het merkbaar betere downstream-nauwkeurigheid.
• Robuustheid: Deze bevindingen houden stand over verschillende modelfamilies, groottes, en zelfs wanneer andere hyperparameters (learning rate, batch size) tijdens zowel pretraining als fine-tuning worden gevarieerd.
• Mechanistische Correlaties:
  • Hogere weight decay correleert met hogere linear probing nauwkeurigheid, wat wijst op meer gestructureerde representaties.
  • Hogere weight decay correleert met een lagere pseudo-rang in attention-matrices (met name $W_{QK}$), wat suggereert dat er een beweging plaatsvindt naar low-rank, robuuste hypothesen.
  • Hogere weight decay correleert met een kleinere Train-Val Gap, wat duidt op verminderde overfitting.

Betekenis en Claims
Het paper betoogt dat het vakgebied de exclusieve afhankelijkheid van pretraining-validatieverlies als enige metriek voor hyperparameteroptimalisatie moet heroverwegen. De auteurs claimen dat modelplasticiteit een onderscheidende en cruciale eigenschap is die vaak wordt opgeofferd wanneer er puur wordt geoptimaliseerd voor pretraining-verlies.

De betekenis van dit werk ligt in:

1. Herwaardering van Defaults: Het suggereert dat standaard hyperparameterkeuzes (bijv. $\lambda=0.1$) mogelijk suboptimaal zijn voor de volledige levenscyclus van LLM-ontwikkeling, met name wanneer downstream-adaptiviteit het uiteindelijke doel is.
2. Beperkingen van Cross-Entropy: Het benadrukt dat cross-entropy verlies een imperfecte proxy is voor downstream-capaciteit, aangezien modellen met vergelijkbare verliezen fundamenteel verschillende plasticiteit kunnen vertonen.
3. Veelzijdige Rol van Hyperparameters: Het onderstreept dat een enkele hyperparameter zoals weight decay een complexe, veelzijdige rol speelt in het vormen van modelgedrag, wat niet alleen generalisatie beïnvloedt, maar ook de structurele eigenschappen (lineariteit, rang) die toekomstig leren faciliteren.

De auteurs blijven bescheiden en merken op dat hoewel hun resultaten robuust zijn over de geteste schalen, de specifieke optimale waarden kunnen verschuiven in extreem grootschalige of zwaar overtrained scenario's buiten hun huidige experimentele reikwijdte. Ze pleiten voor toekomstig werk om de stability-plasticity trade-off en de toepasbaarheid van deze bevindingen op andere fundamentele modelmodaliteiten verder te onderzoeken.