Technische Samenvatting: Het exploiteren van gewichtsruimte-symmetrieën voor het benaderen van kromming

Probleemstelling

Nauwkeurige schatting van de kromming (de Hessiaanse matrix) van een verliesfunctie is fundamenteel voor veel subgebieden van machine learning, waaronder tweede-orde optimalisatie, Bayesiaanse deep learning, continual learning en modelcompressie. Het construeren, opslaan en inverteren van de volledige Hessiaan is echter computationeel en geheugentechnisch onhaalbaar voor moderne diepe netwerken. Hoewel bestaande benaderingstechnieken zoals (E)KFAC, Shampoo en Soap schaalbaarheid hebben bereikt door uit te gaan van specifieke blokdiagonaal- of Kronecker-gefactoreerde structuren, ontbreekt het het vakgebied aan een principieel theoretisch kader dat uitlegt waarom deze specifieke structuren werken of hoe ze systematisch afgeleid kunnen worden op basis van de modelarchitectuur. Bovendien heeft eerder werk de structurele beperkingen op kromming, die voortvloeien uit de bekende gewichtsruimte-symmetrieën in neurale netwerken (bijv. permutatie-invariantie van verborgen eenheden), grotendeels genegeerd.

Methodologie

De auteurs stellen een nieuw kader voor om de Hessiaan te benaderen door analytisch te middelen over de "groep-baan" (group orbit) gegenereerd door gewichtsruimte-symmetrieën. De kerninzicht berust op het feit dat veel verliesfuncties invariant zijn onder specifieke parametertransformaties (groepsacties), wat impliceert dat de gradiënt van de verliesfunctie onder deze transformaties eveneens equivariant is.

1. Baan-middeling en Gradiënt-equivariantie

Gegeven een huidige parameterset $w$ en gradiënt $g$, beschouwen de auteurs de verzameling van alle paren $(w^{(i)}, g^{(i)})$ die bereikbaar zijn door symmetrietransformaties van een groep $G toe te passen. Omdat de verliesfunctie invariant is, is de gradiënt in elk punt in de baan gerelateerd aan de oorspronkelijke gradiënt door de groepsactie. Door deze gradiënten over de gehele baan te middelen, aggregeert de methode kromminginformatie zonder dat meerdere gradiëntberekeningen nodig zijn.

2. Gestructureerde Hessiaanse Benadering

De methode leidt een gestructureerde benadering van de Hessiaan af, $\hat{H}$, door een secant-conditie op te lossen op basis van de baan-gemiddelden van de parametervariaties en gradiëntvariaties.

• Commutant-algebra: De resulterende Hessiaanse benadering behoort tot de commutant-algebra van de symmetriegroep. Dit betekent dat de benadering gerepresenteerd kan worden als een lineaire superpositie van een klein aantal basis-elementen (ijle, binaire tensoren gevormd door Kronecker-producten van identiteits- en alles-één-matrices).
• Factorisatie: In plaats van de volledige matrix op te slaan, slaat de methode alleen de scalaire of matrix "factoren" op die deze basis-elementen wegen. Dit maakt een hanteerbare schatting, opslag en inversie mogelijk.
• Vereenvoudiging: De auteurs stellen een praktische vereenvoudiging voor waarbij wordt aangenomen dat het baan-gemiddelde van de parametervariaties proportioneel is aan de identiteitsmatrix. Dit leidt tot een krommingsschatting $\hat{H}_g = S_g^{1/2}$, waarbij $S_g$ het baan-gemiddelde is van het buitenproduct van de gradiënt.

3. De "Symo" Optimizer

Het kader wordt geïnstantieerd als een nieuwe optimizer genaamd Symo (Symmetry-based optimizer). De update-regel is $w_{t+1} = w_t - \eta (\hat{H}_g)^{-1} g_t$. De auteurs bieden een just-in-time compiler die symbolische specificaties van de symmetrieën van een netwerk neemt en efficiënte PyTorch-routines genereert om de noodzakelijke factoren te berekenen.

Belangrijkste Bijdragen

1. Nieuwe Hessiaanse Benadering: Een methode om een gestructureerde Hessiaanse schatting af te leiden uit een enkele gradiëntberekening door te middelen over de baan van de symmetriegroep. De benadering is bewezen zeer gestructureerd en bevindt zich in de commutant-algebra.
2. Theoretische Unificatie van Shampoo en Muon: De paper bewijst dat specifieke keuzes van symmetriegroepen binnen dit kader exact de update-regels van de Shampoo en Muon optimizers terugbrengen.
   • Door de symmetriegroep te beperken tot specifieke getekende permutatie-subgroepen (bijv. het toepassen van identiteit-symmetrie op input/output-lagen en getekende permutaties op verborgen lagen), wordt de afgeleide krommingmatrix een blokdiagonaal met Kronecker-factoren, wat wiskundig equivalent is aan Shampoo/Muon.
   • Dit biedt een fundamentele verklaring voor de effectiviteit van deze optimizers: ze exploiteren impliciet gewichtsruimte-symmetrieën, waardoor ze een "sweet spot" vinden tussen benaderingsnauwkeurigheid en computationele kosten.
3. Complexiteitscontrole: Het kader toont aan dat de complexiteit van de Hessiaanse benadering direct wordt bepaand door de grootte van de gekozen symmetriegroep. Grotere groepen leveren lagere-dimensionale commutant-algebra'ën op (minder factoren, goedkopere berekening) maar grotere banen (potentieel minder lokale nauwkeurigheid). Kleinere groepen leveren meer termen maar betere lokale getrouwheid op.
4. Empirische Validatie:
   • Hessiaanse Kwaliteit: Experimenten op een kleine MLP tonen aan dat de voorgestelde benaderingen (specifiek de vereenvoudigde $\hat{H}_g$) een hoge cosinus-gelijkenis bereiken met de ware Hessiaan langs de negatieve gradiëntrichting, waarbij ze de baan-gemiddelde parametervariaties overtreffen en de prestaties van Shampoo evenaren in deze metriek.
   • Optimizer Equivalentie: Op MNIST-autoencoders en nanoGPT (een klein taalmodel) produceert de Symo-optimizer trainingscurven die bijna identiek zijn aan Muon, wat de theoretische equivalentie onder specifieke groepsvoorwaarden valideert. Beiden presteren significant beter dan Adam.

Resultaten

• Nauwkeurigheid vs. Kosten Afweging: De paper onderzoekt empirisch verschillende groepkeuzes. Een "grotere groep" (die permutaties van alle lagen inclusief input/output toestaat) presteert nog steeds beter dan Adam, maar is minder effectief dan een "kleinere groep" (die permutaties beperkt tot interne lagen en blokdiagonaal-benaderingen gebruikt), wat overeenkomt met de prestaties van Shampoo. Dit suggereert dat het beperken van de groep om aan te sluiten bij de specifieke architecturale invarianties cruciaal is voor optimale prestaties.
• Rangdeficiëntie: De auteurs merken op dat de volledige baan-gemiddelde matrix $S_w$ rangdeficiënt kan zijn, wat inversie destabiliseert. De voorgestelde vereenvoudiging ($S_w \approx I$) vermijdt dit probleem en levert empirisch betere krommingsschattingen op dan de volledige positief semi-definiete oplossing.
• Schaalbaarheid: De methode schaalt succesvol naar Transformer-architecturen (nanoGPT), wat aantoont dat de symmetrie-bewuste benadering bruikbaar is voor meer dan alleen eenvoudige MLP's.

Betekenis en Claims

De paper claimt een verenigend theoretisch perspectief te bieden voor het bekijken van bestaande tweede-orde optimizers. Door Shampoo en Muon te kaderen als specifieke instanties van symmetrie-gebaseerde Hessiaanse benadering, verplaatst dit werk deze methoden van heuristische engineering naar een fundamentele wiskundige afleiding.

De auteurs stellen dat hun kader:

• Een systematisch recept biedt voor het construeren van kromming-benaderingen die zijn afgestemd op specifieke computationele budgetten door de geschikte symmetriegroepen te selecteren.
• Een fundamentele route biedt van model-symmetrieën naar gestructureerde kromming, wat potentieel nieuwe optimizers voor verschillende modelklassen kan inspireren.
• Toepassingen suggereert buiten optimalisatie, waaronder onzekerheidsschatting, continual learning en modelcompressie, hoewel deze als potentiële toekomstige richtingen worden gepresenteerd in plaats van gedemonstreerde resultaten.

Het paper blijft bescheiden over beperkingen en erkent dat de nauwkeurigheid van de Taylor-expansie rond het baan-gemiddelde niet volledig begrensd is voor getrainde gewichten, en dat de formele garantie van eigenvector-uitlijning tussen de preconditioner en de ware Hessiaan een openstaand probleem blijft. Daarnaast, hoewel het kader de afleiding van factoren automatiseert, blijft het selecteren van de optimale groep voor nieuwe architecturen een taak die afhankelijk is van de gebruiker.