Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning

1. Enunciato del Problema

Il lavoro affronta l'analisi teorica della convergenza dell'ottimizzatore AdamW-style Shampoo, un algoritmo all'avanguardia vincitore della competizione di addestramento di reti neurali AlgoPerf. Sebbene Shampoo (e le sue varianti come SOAP e Muon) abbia dimostrato prestazioni empiriche superiori rispetto ai metodi di precondizionamento diagonale (come Adam) sfruttando precondizionatori a matrice piena, le sue garanzie teoriche di convergenza per l'ottimizzazione non convessa sono state finora assenti o limitate.

Nello specifico, la letteratura esistente presentava tre principali lacune:

1. Non convessità: La maggior parte delle analisi precedenti era limitata a contesti convessi.
2. Ambito del precondizionamento: Le analisi non convesse precedenti (ad es. Xie et al., 2026) erano limitate al precondizionamento a un solo lato, ignorando il precondizionamento a due lati, più complesso ed efficace, utilizzato nelle implementazioni pratiche.
3. Componenti pratiche: Le teorie esistenti spesso omettevano caratteristiche pratiche critiche come il momento e il decoupled weight decay (stile AdamW), essenziali per l'addestramento moderno di modelli linguistici su larga scala (LLM).

Il lavoro mira a stabilire un tasso di convergenza rigoroso per l'algoritmo generale AdamW-style Shampoo, unificando sia il precondizionamento a un solo lato che quello a due lati all'interno di un quadro non convesso.

2. Metodologia

Gli autori analizzano l'Algoritmo 1, che rappresenta l'ottimizzatore AdamW-style Shampoo. La regola di aggiornamento coinvolge:

• Medie Mobili Esponenziali (EMA): Tracciamento dei primi momenti ($M_k$) e dei secondi momenti ($L_k, R_k$) dei gradienti.
• Precondizionamento a Due Lati: Il passo di aggiornamento è $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$, dove $L_k$ e $R_k$ sono prodotti esterni accumulati dei gradienti (che approssimano $E[GG^T]$ e $E[G^TG]$).
• Esponenti Unificati: I parametri $p$ e $q$ (dove $1/p + 1/q = 1$) permettono all'algoritmo di passare dal precondizionamento a due lati ($p, q < \infty$) a quello a un solo lato ($p=1, q=\infty$ o viceversa).

Innovazioni Tecniche Chiave:
Per gestire la complessità del precondizionamento matriciale a due lati in contesti non convessi, gli autori hanno sviluppato nuove tecniche dimostrative:

1. Disuguaglianza di Hölder di Schatten-p: Invece delle disuguaglianze di Hölder basate su vettori standard utilizzate nell'analisi di Adam, gli autori utilizzano la disuguaglianza di Hölder di Schatten-p per limitare la norma nucleare ($\|\cdot\|_*$) del gradiente. Questo è cruciale poiché il precondizionatore agisce sui valori singolari dei parametri matriciali.
2. Disuguaglianza di Cauchy-Schwarz Matriciale: Per limitare la norma spettrale del passo di aggiornamento ($L^{-1/2p} M R^{-1/2q}$), impiegano una disuguaglianza di Cauchy-Schwarz specifica per matrici per gestire l'interazione tra i precondizionatori sinistro e destro.
3. Analisi del Pavimento Spettrale: Introducono una condizione (Condizione 4) riguardante il limite inferiore dei precondizionatori ($\hat{\epsilon}$). Sostengono che, nella pratica, il pavimento spettrale effettivo dei precondizionatori è significativamente più grande della regolarizzazione numerica $\epsilon$ (ad es. $10^{-12}$), allineando i limiti teorici con le osservazioni empiriche.

3. Contributi Chiave

• Tasso di Convergenza Unificato: Il lavoro stabilisce il primo tasso di convergenza per AdamW-style Shampoo nell'ottimizzazione non convessa che copre sia il precondizionamento a un solo lato che quello a due lati.
• Derivazione del Tasso Ottimale: Dimostrano che la norma nucleare attesa media del gradiente converge a un tasso di:
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  dove $m, n$ sono le dimensioni della matrice, $K$ è il numero di iterazioni e $C$ corrisponde alla costante nel tasso ottimale di SGD.
• Giustificazione Teorica per Impostazioni Pratiche: L'analisi spiega perché impostare il regolarizzatore numerico $\epsilon$ a valori estremamente piccoli (ad es. $10^{-12}$) non ostacola la convergenza. La teoria mostra che il pavimento del precondizionatore effettivo è determinato dalla scala del rumore del gradiente, non dal $\epsilon$ numerico.
• Relazione tra Norme: Gli autori dimostrano che il tasso di norma nucleare derivato è analogo al tasso ottimale di norma di Frobenius di SGD ($O(K^{-1/4})$) nella condizione ideale in cui $\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$.

4. Risultati

• Limiti Teorici:
  • Il limite nel caso peggiore dipende da $\sqrt{m+n}$, che è significativamente più stretto della dipendenza da $\sqrt{mn}$ (dimensione totale $d$) trovata in alcuni metodi diagonali o nei limiti inferiori a matrice piena.
  • Il tasso di convergenza corrisponde al limite inferiore noto per l'ottimizzazione stocastica non convessa ($O(K^{-1/4})$) fino a un fattore di $\sqrt{m+n}$.
• Validazione Empirica (Pre-addestramento GPT-2):
  • Rapporto delle Norme del Gradiente: Gli esperimenti sul pre-addestramento di GPT-2 hanno mostrato che il rapporto $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ rimane costantemente vicino al limite superiore teorico $\sqrt{\min(m,n)}$. Questo valida l'assunzione che il tasso di norma nucleare sia effettivamente equivalente al tasso di norma di Frobenius nella pratica.
  • Pavimento Spettrale: I valori propri minimi dei precondizionatori ($L_k, R_k$) sono stati osservati essere ordini di grandezza più grandi del $\epsilon$ numerico ($10^{-12}$) e comparabili alla scala dipendente dal rumore $\hat{\sigma}^2/(m+n)$. Questo conferma che l'assunzione "non pratica" di un pavimento spettrale più grande vale nell'addestramento reale di LLM.

5. Significato

• Colmare il Divario tra Teoria e Pratica: Questo lavoro fornisce la prima fondazione teorica rigorosa per la specifica variante di Shampoo che ha vinto la competizione AlgoPerf, spiegando perché funziona così bene nella pratica.
• Avanzamento rispetto ai Metodi Diagonali: Convalida teoricamente la superiorità del precondizionamento non diagonale (a matrice piena) rispetto ai metodi diagonali (come Adam) per parametri strutturati a matrice, mostrando che le correlazioni tra parametri catturate da Shampoo portano a direzioni di convergenza migliori.
• Guida per gli Iperparametri: L'analisi offre una giustificazione teorica per la scelta degli iperparametri (ad es. la relazione tra momento $\theta$ e decadimento del secondo momento $\beta$, e il ruolo di $\epsilon$), guidando i praticanti nell'ottimizzazione di questi ottimizzatori per modelli su larga scala.
• Fondamento per Futuri Ottimizzatori: Unificando il precondizionamento a un solo lato e a due lati e incorporando momento e decadimento dei pesi, questo quadro stabilisce un nuovo standard per l'analisi di ottimizzatori avanzati di tipo second'ordine nel deep learning.

In sintesi, il lavoro unisce con successo l'analisi teorica di Shampoo, dimostrando che raggiunge un tasso di convergenza ottimale per problemi non convessi catturando al contempo le complesse strutture matriciali intrinseche nei moderni modelli di deep learning.