Riepilogo Tecnico: GEFEN: OTTIMIZZATORE STOCASTICO OTTIMIZZATO

Definizione del Problema

Il deep learning moderno si affida pesantemente all'ottimizzatore AdamW grazie alle sue robuste prestazioni su diverse architetture. Tuttavia, AdamW mantiene medie mobili del primo e del secondo momento per ogni parametro, tipicamente in FP32. Ciò comporta un'impronta di memoria dello stato dell'ottimizzatore di circa due volte la dimensione dei parametri del modello (circa 8 byte per parametro). Per modelli su larga scala (ad es. 13B di parametri), questo stato da solo può consumare circa 97 GiB di memoria, limitando severamente le dimensioni dei modelli addestrabili, le dimensioni massime dei batch e la velocità di addestramento su hardware con memoria limitata.

Le alternative esistenti orientate all'efficienza della memoria affrontano compromessi significativi:

• Ottimizzatori Quantizzati (ad es. Adam8bit, Adam4bit): Riducono la memoria ma spesso si affidano a codebook creati manualmente e subottimali, dimensioni di blocco fisse (che introducono iperparametri) e potrebbero non essere universalmente applicabili (ad es. evitando la quantizzazione degli embedding).
• Ottimizzatori Strutturati (ad es. Adam-mini): Condividono le statistiche del secondo momento basandosi su regole manuali che mimano le strutture block-diagonali dell'Hessiana. Tuttavia, richiedono metadati specifici dell'architettura (ad es. il numero di teste di attenzione) e mancano di una forte giustificazione teorica sul perché il raggruppamento allineato all'Hessiana sia benefico.

L'obiettivo è progettare un ottimizzatore che riduca significativamente l'impronta di memoria (oltre ad Adam-mini) mantenendo prestazioni simili ad AdamW, senza richiedere configurazioni manuali, iperparametri aggiuntivi o ricalibrazione del tasso di apprendimento (learning rate).

Metodologia

Il metodo proposto, Gefen, ottiene la riduzione della memoria attraverso due meccanismi primari: la condivisione automatica del secondo momento e la quantizzazione appresa del primo momento.

1. Motivazione Teorica

L'approccio si basa su un risultato teorico (Teorema 4.1) che dimostra come, per i parametri con grandi voci miste dell'Hessiana, il rapporto tra i loro quadrati dei gradienti sia costretto verso l'unità. Ciò implica che i parametri con alta affinità Hessiana possiedono naturalmente magnitudo di gradienti al quadrato simili, rendendoli candidati ideali per la condivisione delle stime del secondo momento. Sebbene il calcolo diretto dell'Hessiana sia intrattabile su larga scala, gli autori osservano che la struttura dei quadrati dei gradienti nel primo step di addestramento correla con questa struttura dell'Hessiana.

2. Partizionamento Automatico dei Blocchi

Invece di fare affidamento su regole manuali o sul calcolo dell'Hessiana, Gefen inferisce la struttura ottimale dei blocchi dai quadrati dei gradienti del primo step di ottimizzazione ($g_0$).

• Algoritmo: Per un tensore di parametri appiattito, l'algoritmo testa tutti i divisori propri $p$ della dimensione del tensore $n$. Rimodella i quadrati dei gradienti in blocchi di dimensione $p$ e calcola l'eterogeneità intra-blocco $E(p)$ (la varianza dei quadrati dei gradienti all'interno dei blocchi).
• Strategia di Selezione: L'algoritmo seleziona il periodo $p$ che produce il primo calo pronunciato di $E(p)$ all'aumentare della dimensione del blocco, favorendo blocchi più grandi per l'efficienza della memoria pur garantendo la somiglianza dei gradienti. Se non viene trovato alcun calo significativo o se il $p$ ottimale è $< 8$, il sistema ricorre a $p=1$ (nessuna condivisione).
• Utilizzo: Questo singolo schema di partizionamento viene utilizzato sia per la condivisione del secondo momento che per la quantizzazione del primo momento.

3. Quantizzazione Esatta tramite Programmazione Dinamica

Per ridurre la memoria del primo momento (momentum), Gefen impiega un codebook di quantizzazione apprendibile ed esatto, piuttosto che schemi fissi o euristici.

• Apprendimento del Codebook: Utilizzando i gradienti iniziali, il metodo costruisce un istogramma dei valori dei gradienti normalizzati. Successivamente, risolve il problema della quantizzazione scalare ottimale utilizzando la Programmazione Dinamica (DP) (ispirata a Wang & Song, 2011) per trovare il codebook esatto per la distribuzione e la larghezza di bit data.
• Stabilità: A differenza di Lloyd-Max, che può soffrire di "contrazione" (restringimento dell'intervallo di quantizzazione nel tempo) e sensibilità all'inizializzazione, Gefen forza gli endpoint del codebook a essere fissati a $-1$e$1$. Ciò assicura che l'intervallo di quantizzazione sia preservato durante l'intero addestramento.
• Efficienza: Il codebook viene appreso una sola volta all'inizio dell'addestramento (utilizzando CUDA per l'istogramma e CPU/Numba per la DP) e riutilizzato per l'intera sessione di addestramento.

4. Flusso dell'Algoritmo

Gefen opera come un sostituto diretto (drop-in replacement) per AdamW:

1. Inizializzazione: Partiziona i parametri in base a $g_0$ e apprende il codebook di quantizzazione.
2. Step di Aggiornamento:
   • Calcola i gradienti.
   • Aggiorna il primo momento (momentum) dequantizzando i valori del codebook memorizzato, applicando il momentum, ri-quantizzando e aggiornando il valore assoluto massimo per blocco.
   • Aggiorna il secondo momento calcolando la media dei quadrati dei gradienti all'interno di ogni blocco e applicando la media mobile esponenziale.
   • Applica la correzione del bias e aggiorna i parametri utilizzando il secondo momento condiviso e il primo momento quantizzato.

Contributi Chiave

1. Efficienza della Memoria: Gefen riduce l'impronta di memoria dell'ottimizzatore di circa 8× rispetto ad AdamW (riducendo da ~2.0× la dimensione dei parametri a ~0.25×), corrispondente a una riduzione di ~6.5 GiB per miliardo di parametri.
2. Adattamento Automatico: Il metodo non richiede metadati specifici dell'architettura, regole di raggruppamento manuali o iperparametri aggiuntivi oltre ai default di AdamW. Inferisce automaticamente le strutture dei blocchi dai dati.
3. Fondamento Teorico: Fornisce un collegamento teorico tra le voci dell'Hessiana e i rapporti dei quadrati dei gradienti, giustificando la condivisione delle statistiche del secondo momento per i parametri allineati all'Hessiana.
4. Quantizzazione Ottimale: Introduce un approccio esatto basato su istogrammi e programmazione dinamica per la quantizzazione del momentum, evitando la subottimalità e l'instabilità dei metodi iterativi come Lloyd-Max.
5. Parità delle Prestazioni: Dimostra che queste riduzioni non comportano una perdita di qualità del modello, eguagliando le prestazioni di AdamW in vari compiti.

Risultati Sperimentali

Gli autori hanno valutato Gefen rispetto ad AdamW, Adam-mini, Adam8bit, Adam4bit, Adafactor e SM3 su modelli CNN, ResNet, GPT-2 e Llama 3.

• Impronta di Memoria di Picco: Gefen ha costantemente ottenuto l'uso della memoria di picco più basso tra i metodi confrontati. Per un modello da 1.5B di parametri, Gefen ha ridotto lo stato dell'ottimizzatore da 11.2 GiB (AdamW) a 1.5 GiB.
• Prestazioni (Qualità): In tutte le impostazioni (OpenWebText, C4, ImageNet, MNIST), Gefen ha ottenuto perdite di validazione e accuratezze comparabili ad AdamW. Notevolmente, ha superato Adam4bit e Adam-mini in diverse metriche mantenendo lo stesso schema del learning rate.
• Throughput: La ridotta impronta di memoria ha permesso microbatch più grandi nell'addestramento distribuito:
  • FSDP (Llama 3-1.5B): Gefen ha abilitato un microbatch size di 2 (rispetto a 1 per gli altri), risultando in un aumento del throughput del 56% (token/sec).
  • DDP (GPT-2-1B): Gefen ha permesso l'addestramento con un microbatch size di 2, mentre AdamW non riusciva a far stare nemmeno un microbatch size di 1. Ciò ha portato a un miglioramento del throughput del 21% rispetto ad Adam-mini.
• Stabilità: Gli studi di ablazione hanno confermato che il codebook appreso rimane stabile durante l'addestramento, validando la decisione di apprenderlo una sola volta all'inizializzazione.

Significato e Rivendicazioni

Il paper posiziona Gefen come un sostituto pratico e "drop-in" per AdamW che affronta il collo di bottiglia della memoria nell'addestramento su larga scala senza sacrificare le prestazioni o richiedere una regolazione manuale.

• Impatto Pratico: Riducendo lo stato dell'ottimizzatore, Gefen consente di addestrare modelli più grandi o di utilizzare batch size più ampi su hardware fisso, migliorando direttamente il throughput di addestramento e l'efficienza del campionamento.
• Generalità: A differenza di altri metodi che si affidano a euristiche manuali o nomi specifici dei tensori, Gefen è agnostico rispetto all'architettura e si adatta automaticamente ai dati.
• Complementarità: Gli autori notano che la compressione dello stato dell'ottimizzatore di Gefen è ortogonale alla compressione della comunicazione dei gradienti (ad es. 1-bit Adam) e alle strategie di addestramento distribuito (FSDP/DDP), il che significa che può essere combinata con tali tecniche per ulteriori guadagni.

Gli autori forniscono un'implementazione Python completa con kernel CUDA fusi, enfatizzando la riproducibilità e la facilità di integrazione nei pipeline di addestramento esistenti.