Sintesi Tecnica: M+Adam – Ottimizzazione Additiva-Moltiplicativa per l'Addestramento a Bassa Precisione

Definizione del Problema

L'addestramento di grandi modelli linguistici (LLM) in piena precisione (FP32) comporta costi significativi in termini di memoria, calcolo ed energia. Sebbene gli acceleratori moderni supportino formati a precisione ridotta come BF16, FP8 e FP4, le pipeline di addestramento standard mantengono tipicamente "pesi master" ad alta precisione (FP32) per garantire la stabilità. L'addestramento diretto con pesi master a bassa precisione porta spesso a una degradazione dell'accuratezza e all'instabilità.

Gli autori identificano un modo specifico di fallimento nell'ottimizzazione a bassa precisione: la risoluzione di quantizzazione grossolana a grandi magnitudo dei pesi. Nei sistemi in virgola mobile, la spaziatura tra i valori rappresentabili (bin di quantizzazione) cresce con la magnitudo del numero. Di conseguenza, gli ottimizzatori puramente additivi (come Adam) spesso non riescono a progredire quando le magnitudo dei pesi sono elevate, poiché i passi di aggiornamento calcolati cadono al di sotto del passo di quantizzazione locale e vengono arrotondati a zero. Al contrario, gli ottimizzatori puramente moltiplicativi (come Madam), che scalano i pesi rispetto alla loro magnitudo attuale, falliscono vicino allo zero o quando sono necessari cambi di segno, poiché non possono uscire da un'inizializzazione esatta a zero o invertire i segni efficacemente.

Metodologia: M+Adam

Il documento propone M+Adam, un nuovo ottimizzatore che combina i punti di forza complementari dei meccanismi di aggiornamento additivo e moltiplicativo per superare i limiti dell'aritmetica a bassa precisione.

Meccanismo Centrale

M+Adam mantiene due rami di aggiornamento paralleli che vengono combinati ad ogni passo:

1. Ramo Additivo (stile Adam): Opera direttamente sullo spazio dei pesi. Calcola i momenti standard del gradiente ($g$) per produrre un aggiornamento additivo $u^{(a)}$. Questo ramo è efficace per gestire piccole magnitudo, uscire dall'inizializzazione a zero e facilitare i cambi di segno.
2. Ramo Moltiplicativo (stile Madam): Opera sulla scala logaritmica (spazio degli esponenti). Calcola il gradiente rispetto all'esponente, $g^{(e)} = (\ln 2) w_t g$, e utilizza momenti adattivi per produrre un aggiornamento moltiplicativo relativo $u^{(m)}$. Questo ramo garantisce il progresso ad alte magnitudo dove gli aggiornamenti additivi verrebbero arrotondati a zero.

La regola di aggiornamento del peso è definita come:
$$w_{t+1} = w_t + w_t u^{(m)}_t + u^{(a)}_t$$
Questa struttura consente all'ottimizzatore di eseguire aggiornamenti relativi consapevoli della scala quando le magnitudo sono elevate (tramite il termine moltiplicativo), mantenendo al contempo la flessibilità di cambiare segno e aggiornare valori vicini allo zero (tramite il termine additivo).

Analisi Teorica

Gli autori forniscono una garanzia teorica per M+Adam sotto standard assunzioni di regolarità (smoothness). Modellano l'aggiornamento come un'operazione "scala-poi-aggiungi" in coordinate locali. Il Teorema 4.1 stabilisce che se i tassi di apprendimento e le costanti di regolarità soddisfano specifiche condizioni, il metodo garantisce una discesa monotona ($L(w_{t+1}) < L(w_t)$) a meno che il gradiente non sia zero. Ciò dimostra che la geometria combinata è una strategia di ottimizzazione valida.

Dettagli di Implementazione

• Supporto alla Precisione: Il metodo è progettato per l'archiviazione dei pesi master in BF16, FP8 e NVFP4.
• Gestione dello Stato: M+Adam introduce uno stato aggiuntivo dell'ottimizzatore (l'accumulatore del secondo momento per il ramo moltiplicativo). Gli autori dimostrano che questo stato può essere compresso utilizzando approssimazioni a basso rango (specificamente Apollo-E, un'approssimazione di rango 4) per minimizzare l'overhead di memoria, mantenendo l'impronta simile a quella di AdamW con stato BF16 standard.
• Nessun Arrotondamento Stocastico: A differenza di alcuni approcci precedenti, M+Adam raggiunge la stabilità senza richiedere l'arrotondamento stocastico (SR) o la compensazione di tipo Kahan, sebbene sia compatibile con essi.

Contributi Chiave

1. Identificazione di Modi di Fallimento Complementari: Il documento caratterizza formalmente i modi di fallimento degli ottimizzatori additivi (stallo ad alte magnitudo dovuto all'arrotondamento) e degli ottimizzatori moltiplicativi (stallo vicino allo zero e incapacità di invertire i segni), mostrando che sono complementari.
2. Algoritmo M+Adam: Un nuovo ottimizzatore che unifica queste due geometrie, fornendo un approccio rigoroso all'addestramento a bassa precisione senza pesi master FP32.
3. Garanzia di Discesa Teorica: Una prova di discesa monotona per la regola di aggiornamento additivo-moltiplicativo combinata sotto standard assunzioni di regolarità.
4. Validazione Empirica su Scala: Esperimenti estesi su modelli di tipo LLaMA che vanno da 60M a 1B di parametri attraverso budget di addestramento da 1 a 8× Chinchilla.

Risultati Sperimentali

Gli autori hanno valutato M+Adam in vari regimi di precisione: BF16/BF16, BF16/FP8, FP8/FP8 e NVFP4/FP8.

• Miglioramento Consistente: M+Adam ha superato costantemente i baseline AdamW ottimizzati in tutti i modelli e regimi di precisione.
• Riduzioni della Perplessità: Nel regime di bassa precisione più aggressivo (NVFP4/FP8), M+Adam ha ridotto la perplessità di margini significativi rispetto ad AdamW (ad esempio, un miglioramento del 16,6% per il modello da 350M). Anche nell'addestramento BF16/BF16, M+Adam ha ridotto la perplessità dal 2,8% al 13,3% a seconda della dimensione del modello.
• Robustezza: Il metodo è rimasto stabile senza l'uso dell'arrotondamento stocastico. Ha inoltre superato AdamW combinato con l'arrotondamento stocastico (AdamW+SR) e la compensazione di Kahan.
• Scalabilità: I guadagni di prestazione sono persistiti all'aumentare del budget di addestramento (fino a 8× Chinchilla) e all'aumentare della dimensione del modello fino a 1 miliardo di parametri.
• Efficienza della Memoria: Utilizzando la tecnica di compressione Apollo-E, lo stato di memoria aggiuntivo richiesto per il ramo moltiplicativo è stato ridotto, risultando in un'impronta dello stato dell'ottimizzatore solo leggermente superiore a quella di un AdamW BF16 standard, pur mantenendo la maggior parte dei guadagni in termini di perplessità.

Significato e Rivendicazioni

Il documento afferma che gli aggiornamenti additivi-moltiplicativi combinati sono particolarmente utili quando i pesi master a bassa precisione diventano una fonte dominante di errore di ottimizzazione.

Il significato di questo lavoro risiede nel fornire una soluzione rigorosa all'instabilità numerica dell'addestramento a bassa precisione. Allineando la geometria di aggiornamento dell'ottimizzatore con la struttura non uniforme dei numeri in virgola mobile, M+Adam mitiga gli errori di arrotondamento che affliggono gli standard ottimizzatori additivi ad alte magnitudo e i problemi di stagnazione degli ottimizzatori moltiplicativi vicino allo zero.

Gli autori concludono che M+Adam rappresenta un passo verso l'addestramento end-to-end a bassa precisione che non dipende dai pesi master FP32. Questa capacità è cruciale per sfruttare l'hardware IA moderno (tensor core) ottimizzato per FP8 e FP4, riducendo potenzialmente l'overhead di memoria e i costi energetici per l'addestramento di modelli su larga scala. Il lavoro suggerisce che i futuri ottimizzatori dovrebbero essere progettati tenendo conto dei vincoli specifici della rappresentazione in virgola mobile, piuttosto che trattare la bassa precisione semplicemente come un problema di quantizzazione da risolvere con tecniche di arrotondamento.