Technische Samenvatting: AMUSE (Anytime Muon met Stabiele Gradiëntevaluatie)

Probleemstelling

Moderne deep learning-optimalisatie vertrouwt doorgaans op de AdamW-optimizer in combinatie met voorgeschreven leerplanschema's (bijvoorbeeld cosinusdecay). Recente vooruitgang heeft deze componenten afzonderlijk uitgedaagd: Schedule-Free (SF)-optimalisatie verwijdert expliciete schema's door gradiënten te evalueren op geïnterpoleerde punten tussen de huidige en de gemiddelde iteraten, terwijl Muon de update-geometrie voor matrixwaardige parameters verbetert door momentum te orthogonaliseren.

Ondanks Muons empirische succes bij het voortrainen van grote taalmodellen (LLM's), is het onderliggende mechanisme niet volledig begrepen. Het artikel identificeert een specifieke instabiliteit in Muon: terwijl de orthogonalisatie vooruitgang versnelt langs "vlakke" richtingen (de "rivier" van het verlieslandschap), versterkt het onbedoeld ruis in "steile" richtingen met hoge kromming (de "valleiwanden"), wat leidt tot oscillerende trajecten. Omgekeerd stabiliseren Schedule-Free-methoden trajecten door ze dicht bij de rivier te houden, maar missen ze mogelijk de snelle bulkvooruitgang die Muons orthogonalisatie biedt. Het kernprobleem is het ontwerpen van een optimizer die Muons snelle convergentie combineert met de stabiliteit van Schedule-Free-gemiddelden, zonder dat handmatige leerplanschema's nodig zijn.

Methodologie: AMUSE

De auteurs stellen AMUSE (Anytime Muon met Stabiele Gradiëntevaluatie) voor, wat Muons georthogonaliseerde matrixupdates integreert met het Schedule-Free-raamwerk met behulp van een tijdsvariërend interpolatiecoëfficiënt.

Theoretische Grondslag: Rivier-Vallei Landschap

De analyse is gebaseerd op het perspectief van het rivier-vallei verlieslandschap:

• De Rivier: Een hoogdimensionale, laagkrommende bulk-onderruimte waar nuttige trainingsvooruitgang plaatsvindt.
• De Valleiwanden: Een laagdimensionale, hoogkrommende dominante onderruimte die oscillaties induceert.
• Muons Mechanisme: Het orthogonaliseren van de momentummatrix verhoogt het relatieve bulkcomponent van updates, wat vooruitgang langs de rivier versnelt. Het voorkomt echter ook dat dominante componenten vervagen, waardoor kleine ruis in dominante richtingen wordt versterkt tot oscillaties over de valleiwanden.
• Schedule-Free Mechanisme: Het evalueren van gradiënten op een gemiddeld traject (dichter bij de rivier) filtert fluctuaties met hoge kromming uit voordat de updatestap plaatsvindt.

Algoritmeontwerp

AMUSE maakt gebruik van een tijdsvariërend interpolatiecoëfficiënt $\beta_t$ om het gradiëntevaluatiepunt $Y_t$ te bepalen tussen de snelle basissequentie $Z_t$ (Muon-updates) en de stabiele gemiddelde sequentie $X_t$:
$$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$$

Het schema voor $\beta_t$ is ontworpen om vroege adaptatie en stabiliteit in een later stadium in evenwicht te brengen:

1. Warmup-fase ($t \le T_0$): $\beta_t$ is vastgesteld op een kleine waarde $\beta_1$. Dit stelt de optimizer in staat gradiënten dicht bij de snelle Muon-sequentie ($Z_t$) te evalueren voor snelle initiële adaptatie.
2. Na Warmup ($t > T_0$): $\beta_t$ neemt geleidelijk toe tot 1 volgens:
   $$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$$
   Hierbij regelt $\rho \in [0, 1]$ de snelheid van de verschuiving. Naarmate $\beta_t$ toeneemt, verschuift het gradiëntevaluatiepunt naar de gemiddelde sequentie $X_t$. Dit onderdrukt de componenten met hoge kromming (dominante componenten) van de gradiënt voordat ze Muons orthogonalisatie ondergaan, waardoor het traject wordt gestabiliseerd.

Cruciaal is dat dit schema onafhankelijk is van de totale trainingshorizon, wat anytime training mogelijk maakt (training kan op elk moment worden gestopt zonder prestatieverlies). AMUSE vereist geen schema voor leerplandecay.

Belangrijkste Bijdragen

1. Geometrische Analyse van Muon: Het artikel biedt een geometrische verklaring voor Muons efficiëntie en instabiliteit. Het toont aan dat orthogonalisatie bulk-gerichte updates bevordert (wat vooruitgang langs de rivier versnelt), maar ruis in dominante richtingen versterkt (wat leidt tot oscillaties langs de valleiwanden).
2. Stabilisatie via Interpolatie: De auteurs tonen aan dat het evalueren van gradiënten dichter bij het gemiddelde traject (via Schedule-Free) dominante componenten vóór orthogonalisatie reduceert, waardoor oscillaties worden verminderd.
3. AMUSE Algoritme: Een nieuwe optimizer die dynamisch het gradiëntevaluatiepunt verschuift van het snelle Muon-traject naar het stabiele gemiddelde traject. Dit combineert de snelle convergentie van Muon met de stabiliteit van Schedule-Free-methoden.
4. Ondersteuning voor Anytime Training: Door gebruik te maken van een tijdsvariërend schema dat onafhankelijk is van het totale aantal stappen en het vermijden van leerplandecay, ondersteunt AMUSE flexibele trainingshorizons.

Experimentele Resultaten

AMUSE werd geëvalueerd op diverse benchmarks, waaronder beeldclassificatie (CIFAR-10/100, SVHN, ImageNet-1k), beeldsegmentatie (ISIC 2018) en LLM-voortraining (124M, 720M en 1B parameter Llama-achtige modellen op FineWeb).

• Pareto-grens Prestatie-Iteratie: AMUSE presteert consistent beter dan (Schedule-Free) AdamW en Muon op alle taken.
• Efficiëntiewinst:
  • Bij 720M LLM-voortraining bereikt AMUSE Muons eindprestatie met 1,51$\times$ minder stappen.
  • Op ImageNet met ResNet-50 bereikt het dezelfde prestatie met 1,12$\times$ minder stappen.
  • Op ImageNet met ViT-finetuning zijn 3,08$\times$ minder stappen vereist.
• Stabiliteit: In tegenstelling tot standaard Muon, dat aanzienlijk profiteert van post-hoc exponentiële gewichtsgemiddelden (EWA) of leerplandecay, behoudt AMUSE een stabiel traject gedurende de training, waarbij minimale verbetering wordt waargenomen wanneer deze post-hoc technieken worden toegepast.
• Schaalbaarheid: AMUSE schaalt effectief naar grote batchgroottes (tot 2048) en grote modelgroottes (1B parameters), waar andere Schedule-Free-varianten soms moeite hebben.

Betekenis en Claims

Het artikel stelt dat AMUSE een principiële aanpak voor optimizerontwerp vertegenwoordigt, afgeleid van een systematisch begrip van trainingsdynamiek en de geometrie van het verlieslandschap. Door expliciet in te spelen op de afweging tussen snelle bulkvooruitgang en stabiliteit langs de valleiwanden, bereikt AMUSE superieure prestaties zonder de noodzaak van handmatig leerplanschema's.

De auteurs merken bescheiden op dat hoewel AMUSE bestaande optimizers overtreft, het een lichte geheugenoverhead met zich meebrengt (een extra staatkopie voor de basissequentie $Z_t$) in vergelijking met standaard Muon, hoewel dit vergelijkbaar blijft met AdamW en Schedule-Free AdamW. Zij suggereren dat toekomstig werk zich zou kunnen richten op het bereiken van vergelijkbare stabiliteit en prestaties zonder deze extra geheugenkosten. Het artikel benadrukt dat het "rivier-vallei"-perspectief een robuust kader biedt voor het ontwerpen van optimizers die zowel efficiënt als stabiel zijn.