Technische Samenvatting: Het verbeteren van neurale netwerktraining door de magnitude en richting van gewichtsvectoren te ontkoppelen

1. Probleemstelling

Moderne training van neurale netwerken vertrouwt op optimizers zoals Adam en Muon die gewichtsmatrices behandelen als één enkel, monolithisch object. Echter, elke gewichtsmatrix $W$ bevat inherent twee afzonderlijke grootheden: een magnitude ($\|W\|$) en een richting ($\hat{W} = W/\|W\|$). In standaardoptimalisatie zijn deze twee grootheden aan elkaar gekoppeld, wat leidt tot instabiele en onvoorspelbare trainingsdynamiek:

• Afhankelijkheid van directionele verandering: De hoekverandering die wordt veroorzaakt door een update hangt invers af van de huidige magnitude. Bijgevolg controleert de leersnelheid (LR) niet direct de snelheid van directionele verandering, die kan variëren tussen lagen en over de tijd.
• Onbedoelde magnitude-drift: Updates neigen loodrecht op de huidige gewichten te staan (door schaalinvariantie of ruis), waardoor de magnitude zelfs in de afwezigheid van een radiale gradiënt groeit. Deze "ongecontroleerde groei" dwingt beoefenaars om gebruik te maken van indirecte mechanismen zoals weight decay om de training te stabiliseren en warmup om vroege instabiliteit te voorkomen.
• Schaalbeperkingen: Huidige methoden voor het schalen van training (bijv. $\mu$P) vereisen vaak het opnieuw afstemmen van hyperparameters wanneer de breedte of diepte van een model verandert. Bovendien missen recente benaderingen die gewichten op een vaste-norm sfeer beperken vaak leerbare magnitude-parameters, waardoor de controle over de schaal uitsluitend bij normalisatielagen ligt.

2. Methodologie: Magnitude–Direction (MD) Ontkoppeling

De auteurs stellen Magnitude–Direction (MD) Ontkoppeling voor, een modificatie van de optimizer die elke gewichtsmatrix factoriseert in een richting met een vaste norm en leerbare magnitude-gains, die met afzonderlijke leersnelheden worden bijgewerkt.

Kernfactorisatie

Elke gewichtsmatrix $W$ wordt herparameteriseerd als:
$$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$$
Waarbij:

• $\hat{W}$ de richting is, beperkt tot een vaste hypersfeer (bijv. Frobenius-norm).
• $\gamma_{\text{row}}$ en $\gamma_{\text{col}}$ leerbare magnitude-gains zijn (vectoren die respectievelijk op rijen en kolommen inwerken).

Optimalisatieprocedure

De methode werkt volledig binnen de optimizer-stap, waardoor de forward pass van het model ongewijzigd blijft (het model ziet één gefuseerde tensor $W$):

1. Herstel: Herstel $\hat{W}$ en de gains $\gamma$ uit de gefuseerde gewichtsmatrix $W$.
2. Gradiëntsplitsing: Splits de gradiënt $G = \partial L / \partial W$ in een richting-gradiënt ($\partial L / \partial \hat{W}$) en gain-gradiënten (gesommeerd over de respectievelijke assen).
3. Richting-update: Update $\hat{W}$ met de basisoptimizer (bijv. Adam, Muon) met leersnelheid $\eta_W$, en projecteer deze vervolgens terug op de sfeer om de vaste norm te behouden.
4. Gain-update: Update de gains $\gamma$ (optioneel via een gladde functie zoals softplus om positiviteit te waarborgen) met een afzonderlijke leersnelheid $\eta_\gamma$.
5. Reassemblage: Reconstrueer $W$ voor de volgende forward pass.

Belangrijke Eigenschappen

• Optimizer Agnostisch: De methode fungeert als een wrapper rond elke basisoptimizer (Adam, Muon, Shampoo, etc.).
• Directe LR-controle: Door de richting op een sfeer te fixeren, stelt de leersnelheid direct de relatieve hoekverandering in, wat de afhankelijkheid van de huidige gewichtsmagnitude wegneemt.
• Eliminatie van Heuristieken: De methode elimineert de noodzaak voor weight decay (omdat gewichten al norm-beperkt zijn) en warmup (omdat grote, destabiliserende vroege updates vanaf stap één worden gereguleerd).
• Fijnmazige Schaalcontrole: In tegenstelling tot methoden die uitsluitend vertrouwen op de gains van normalisatielagen, biedt MD Ontkoppeling expliciete, geleerde controle over de schaal van individuele rijen en kolommen, waardoor het model specifieke kenmerken kan versterken of dempen.

3. Belangrijkste Bijdragen en Bevindingen

Empirische Resultaten

De auteurs valideren MD Ontkoppeling bij dense modellen (181M tot 1.29B parameters) en grote Mixture-of-Experts (MoE) modellen (tot 6.7B totale parameters).

• Prestatiewinst: MD Ontkoppeling presteert beter dan goed afgestemde baselines (AdamW en Muon). Op grote MoE-modellen bereikt MuonMD (Muon met MD Ontkoppeling) de loss van AdamW met ongeveer 2$\times$ minder rekenkracht.
• Leersnelheid Transfer: De methode maakt zero-shot leersnelheid transfer mogelijk over de breedte van het model. De optimale matrix-LR gevonden op een klein model blijft optimaal voor aanzienlijk grotere modellen zonder hernieuwde afstemming. Dit geldt voor zowel breedte- als dieptescaling.
• Warmup-vrije Training: Experimenten tonen aan dat het schrappen van warmup niet alleen de stabiliteit behoudt, maar ook de uiteindelijke loss verbetert, aangezien de vroege trainingsstappen niet langer worden verspild aan gereduceerde leersnelheden.
• Ablatie Inzichten:
  • Normalisatie-as: Het beperken van de richting via de Frobenius-norm (gehele matrix) blijkt de meest flexibele methode te zijn en presteert vergelijkbaar met rij/kolom-beperkingen.
  • Gain Parametrisatie: Een gecombineerde per-rij en per-kolom gain ($\gamma_{\text{row}} + \gamma_{\text{col}}$) presteert beter dan scalaire of single-axis gains. De softplus parametrisatie voor gains biedt een lichte, consistente voorsprong boven directe of exponentiële updates.
  • Embeddings: Het beperken van embedding-vectoren tot een eenheids L2-norm presteert iets beter dan ze onbeperkt laten.

Vergelijking met Gerelateerd Werk

Het artikel plaatst MD Ontkoppeling binnen een groeiend corpus van werk over sferische training en herparametrisatie:

• In tegenstelling tot nGPT, dat sferische beperkingen bundelt met specifieke architecturale wijzigingen (bijv. L2-normalisatie, geïnterpoleerde residuen), is MD Ontkoppeling een modificatie aan de optimizer-zijde die kan worden toegepast op standaardarchitecturen. Wanneer toegepast op de architectuur van nGPT, verbetert MD Ontkoppeling de prestaties verder.
• In tegenstelling tot Muown of HyperP, die mogelijk normen vastzetten zonder leerbare magnitudes of de methode koppelen aan specifieke optimizers, biedt MD Ontkoppeling expliciete, geleerde magnitudes en is het optimizer-agnostisch.
• Het adresseert de "magnitude-direction interferentie" die is geïdentificeerd in eerder werk (bijv. Kosson et al., 2024a) door de dynamiek van de twee grootheden te scheiden.

4. Betekenis en Claims

Het artikel claimt dat het behandelen van magnitude en richting als afzonderlijk gecontroleerde grootheden leidt tot voorspelbaardere trainingsdynamiek en een eenvoudige, breed toepasbare verbetering van moderne optimizers.

• Eenvoud: De methode vereist geen architecturale wijzigingen en voegt minimale computationele overhead toe (elementgewijze operaties tijdens de optimizer-stap).
• Robuustheid: Het vereenvoudigt het trainingsrecept door de noodzaak voor weight decay en warmup te verwijderen, waardoor de hyperparameter-zoekruimte wordt verkleind.
• Schaalbaarheid: Het biedt een principieel mechanisme voor leersnelheid transfer over modelbreedte, wat een belangrijke bottleneck bij het schalen van large language models aanpakt.

De auteurs concluderen dat hoewel de extra logica in de optimizer-stap een bescheiden overhead introduceert, de voordelen in trainingsstabiliteit, rekenefficiëntie en hyperparameter transfer MD Ontkoppeling tot een significante vooruitgang maken voor het trainen van grootschalige neurale netwerken.