ZO-Act: Efficient Zeroth-Order Fine-Tuning via One-Shot Activation-Informed Low-Rank Subspaces

Riepilogo Tecnico: ZO-Act

Problema

Il fine-tuning dei Large Language Models (LLM) si basa tipicamente sull'ottimizzazione del primo ordine (backpropagation), che diventa proibitiva in termini di memoria man mano che i modelli scalano a miliardi di parametri a causa dei requisiti di archiviazione per attivazioni, stati dell'ottimizzatore e computazioni del gradiente. L'ottimizzazione del ordine zero (Zeroth-order, ZO) offre un'alternativa basata solo sul forward pass, stimando le direzioni di aggiornamento utilizzando solo le valutazioni della perdita (loss), evitando così la backpropagation. Tuttavia, i metodi ZO esistenti affrontano sfide significative:

1. Alta Varianza: Perturbare l'intero spazio dei parametri (come in MeZO) risulta in stime del gradiente ad alta varianza, portando a una convergenza lenta e instabilità.
2. Sottospazi Subottimali: Sebbene alcuni metodi limitino le perturbazioni a sottospazi a bassa dimensione (ad es., LOZO, SubZero), essi spesso si affidano a sottospazi costruiti casualmente o a perturbazioni casuali che potrebbero non allinearsi con il reale flusso del gradiente.
3. Compatibilità con l'Ottimizzatore: Molti metodi ZO faticano a utilizzare efficacemente gli standard ottimizzatori basati sul momentum (come Adam), poiché la struttura della perturbazione non espone variabili addestrabili esplicite e leggere.
4. Vincoli di Quantizzazione: Il fine-tuning di LLM quantizzati (ad es., INT4) è difficile con metodi che richiedono la materializzazione di perturbazioni dell'intero peso.

Metodologia: ZO-Act

Gli autori propongono ZO-Act, un metodo di fine-tuning ZO a colpo singolo (one-shot) informato dalle attivazioni. L'intuizione centrale è che le componenti dominanti dei gradienti dei pesi nei layer lineari risiedono nel sottospazio spaziato dai vettori singolari destri superiori della matrice di attivazione di input.

Meccanismo Centrale

1. Inizializzazione One-Shot: Per ogni layer lineare, ZO-Act esegue un singolo passaggio forward su un batch di calibrazione per calcolare la matrice di attivazione di input $X$. Successivamente, calcola i primi $r$ vettori singolari destri ($V_r$) di $X$ tramite una SVD (Singular Value Decomposition) sottile.
2. Parametrizzazione del Sottospazio Fisso: L'aggiornamento del peso è limitato a un sottospazio a basso rango definito da questi vettori di attivazione congelati. L'aggiornamento effettivo del peso è parametrizzato come:
   $$\Delta W = V_r B$$
   dove $V_r \in \mathbb{R}^{m \times r}$ è la base di attivazione congelata, e $B \in \mathbb{R}^{r \times n}$ è una matrice di coefficienti leggera e addestrabile. I pesi pre-addestrati originali $W$ e la base $V_r$ rimangono congelati durante tutto l'addestramento.
3. Ottimizzazione nello Spazio dei Coefficienti: Invece di perturbare l'intera matrice dei pesi $W$, ZO-Act campiona perturbazioni casuali $Z$ nello spazio dei coefficienti a bassa dimensione ($Z \sim \mathcal{N}(0, I)$). Il peso effettivo perturbato diventa $W_{eff} = W + V_r(B + \mu Z)$.
4. Stima Solo Forward: Il gradiente per $B$ viene stimato utilizzando valutazioni della perdita a differenza finita (forward-difference):
   $$\hat{g}_B = \frac{L(B + \mu Z) - L(B)}{\mu} Z$$
   Questo gradiente stimato viene poi utilizzato per aggiornare $B$ usando ottimizzatori standard del primo ordine come Adam.

Analisi Teorica

Il paper analizza ZO-Act come un'ottimizzazione di ordine zero su uno spazio di coefficienti ristretto $\beta$ (dimensione $k$) piuttosto che sull'intero spazio dei pesi (dimensione $d$).

• Riduzione della Varianza: La varianza dello stimatore ZO scala linearmente con la dimensione della perturbazione. Riducendo la dimensione da $d$ (pesi completi) a $k = \sum r n_\ell$ (coefficienti, dove $r \ll m$), ZO-Act riduce significativamente il termine di convergenza dipendente dalla varianza e l'errore di differenza finita.
• Trade-off Bias-Varianza: Limitare gli aggiornamenti al sottospazio informato dalle attivazioni introduce un bias di approssimazione del sottospazio. Tuttavia, gli autori sostengono che le attivazioni e i gradienti degli LLM presentano strutture a basso rango, il che significa che le direzioni di aggiornamento dominanti sono ben catturate dai vettori singolari superiori delle attivazioni di input. Pertanto, il bias è controllato e mitigato dalla natura dei dati del sottospazio.

Contributi Chiave

1. Sottospazio Informato dalle Attivazioni: A differenza dei precedenti metodi di sottospazio che utilizzano matrici a basso rango casuali o statiche, ZO-Act deriva il sottospazio di perturbazione direttamente dalle attivazioni di input, garantendo che il sottospazio sia allineato con le direzioni di attivazione dominanti del layer.
2. Ottimizzazione Esplicita dei Coefficienti: Congelando la base e ottimizzando solo la matrice dei coefficienti $B$, ZO-Act trasforma il fine-tuning ZO in un problema di ottimizzazione esplicito su uno spazio a bassa dimensione. Ciò consente l'applicazione diretta di ottimizzatori basati sul momentum (ad es., Adam) ed evita la materializzazione di perturbazioni dell'intero peso.
3. Compatibilità con la Quantizzazione: Il metodo supporta naturalmente gli LLM quantizzati (ad es., INT4) perché i pesi a basso bit rimangono congelati e vengono aggiornati solo i coefficienti a basso bit.
4. Garanzie Teoriche: Il paper fornisce un'analisi di convergenza che mostra come ZO-Act riduca sia la varianza dello stimatore del gradiente che l'errore di differenza finita rispetto al ZO a pieno peso, a costo di un bias di sottospazio gestibile.

Risultati Sperimentali

Gli autori hanno valutato ZO-act su Llama-3-8B, OPT-13B e un modello Llama-3-8B quantizzato in INT4 attraverso compiti di comprensione del linguaggio, risposta a domande e ragionamento di senso comune.

• Performance: ZO-Act ha superato costantemente i forti baseline ZO (inclusi MeZO, LOZO, SubZero, AGZO e ZO-Muon).
  • Su Llama-3-8B, ha ottenuto i migliori risultati ZO a precisione intera su SST-2, RTE, BoolQ e WiC. In particolare, ha migliorato le prestazioni su RTE da 81.2 (ZO-Muon) a 87.0 e su CB da 69.6 a 89.3.
  • Su OPT-13B, ha ottenuto risultati simili ai vertici su SST-2, RTE, CB e BoolQ.
  • Sui Modelli Quantizzati in INT4, ZO-Act è rimasto competitivo con i baseline ZO a precisione intera, dimostrando la sua efficacia per il fine-tuning con vincoli di memoria.
• Efficienza: ZO-Act ha eguagliato l'uso di memoria più basso tra i metodi confrontati (circa 16.1 GB per Llama-3-8B, scendendo a 5.8 GB per INT4) e ha raggiunto il tempo di esecuzione più veloce (42.0 minuti rispetto ai 44.5+ minuti degli altri) sotto un budget fisso di query forward.
• Sensibilità al Rango: Gli esperimenti hanno mostrato che un rango di $r=1$ era spesso sufficiente per i modelli a precisione intera, fornendo i migliori risultati di accuratezza e minimizzando la dimensione della perturbazione. Rangi più alti ($r=32$) sono stati benefici per i modelli quantizzati per compensare la ridotta capacità.
• Stabilità: Studi diagnostici hanno confermato che il sottospazio informato dalle attivazioni rimane stabile durante il fine-tuning e cattura una parte significativa dell'energia del gradiente completo, validando il design di inizializzazione "one-shot".

Significato e Rivendicazioni

Il paper sostiene che ZO-Act rappresenti un passo avanti significativo nel rendere il fine-tuning di ordine zero un'alternativa pratica e ad alte prestazioni per i grandi LLM e i modelli quantizzati. La sua importanza risiede nel:

• Colmare il Divario: Riduce il divario di performance tra il fine-tuning di ordine zero e quello del primo ordine riducendo la varianza dello stimatore attraverso sottospazi informati dai dati.
• Abilitare gli Ottimizzatori: Risolve la difficoltà di applicare ottimizzatori standard come Adam ai metodi ZO esponendo variabili addestrabili esplicite a bassa dimensione.
• Scalabilità: Fornisce una soluzione efficiente dal punto di vista della memoria per il fine-tuning di modelli quantizzati, fondamentale per il deployment su hardware con risorse limitate.

Gli autori riconoscono i limiti, notando che ZO-Act non colma ancora completamente il divario di performance con i metodi del primo ordine a causa del rumore intrinseco della stima ZO e del bias del sottospazio. Inoltre, ZO-Act rimane più lento in termini di tempo di esecuzione (wall-clock time) rispetto ai metodi basati sulla backpropagation a causa dei molteplici passaggi forward richiesti per ogni aggiornamento. Tuttavia, il metodo si posiziona come una scelta superiore per scenari in cui la backpropagation è proibitiva o non disponibile in termini di memoria.