Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration

Riepilogo Tecnico: Accelerazione dell'Ottimizzazione Spettrale di Ordine Zero con Ortogonalizzazione Parziale dall'Iterazione di Potenza

Enunciato del Problema
L'ottimizzazione di ordine zero (ZO) si è affermata come una tecnica critica per il fine-tuning dei Modelli Linguistici di Grande Dimensione (LLM) su dispositivi edge con vincoli di memoria, poiché evita la retropropagazione intensiva in termini di memoria richiesta dai metodi di primo ordine (FO). Tuttavia, i metodi ZO soffrono di un'alta varianza nelle stime del gradiente a causa delle perturbazioni Monte Carlo, portando a una convergenza significativamente più lenta rispetto ai metodi FO.

I recenti tentativi di accelerare l'ottimizzazione ZO si sono concentrati sulla ricerca in sottospazi a bassa dimensionalità (ad es. LOZO, SubZero) e su ottimizzatori spettrali come Muon, che sfruttano la geometria della matrice attraverso l'ortogonalizzazione. Sebbene Muon superi AdamW negli scenari FO ortogonalizzando i gradienti per sfruttare le direzioni spettrali deboli, gli autori osservano che l'applicazione diretta di Muon agli scenari ZO produce guadagni limitati. Il problema centrale identificato è che le stime del gradiente ZO sono intrinsecamente rumorose e inaffidabili. A differenza dei gradienti FO, che sono concentrati nelle direzioni singolari dominanti, le stime ZO presentano una "coda spettrale più pesante e piatta" (Figura 1). Di conseguenza, l'ortogonalizzazione Newton-Schulz (NS) standard di Muon, che equalizza l'energia su tutto lo spettro attivo, amplifica involontariamente queste direzioni rumorose e inaccurate, trasformando gli errori di stima in componenti di aggiornamento dannose.

Metodologia: ZO-MOPI
Gli autori propongono ZO-MOPI (Momentum di Ordine Zero con Ortogonalizzazione Parziale), un metodo progettato per accelerare l'ottimizzazione ZO estraendo selettivamente le direzioni spettrali dominanti mentre sopprime il rumore. L'approccio integra tre componenti chiave:

1. Ortogonalizzazione Parziale tramite Iterazione di Potenza in Streaming (SPI):
   Invece di eseguire un'ortogonalizzazione a spettro completo (come le iterazioni NS), ZO-MOPI adotta una strategia di "ortogonalizzazione parziale". Si concentra esclusivamente sulle top-$k$ direzioni singolari dominanti, scartando la coda spettrale rumorosa.

   • Meccanismo: Il metodo sostituisce la procedura NS con una Iterazione di Potenza in Streaming (SPI) più rapida e concentrata. SPI traccia il sottospazio di dimensione $k$ principale durante i passaggi di addestramento.
   • Efficienza: Memorizzando il sottospazio singolare destro ($V_{t-1}$) e utilizzandolo come avvio caldo per l'iterazione successiva, SPI evita la SVD completa o costose moltiplicazioni di matrici. Il costo computazionale è ridotto a $O(rnk + nk^2)$, dove $k \ll r \ll \min(m, n)$, rispetto al costo più elevato dei metodi a spettro completo.
   • Stabilità: Il metodo si basa sull'osservazione che le direzioni spettrali dominanti evolvono in modo regolare sotto l'effetto della smoothing del momentum, rendendo il sottospazio precedente un'inizializzazione affidabile per il passo corrente.

2. Stima ZO nel Sottospazio:
   Per ridurre ulteriormente la varianza, la stima del gradiente è vincolata a un sottospazio a rango basso. La matrice di perturbazione casuale $Z$ è decomposta in $AB$, dove $A$ è una matrice di sottospazio fissa e $B$ è una perturbazione casuale all'interno di tale sottospazio. Ciò riduce la dimensione di ricerca effettiva da $mn$a$rn$, abbassando significativamente la varianza dello stimatore del gradiente.

3. Momentum Proiettato:
   Il momentum è mantenuto all'interno del sistema di coordinate ridotto $r \times n$ invece che nello spazio dei parametri completo per risparmiare memoria. Una componente critica del metodo è la proiezione del momentum quando il sottospazio $A$ viene aggiornato periodicamente (ogni $\nu$ iterazioni). Il momentum $M_{old}$ viene proiettato sul nuovo sottospazio $A_{new}$ tramite $M_{new} \leftarrow A_{new}^\top A_{old} M_{old}$.

   • Necessità: Gli autori dimostrano che il momentum non serve meramente alla riduzione della varianza, ma è essenziale per la stabilità della SPI. Senza momentum, il sottospazio dominante fluttua drasticamente a causa delle stime ZO rumorose, causando l'instabilità dell'iterazione di potenza.

Contributi Chiave

• Identificazione della Disallineamento Spettrale: Il documento evidenzia che l'ortogonalizzazione a spettro completo (Newton-Schulz) è dannosa negli scenari ZO perché amplifica il rumore intrinseco nelle stime del gradiente ZO, mentre i gradienti FO sono naturalmente concentrati.
• Ortogonalizzazione Parziale: L'introduzione di una strategia di ortogonalizzazione selettiva che preserva solo i componenti spettrali più affidabili (top-$k$) ignorando la coda rumorosa.
• Iterazione di Potenza in Streaming (SPI): Un algoritmo leggero ed efficiente per tracciare le direzioni singolari dominanti che sfrutta la continuità temporale e l'avvio caldo per evitare calcoli costosi.
• Integrazione del Momentum Proiettato: Un meccanismo per mantenere la stabilità del momentum attraverso gli aggiornamenti del sottospazio, dimostrato essere critico per l'efficacia del tracciamento spettrale in ambienti rumorosi.

Risultati Sperimentali
Gli autori valutano ZO-MOPI su compiti di fine-tuning di LLM utilizzando i modelli OPT-13B, LLaMA3-8B e Gemma2-2B su benchmark SuperGLUE (SST-2, RTE, BoolQ, SQuAD).

• Velocità di Convergenza: ZO-MOPI raggiunge una convergenza da 1,5 a 4 volte più rapida rispetto allo stato dell'arte attuale ZO-Muon. Ad esempio, sui compiti RTE e SQuAD, ZO-Muon richiede circa 3 volte il tempo reale per raggiungere la stessa accuratezza di ZO-MOPI.
• Accuratezza Finale: Il metodo raggiunge accuratezze finali competitive o superiori rispetto a baseline solide come MeZO, LOZO e ZO-Muon. In diversi casi (ad es. OPT-13B su SST-2 e RTE), ZO-MOPI stabilisce nuovi record di accuratezza ZO.
• Efficienza: Il metodo mantiene l'efficienza di memoria del fine-tuning ZO, con un sovraccarico di memoria aggiuntivo trascurabile per la memorizzazione del sottospazio e degli stati dell'iterazione di potenza.
• Studi di Ablazione: Gli esperimenti confermano che la rimozione del momentum o l'uso di un intervallo di aggiornamento del sottospazio troppo breve degrada significativamente le prestazioni, validando la necessità della proiezione del momentum proposta e delle strategie di campionamento pigro del sottospazio.

Significato e Affermazioni
Il documento afferma che ZO-MOPI colma con successo il divario di prestazioni tra l'ottimizzazione ZO e FO per il fine-tuning degli LLM affrontando le caratteristiche spettrali specifiche delle stime di gradiente rumorose. Gli autori sostengono che la chiave per accelerare l'ottimizzazione ZO non risiede semplicemente nella riduzione della varianza tramite proiezione del sottospazio, ma anche nell'adozione di un allineamento spettrale selettivo che eviti di amplificare il rumore. Sostituendo l'ortogonalizzazione aggressiva a spettro completo con un'ortogonalizzazione parziale tramite iterazione di potenza, il metodo consente agli ottimizzatori spettrali di funzionare efficacemente nel regime ZO ad alta varianza. Il lavoro suggerisce che la struttura spettrale può essere meglio sfruttata anche sotto stime rumorose, offrendo una via verso un fine-tuning efficiente e vincolato alla memoria di modelli di grandi dimensioni senza fare affidamento sulla retropropagazione.