How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks

Technische Samenvatting: Hoe de Optimizer de Geleerde Oplossingen in Equivariante Neuronale Netwerken Vormgeeft

Probleemstelling
Equivariante neuronale netwerken coderen geometrische symmetrieën direct in hun architectuur, wat een natuurlijke inductieve bias biedt voor data met bekende geometrische symmetrieën. Echter maken harde equivariantie-beperkingen optimalisatie vaak moeilijk, wat leidt tot verlieslandschappen met kritieke punten en schijnbare minima. Bijgevolg kunnen deze modellen op grote schaal minder goed presteren dan minder beperkte alternatieven. De dominante reactie in recente literatuur is geweest om architecturale beperkingen te verzachten via benaderende of ontspannen equivariantiemethoden. Daarentegen blijft de rol van de optimizer bij het vormgeven van de oplossingen van equivariante modellen onderbelicht. Dit artikel onderzoekt of de keuze van de optimizer, onafhankelijk van architecturale aanpassingen, de tijdens training bereikte oplossingen significant kan veranderen en de prestaties kan verbeteren.

Methodologie
De auteurs voeren een vergelijkende studie uit tussen Adam (Kingma & Ba, 2015) en Muon (Jordan et al., 2024), een onlangs voorgestelde optimizer die zijn momentumbuffer orthogonaliseert via Newton-Schulz-iteraties. De evaluatie beslaat drie verschillende settings:

1. 3D-Puntenwolkclassificatie: Met behulp van de ModelNet40-dataset (en de beschadigde variant, ModelNet40-C) met drie architecturen die verschillende geometrische inductieve biases bezitten:
   • EGNN: Een expliciet $E(3)$-equivariant grafisch neuronale netwerk.
   • DGCNN: Een permutatie-equivariant grafisch CNN met dynamische $k$-NN-connectiviteit.
   • PointNet: Een permutatie-invariante architectuur via symmetrische pooling.
2. Moleculair Leren: Met behulp van de QM9-dataset met GotenNet ($E(3)$-equivariant) en GINE (permutatie-equivariant) op peptide- en ZINC-datasets.
3. Geometrische Analyse: Om te begrijpen waarom de prestaties verschillen, analyseren de auteurs de getrainde checkpoints met behulp van:
   • Lokale Verliesgeometrie: Visualisatie van 2D-verliesplakken en 3D-oppervlakken rond checkpoints, en schatting van Hessiaan-samenvattingen (top-eigenwaarde en trace) via power-iteratie en Hutchinson-schatters.
   • Spectrale Structuur: Berekening van de stabiele rang en effectieve rang (gebaseerd op Shannon-entropie van singuliere waarden) voor zowel gewichtsmatrices als representaties van tussenliggende lagen.

Belangrijkste Resultaten

• Prestatiewinst: Op ModelNet40 presteert Muon consistent beter dan Adam over alle architecturen (EGNN, PointNet, DGCNN) in zowel schone als beschadigde settings. Bijvoorbeeld, op schone ModelNet40 verbeterde Muon de EGNN-nauwkeurigheid met +5,17%, PointNet met +2,67% en DGCNN met +1,96%. Vergelijkbare trends werden waargenomen op de QM9-moleculaire dataset, waar Muon de meeste doelstellingen verbeterde.
• Kenmerken van het Verlieslandschap: Visualisaties van het verliesoppervlak rond getrainde checkpoints onthullen dat Muon-oplossingen een gladdere, regelmatiger lokale geometrie vertonen in vergelijking met de onregelmatige contouren die vaak bij Adam worden gezien. Echter, Hessiaan-schattingen tegenspreken de intuïtie dat gladdere geometrie impliceert lagere kromming; Muon-checkpoints hebben daadwerkelijk grotere top-eigenwaarden en traceschattingen (wat hogere kromming aangeeft) dan Adam-checkpoints.
• Spectrale Eigenschappen: Muon levert consistent hogere stabiele en effectieve rangen op voor zowel gewichten als tussenliggende representaties over alle lagen en architecturen. Dit suggereert dat Muon de capaciteit over singuliere richtingen gelijkmatiger verdeelt, waardoor concentratie van spectra in enkele dominante richtingen, die vaak geassocieerd wordt met op gradiënten gebaseerde methoden, wordt vermeden.

Betekenis en Beweringen
Het artikel stelt dat de optimizer een betekenisvolle bron van inductieve bias is die interageert met geometrische beperkingen. De auteurs beweren dat:

1. Keuze van Optimizer Maakt Uit: Het veranderen van de optimizer van Adam naar Muon de prestaties in equivariante modellen kan verbeteren zonder de architectuur aan te passen of beperkingen te verzachten.
2. Verschillende Oplossingsruimtes: Muon en Adam convergeren naar kwalitatief verschillende oplossingen, gekenmerkt door verschillende lokale verliesgeometrieën en spectrale structuren.
3. Mechanisme-hypothese: De georthogonaliseerde momentumupdates in Muon lijken een traject-niveau bias te bieden naar gunstigere regio's van het verlieslandschap, waarbij singuliere richtingen van updates gelijkmatiger worden behandeld.
4. Ontwerpprincipe: De resultaten suggereren een potentieel ontwerpprincipe voor equivariante optimizers: update-regels die voorkomen dat de singuliere spectra van gewichten en representaties concentreren in enkele richtingen, kunnen training verbeteren.

De auteurs houden een bescheiden standpunt aan, opmerkend dat hoewel de correlatie tussen spectrale eigenschappen en nauwkeurigheidswinsten consistent is, een causaal verband niet is vastgesteld. Zij concluderen dat de interactie tussen optimizerontwerp en geometrische inductieve bias nader onderzoek vereist, met name met betrekking tot hoe optimizers de structuur van symmetriegroepen kunnen benutten in plaats van equivariantie te behandelen als een black-box beperking.