Rethinking Bregman Divergences in Kronecker-Factored Optimizers

Sintesi Tecnica: Ripensare le Divergenze di Bregman negli Ottimizzatori Fattorizzati di Kronecker

Definizione del Problema
Gli ottimizzatori in stile Shampoo approssimano la matrice di covarianza del gradiente $C = E[\text{vec}(G)\text{vec}(G)^\top]$ utilizzando una struttura fattorizzata di Kronecker $L \otimes R$ per ridurre la complessità computazionale. Un lavoro recente (Lin et al., 2026) ha unificato varie varianti di Shampoo inquadrandole come problemi di minimizzazione di divergenze di matrici di Bregman tra $C$ e $L \otimes R$. Tuttavia, rimane un vuoto critico: per matrici di covarianza dense generiche, la fattorizzazione esatta di Kronecker è impossibile, il che significa che l'errore di approssimazione è inevitabile. Non è chiaro come diverse divergenze di Bregman distribuiscano questo errore attraverso lo spettro della covarianza e se i precondizionatori risultanti si fidino equamente di tutte le parti dello spettro. Inoltre, poiché la covarianza del gradiente spesso differisce significativamente dall'effettiva Hessiana, non tutte le componenti spettrali sono adatte al precondizionamento. La domanda centrale affrontata è: Quando l'approssimazione esatta di Kronecker è impossibile, quali parti dello spettro della covarianza dovrebbero essere approssimate e seguite da un ottimizzatore strutturato?

Metodologia e Analisi
Gli autori affrontano questo problema attraverso due lenti analitiche principali: la geometria spettrale delle approssimazioni di Bregman e l'affidabilità empirica degli eigenspazi della covarianza.

1. Geometria Spettrale dell'Approssimazione di Bregman:
   Il documento caratterizza come le divergenze di Frobenius, von Neumann e LogDet penalizzino diversamente gli errori di approssimazione. Decomponendo la divergenza in termini di spettro marginale e un termine di accoppiamento cross-spettrale $g_F(\lambda_i, \omega_j)$, gli autori dimostrano che:

   • Divergenza di Frobenius: Il termine di accoppiamento $-\lambda_i \omega_j$ impone una penalità di magnitudo bilineare, enfatizzando fortemente lo spettro superiore (grandi autovalori).
   • Divergenza di von Neumann: Il termine $-\lambda_i \log \omega_j$ mantiene una preferenza per lo spettro superiore, ma la smorza logaritmicamente.
   • Divergenza LogDet: Il termine $\lambda_i / \omega_j$ crea una distorsione del rapporto relativo, rendendo l'approssimazione altamente sensibile ai piccoli autovalori (lo spettro della coda).

   Il Teorema 2 stabilisce che i punti stazionari di questi problemi di minimizzazione non sono determinati dal residuo grezzo $(S^* - C)$, ma da un residuo pesato dalla divergenza $\nabla^2 F(S^*)[S^* - C]$. Di conseguenza, i fattori di Kronecker risultanti si allineano alle specifiche preferenze spettrali della divergenza scelta (ad esempio, i fattori LogDet sono più sensibili ai piccoli autovalori, mentre i fattori di Frobenius si concentrano sui sottospazi dominanti).

2. Affidabilità Spettrale del Precondizionamento della Covarianza:
   Gli autori investigano empiricamente l'allineamento tra la covarianza del gradiente e gli eigenspazi dell'Hessiana su un modello Transformer addestrato su SST-2. I risultati (Tabella 3) indicano che il top eigenspazio della covarianza è sostanzialmente meglio allineato con l'Hessiana rispetto alle componenti a autovalore inferiore. Si osserva che lo spettro della coda è più rumoroso e meno affidabile per il precondizionamento, in linea con i risultati precedenti (Kunstner et al., 2019; Thomas et al., 2020).

3. Soluzione Proposta: Ottimizzatore di Kronecker Consapevole del Sottospazio (BregTop):
   Motivati dal disallineamento tra le preferenze spettrali indotte dalla divergenza e l'effettiva affidabilità dello spettro, gli autori propongono BregTop. Questo ottimizzatore combina i punti di forza di entrambe le osservazioni:

   • Seleziona i fattori di Kronecker utilizzando una divergenza sensibile al "top" (ad esempio, von Neumann o Frobenius).
   • Applica il precondizionamento basato sugli autovalori solo all'interno del top sottospazio (le $\lceil \rho mn \rceil$ voci più grandi dello spettro congiunto di Kronecker).
   • Tratta il sottospazio complementare inferiore con una costante di accelerazione isotropa adattiva ($\chi_{\rho, q}$), invece di tentare di precondizionare la coda rumorosa con informazioni di curvatura potenzialmente fuorvianti.

   Teoricamente, ciò decompone il precondizionatore in una componente superiore scalata dalla curvatura e una componente di accelerazione uniforme per lo spazio inferiore (Teorema 4).

Contributi Chiave

• Caratterizzazione della Geometria Spettrale: Il documento fornisce una caratterizzazione formale di come le divergenze di Frobenius, von Neumann e LogDet inducano diverse penalità dipendenti dagli autovalori, mostrando che i fattori di Kronecker stazionari sono governati da residui pesati dalla divergenza piuttosto che da errori grezzi.
• Affidabilità Spettrale Empirica: Il lavoro dimostra che, sebbene l'intero spettro della covarianza sia inaffidabile per il precondizionamento, il top eigenspazio mantiene un forte allineamento con l'Hessiana, giustificando un approccio consapevole del sottospazio.
• Design di un Ottimizzatore Consapevole del Sottospazio: Gli autori introducono BregTop, un algoritmo che passa adattivamente tra lo scaling basato sugli autovalori nel top sottospazio affidabile e l'accelerazione isotropa nel bottom sottospazio rumoroso.

Risultati Sperimentali
Gli autori hanno valutato BregTop contro Shampoo standard, KL-Shampoo, VN-Shampoo e F-Shampoo su un Transformer a 4 strati addestrato sull'intero dataset SST-2. La metrica utilizzata è stata il numero di step necessari affinché la perdita di addestramento con smoothing EMA scendesse sotto lo 0,05.

• BregTop-VN-v2 ha raggiunto la perdita target in 1003 step, superando lo Shampoo standard (1119 step) e il baseline VN-Shampoo (1127 step).
• Le varianti BregTop hanno costantemente migliorato i rispettivi baseline di divergenza (ad esempio, BregTop-VN ha migliorato rispetto a VN-Shampoo; BregTop-F ha migliorato rispetto a F-Shampoo), validando il beneficio di isolare il top sottospazio per il precondizionamento.

Significatività e Rivendicazioni
Il documento sostiene che il design degli ottimizzatori strutturati debba tenere conto dell'inevitabile errore di approssimazione nella fattorizzazione di Kronecker e della variabile affidabilità dello spettro della covarianza. Riconoscendo che diverse divergenze di Bregman intrinsecamente danno priorità a diverse regioni spettrali, e che solo il top spettro è affidabilmente allineato con l'Hessiana, l'approccio consapevole del sottospazio proposto offre una strategia di precondizionamento più robusta. Gli autori osservano che il loro design è coerente con la dinamica di addestramento osservata in Song et al. (2025), dove l'apprendimento avviene nel bottom sottospazio; pertanto, l'applicazione dell'accelerazione isotropa in tale spazio preserva il flusso del gradiente necessario pur smorzando gli aggiornamenti rumorosi nelle direzioni ad alta curvatura tramite il precondizionamento del top sottospazio. I risultati suggeriscono che questo approccio ibrido migliora l'efficienza dei passi senza richiedere informazioni di secondo ordine complete.