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

Titolo

Pre-training di LLM senza decadimento del Learning Rate (LR) migliora il Supervised Fine-Tuning (SFT)

1. Il Problema

Nell'addestramento su larga scala dei Large Language Models (LLM), la pianificazione del Learning Rate (LR) è un aspetto critico. Attualmente, la pratica standard prevede l'uso di schedulers basati sul decadimento (come Cosine decay, Linear decay o Warmup-Stable-Decay - WSD), che riducono progressivamente il LR verso la fine della fase di pre-training.

• Obiettivo attuale: Questi schedulers sono ottimizzati per minimizzare la loss di pre-training e massimizzare le prestazioni sui benchmark zero-shot immediati.
• Il limite: La ricerca mostra che un modello con ottime prestazioni di pre-training non garantisce necessariamente prestazioni superiori dopo il Supervised Fine-Tuning (SFT), che è la fase cruciale per l'adattamento a compiti specifici e il seguire le istruzioni umane.
• Domanda di ricerca: È possibile che il decadimento del LR, sebbene utile per le metriche di pre-training, comprometta l'adattabilità del modello durante le fasi successive di post-training?

2. Metodologia

Gli autori hanno condotto un'indagine empirica sistematica confrontando schedulers di LR tradizionali con una variante proposta chiamata Warmup-Stable-Only (WSO).

• Approccio WSO: Questo scheduler rimuove completamente la fase di decadimento. Dopo una fase di warmup, il LR rimane costante fino alla fine del pre-training (e potenzialmente del mid-training), mantenendo un valore pari al massimo ($\eta_{max}$).
• Confronto: Sono stati confrontati quattro schedulers:
  1. WSO: Warmup + Stabile (nessun decadimento).
  2. WSD: Warmup + Stabile + Decadimento (alla fine).
  3. Cosine: Decadimento Cosine.
  4. Linear: Decadimento Lineare.
• Configurazioni Sperimentali:
  • Modelli: Architetture simili a Llama 3 con 1B e 8B parametri.
  • Fasi di Addestramento:
    1. Pre-training: Addestramento su dataset come FineWeb-Edu e OLMo-mix.
    2. Mid-training: Una fase intermedia di addestramento su dati di alta qualità (es. dolmino-mix) per espandere il dominio o il contesto.
    3. Over-training: Addestramento su un volume di token molto superiore all'ottimo di Chinchilla (fino a 2T di token).
  • Post-training: Tutti i modelli sono stati sottoposti a SFT (Supervised Fine-Tuning) utilizzando il mix Tulu-3, con una ricerca esaustiva del Learning Rate ottimale per ogni modello pre-addestrato.
• Analisi Teorica: È stata effettuata un'analisi della geometria del loss landscape misurando la sharpness (nitidezza) dei minimi utilizzando l'estimatore di traccia di Hutchinson per l'ematrice Hessiana.

3. Risultati Chiave

A. Inversione delle Prestazioni (PT vs SFT)

I risultati mostrano una chiara inversione tra le prestazioni di pre-training e quelle post-SFT:

• Pre-training: Gli schedulers con decadimento (specialmente Linear e WSD con $\alpha_{pre}=0.0$) ottengono i migliori risultati in termini di validation loss e punteggi zero-shot sui benchmark standard.
• Post-SFT: I modelli addestrati con WSO (nessun decadimento) superano costantemente tutti gli altri scheduler nelle prestazioni dopo il fine-tuning, sia per modelli da 1B che da 8B parametri.
  • Esempio: Nel caso da 8B parametri, WSO ha ottenuto un miglioramento di +1.1 nel punteggio medio SFT rispetto al miglior scheduler con decadimento, nonostante avesse prestazioni inferiori durante il pre-training.

B. Generalità in Mid-training e Over-training

Il vantaggio di WSO si è dimostrato robusto in scenari complessi:

• Mid-training: Anche quando applicato in una pipeline a tre stadi (Pre-training -> Mid-training -> SFT), l'uso di WSO in entrambe le fasi (o anche solo nel mid-training) ha portato a migliori risultati finali rispetto all'uso di decadimento.
• Over-training: In esperimenti con 2 trilioni di token (100x l'ottimo di Chinchilla), WSO ha mantenuto il suo vantaggio, dimostrando che il beneficio non è limitato a specifici regimi di calcolo.

C. Analisi del Loss Landscape

L'analisi della sharpness rivela il motivo teorico del successo di WSO:

• Decadimento: Gli schedulers con decadimento tendono a portare il modello in minimi acuti (sharp minima), caratterizzati da una alta curvatura del loss landscape.
• WSO: Mantiene il modello in minimi piatti (flat minima).
• Correlazione: È stata osservata una forte correlazione negativa (Pearson $r = -0.709$) tra la sharpness misurata alla fine del pre-training e le prestazioni finali dopo SFT. Minimi più piatti implicano una maggiore adattabilità ai nuovi dati durante il fine-tuning, poiché le prestazioni del modello sono più stabili di fronte agli aggiornamenti dei parametri.

4. Contributi Principali

1. Dimostrazione Sistematica: Fornisce prove empiriche che WSO supera gli scheduler con decadimento nei compiti downstream dopo SFT, su modelli di diverse scale (1B, 8B) e in diverse configurazioni di addestramento.
2. Validazione in Regimi Moderni: Estende la validità di questo approccio ai paradigmi moderni di addestramento, inclusi il mid-training e l'over-training, sfidando la convenzione di ridurre il LR alla fine.
3. Spiegazione Geometrica: Collega esplicitamente la scelta dello scheduler LR alla geometria del loss landscape, dimostrando che WSO preserva minimi piatti che favoriscono l'adattabilità, mentre il decadimento porta a minimi acuti che limitano la flessibilità futura.

5. Significato e Implicazioni

• Cambiamento di Paradigma: Il lavoro suggerisce che l'ottimizzazione delle metriche di pre-training (come la loss minima) non è l'obiettivo finale per la costruzione di LLM utili. Ottimizzare per il pre-training può essere controproducente per l'adattabilità finale.
• Guida Pratica: Per gli sviluppatori di LLM, l'adozione di WSO (mantenendo il LR costante fino alla fine) è una strategia semplice ma potente per rilasciare modelli più "portabili" e facili da adattare a compiti specifici tramite SFT.
• Strategia di Rilascio: Gli autori raccomandano di rilasciare modelli pre-addestrati con WSO, poiché offrono una base superiore per le fasi successive di allineamento e specializzazione, migliorando l'efficienza complessiva del ciclo di vita del modello.

In sintesi, il paper dimostra che evitare il decadimento del Learning Rate durante il pre-training, controintuitivamente, produce modelli che, sebbene possano sembrare leggermente meno performanti "grezzi", sono significativamente più capaci di adattarsi e eccellere dopo il fine-tuning supervisionato.