NuMuon: Nuclear-Norm-Constrained Muon for Compressible LLM Training

Dit paper introduceert NuMuon, een optimizer die de kern-norm-beperking toevoegt aan Muon om de compressibiliteit van LLM-weights te vergroten en de kwaliteit na compressie te verbeteren, terwijl de gunstige convergentie van Muon behouden blijft.

De kern

NuMuon: De Kunst van het Slimme Opslaan van AI-Modellen

Stel je voor dat je een gigantische bibliotheek bouwt met miljoenen boeken (de kennis van een AI). Deze bibliotheek is zo groot dat hij een heel gebouw nodig heeft om in te passen, en het kost een fortuin om het te onderhouden. Dit is het probleem met moderne "Large Language Models" (LLM's): ze zijn slim, maar ze zijn ook enorm zwaar en duur om te draaien.

De onderzoekers van Pluralis Research hebben een nieuwe manier bedacht om deze bibliotheken te bouwen die niet alleen slim zijn, maar ook makkelijk in een klein kofferje passen zonder dat de inhoud verandert. Ze noemen hun methode NuMuon.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Muon" Bibliotheek

Er was al een nieuwe manier om deze bibliotheken te bouwen, genaamd Muon. Muon is een slimme "architect" (een optimizer) die de boeken op een heel efficiënte manier rangschikt. Het werkt sneller dan de oude methoden (zoals AdamW).

Maar er was een verrassing: toen de onderzoekers keken hoe Muon de boeken rangschikte, zagen ze iets raars. Hoewel Muon beloofde om alles op te slaan (elk detail, elke pagina), bleek dat de bibliotheek eigenlijk vol zat met herhalingen. Veel boeken waren bijna identiek aan elkaar.

• De Analogie: Stel je voor dat je een muur bouwt met bakstenen. Muon zegt: "Ik gebruik elke baksteen die er is!" Maar als je kijkt, zie je dat de muur eigenlijk alleen maar uit 10% unieke bakstenen bestaat en de rest is gewoon dezelfde steen die 100 keer wordt gebruikt.
• Het Resultaat: Omdat er zoveel herhalingen zijn, kun je de muur makkelijk kleiner maken (compresseren). Maar als je te hard knijpt (te veel comprimeren), valt de muur in elkaar. De Muon-muur is kwetsbaar als je hem te veel wilt verkleinen.

2. De Oplossing: NuMuon (De "Nucleaire" Architect)

De onderzoekers dachten: "Waarom wachten tot de herhalingen vanzelf ontstaan? Laten we de architect dwingen om vanaf dag één een compacte muur te bouwen."

Ze introduceerden NuMuon. Dit is een升级版 (upgrade) van Muon.

• De Analogie: Stel je voor dat Muon een schilder is die een groot doek vult met verf, en later pas ziet dat hij veel dezelfde kleuren heeft gebruikt. NuMuon is een schilder die direct een beperkt palet aan kleuren kiest. Hij zegt: "Ik ga alleen de 10 belangrijkste kleuren gebruiken om dit schilderij te maken."
• De Wiskundige Term (Nucleaire Norm): In de paper wordt dit een "nucleaire norm-beperking" genoemd. In het gewone leven betekent dit: "Houd het aantal unieke dingen dat je gebruikt, laag."

3. Hoe werkt het in de praktijk?

NuMuon gebruikt een slimme truc. In plaats van te proberen alles te leren, leert het model alleen de belangrijkste patronen.

• De "Top-k" Truc: Stel je voor dat je een lijst hebt met 1000 woorden die je gebruikt. Muon zou proberen ze allemaal te onthouden. NuMuon zegt: "Weet je wat? We houden alleen de top 100 woorden bij. Die zijn het belangrijkst. De rest is ruis."
• De Scheduling: De onderzoekers ontdekten dat je niet direct naar 10 woorden moet gaan. Je begint met 1000, en langzaam, naarmate het model "slimmer" wordt, schroef je het aantal unieke woorden terug naar 100. Dit is als een leerling die eerst alles opschrijft, en later alleen de samenvatting onthoudt.

4. Waarom is dit geweldig?

Het resultaat is een AI-model dat twee dingen doet die normaal gesproken tegenstrijdig zijn:

1. Het is net zo slim als Muon: Het leert even goed en snel tijdens het trainen.
2. Het is super-compressibel: Omdat het model van nature al is opgebouwd uit "herhalende patronen" (lage rang), kun je het later extreem klein maken zonder dat het dom wordt.

De Vergelijking:

• Oude methoden (AdamW): Bouwen een zware, rommelige bibliotheek. Zwaar om te dragen, en als je hem verkleint, verdwijnt de kennis.
• Muon: Bouwt een snelle bibliotheek, maar hij is nog steeds wat rommelig. Als je hem te veel verkleint, valt hij in elkaar.
• NuMuon: Bouwt een bibliotheek die van nature al compact is. Je kunt hem in een rugzak stoppen en hij werkt nog steeds perfect.

Conclusie

Met NuMuon kunnen bedrijven goedkopere, snellere AI-modellen maken die op gewone telefoons of servers draaien, zonder dat ze hun intelligentie verliezen. Het is alsof je een gigantische, zware koffer vervangt door een slimme, opvouwbare tas die precies past bij wat je nodig hebt.

Kortom: NuMuon zorgt ervoor dat AI niet alleen slim is, maar ook efficiënt en betaalbaar voor iedereen.

Technische samenvatting

Probleemstelling

De snelle vooruitgang van Large Language Models (LLMs) wordt steeds meer beperkt door de hoge kosten voor geheugen en implementatie. Om dit op te lossen, worden compressietechnieken gebruikt die vaak vertrouwen op de laag-rang (low-rank) structuur van trainingsgewichten.

• De uitdaging: Bestaande optimalisatoren zoals AdamW vertonen een impliciete bias naar laag-rang structuren, maar recente optimalisatoren zoals Muon (die volledige rang-updates toepassen via orthogonalisatie) zijn ontworpen voor trainings-efficiëntie zonder expliciete rangbeperking.
• De observatie: Hoewel Muon volledige rang-updates gebruikt, bleek uit empirisch onderzoek dat Muon-getrainde modellen toch een opvallende laag-rang structuur vertonen. Echter, deze emergente structuur is niet robuust genoeg voor agressieve compressie; bij hoge compressiepercentages degradeert de prestatie van Muon-modellen snel.
• De vraag: Hoe kunnen we de trainingsdynamiek van Muon zo sturen dat de gewichten expliciet een robuuste laag-rang structuur aannemen, zonder de gunstige convergentie-eigenschappen van Muon te verliezen?

Methodologie: NuMuon

De auteurs introduceren NuMuon, een variant van de Muon-optimalisator die een kern-norm (nuclear-norm) beperking toevoegt aan de update-richting.

1. Theoretisch Kader (LMO):

   • Muon wordt geïnterpreteerd als een Linear Minimization Oracle (LMO) over een spectrale-norm-bol. Dit betekent dat Muon de update-richting kiest die de lineaire geobjectiveerde functie minimaliseert binnen een bol van spectrale norm.
   • NuMuon breidt dit uit door de update-richting te beperken tot de doorsnede van een spectrale-norm-bol en een kern-norm-bol (de som van de singuliere waarden). De kern-norm fungeert als een convexe proxy voor rang.

2. De NuMuon Update:

   • Het probleem wordt gereduceerd tot een lineair programma over de singuliere waarden van de momentum-buffer.
   • De optimale oplossing blijkt een top-$k$ singuliere vector truncatie te zijn. In plaats van een volledige orthogonalisatie (zoals bij Muon), selecteert NuMuon alleen de $k$ belangrijkste singuliere richtingen.
   • De update wordt gegeven door: $\Delta W = -\rho \sum_{i=1}^k u_i v_i^\top$, waarbij $k$ de gewenste rang is.

3. Praktische Implementatie:

   • Efficiëntie: Om de top-$k$ singuliere vectoren efficiënt te berekenen op grote schaal, gebruiken de auteurs een gerandomiseerde Block Krylov-methode in plaats van een volledige SVD.
   • Rank Scheduler: Omdat training vroeg in het proces vaak een hogere effectieve rang vereist, gebruiken ze een cosine-rangplanner. De rang begint hoog en neemt geleidelijk af naar een lager doelwaarde tijdens de training.

4. Convergentie:

   • De auteurs bewijzen convergentiegaranties voor NuMuon in niet-convexe settings. Ze tonen aan dat de convergentiebound afhangt van de gekozen rang $k$ en de "tail energy" (de energie van de gradiënt buiten de top-$k$ componenten), die empirisch klein blijkt te zijn.

Belangrijkste Bijdragen

1. Empirische Inzicht: Het paper onthult dat Muon-modellen, ondanks hun volledige rang-updates, van nature een laag-rang structuur ontwikkelen, maar dat deze onvoldoende is voor extreme compressie.
2. NuMuon Optimalisator: Een nieuwe optimalisator die de spectrale dynamiek van Muon behoudt maar de update-rang expliciet controleert via een kern-norm budget.
3. Wiskundige Formulering: Een afleiding die laat zien dat de NuMuon-stap reduceert tot een lineair programma met een gesloten vorm oplossing (top-$k$ truncatie).
4. Convergentieanalyse: Theoretische garanties voor de convergentie van NuMuon onder aannames van gladheid en begrenste gradiëntvariatie.

Resultaten

De auteurs hebben NuMuon getest op modellen van 0,6B tot 1,8B parameters (Qwen3, Olmo2, Llama3) en vergeleken met AdamW en Muon.

• Convergentie: NuMuon convergeert vergelijkbaar met Muon en presteert beter dan AdamW in termen van trainings- en validatie-verwarring (perplexity).
• Stabiele Rang: NuMuon produceert gewichtsmatrices met een aanzienlijk lagere stabiele rang dan Muon en AdamW, wat aangeeft dat de informatie sterker geconcentreerd is in minder dimensies.
• Compressie-Prestaties:
  • Bij gebruik van state-of-the-art compressietechnieken (ASVD, SVD-LLM, Dobi-SVD) behouden NuMuon-modellen hun prestaties veel beter bij hoge compressiepercentages (tot 80%).
  • Bij 40-80% compressie vertonen NuMuon-modellen tot 55,9% betere trade-offs tussen compressie en kwaliteit (lagere perplexity) vergeleken met Muon-baselines.
  • In termen van inferentie-throughput (tokens per seconde) bij een vaste perplexiteit, bieden NuMuon-modellen de snelste inferentie bij extreme compressie.
• Subspace Alignment: Analyses tonen aan dat de updates van NuMuon beter aligned zijn met de dominante spectrale subruimte van de gewichten dan die van Muon, wat de robuustheid tegen compressie verklaart.

Betekenis en Impact

Dit werk is significant voor de implementatie van LLMs in productieomgevingen met beperkt geheugen:

• Efficiëntie: Het biedt een manier om LLMs te trainen die "compressie-vriendelijk" zijn, waardoor ze later efficiënter kunnen worden gedeployed zonder dat er een aparte, kostbare fine-tuning-fase nodig is om compressie-veerkracht te herstellen.
• Optimalisatie-Design: Het illustreert dat de keuze van de optimalisator directe invloed heeft op de geometrie van de gewichten en de daaropvolgende compressibiliteit.
• Toekomstige Richtingen: De methode is ook relevant voor gedistribueerde training in bandbreedte-beperkte omgevingen, aangezien de gefactoriseerde vorm van de updates minder communicatie vereist.

Kortom, NuMuon combineert de trainings-efficiëntie van Muon met de praktische voordelen van laag-rang structuren, waardoor het een ideale kandidaat is voor het trainen van schaalbare en kosteneffectieve Large Language Models.