Weight Decay Improves Language Model Plasticity

Sintesi Tecnica: Il Weight Decay Migliora la Plasticità dei Modelli Linguistici

Problematica
Le pratiche attuali per l'addestramento dei Large Language Models (LLM) tipicamente scindono il pre-addestramento dal post-addestramento (fine-tuning, allineamento). L'ottimizzazione degli iperparametri e le leggi di scala durante il pre-addestramento vengono valutate prevalentemente sulla base della perdita di validazione del modello base (cross-entropy), partendo dal presupposto che una perdita di pre-addestramento inferiore correli con un modello più capace nelle fasi successive. Tuttavia, questo approccio trascura la plasticità del modello — ovvero la capacità di un modello addestrato di adattarsi con successo a nuovi dati e compiti durante l'ulteriore addestramento senza riinizializzazione. Il paper indaga se l'ottimizzazione degli iperparametri di pre-addestramento esclusivamente per la performance di pre-addestramento comprometta involontariamente l'adattabilità del modello durante il successivo fine-tuning. Nello specifico, gli autori si concentrano sul weight decay, un comune iperparametro di regolarizzazione, per determinare se il suo valore ottimale per minimizzare la perdita di pre-addestramento differisca dal suo valore ottimale per massimizzare l'adattabilità a valle.

Metodologia
Gli autori conducono uno studio empirico sistematico su due famiglie di modelli (Llama-2 e OLMo-2) e molteplici dimensioni di modello (fino a 4 miliardi di parametri). Il design sperimentale varia l'iperparametro di weight decay ($\lambda$) durante il pre-addestramento, controllando al contempo altri fattori.

1. Regimi di Addestramento: I modelli sono addestrati in due regimi distinti:
   • Compute-optimal: 20 token-per-parametro (TPP).
   • Overtrained (Sovra-addestrato): 140 TPP.
2. Valutazione del Fine-tuning: I modelli pre-addestrati vengono sottoposti a fine-tuning su una suite diversificata di compiti downstream, tra cui:
   • Sei dataset di ragionamento Chain-of-Thought (CoT) (es. MetaMathQA, MMLUProCoT).
   • Cinque benchmark di comprensione del linguaggio e ragionamento del senso comune (es. HellaSwag, WinoGrande).
   • Un compito di allineamento della sicurezza.
3. Metriche: La performance downstream è valutata utilizzando molteplici metriche, inclusi Greedy (Pass@1), Pass@16, Majority Vote (Maj@16), Reward Model scoring (RM@16) e Correct Ratio.
4. Analisi Meccanicistica: Per comprendere perché il weight decay influenzi la plasticità, gli autori analizzano:
   • Linear Probing: Misurazione della separabilità lineare delle rappresentazioni interne (classificazione del sentiment e del topic).
   • Attention Rank: Analisi del pseudo-rank delle matrici di proiezione query-key ($W_{QK}$) e value-projection ($W_{VP}$).
   • Overfitting: Calcolo del divario tra la perdita di addestramento e la perdita di validazione (Train-Val Gap).

Contributi Chiave

1. Il Weight Decay come Motore di Plasticità: Il paper dimostra che il weight decay è un fattore critico nel plasmare la plasticità del modello. Contrariamente ai default standard, valori di weight decay più elevati durante il pre-addestramento portano costantemente a una migliore performance downstream dopo il fine-tuning.
2. Disaccoppiamento degli Obiettivi: Lo studio fornisce prove che il valore dell'iperparametro che minimizza la perdita di validazione del pre-addestramento non è necessariamente il valore che massimizza la performance downstream. In diversi regimi, modelli con una perdita di pre-addestramento più alta (dovuta a un weight decay maggiore) superano i modelli con una perdita di pre-addestramento inferiore dopo il fine-tuning.
3. Intuizioni Meccanicistiche: Gli autori identificano tre meccanismi specifici attraverso i quali il weight decay migliora la plasticità:
   • Incoraggia rappresentazioni linearmente separabili, rendendo le caratteristiche interne più accessibili per i compiti downstream.
   • Regolarizza le matrici di attenzione, in particolare riducendo il rango del prodotto query-key ($W_{QK}$), portando a strutture più compresse e robuste.
   • Riduce l'overfitting sui dati di pre-addestramento, "dimenticando" efficacementamente il rumore specifico dell'addestramento e preservando la capacità del modello di apprendere nuove informazioni.

Risultati Chiave

• Spostamento del Weight Decay Ottimale: Per l'addestramento compute-optimal (20 TPP), il weight decay ottimale per la performance downstream è significativamente maggiore dello standard di default di 0.1. Ad esempio, i valori ottimali sono stati individuati a 0.5–1.0 per i modelli Llama-2 e OLMo-2 in questo regime. Nel regime overtrained (140 TPP), il valore ottimale per la performance downstream (0.3) rimane superiore al valore che minimizza la perdita di pre-addestramento (0.1).
• Trade-off Controintuitivi: Il paper evidenzia un trade-off in cui modelli base "peggiori" (perdita di pre-addestramento più alta) possono diventare modelli finali "migliori" dopo il fine-tuning. Ad esempio, un modello OLMo-2-1B-140x pre-addestrato con $\lambda=0.3$ o $1.0$ aveva una perdita di pre-addestramento superiore rispetto al baseline $\lambda=0.1$, ma ha ottenuto una accuratezza downstream sensibilmente migliore.
• Robustezza: Queste scoperte sono valide attraverso diverse famiglie di modelli, dimensioni e anche quando altri iperparametri (learning rate, batch size) vengono variati sia durante il pre-addestramento che durante il fine-tuning.
• Correlazioni Meccanicistiche:
  • Un weight decay più elevato correla con una maggiore accuratezza del linear probing, indicando rappresentazioni più strutturate.
  • Un weight decay più elevato correla con un pseudo-rank inferiore nelle matrici di attenzione (particolarmente in $W_{QK}$), suggerendo un movimento verso ipotesi a basso rango e robuste.
  • Un weight decay più elevato correla con un Train-Val Gap minore, indicando un ridotto overfitting.

Significatività e Rivendicazioni
Il paper sostiene che il settore debba riconsiderare l'uso esclusivo della perdita di validazione del pre-addestramento come unico metrica per l'ottimizzazione degli iperparametri. Gli autori affermano che la plasticità del modello è una proprietà distinta e cruciale che viene spesso sacrificata quando si ottimizza puramente per la perdita di pre-addestramento.

La significatività di questo lavoro risiede nel fatto che:

1. Rivalutazione dei Default: Suggerisce che le scelte standard degli iperparametri (es. $\lambda=0.1$) potrebbero essere subottimali per l'intero ciclo di vita dello sviluppo di un LLM, specialmente quando l'adattabilità downstream è l'obiettivo finale.
2. Limitazioni della Cross-Entropy: Evidenzia come la perdita di cross-entropy sia un proxy imperfetto per la capacità downstream, poiché modelli con perdite simili possono esibire una plasticità vastamente differente.
3. Ruolo Multifaccettato degli Iperparametri: Sottolinea come un singolo iperparametro come il weight decay svolga un ruolo complesso e multifaccettato nel plasmare il comportamento del modello, influenzando non solo la generalizzazione ma anche le proprietà strutturali (linearità, rango) che facilitano l'apprendimento futuro.

Gli autori mantengono un tono modesto, osservando che, sebbene i loro risultati siano robusti attraverso le scale testate, i valori ottimali specifici potrebbero spostarsi in scenari di scala estremamente grande o pesantemente overtrained oltre l'attuale ambito sperimentale. Chiedono che il lavoro futuro approfondisca ulteriormente il trade-off tra stabilità e plasticità e l'applicabilità di queste scoperte ad altre modalità di foundation model.