BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization

Sintesi Tecnica: BROS: Sottospazi Randomizzati Corretti per Bias per l'Ottimizzazione Bilevel a Ciclo Singolo Efficiente in Memoria

Enunciato del Problema
L'ottimizzazione bilevel stocastica (SBO) è un framework fondamentale per applicazioni di deep learning come l'ottimizzazione degli iperparametri, la ricampionatura dei dati, l'apprendimento di rappresentazioni e la pulizia degli iper-dati. La formulazione standard coinvolge una variabile di livello superiore $x$ e una variabile di livello inferiore $Y$ (tipicamente i parametri di una rete neurale), dove l'obiettivo di livello superiore dipende dalla soluzione del problema di ottimizzazione di livello inferiore.

I metodi SBO a ciclo singolo esatti esistenti (ad esempio MA-SOBA) forniscono forti garanzie di convergenza (complessità del campione $O(\epsilon^{-2})$) ma soffrono di gravi colli di bottiglia nella memoria quando la rete di livello inferiore è grande. Questi metodi si basano su prodotti vettore-Hessiano (HVP) e prodotti vettore-Jacobiano (JVP) a spazio completo per mantenere variabili ausiliarie, richiedendo una memoria proporzionale al conteggio completo dei parametri. Al contrario, i metodi surrogati efficienti in memoria (ad esempio metodi di penalità, approcci di ordine zero) riducono l'utilizzo della memoria ma spesso sacrificano garanzie di convergenza esatta, richiedono assunzioni più forti o ottimizzano un problema surrogato anziché l'obiettivo bilevel originale.

Metodologia: BROS
Gli autori propongono BROS (Bilevel optimization via RandOmized Subspaces), un metodo SBO a ciclo singolo efficiente in memoria progettato per problemi parametrizzati da matrici multistrato. BROS affronta il collo di bottiglia della memoria eseguendo aggiornamenti di livello inferiore e di livello ausiliario all'interno di sottospazi randomizzati a bassa dimensionalità, preservando al contempo la struttura esatta del ipergradiente a ciclo singolo.

I componenti fondamentali di BROS sono:

1. Decomposizione in Sottospazi Randomizzati:
   Invece di aggiornare i parametri completi di livello inferiore $Y \in \mathbb{R}^{m \times n}$, BROS campiona un proiettore Haar scalato $P \in \mathbb{R}^{m \times r}$ (dove $r \ll m$) ad ogni iterazione. L'aggiornamento di livello inferiore è calcolato nel sottospazio ridotto $B \in \mathbb{R}^{r \times n}$ tramite l'obiettivo proiettato $\tilde{g}(x, B) = g(x, Y + PB)$. La variabile a spazio completo viene aggiornata come $Y^{k+1} = Y^k + P B^k$. Ciò riduce l'impronta di memoria di gradienti, attivazioni e direzioni ausiliarie da $O(mn)$a$O(rn)$.

2. Aggiornamenti Ausiliari Corretti per Bias:
   Un approccio ingenuo "schizza-poi-alza" (sketch-then-lift) per la variabile ausiliaria $Z$ (che risolve il sistema lineare che coinvolge l'Hessiano) introduce un bias perché l'azione dell'Hessiano alzata $P \tilde{H} [P^\top Z]$ non è uno stimatore non distorto dell'azione dell'Hessiano a spazio completo $H[Z]$. Questo bias distorce l'ipergradiente e può portare a limiti non stazionari.
   Per risolvere ciò, BROS introduce una correzione bi-sonda Rademacher. L'algoritmo campiona vettori Rademacher indipendenti (vettori di segno) $u$ e $\xi$ per stimare termini di disturbo simili alla traccia e simili alla trasposta indotti dalla proiezione. Queste sonde permettono la costruzione di uno stimatore HVP corretto, indicato come $\widehat{H}[Z]$, che è non distorto condizionatamente per l'azione dell'Hessiano a spazio completo senza richiedere la formazione di matrici a spazio completo.

3. Ricorsione a Ciclo Singolo:
   BROS mantiene una ricorsione accoppiata a ciclo singolo sulla variabile superiore $x$, la variabile inferiore $Y$, la variabile ausiliaria $Z$ e uno stimatore di ipergradiente a media mobile $h$. Gli aggiornamenti utilizzano il gradiente inferiore proiettato e la direzione ausiliaria corretta per bias.

Contributi Chiave

• Ricorsione Efficiente in Memoria: BROS scompone le ricorsioni di livello inferiore e ausiliarie in sottospazi randomizzati, riducendo significativamente la dimensionalità dei gradienti memorizzati, delle attivazioni e dei calcoli HVP/JVP.
• Recupero Non Distorto dell'Ipergradiente: Il documento identifica il bias specifico indotto dalla proiezione nelle azioni dell'Hessiano ausiliario alzato e propone una correzione bi-sonda Rademacher a livello di strato. Ciò garantisce che lo stimatore rimanga non distorto per l'azione completa dell'Hessiano multistrato senza imporre assunzioni strutturali aggiuntive (ad esempio Hessiani a blocchi diagonali).
• Garanzie Teoriche: Sotto assunzioni standard (liscezza, convessità forte del livello inferiore), gli autori dimostrano che BROS preserva la complessità del campione $O(\epsilon^{-2})$ dei metodi esatti a ciclo singolo (specificamente MA-SOBA) per trovare un punto stazionario $\epsilon$. Il tasso di convergenza dipende dal rapporto di proiezione $\omega_P = \max(m_\ell/r_\ell)$, recuperando il tasso a spazio completo quando $r=m$.

Risultati Sperimentali
Gli autori valutano BROS su quattro compiti di apprendimento bilevel:

1. Pulizia degli Iper-Dati MNIST: BROS raggiunge un'accuratezza di test dell'84,30%, corrispondente da vicino alla linea di base MA-SOBA a spazio completo (84,46%) e superando i metodi basati su penalità e di ordine zero.
2. Apprendimento di Miscela di Dati (modelli stile GPT da 280M): BROS corrisponde alle prestazioni di MA-SOBA (log-perplessità finale 2,8109 contro 2,8098) riducendo la memoria di picco per l'addestramento proxy di circa 44,9% rispetto a MA-SOBA e del 27% rispetto ai metodi di Penalità.
3. Apprendimento di Iper-Rappresentazione CIFAR-10: BROS raggiunge un'accuratezza di test finale del 78,03%, comparabile a MA-SOBA (78,57%).
4. Ricampionamento di Campioni ViT CIFAR-100: BROS raggiunge la migliore accuratezza di validazione (49,88%) e di test (49,05%) tra i metodi confrontati, utilizzando solo 6,42 GB di memoria di picco rispetto a oltre 23 GB per le linee di base a spazio completo.

Significato e Affermazioni
Il documento afferma che BROS colma con successo il divario tra efficienza in memoria e convergenza teorica nell'ottimizzazione bilevel stocastica. A differenza dei metodi surrogati che scambiano esattezza per efficienza, BROS mantiene la struttura esatta di aggiornamento dell'ipergradiente a ciclo singolo e il suo ordine di convergenza associato $O(\epsilon^{-2})$. Utilizzando sottospazi randomizzati e un nuovo meccanismo di correzione del bias, BROS abilita compiti di ottimizzazione bilevel (come la ricampionatura dei dati e la regolazione degli iperparametri) su grandi reti neurali che altrimenti sarebbero non fattibili a causa di vincoli di memoria, riducendo l'utilizzo di memoria di picco fino al 44,9% mentre corrisponde da vicino alle prestazioni a spazio completo.