Technische Samenvatting: Zeta – Dual Whitening voor Matrixoptimalisatie via Coördinaat-Adaptieve Preconditionering

Probleemstelling

De training van grootschalige neurale netwerken is steeds meer verschoven naar matrix-bewuste optimizers die de structurele eigenschappen van gewichtsparameters benutten, waarbij ze verder gaan dan element-gewijze adaptatie. Een leidend voorbeeld is Muon, dat orthogonale restricties op gewichtsupdates afdwingt via spectrale whitening (specifiek, Newton–Schulz iteratie) om singuliere waarden te egaliseren en te voorkomen dat één enkele richting de update domineert.

Echter, dit paper identificeert een kritieke, onderbelichte kwetsbaarheid in Muon en vergelijkbare spectrale methoden: de Newton–Schulz iteratie is zeer gevoelig voor de conditionering van de input. Hoewel de iteratie goed functioneert wanneer singuliere waarden goed verdeeld zijn, vertonen ruwe momentummatrices in de diepe netwerktraining vaak een ernstige coördinaat-gewijze schaalheterogeniteit. In deze matrices bezitten enkele coördinaten disproportioneel grote magnituden, wat het singuliere spectrum scheef trekt. Wanneer dergelijke ongeconditioneerde matrices de Newton–Schulz stap binnenkomen, degradeert het orthogonalisatieproces, wat leidt tot verstoorde update-richtingen die de gradiëntgeometrie niet langer getrouw reflecteren. De auteurs stellen dat spectrale whitening impliciet een schaal-genormaliseerde input vereist, een conditie die ruw momentum systematisch schendt.

Methodologie: Zeta Optimizer

Om dit structurele gat te overbruggen, stellen de auteurs Zeta voor, een dual whitening optimizer die coördinaat-whitening en spectrale whitening organiseert in een strikt geordende, niet-instelbare pipeline. De kerninzicht is dat spectrale whitening moet worden gebouwd op statistische stabilisatie; coördinaat-whitening vestigt de statistische isotropie die nodig is om spectrale whitening betrouwbaar te laten functioneren.

De Zeta-pipeline werkt als volgt:

1. Coördinaat-Whitening (Preconditionering):
   Voordat het de spectrale fase ingaat, past Zeta een coördinaat-gewijze normalisatie toe op de momentummatrix. Het houdt exponentiële voortschrijdende gemiddelden bij van het eerste moment ($M_t$) en het tweede moment ($V_t$), vergelijkbaar met Adam. De momentummatrix wordt vervolgens element-gewijs genormaliseerd door het voortschrijdende tweede moment:
   $$\tilde{G}_t = \frac{M_t}{\sqrt{V_t} + \epsilon}$$
   Deze stap homogeniseert de schaalverschillen tussen coördinaten en transformeert de matrix naar een statistisch isotrope staat. Theoretisch fungeert dit als een benaderende diagonale preconditioner, die de schaalimbalans aanpakt die anders het singuliere spectrum zou scheef trekken.

2. Spectrale Whitening (Orthogonalisatie):
   De gehomogeniseerde matrix $\tilde{G}_t$ wordt vervolgens ingevoerd in de Newton–Schulz iteratie. Omdat de input nu vrij is van extreme entry-wise schaalverschillen, is het singuliere spectrum beter geconditioneerd. De iteratie convergeert naar een getrouwe orthogonale benadering, waardoor de singuliere waarden worden geëgaliseerd om een gebalanceerde update-richting te produceren.

3. Update Constructie:
   De resulterende orthogonale matrix $U_t$ wordt geschaald met een RMS-genormaliseerde factor om ervoor te zorgen dat de update-magnitude compatibel is met standaard optimizer-configuraties (bijv. AdamW takken) zonder de spectrale richting te veranderen.

4. Dual-Path Executie:
   Voor matrix-waardige parameters (bijv. attention en feed-forward projecties) wordt de volledige duale pipeline toegepast. Voor niet-matrix parameters (bijv. biases, LayerNorm schalen), wordt de spectrale stap overgeslagen en valt de optimizer terug op een AdamW-stijl update.

Theoretische Bijdragen

Het paper biedt drie belangrijke theoretische garanties ter ondersteuning van het ontwerp:

• Benaderende Whitening: De coördinaat-normalisatiestap wordt bewezen als een efficiënte diagonale benadering van de ideale whitening-transformatie (Theorem 1), wat verder gaat dan heuristische schaling naar een principiële preconditioneringstap.
• Verminderde Orthogonalisatiefout: De auteurs bewijzen dat coördinaat-whitening de condition number van de input voor de Newton–Schulz iteratie strikt verbetert. Bijgevolg wordt de orthogonalisatiefout van de spectrale fase strikt verminderd vergeleken met opereren op ruw momentum (Theorem 2).
• Schaalbehoud: De uiteindelijke RMS-schaling behoudt de spectrale richting van de update terwijl het de globale magnitude synchroniseert met de rest van het netwerk (Theorem 3).

Empirische Resultaten

Zeta werd geëvalueerd over taalmodellen (0.6B tot 8B parameters), mixture-of-experts (MoE) architecturen en visietaken, waarbij het consistent de sterke baselines (AdamW, Muon, en AdaMuon) evenaarde of overtrof.

• Convergentiesnelheid: In pretraining-taken demonstreerde Zeta snellere convergentie dan Muon en AdamW. Bijvoorbeeld, op een 1.7B parameter model bereikte Zeta een 1.67× versnelling (relatief aan AdamW) om hetzelfde verliesniveau te bereiken, waarmee het Muon (1.52×) en AdaMuon (1.61×) versloeg. Op een 8B model bereikte Zeta een 1.25× versnelling versus Muon's 1.19×.
• Efficiëntie: De per-iteratie tijdskosten van Zeta bleven vergelijkbaar met AdamW (bijv. 25.47s versus 24.46s voor het 8B model op 16 GPU's), wat aangeeft dat het dual whitening-proces geen significante runtime-penalties met zich meebrengt.
• MoE Prestaties: In MoE-architecturen, waar sparse expert activatie zorgt voor substantiële gradiëntheterogeniteit, behield Zeta een duidelijk voordeel en convergeerde het naar een lager verlies met minder fluctuaties dan concurrenten.
• Downstream Taken: Op 12 taalbegrip-benchmarks (bijv. MMLU, GSM8k) behaalden met Zeta getrainde modellen de hoogste gemiddelde nauwkeurigheid over zowel de 1.7B als 8B schalen. Eveneens op CIFAR-100 visietaken behaalde Zeta de hoogste Top-1 validatie-nauwkeurigheid voor zowel ViT-Tiny als ViT-Base.
• Hyperparameter Robuustheid: Een grid search op momentumcoëfficiënten ($\beta_1, \beta_2$) toonde aan dat Zeta robuust is tegen hyperparameterkeuzes, waarbij het uiteindelijke trainingverlies met minder dan 1.2% varieerde over het geteste bereik.

Betekenis en Claims

Het paper claimt dat Zeta een fundamentele mismatch in bestaande matrix-bewuste optimizers oplost: de toepassing van krachtige geometrische operaties (spectrale whitening) op signalen met ongecontroleerde coördinaat-gewijze schalen. Door de prevalentie van intra-matrix schaalimbalans te verifiëren via chi-kwadraat uniformiteitstesten en aan te tonen dat coördinaat-whitening dit effectief corrigeert, stellen de auteurs vast dat het oplossen van schaalimbalans vóór orthogonalisatie een vereiste is voor stabiele spectrale methoden.

De betekenis van Zeta ligt in de unificatie van twee voorheen concurrerende optimalisatieparadigma's (coördinaat-gewijze adaptieve schaling en spectrale orthogonalisatie) in een enkele, logisch noodzakelijke pipeline. De volgorde van operaties is geen instelbare hyperparameter maar een wiskundige afhankelijkheid: coördinaat-whitening moet voorafgaan aan spectrale whitening om te garanderen dat de input goed geconditioneerd is. Deze aanpak leidt tot snellere convergentie en betere generalisatie zonder de geometrische voordelen van matrix-niveau optimalisatie op te offeren.