A Convergence Analysis of Adaptive Optimizers under Floating-point Quantization

Deze paper introduceert het eerste theoretische raamwerk voor de convergentieanalyse van adaptieve optimalisatoren zoals Adam en Muon onder floating-point kwantisatie, waarbij wordt aangetoond dat ze bij voldoende mantissenlengte convergentieraten behouden die dicht bij die van full-precision liggen, met een grotere gevoeligheid voor kwantisatiefouten bij Adam dan bij Muon.

De kern

Titel: Waarom werken slimme trainingsmethoden ook met "ruwe" getallen? Een uitleg in gewoon Nederlands.

Stel je voor dat je een gigantische kunstwerk aan het bouwen bent: een Groot Taalmodel (zoals de slimme AI's die we vandaag de dag kennen). Om dit te doen, moet je een enorm complex puzzelstuk oplossen. Dit proces heet "trainen".

Normaal gesproken gebruiken computers hiervoor zeer precieze getallen (zoals een chirurg die een microscoop gebruikt). Maar deze modellen zijn zo groot dat ze de geheugenruimte van elke computer op aarde zouden vullen. Daarom proberen ingenieurs het "ruwer" te doen: ze gebruiken minder precieze getallen (zoals FP8 of BF16). Het is alsof je van een microscoop overschakelt naar een vergrootglas.

Het probleem:
Wetenschappers wisten al lang dat dit in de praktijk werkt (AI's worden er niet dom van), maar ze hadden geen goed wiskundig bewijs waarom dit zo is. Bestaande theorieën gingen ervan uit dat alle getallen perfect waren. Ze konden niet verklaren waarom het trainen met "ruwe" getallen (waarbij kleine afrondingsfoutjes optreden) toch succesvol blijft.

De oplossing van dit papier:
De auteurs (Xuan Tang, Jichu Li en Difan Zou) hebben een nieuwe theorie bedacht. Ze kijken niet alleen naar de foutjes in de berekeningen, maar naar alles: de getallen die de AI leert (gewichten), de signalen die ze ontvangt (gradiënten) en de "herinneringen" die de AI heeft (momentum).

Ze vergelijken twee populaire methoden om dit puzzelstuk op te lossen:

1. Adam: De oude, bewezen kampioen.
2. Muon: De nieuwe, snelle uitdager.

Hier is de uitleg met een creatieve analogie:

De Analogie: De Bergbeklimmers

Stel je voor dat je twee bergbeklimmers hebt die een berg moeten beklimmen (de berg is het probleem dat ze moeten oplossen). Ze hebben een kaart en een kompas, maar hun apparatuur is niet perfect.

1. De "Ruwheid" (Quantisatie)
In plaats van een perfect kompas met een naald die precies op Noord wijst, gebruiken ze een kompas dat soms een beetje aarzelt of een kaart met wat vlekken. Dit zijn de afrondingsfoutjes door de lage precisie.

• Vroeger dachten wetenschappers: "Als je kompas maar een beetje schuift, ga je de berg niet op."
• Deze paper zegt: "Nee, zolang de schuif niet te groot is, komen ze er toch."

2. De Twee Klimmers (Adam vs. Muon)

• Adam (De voorzichtige klimmer):
  Adam is slim, maar hij is erg afhankelijk van zijn "herinnering" aan hoe steil de berg was in het verleden. Hij kijkt naar een gemiddelde van de afgelopen stappen.

  • Het probleem: Als je kompas (de data) een beetje vlekken heeft, en Adam kijkt naar een hele lange geschiedenis (wat hij doet), dan stapelen die kleine vlekjes zich op. Het is alsof je een foutje in je notities maakt, en elke dag dat je terugkijkt, wordt dat foutje groter.
  • De conclusie van de paper: Adam is erg gevoelig voor ruwe getallen, vooral als hij heel voorzichtig is (wat hij vaak is). Hij heeft een heel scherp kompas nodig om niet de weg kwijt te raken. Als je hem te ruwe getallen geeft, begint hij te trillen en stopt hij met klimmen.

• Muon (De flexibele klimmer):
  Muon is een nieuwere methode. Hij gebruikt een andere techniek (gebaseerd op wiskundige rotaties, ofwel SVD) om zijn richting te bepalen.

  • Het voordeel: Muon is minder afhankelijk van die lange, gevoelige geschiedenis. Hij kijkt meer naar de huidige situatie en past zich flexibeler aan.
  • De conclusie van de paper: Muon is veel robuuster. Hij kan prima overweg met een kompas dat wat vlekken heeft. Zelfs als de getallen niet perfect zijn, blijft hij stabiel omhoog klimmen.

Wat betekent dit voor de toekomst?

De auteurs hebben bewezen dat:

1. Het werkt: Je kunt deze slimme methoden gebruiken met "ruwe" getallen zonder dat de AI faalt.
2. De regels zijn duidelijk: Je moet alleen zorgen dat de "ruwheid" (het aantal bits) niet te laag is. Als je maar een paar bits hebt, wordt het te rommelig. Maar met een beetje meer (zoals 8 bits), is het prima.
3. Muon is de toekomst voor lage precisie: Omdat Muon minder gevoelig is voor die kleine foutjes, is hij waarschijnlijk de betere keuze voor de toekomst, waar we steeds meer AI's op goedkope, snelle hardware willen draaien.

Samenvattend in één zin:
Deze paper legt uit waarom het trainen van super-slimme AI's met "goedkope" en minder precieze rekenkracht toch werkt, en onthult dat de nieuwe methode Muon hier veel beter tegen kan dan de oude standaard Adam, omdat hij minder snel in paniek raakt door kleine rekenfoutjes.

Dit is een grote stap om AI goedkoper en sneller te maken voor iedereen!

Technische samenvatting

Probleemstelling

De snelle schaalvergroting van Large Language Models (LLMs) heeft low-precision training (bijvoorbeeld met BF16 of FP8) onmisbaar gemaakt om het geheugengebruik te verminderen en de rekenefficiëntie te verhogen. Hoewel deze methoden empirisch zeer succesvol zijn, ontbreekt er een strikte theoretische onderbouwing voor het gebruik van adaptieve optimalisatoren (zoals Adam en Muon) onder volledige kwantisatie.

Bestaande theoretische werken focussen voornamelijk op:

1. SGD met gekwantiseerde gradiënten (QSGD): Vaak gebaseerd op aannames van onbevooroordeelde kwantisatie of foutfeedback-mechanismen.
2. Beperkte componenten: Veel analyses kwantiseren alleen de gradiënten of weights, maar negeren de kwantisatie van de optimizer states (zoals momenten en tweede-moment schattingen), wat cruciaal is in moderne low-bit LLM-training.
3. Onrealistische aannames: Veel theorieën gaan uit van onbevooroordeelde kwantisatie, terwijl hardware-gebaseerde floating-point kwantisatie (zoals FP32 naar BF16) inherent een relatieve fout introduceert en niet noodzakelijk onbevooroordeeld is.

Dit creëert een kloof tussen de empirische succesvolle toepassing van volledig gekwantiseerde training en het theoretisch inzicht in waarom en wanneer deze convergeren.

Methodologie

De auteurs introduceren het eerste theoretische kader voor het analyseren van de convergentie van adaptieve optimalisatoren onder een hardware-bewust floating-point kwantisatiemodel.

1. Kwantisatiemodel:

   • In plaats van onbevooroordeelde kwantisatie of foutfeedback, gebruiken ze een relatieve foutmodel (Assumptie 3.1).
   • Voor een variabele $x$ en zijn gekwantiseerde versie $x_Q$ geldt: $|x_Q - x| \leq q|x|$, waarbij $q = \Theta(2^{-M})$ en $M$ de lengte van de mantisse is. Dit model weerspiegelt nauwkeurig het gedrag van floating-point formaten (waarbij de exponent behouden blijft en alleen de mantisse wordt afgekapt).
   • Het kader modelleert kwantisatie voor alle componenten: weights ($W$), gradiënten ($G$), eerste momenten ($M$) en tweede momenten ($V$).

2. Onderzochte Optimalisatoren:

   • Adam: De standaard adaptieve optimizer.
   • Muon: Een nieuwere optimizer die gebruikmaakt van een SVD-based sign-operator (Singular Value Decomposition) voor matrix-gebaseerde optimalisatie, wat populair wordt in LLM-training.

3. Analytische Benadering:

   • De auteurs leiden convergentiegaranties af voor gladde, niet-convexe objectief functies onder standaard stochastische gradiëntveronderstellingen.
   • Ze analyseren hoe de relatieve fouten ($q_W, q_G, q_M, q_V$) de convergentiesnelheid beïnvloeden en hoe deze fouten zich voortplanten door de iteraties.

Belangrijkste Bijdragen

1. Eerste Theoretisch Kader: Het paper biedt de eerste rigoureuze convergentieanalyse voor adaptieve optimalisatoren (Adam en Muon) waarbij gradiënten, weights én optimizer states allemaal onderhevig zijn aan floating-point kwantisatie, zonder reliance op onrealistische foutfeedback-mechanismen.
2. Convergentiegaranties:
   • Ze bewijzen dat zowel Adam als Muon convergentiesnelheden behalen die dicht bij hun full-precision tegenhangers liggen, mits de mantisse-lengte $M$ logaritmisch schaalt met het aantal iteraties ($M = \Omega(\log T)$).
   • De verwachte norm van de gradiënt convergeert met een snelheid van $\tilde{O}(T^{-1/4})$, wat overeenkomt met de bekende rates voor full-precision Adam.
3. Differentiatie in Sensitiviteit:
   • Adam: De analyse toont aan dat Adam extreem gevoelig is voor kwantisatie van weights en tweede momenten. Dit komt door de afhankelijkheid van de parameter $\beta_2$ (die vaak dicht bij 1 wordt ingesteld). De inverse vierkantswortel van de historische gradiëntvariantie versterkt kwantisatiefouten in de tweede momenten.
   • Muon: De analyse toont aan dat Muon weinig gevoelig is voor kwantisatie. De SVD-based sign-operator voorkomt de versterking van fouten door de inverse vierkantswortel van de variantie. Muon vereist minder strikte foutcontrole (relatieve fouten van orde $O(T^{-1/2})$ zijn voldoende, terwijl Adam soms $O(T^{-2})$ vereist voor weights en tweede momenten).

Resultaten

1. Theoretische Resultaten:

   • Voor Adam: Om de $\tilde{O}(T^{-1/4})$ snelheid te behouden, moeten de relatieve fouten voor weights en tweede momenten zeer snel afnemen (bijv. $q_W, q_V = O(T^{-2})$). Dit verklaart waarom in de praktijk vaak hogere precisie nodig is voor deze componenten.
   • Voor Muon: De optimizer behoudt zijn convergentie met minder strikte voorwaarden ($q = O(T^{-1/2})$ voor alle componenten), wat theoretisch onderbouwt waarom Muon robuuster is in low-precision omgevingen.

2. Empirische Validatie:

   • De theorie werd getest op synthetische data (Rosenbrock-functie), CIFAR-10 en nanoGPT (OpenWebText).
   • Observaties:
     • Lagere mantisse-bits (bijv. M=4) leiden tot significante degradatie in convergentie, vooral bij Adam.
     • Bij matige mantisse-lengtes (bijv. M=10 of hoger) presteren beide optimalisatoren bijna even goed als full-precision.
     • Muon toont consistent betere robustheid dan Adam bij zeer lage precisie (bijv. M=2), wat de theoretische voorspelling bevestigt dat Adam gevoeliger is voor foutversterking door $\beta_2 \to 1$.
     • Experimenten tonen aan dat de kwantisatiefout direct correleert met de convergentieprestaties: minder bits = grotere fout = langzamere convergentie of stagnatie.

Betekenis en Impact

Dit paper sluit een kritieke kloof tussen de empirische praktijk en de theoretische wetenschap van low-precision training:

• Verklaring van Empirisch Succes: Het legt uit waarom low-precision training werkt: zolang de mantisse-lengte logaritmisch schaalt met de trainingsduur, blijven de kwantisatiefouten beheersbaar en convergentie gegarandeerd.
• Ontwerprichting voor Optimizers: De analyse suggereert dat toekomstige low-precision optimalisatoren (zoals Muon) die de versterking van kwantisatiefouten vermijden (bijv. door het vermijden van de inverse vierkantswortel van varianties), inherent robuuster zullen zijn dan traditionele methoden zoals Adam.
• Praktische Richtlijnen: Het biedt theoretische onderbouwing voor de praktijk dat weights en tweede momenten in Adam mogelijk iets hogere precisie nodig hebben dan gradiënten of eerste momenten, en bevestigt dat Muon een superieur alternatief kan zijn voor het trainen van grote modellen met beperkte hardware-resources.

Kortom, dit werk biedt de eerste rigoureuze wiskundige basis voor het gebruik van volledig gekwantiseerde adaptieve optimalisatoren in de moderne deep learning-pipeline.