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

Technische Samenvatting: BROS: Bias-gecorrigeerde Gerandomiseerde deelruimten voor geheugenefficiënte Single-Loop Bilevel Optimalisatie

Probleemstelling
Stochastische bilevel-optimalisatie (SBO) is een fundamenteel raamwerk voor deep learning-toepassingen zoals hyperparameteroptimalisatie, data herwaardering, representatieleren en hyper-data-cleaning. De standaardformulering omvat een bovenliggende variabele $x$ en een onderliggende variabele $Y$ (doorgaans de parameters van een neurale netwerk), waarbij het bovenliggende doel afhankelijk is van de oplossing van het onderliggende optimalisatieprobleem.

Bestaande exacte single-loop SBO-methoden (bijv. MA-SOBA) bieden sterke convergentiegaranties ($O(\epsilon^{-2})$ steekproefcomplexiteit), maar lijden onder ernstige geheugenflessenhals wanneer het onderliggende netwerk groot is. Deze methoden vertrouwen op Hessian-vectorproducten (HVP's) en Jacobiaan-vectorproducten (JVP's) in de volledige ruimte om hulpvariabelen bij te houden, waarbij het geheugenproportioneel is aan het totale aantal parameters. Omgekeerd verminderen geheugenefficiënte surrogate-methoden (bijv. straffingsmethoden, zeroth-order-benaderingen) het geheugengebruik, maar offeren vaak exacte convergentiegaranties op, vereisen sterkere aannames, of optimaliseren een surrogate-probleem in plaats van de oorspronkelijke bilevel-doelfunctie.

Methodologie: BROS
De auteurs stellen BROS (Bilevel-optimalisatie via GeRandomiseerde deelruimten) voor, een geheugenefficiënte single-loop SBO-methode ontworpen voor problemen met gelaagde matrix-geparametriseerde structuren. BROS adresseert de geheugenflessenhals door onderliggende en hulp-niveau-updates uit te voeren binnen randomiseerde laag-dimensionale deelruimten, terwijl de exacte single-loop hypergradiëntstructuur behouden blijft.

De kerncomponenten van BROS zijn:

1. Grandomiseerde deelruimte-decompositie:
   In plaats van de volledige onderliggende parameters $Y \in \mathbb{R}^{m \times n}$ bij te werken, steekt BROS bij elke iteratie een geschaalde Haar-projector $P \in \mathbb{R}^{m \times r}$ (waarbij $r \ll m$). De onderliggende update wordt berekend in de gereduceerde deelruimte $B \in \mathbb{R}^{r \times n}$ via het geprojecteerde doel $\tilde{g}(x, B) = g(x, Y + PB)$. De variabele in de volledige ruimte wordt bijgewerkt als $Y^{k+1} = Y^k + P B^k$. Dit vermindert het geheugenvoetafdruk van gradiënten, activaties en hulp-richtingen van $O(mn)$ naar $O(rn)$.

2. Bias-gecorrigeerde hulp-updates:
   Een naïeve "schets-en-heffen"-benadering voor de hulpvariabele $Z$ (die het lineaire systeem met de Hessian oplost) introduceert bias, omdat de opgeheven Hessian-actie $P \tilde{H} [P^\top Z]$ geen onbevooroordeelde schatter is van de Hessian-actie in de volledige ruimte $H[Z]$. Deze bias vervormt de hypergradiënt en kan leiden tot niet-stationaire limieten.
   Om dit op te lossen, introduceert BROS een Rademacher bi-probe correctie. Het algoritme steekt onafhankelijke Rademacher-vectoren (tekenvectoren) $u$ en $\xi$ om trace-achtige en transposiet-achtige storende termen geïnduceerd door de projectie te schatten. Deze probes maken het mogelijk een gecorrigeerde HVP-schatter te construeren, aangeduid als $\widehat{H}[Z]$, die conditioneel onbevooroordeeld is voor de Hessian-actie in de volledige ruimte zonder dat er een matrix in de volledige ruimte hoeft te worden gevormd.

3. Single-loop recursie:
   BROS handhaaft een gekoppelde single-loop recursie over de bovenliggende variabele $x$, onderliggende variabele $Y$, hulpvariabele $Z$ en een moving-average hypergradiënt-schatter $h$. De updates maken gebruik van de geprojecteerde onderliggende gradiënt en de bias-gecorrigeerde hulp-richting.

Belangrijkste bijdragen

• Geheugenefficiënte recursie: BROS decomposeert de onderliggende en hulp-recursies naar randomiseerde deelruimten, wat de dimensie van opgeslagen gradiënten, activaties en HVP/JVP-berekeningen aanzienlijk vermindert.
• Onbevooroordeelde herstel van hypergradiënt: Het artikel identificeert de specifieke projectie-geïnduceerde bias in opgeheven hulp-Hessian-acties en stelt een gelaagde Rademacher bi-probe correctie voor. Dit zorgt ervoor dat de schatter onbevooroordeeld blijft voor de volledige multilaag Hessian-actie zonder extra structurele aannames op te leggen (bijv. blok-diagonale Hessians).
• Theoretische garanties: Onder standaard aannames (gladheid, sterke convexiteit van het onderliggende niveau) bewijzen de auteurs dat BROS de $O(\epsilon^{-2})$ steekproefcomplexiteit van exacte single-loop-methoden (specifiek MA-SOBA) behoudt voor het vinden van een $\epsilon$-stationair punt. De convergentiesnelheid hangt af van de projectieverhouding $\omega_P = \max(m_\ell/r_\ell)$, waarbij de snelheid in de volledige ruimte wordt hersteld wanneer $r=m$.

Experimentele resultaten
De auteurs evalueren BROS op vier bilevel-leertaken:

1. MNIST Hyper-Data Cleaning: BROS bereikt 84,30% testnauwkeurigheid, wat nauw overeenkomt met de baseline MA-SOBA in de volledige ruimte (84,46%) en beter presteert dan op straffing gebaseerde en zeroth-order-methoden.
2. Data-Mixture Learning (280M GPT-achtige modellen): BROS komt overeen met de prestaties van MA-SOBA (eind log-perplexiteit 2,8109 versus 2,8098) terwijl het piekgeheugen voor proxy-training met ongeveer 44,9% wordt verlaagd ten opzichte van MA-SOBA en met 27% ten opzichte van straffingsmethoden.
3. CIFAR-10 Hyper-Representation Learning: BROS bereikt 78,03% eindtestnauwkeurigheid, vergelijkbaar met MA-SOBA (78,57%).
4. CIFAR-100 ViT Sample Reweighting: BROS bereikt de beste validatie (49,88%) en test (49,05%) nauwkeurigheid onder vergeleken methoden, met slechts 6,42 GB piekgeheugen in vergelijking met meer dan 23 GB voor baselines in de volledige ruimte.

Betekenis en claims
Het artikel claimt dat BROS succesvol de kloof overbrugt tussen geheugenefficiëntie en theoretische convergentie in stochastische bilevel-optimalisatie. In tegenstelling tot surrogate-methoden die exactheid opofferen voor efficiëntie, behoudt BROS de exacte single-loop hypergradiënt-updatestructuur en de bijbehorende $O(\epsilon^{-2})$ convergentieorde. Door gebruik te maken van randomiseerde deelruimten en een nieuw bias-correctiemechanisme, stelt BROS bilevel-optimalisatietaken (zoals data herwaardering en hyperparameter-tuning) op grote neurale netwerken mogelijk die anders onuitvoerbaar zouden zijn vanwege geheugenbeperkingen, met een vermindering van het piekgeheugengebruik tot 44,9% terwijl de prestaties in de volledige ruimte nauw worden benaderd.