The Loss Does Not See the Basis, but Adam Does

Technische Samenvatting: "De verliesfunctie ziet de basis niet, maar Adam wel"

Probleemstelling
Gradient descent (GD) op gefactoreerde modellen ($W = UV^\top$) staat erom bekend een impliciete bias te vertonen naar laag-rang oplossingen, waarbij het vaak convergeert naar minimum-nucleaire-norm interpolanten in ondergedetermineerde matrix-sensing taken. Echter, adaptieve optimizers zoals Adam, ondanks dat ze vergelijkbare trainingsverliezen bereiken, falen vaak in het terugvinden van deze laag-rang structuren, wat leidt tot oplossingen met een hogere effectieve rang en slechtere generalisatie. Hoewel recent werk heeft geïdentificeerd dat adaptieve methoden basis-afhankelijk zijn (gevoelig voor coördinatenrotaties), is het specifieke mechanisme waarmee deze afhankelijkheid de laag-rang inductieve bias in gefactoreerde modellen doorbreekt, nog niet gekarakteriseerd. Dit artikel adresseert de vraag: Is er een criterium om te bepalen of een optimizer de laag-rang bias behoudt of doorbreekt, en kan dit criterium de ingezette methoden sorteren op basis van hun herstelprestaties?

Methodologie en Theoretisch Kader
De auteur introduceert gauge-equivariantie als het centrale criterium voor het classificeren van optimizers. De verliesfunctie $L(U,V) = f(UV^\top)$ bezit een gauge-symmetrie: deze is invariant onder de transformatie $(U,V) \mapsto (UQ, VQ)$ voor elke orthogonale matrix $Q \in O(k)$. Deze symmetrie impliceert dat het roteren van de latente factor-bases de voorspellingen van het model niet verandert.

1. Classificatie van Update-regels:

   • Equivariante Optimizers: Een optimizer is gauge-equivariant als de update-regel commuteert met de gauge-transformatie. Het paper bewijst dat geheugenloze update-regels equivariant zijn als en slechts als ze de vorm hebben van een Gram-bepaalde linker preconditioner: $\Delta = H(GG^\top)G$, waarbij $H$ alleen afhangt van de Gram-matrix $GG^\top$.
   • Niet-Equivariante Optimizers: Optimizers die vaste entrywise nonlineariteiten toepassen op de gradiënt of de optimizer-toestand bij initialisatie (bijv. standaard Adam, RMSProp, Lion, SignSGD), breken deze symmetrie. Specifiek faalt de coördinatie-gewijze tweede-moment adaptatie van Adam om te commueren met orthogonale rotaties, wat leidt tot een "eerste-stap defect" waarbij de productmatrix $W$ onmiddellijk verandert bij het roteren van de initialisatie.

2. Transferstelling:
   De auteur stelt een transferstelling vast die aantoont dat optimizers met een gedeelde scalaire preconditioner (een enkele scalaire noemer voor alle coördinaten) overeenkomen met gradient flow onder een geschaalde tijdparameterisering. Bijgevolg zijn theoretische resultaten over de laag-rang bias van gradient flow (bijv. Gunasekar et al., 2017; Li et al., 2021) direct van toepassing op deze equivariante methoden.

3. Spectrale Schedule:
   Hoewel equivariantie noodzakelijk is om toegang te krijgen tot de laag-rang manifold, is het niet voldoende voor laag-rang herstel. Het specifieke gedrag binnen de equivariante klasse wordt bepaald door de spectrale schedule (hoe singular-waarden groeien). Agressieve gelijke-snelheid schedules (zoals Muon) kunnen exacte laag-rang targets herstellen, maar kunnen falen als het target een significante "spectrale staart" heeft (energie buiten het laag-rang deel), terwijl gulzige (greedy) schedules (zoals GD) deze staarten onderdrukken.

Belangrijkste Bijdragen en Experimenten

• Vergelijking van de Optimizer Zoo: De auteur evalueerde negen optimizers op een ondergedetermineerde matrix-sensing taak ($40 \times 40$, rang 3, $m=2 \times \text{dof}$) tot volledige interpolatie.

  • Resultaat: Er ontstond een duidelijke scheiding. Equivariante methoden (GD, Scalar-Adam, Muon, Shampoo) bereikten herstelfouten $\le 0.286$. Coördinatie-gewijze methoden (Adam, RMSProp, Lion, SignSGD, Adafactor) slaagden er niet in de laag-rang attractor te bereiken, met fouten $\ge 0.42$.
  • Prestaties van Muon: Muon bereikte bijna exact herstel ($6.8 \times 10^{-6}$ fout) op zuivere laag-rang targets, maar degradeerde snel naarmate de spectrale staart van het target toenam, waarbij het de prestaties van GD kruiste bij ongeveer 4% staartenergie.

• Het "Dial" Experiment: Om de oorzaak van het falen van Adam te isoleren, construeerde de auteur een een-parameter familie van optimizers (Adam-$p$) die interporeert tussen standaard Adam ($p=1$, coördinatie-gewijs) en Scalar-Adam ($p=0$, gedeelde scalaire).

  • Resultaat: Naarmate $p$ afnam van 1 naar 0, verbeterden de herstelfout en de effectieve rang monotoon. Dit bevestigde dat coördinatie-gewijze anisotropie in de preconditioner de specifieke methode is die de laag-rang bias vernietigt, onafhankelijk van momentum of adaptiviteit.

• Analyse van Attention Heads: Het paper demonstreert dat individuele attention heads in transformers dezelfde gauge-symmetrie delen ($W_Q^\top W_K$).

  • Resultaat: Het trainen van twee gauge-equivalente transformer kopieën (identieke functie, verschillende bases) met Adam zorgde ervoor dat ze structureel uiteenliepen na een enkele stap, wat resulteerde in een relatieve Frobenius-afstand van 56% in hun head-invarianten ($W_Q^\top W_K$). Equivariante optimizers (SGD, Scalar-Adam) bleven ononderscheidbaar tot op machineprecisie. De auteur merkt op dat hoewel alignment-technieken (zoals permutation re-basin) de gauge kunnen verzoenen, zij de invarianten zelf niet kunnen verzoenen als deze verschillen; dus verwijdert basis-afhankelijke training de overeenkomst die vereist is voor model merging, omdat geen enkele per-head rotatie de kloof tussen de uiteenlopende invarianten kan dichten.

• Validatie in de Praktijk: Op hyperspectrale beeld-completie taken (Indian Pines, Pavia University) bij een gelijkwaardig trainingsverlies:

  • GD behaalde 43–44% lagere helderheid-fouten dan Adam bij lage bemonsteringsdensiteiten, wat overeenkomt met een significant lagere effectieve rang.
  • De testfout van Adam nam toe naarmate deze dieper paste (overfitting op de spectrale staart), terwijl de fout van GD afnam.

• Algoritme Reparatie: De auteur paste hun criterium toe op hun eerdere werk, FlowAdam. Ze identificeerden dat coördinatie-gewijze velocity clipping de equivariantie doorbrak. Het vervangen hiervan door een globale-norm clip herstelde de bias, wat de herstelfout met 37% verbeterde en de methode dichter bij de prestaties van GD bracht.

Betekenis en Claims
Het paper claimt dat de keuze van de basis geen tuning-detail is, maar een fundamentele beslissing over welke interpolant de optimizer selecteert.

• Mechanistische Verklaring: Het biedt een testbaar, constructief criterium (gauge-equivariantie) om te voorspellen of een optimizer de laag-rang inductieve bias van gefactoreerde modellen behoudt.
• Beperkingen: De auteur is bescheiden over de reikwijdte. De bias is mild en voordelig primair wanneer de optimale regularisatie laag is en uitgelijnd met de laag-rang structuur. In regimes die sterke regularisatie of per-coördinaat adaptatie vereisen (bijv. stijve problemen in PINNs), kunnen coördinatie-gewijze methoden nog steeds beter presteren dan equivariante methoden.
• Geen Benchmark Claims: Het werk pleit niet voor een nieuwe state-of-the-art optimizer, maar biedt eerder een mechanistische verklaring voor geobserveerde fenomenen, waardoor duidelijk wordt waarom bepaalde methoden (zoals Adam) falen in het herstellen van laag-rang structuren in specifieke settings, terwijl anderen (zoals GD of Muon) dat wel doen.

Samenvattend betoogt het paper dat "de verliesfunctie de basis niet ziet" (het is gauge-invariant), maar "Adam dat wel doet" (zijn updates zijn basis-afhankelijk), en dat deze afhankelijkheid de wortel is van het falen om laag-rang interpolanten te selecteren in gefactoreerde modellen.