Riassunto Tecnico: LoRDO – Ottimizzazione a Basso Rango Distribuita con Comunicazione Infrequente

1. Definizione del Problema

L'addestramento distribuito di modelli di fondazione tramite Distributed Data Parallel (DDP) è sempre più limitato dalla larghezza di banda dell'interconnessione e dall'overhead di memoria richiesto per memorizzare gli stati dell'ottimizzatore (ad esempio, i momenti di Adam). Sebbene le strategie che impiegano una comunicazione infrequente (aggiornamenti locali) riducano la frequenza di sincronizzazione, esse rimangono strozzate dai costi di memoria e di comunicazione degli stati dell'ottimizzatore a pieno rango.

Gli ottimizzatori adattivi a basso rango (ad esempio, GaLore, LDAdam) offrono una via per alleviare questi vincoli proiettando i gradienti in uno spazio sottospazio a basso rango. Tuttavia, generalizzare questi metodi al regime di sincronizzazione infrequente presenta una sfida di ottimizzazione critica:

1. Proiezioni Locali: Calcolare le proiezioni a basso rango localmente su singoli worker utilizzando gradienti stocastici introduce un rumore significativo dovuto alle piccole dimensioni effettive del batch per ogni worker. Ciò porta a sottospazi di proiezione subottimali e al degrado delle prestazioni.
2. Proiezioni Globali: L'uso di una matrice di proiezione globale derivata da pseudo-gradienti aggregati (come in DDP) fornisce una base più stabile. Tuttavia, nel regime di aggiornamento locale, questo approccio restringe permanentemente la traiettoria di ottimizzazione a un sottospazio a basso rango fisso. Senza segnali a pieno rango, l'ottimizzazione ristagna perché la base di proiezione non riesce a "rigenerarsi" o a esplorare nuove direzioni al di fuori dello spazio iniziale.

La domanda centrale affrontata è: Come possiamo progettare ottimizzatori a basso rango ad alte prestazioni per l'addestramento con comunicazione efficiente che evitino la stagnazione del sottospazio mantenendo al contempo bassi carichi di memoria e di comunicazione?

2. Metodologia: LoRDO

Gli autori propongono LoRDO (Low-Rank Distributed Optimization), un framework che unifica l'ottimizzazione a basso rango con la sincronizzazione infrequente. La metodologia affronta la tensione tra stabilità della proiezione ed esplorazione del sottospazio attraverso tre meccanismi chiave:

A. Matrici di Proiezione Globali

Invece di far determinare a ogni worker la propria matrice di proiezione $Q_m$ basata su gradienti stocastici locali (che soffre di alta varianza), LoRDO impiega una matrice di proiezione globale $Q_t$. Questa matrice viene calcolata al confine di sincronizzazione utilizzando lo pseudo-gradiente aggregato ($\Delta_t$), che rappresenta il cambiamento totale dei parametri del modello su $K$ passi locali.

• Beneficio: Lo pseudo-gradiente ha una dimensione effettiva del batch di $M \times B$ (dove $M$ è il numero di worker), fornendo un segnale più stabile con varianza ridotta rispetto ai gradienti locali. Incorpora inoltre l'informazione di curvatura racchiusa nello pseudo-gradiente, portando a una stima della base di qualità superiore.

B. Momento Quasi-Iperbolico a Pieno Rango (QHM)

Per risolvere il problema della "stagnazione del sottospazio" inerente alle proiezioni globali, LoRDO inietta un termine di momento quasi-iperbolico a pieno rango nell'aggiornamento locale.

• Meccanismo: La regola di aggiornamento combina il momento a basso rango con un segnale di gradiente a pieno rango scalato dal secondo momento.
  $$\text{Update} \propto (1-\omega)\frac{\hat{G}_t}{\mu(\sqrt{\hat{v}_t} + \epsilon)} + \omega Q_t \left[ \frac{u_t}{\sqrt{v_t} + \epsilon} \right]$$
• Effetto: Questa iniezione impedisce allo pseudo-gradiente aggregato di rimanere intrappolato in un sottospazio fisso di rango-$r$. Permette alla procedura di ottimizzazione di aggregare temporalmente sottospazi di rango-$r$ provenienti da diversi worker, consentendo un'eventuale esplorazione completa dello spazio. Ciò ripristina la capacità di esplorare l'intero spazio delle soluzioni senza aumentare l'overhead di comunicazione.

C. Feedback dell'Errore e Rotazione del Momento

LoRDO incorpora il feedback dell'errore localmente per tenere conto delle informazioni perse durante la compressione a basso rango. Inoltre, dopo il calcolo di una nuova matrice di proiezione, gli stati dell'ottimizzatore (momenti) vengono ruotati per garantire che rimangano accumulati sullo stesso sottospazio, una tecnica adattata da LDAdam.

3. Contributi Chiave

1. Analisi dei Modelli di Fallimento della Proiezione: Gli autori dimostrano che le proiezioni locali danneggiano le prestazioni a causa dell'alta varianza derivante dalle piccole dimensioni del batch dei worker, mentre le proiezioni globali senza segnali a pieno rango inducono la stagnazione del sottospazio.
2. Ripristino dell'Esplorazione del Sottospazio Completo: LoRDO introduce un ottimizzatore a basso rango che inietta un segnale di gradiente a pieno rango (tramite QHM) negli aggiornamenti locali, mantenendo al contempo una proiezione globale derivata da pseudo-gradienti aggregati. Questo evita la stagnazione senza aumentare l'overhead di comunicazione o di memoria.
3. Parità con DDP ed Efficienza: Il framework raggiunge una quasi-parità con il DDP a basso rango sincrono (gap di perplexity < 1%) pur riducendo l'overhead di memoria e comunicazione fino a $8\times$–$12\times$. Sotto forti vincoli di memoria che richiedono bassi ranghi, LoRDO supera persino DDP in termini di perplexity.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su transformer decoder-only scalati a 16M, 125M e 720M di parametri utilizzando la miscela di dati SmolLM2.

• Stagnazione vs. Esplorazione: Nei modelli da 16M, le proiezioni globali senza QHM non sono riusciti ad apprendere (perplexity 46.8) rispetto a DDP (36.6), confermando la stagnazione del sottospazio. L'aggiunta di QHM a pieno rango ha ripristinato le prestazioni, permettendo a LoRDO-Global di superare le varianti locali.
• Qualità della Proiezione: LoRDO-Global ha dimostrato una stabilità e prestazioni superiori, in particolare ai ranghi più bassi ($r=8, 16$), grazie alla base di proiezione di qualità superiore derivata dallo pseudo-gradiente aggregato. Ha mostrato una minore sensibilità alle variazioni della dimensione del batch locale rispetto ai metodi di proiezione locale.
• Scalabilità a 720M: Alla scala di 720M, LoRDO-Global ($r=256$) ha raggiunto la quasi-parità con il DDP a basso rango ($r=256$) con un gap di perplexity inferiore all'1% (10.41 vs 10.34). Ha inoltre eguagliato l'accuratezza nei task downstream (ARC, HellaSwag, MMLU, PIQA) con DDP.
• Guadagni di Efficienza:
  • Comunicazione: LoRDO riduce l'overhead di comunicazione di circa 10$\times$ rispetto al DDP a basso rango e di 25$\times$ rispetto al DDP a pieno rango.
  • Memoria: Riduce l'overhead di memoria dello stato dell'ottimizzatore di 8$\times$ rispetto alle controparti a pieno rango.
  • Tempo di esecuzione (Wall-Clock Time): In ambienti a bassa larghezza di banda (modellati a 10 Gbps), LoRDO supera significativamente DDP in tempo di esecuzione grazie alla riduzione della frequenza di sincronizzazione.

5. Significato e Rivendicazioni

Il documento afferma che LoRDO risolve la tensione fondamentale tra stabilità della proiezione ed esplorazione del sottospazio nell'ottimizzazione efficiente della comunicazione a basso rango. Utilizzando pseudo-gradienti aggregati per la stabilità della proiezione e un momento quasi-iperbolico a pieno rango per l'esplorazione, LoRDO consente:

• Addestramento Scalabile: Prestazioni competitive su larga scala (fino a 720M di parametri) con costi di comunicazione e memoria significativamente ridotti.
• Efficienza delle Risorse: Fattibilità dell'addestramento su hardware con risorse limitate (memoria/banda ridotte), dove gli ottimizzatori a pieno rango sono intrattabili.
• Apprendimento Decentralizzato: Aprire la strada a un apprendimento efficiente in ambienti decentralizzati, estendendo l'addestramento del modello oltre i singoli data center.

Gli autori segnalano alcune limitazioni, tra cui la mancanza di una prova formale di convergenza e l'attuale focalizzazione sugli ottimizzatori adattivi (Adam) piuttosto che su quelli basati su matrici (ad esempio, Muon), lasciando queste questioni per il lavoro futuro. Tuttavia, sostengono che le tendenze osservate nei modelli più piccoli si applichino anche a scale maggiori.