Riepilogo Tecnico: Cosine-Gated Adam-Decay (CGAD) per DiLoCo Decoupled

1. Enunciato del Problema

Nei sistemi DiLoCo Decoupled (DDiLoCo) asincroni, i worker eseguono molteplici passi di ottimizzazione locale prima di scambiare pseudo-gradienti con un server di parametri globale. Una sfida critica in questa architettura è la stalezza (staleness): gli pseudo-gradienti applicati al round esterno $t$ possono essere stati calcolati $\tau$ round prima a causa della latenza di rete, di worker lenti (straggler) o di budget di asincronia.

L'ottimizzatore esterno standard per DiLoCo è SGD con momento di Nesterov. Questo metodo applica un passo di aggiornamento uniforme a ogni pseudo-gradiete in arrivo senza condizionarlo alla sua età ($\tau$). All'aumentare di $\tau$, il buffer del momento accumula una visione obsoleta del paesaggio della loss. Di conseguenza, l'ottimizzatore compie passi che superano le direzioni che non scendono più la loss corrente, portando a instabilità o divergenza, in particolare all'aumentare della scala del modello. Le analisi teoriche esistenti sul momentum-SGD asincrono suggeriscono che i tassi di convergenza degradano con un fattore proporzionale a $\tau_{max}^2$ (il quadrato del ritardo massimo).

2. Metodologia: Cosine-Gated Adam-Decay (CGAD)

Gli autori propongono Cosine-Gated Adam-Decay (CGAD), una sostituzione diretta ("drop-in") per l'ottimizzatore esterno Nesterov che condiziona esplicitamente gli aggiornamenti all'età dello pseudo-gradiete.

Meccanismo Principale

CGAD modifica l'aggiornamento Adam standard scalando ogni pseudo-gradiete in arrivo $g_f$ per il frammento $f$ mediante un fattore $\sigma(\tau)$ prima che entri nei buffer del primo e del secondo momento di Adam:
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

Dove:

• Decadimento Esponenziale ($e^{-\alpha \tau}$): Modella il decadimento informativo dei gradienti obsoleti. All'aumentare di $\tau$, il contributo del gradiente ai buffer del momento viene ridotto esponenzialmente.
• Porta Cosine ($\gamma(\tau)$): Una funzione di taglio liscia definita come:
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  Questa porta è uguale a 1 per i gradienti freschi, decade dolcemente a 0 a un taglio scelto $\tau_{cut}$, e annulla rigorosamente i contributi per $\tau > \tau_{cut}$.

Scelte Chiave di Progettazione

1. Porta Pre-Buffer: La scalatura $\sigma(\tau)$ viene applicata al gradiente prima che aggiorni il primo momento ($m$) e il secondo momento ($v$). Questo impedisce ai gradienti obsoleti di gonfiare il secondo momento ($v$), il quale altrimenti attenuerebbe i successivi aggiornamenti freschi.
2. Variante per Frammento (PA-CGAD): Per scheduler di sincronizzazione parziale (es. Streaming DiLoCo), il metodo traccia l'età dei singoli frammenti ($a_f$) e porta gli aggiornamenti in base a $\max(\tau, a_f)$, consentendo un modello di stalezza unificato sia per il ritardo di rete che per i ritardi di scheduling.
3. Iperparametri: Il metodo introduce due iperparametri, $\alpha$ (tasso di decadimento) e $\tau_{cut}$ (taglio), con valori predefiniti ($\alpha=0.2, \tau_{cut}=32$) che si trasferiscono attraverso le scale dei modelli.

3. Contributi Teorici

Il documento fornisce un limite di convergenza non asintotico per un aggiornamento adattivo con porta idealizzato su obiettivi non convessi lisci.

• Indipendenza Stalezza-Bias: Il termine di stalezza-bias del limite di convergenza dipende dall'iperparametro di porta $\alpha$ (specificamente $O(1/e^\alpha)$) piuttosto che dal ritardo massimo realizzato $\tau_{max}$.
• Contrasto con l'Analisi Standard: Le analisi standard del momentum-SGD asincrono producono un termine di complessità che scala con $\tau_{max}^2$. Sostituendo questo con un termine dipendente da $\alpha$, il meccanismo di porta disaccoppia teoricamente la penalità di convergenza dalla distribuzione del ritardo nel caso peggiore, a condizione che $\tau_{cut} \ge 1/\alpha$.
• Ambito: Gli autori notano che questo è un motivo per il meccanismo di porta piuttosto che un teorema di convergenza completo per Adam empirico, poiché assume passi normalizzati limitati piuttosto che derivarli da statistiche empiriche.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul preaddestramento di modelli linguistici in stile Llama a tre scale: 25 M, 1 B e 7 B parametri, utilizzando ritardi controllati fino a $\tau=16$.

Stabilità e Convergenza

• Scala 1 B: A $\tau=16$, la ricetta Nesterov standard ha diverso (loss $\approx 319$), mentre CGAD ha raggiunto una loss di 7.44, comparabile alle sue prestazioni a $\tau=0$ (7.48). Il metodo Nesterov ha perso completamente la convergenza a $\tau \ge 8$.
• Scala 25 M: CGAD e altre baseline consapevoli della stalezza (Adam-Decay, SDM, Nesterov Ritardato) sono rimaste tutte stabili, mentre Nesterov ha diverso a $\tau \ge 4$. A questa scala, le differenze tra i metodi consapevoli della stalezza sono state marginali.

Assicurazione sulla Scala e Robustezza

La scoperta più significativa emerge alla scala 7 B in condizioni di vincolo di memoria (pesi bf16, AdamW a 8 bit):

• Varianza Seed-to-Seed: La baseline più vicina, Adam-Decay (CGAD senza il taglio cosine), ha mostrato un aumento di 27 volte della deviazione standard seed-to-seed ($\sigma$) mentre il modello scalava da 25 M a 7 B (da 0.07 a 1.89 a $\tau=8$). Al contrario, il $\sigma$ di CGAD è rimasto piatto (circa 0.38).
• Rischio Singolo Shot: Utilizzando "media + $\sigma$" come proxy per il rischio di distribuzione singola, il rischio di Adam-Decay a 7 B / $\tau=8$ (13.19) ha superato la loss di livello casuale (10.83) di 2.4 nat, implicando un'alta probabilità di fallimento in una singola esecuzione di produzione. Il rischio di CGAD (11.28) è rimasto ben al di sotto di questa soglia.
• Meccanismo di Fallimento: Gli autori attribuiscono il fallimento di Adam-Decay a 7 B all'accumulo di rumore di arrotondamento al quadrato nel buffer del secondo momento ($v$) derivante da gradienti obsoleti che non vengono mai completamente azzerati. Il taglio cosine in CGAD previene ciò azzerando rigorosamente i contributi oltre $\tau_{cut}$.

Sensibilità alla Sintonizzazione

• Nesterov vs. CGAD: Anche quando Nesterov è stato risintonizzato (rimuovendo completamente il momento), non è stato in grado di eguagliare le prestazioni di una configurazione CGAD sintonizzata ad alti ritardi. Nesterov sintonizzato richiedeva iperparametri diversi per diversi regimi di ritardo, mentre CGAD manteneva la robustezza su tutto l'intervallo.
• Sensibilità agli Iperparametri: Il $\alpha$ ottimale si sposta leggermente con la scala (0.2 per 25 M, 0.4 per 1 B a $\tau=8$), ma il valore predefinito $\alpha=0.2$ rimane robusto su tutto lo spettro dei ritardi.

5. Significato e Affermazioni

Il documento afferma che CGAD è la scelta di scaling più robusta per l'addestramento DiLoCo asincrono, affrontando specificamente la fragilità del momento Nesterov sotto grandi ritardi controllati.

• Assicurazione sulla Scala: Il contributo principale è il taglio cosine, che agisce come "assicurazione sulla scala". Mentre il solo decadimento esponenziale (Adam-Decay) è sufficiente per modelli più piccoli o impostazioni in precisione completa, il taglio rigido è critico per modelli di frontiera (7 B+) che operano con stati dell'ottimizzatore compressi (8 bit/bf16). Previene l'accumulo di varianza obsoleta che porta a varianza catastrofica seed-to-seed e al fallimento singolo-shot.
• Compatibilità Drop-in: CGAD richiede solo una modifica di una riga alle implementazioni Adam esistenti e aggiunge due iperparametri con valori predefiniti che si trasferiscono attraverso le scale.
• Impatto Pratico: Per l'addestramento asincrono di frontiera, il metodo offre la differenza tra una singola esecuzione di addestramento che viene rilasciata con successo e una che fallisce (si posiziona sopra la loss casuale).

Gli autori rimangono modesti riguardo all'ambito teorico, riconoscendo che il limite di convergenza si applica a un aggiornamento idealizzato e che la superiorità empirica è dimostrata specificamente all'interno della famiglia DiLoCo e nei contesti di preaddestramento in stile Llama. Notano che il conteggio dei passi interni ($H=8$) utilizzato negli esperimenti è inferiore alle impostazioni DiLoCo di produzione ($H=100-500$), sebbene gli studi di sensibilità suggeriscano che la classificazione persista a valori di $H$ più elevati.