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

Technische Samenvatting: Versnelling van Zeroth-Order Spectrale Optimalisatie met Partiële Orthogonalisatie via Power Iteration

Probleemstelling
Zeroth-order (ZO) optimalisatie is een cruciale techniek geworden voor het fijnafstemmen van Large Language Models (LLM's) op geheugenbeperkte randapparaten, omdat het de geheugenintensieve backpropagation vermijdt die nodig is bij first-order (FO) methoden. Echter, ZO-methoden lijden onder een hoge variantie in gradiëntschattingen door Monte Carlo-perturbaties, wat leidt tot aanzienlijk langzamere convergentie in vergelijking met FO-methoden.

Recente pogingen om ZO-optimalisatie te versnellen hebben zich gericht op zoektochten in laagdimensionale deelruimten (bijv. LOZO, SubZero) en spectrale optimalisatoren zoals Muon, die matrixgeometrie benutten via orthogonalisatie. Hoewel Muon AdamW in FO-omgevingen overtreft door gradiënten te orthogonaliseren om zwakke spectrale richtingen te benutten, stellen de auteurs vast dat het direct toepassen van Muon op ZO-omgevingen beperkte winst oplevert. Het kernprobleem dat is geïdentificeerd, is dat ZO-gradiëntschattingen inherent ruisachtig en onbetrouwbaar zijn. In tegenstelling tot FO-gradiënten, die geconcentreerd zijn in dominante singuliere richtingen, vertonen ZO-schattingen een "zwaardere en plattere spectrale staart" (Figuur 1). Bijgevolg versterkt de standaard Newton-Schulz (NS) orthogonalisatie van Muon, die energie over het volledige actieve spectrum egaliseert, per ongeluk deze ruisachtige, onnauwkeurige richtingen, waardoor schattingsfouten worden omgezet in schadelijke updatecomponenten.

Methodologie: ZO-MOPI
De auteurs stellen ZO-MOPI (Zeroth-Order Momentum met Partiële Orthogonalisatie) voor, een methode die is ontworpen om ZO-optimalisatie te versnellen door selectief dominante spectrale richtingen te extraheren en ruis te onderdrukken. De aanpak integreert drie kerncomponenten:

1. Partiële Orthogonalisatie via Streaming Power Iteration (SPI):
   In plaats van volledige-spectrum orthogonalisatie (zoals NS-iteraties) uit te voeren, hanteert ZO-MOPI een strategie van "partiële orthogonalisatie". Het richt zich uitsluitend op de top-$k$ dominante singuliere richtingen en verwijdert de ruisachtige spectrale staart.

   • Mechanisme: De methode vervangt de NS-procedure door een snellere, meer geconcentreerde Streaming Power Iteration (SPI). SPI volgt de leidende $k$-dimensionale deelruimte over trainingsstappen heen.
   • Efficiëntie: Door het opslaan van de rechter singuliere deelruimte ($V_{t-1}$) en het gebruik daarvan als warm start voor de volgende iteratie, vermijdt SPI volledige SVD of dure matrixvermenigvuldigingen. De rekenkosten worden gereduceerd tot $O(rnk + nk^2)$, waarbij $k \ll r \ll \min(m, n)$, in vergelijking met de hogere kosten van volledige-spectrum methoden.
   • Stabiliteit: De methode vertrouwt op de observatie dat dominante spectrale richtingen zich soepel ontwikkelen onder momentum-smoothing, waardoor de vorige deelruimte een betrouwbare initialisatie vormt voor de huidige stap.

2. Subspace ZO-schatting:
   Om de variantie verder te verlagen, wordt de gradiëntschatting beperkt tot een laag-rang deelruimte. De willekeurige perturbatiematrix $Z$ wordt ontbonden in $AB$, waarbij $A$ een vaste deelruimtematrix is en $B$ een willekeurige perturbatie binnen die deelruimte. Dit verlaagt de effectieve zoekdimensionaliteit van $mn$ naar $rn$, wat de variantie van de gradiëntschatter aanzienlijk verlaagt.

3. Geprojecteerd Momentum:
   Momentum wordt onderhouden binnen het gereduceerde $r \times n$ coördinatenstelsel in plaats van de volledige parameter ruimte om geheugen te besparen. Een kritiek onderdeel van de methode is de projectie van momentum wanneer de deelruimte $A$ periodiek wordt ververst (elke $\nu$ iteraties). Het momentum $M_{old}$ wordt geprojecteerd op de nieuwe deelruimte $A_{new}$ via $M_{new} \leftarrow A_{new}^\top A_{old} M_{old}$.

   • Noodzaak: De auteurs tonen aan dat momentum niet alleen dient voor variantiereductie, maar essentieel is voor de stabiliteit van SPI. Zonder momentum fluctueert de dominante deelruimte drastisch door de ruisachtige ZO-schattingen, waardoor de power iteration instabiel wordt.

Belangrijkste Bijdragen

• Identificatie van Spectrale Mismatch: Het artikel benadrukt dat volledige-spectrum orthogonalisatie (Newton-Schulz) schadelijk is in ZO-omgevingen omdat het de ruis inherent aan ZO-gradiëntschattingen versterkt, terwijl FO-gradiënten van nature geconcentreerd zijn.
• Partiële Orthogonalisatie: De introductie van een selectieve orthogonalisatiestrategie die alleen de meest betrouwbare spectrale componenten (top-$k$) behoudt en de ruisachtige staart negeert.
• Streaming Power Iteration (SPI): Een lichtgewicht, efficiënt algoritme voor het volgen van dominante singuliere richtingen dat gebruikmaakt van temporele continuïteit en warm-starting om dure berekeningen te vermijden.
• Integratie van Geprojecteerd Momentum: Een mechanisme om momentumstabiliteit te handhaven tijdens deelruimteverversingen, wat blijkt cruciaal te zijn voor de effectiviteit van spectrale tracking in ruisachtige omgevingen.

Experimentele Resultaten
De auteurs evalueren ZO-MOPI op LLM-fijneinstellingstaken met behulp van OPT-13B, LLaMA3-8B en Gemma2-2B-modellen over SuperGLUE-benchmarks (SST-2, RTE, BoolQ, SQuAD).

• Convergentiesnelheid: ZO-MOPI bereikt 1,5x tot 4x snellere convergentie in vergelijking met de huidige state-of-the-art ZO-Muon. Bijvoorbeeld, op de RTE- en SQuAD-taken vereist ZO-Muon ongeveer 3x de wandkloktijd om dezelfde nauwkeurigheid te bereiken als ZO-MOPI.
• Eindnauwkeurigheid: De methode bereikt concurrerende of superieure eindnauwkeurigheden in vergelijking met sterke baselines zoals MeZO, LOZO en ZO-Muon. In verschillende gevallen (bijv. OPT-13B op SST-2 en RTE) stelt ZO-MOPI nieuwe beste ZO-nauwkeurigheidsrecords.
• Efficiëntie: De methode behoudt de geheugenefficiëntie van ZO-fijneinstelling, met slechts verwaarloosbare extra geheugenoverhead voor het opslaan van de deelruimte en power iteration-statussen.
• Ablatiestudies: Experimenten bevestigen dat het verwijderen van momentum of het gebruik van een te korte deelruimteverversingsinterval de prestaties aanzienlijk verslechtert, wat de noodzaak bevestigt van de voorgestelde momentumprojectie en strategieën voor lazy subspace sampling.

Betekenis en Claims
Het artikel claimt dat ZO-MOPI succesvol de prestatiekloof tussen ZO- en FO-optimalisatie voor LLM-fijneinstelling overbrugt door de specifieke spectrale kenmerken van ruisachtige gradiëntschattingen aan te pakken. De auteurs betogen dat de sleutel tot het versnellen van ZO-optimalisatie niet alleen het verminderen van variantie door middel van deelruimteprojectie is, maar ook het aannemen van een selectieve spectrale uitlijning die het versterken van ruis vermijdt. Door agressieve volledige-spectrum orthogonalisatie te vervangen door partiële orthogonalisatie via power iteration, maakt de methode het mogelijk dat spectrale optimalisatoren effectief functioneren in het hoog-variantie ZO-regime. Het werk suggereert dat spectrale structuur zelfs onder ruisachtige schattingen beter kan worden benut, wat een weg biedt naar efficiënte, geheugenbeperkte fijneinstelling van grote modellen zonder afhankelijkheid van backpropagation.