Technische Samenvatting: Pro-KLShampoo

Probleemstelling

Moderne vooropleiding van Large Language Models (LLM) is sterk afhankelijk van optimalisatiealgoritmen die de matrixstructuur van gradiënten benutten. Er hebben zich twee verschillende fronten ontwikkeld:

1. Expliciete Kronecker-gefactoreerde preconditionering: Vertegenwoordigd door KL-Shampoo, dat de preconditioner schat door de Kullback-Leibler (KL)-divergentie te minimaliseren tussen de ware gradiëntcovariantie en een Kronecker-gestructureerde benadering.
2. Orthogonalisatie van gradiëntmomentum: Geïllustreerd door Muon, dat het momentum orthogonaliseert om de steilste afdaling onder de spectrale norm te realiseren.

Deze benaderingen worden doorgaans geïsoleerd ontwikkeld. Hoewel KL-Shampoo sterke theoretische onderbouwing biedt via KL-minimalisatie, brengt het hoge rekenkundige en geheugenkosten met zich mee vanwege het handhaven van volledige Kronecker-factoren. Omgekeerd zijn orthogonalisatiemethoden zoals Muon efficiënter, maar werken ze binnen een ander geometrisch raamwerk (steilste afdaling onder de spectrale norm) zonder het perspectief van KL-schatting. Het artikel adresseert de noodzaak om deze benaderingen te verenigen om de statistische efficiëntie van KL-Shampoo te bereiken met de rekenkundige efficiëntie van orthogonalisatiemethoden.

Methodologie: Pro-KLShampoo

De auteurs stellen Pro-KLShampoo voor (Projected KL-Shampoo met Whitening Hersteld door Orthogonalisatie). De kerninnovatie vloeit voort uit een structurele observatie van de preconditioners van KL-Shampoo: hun eigenwaarde-spectra vertonen een "spike-and-flat" (piek-en-vlak) vorm. Dit bestaat uit een paar dominante eigenwaarden (de "piek") gevolgd door een ongeveer uniforme staart (het "vlak"). Deze structuur wordt empirisch waargenomen over lagen en trainingsfasen heen en is theoretisch exact onder een gradiëntmodel van signaal-plus-ruis met rang $\rho$.

1. Structurele Restrictie

Pro-KLShampoo beperkt een van de Kronecker-factoren van KL-Shampoo (specifiek de rechterfactor $R$ voor rechthoekige lagen waar $m \le n$) tot een parametrische familie die is afgestemd op de spike-and-flat-structuur:

• Projectie-ruimte ($r$-dimensionaal): De preconditioner behoudt de volledige spectrale structuur op een bijgehouden deelruimte opgespannen door $U \in \text{St}(n, r)$.
• Complement-ruimte ($n-r$ dimensies): De overige richtingen worden samengevoegd tot één gedeelde eigenwaarde $\mu_\perp$.

Deze restrictie verlaagt de opslag van de rechterfactor van $O(n^2)$ naar $O(nr)$ en de kosten van de QR-decompositie van $O(n^3)$ naar $O(r^3)$.

2. Herstel via Orthogonalisatie

Een naïeve restrictie tot een scalaire staart verliest informatie als de staart niet perfect uniform is. Om de algebraïsche vorm van de volledige KL-Shampoo-preconditioner op de complement-ruimte te herstellen, past Pro-KLShampoo orthogonalisatie (polair decompositie) toe op het complement van de gradiëntcomponent.

De updateregel is:
$$\Delta W = -\alpha_{kl} L^{-1/2} G U S^{-1/2} U^\top - c_a \text{polar}\left(\mu_\perp^{-1/2} L^{-1/2} G P_\perp\right)$$
Waarbij:

• De eerste term de volledige KL-Shampoo-update binnen de bijgehouden deelruimte vertegenwoordigt.
• De tweede term orthogonalisatie toepast op het complement.
• De auteurs een identiteit bewijzen (Eq. 12) die aantoont dat deze orthogonalisatie de algebraïsche vorm van de whitening per richting van de volledige KL-Shampoo op het complement herstelt, en zo Muon-stijl orthogonalisatie effectief verbindt met het KL-schattingkader.

3. Deelruimte-tracking

De optimale deelruimte $U$ wordt geïdentificeerd als het eigenruimte van de matrix $\Phi_L = E[G^\top L^{-1} G]$. Onder de spike-and-flat-conjectuur is de eigenruimte van de top-$r$ van $\Phi_L$ de optimale keuze. Het algoritme houdt deze deelruimte bij met behulp van één stap van machtsiteratie (via QR-decompositie) op een lopende schatting van $\Phi_L$.

Belangrijkste Bijdragen

1. Structurele Observatie & Parametrische Restrictie: Het artikel identificeert de "spike-and-flat" eigenwaarde-structuur in de preconditioners van KL-Shampoo (empirisch robuust en theoretisch exact onder gradiëntmodellen met lage rang) en benut dit door het KL-doel te beperken tot een parametrische familie.
2. Algebraïsch Herstel via Orthogonalisatie: Het toont aan dat het toepassen van orthogonalisatie op de complement-ruimte de algebraïsche vorm van de volledige KL-Shampoo-preconditioner herstelt. Dit biedt een theoretische brug tussen Muon-stijl orthogonalisatie en KL-schatting.
3. Convergentiegarantie: De auteurs bieden een convergentiegarantie voor niet-convexe optimalisatie met een snelheid van $O(1/\sqrt{T})$ voor het geïdealiseerde algoritme onder operator-norm gladheid.
4. Empirische Prestaties: Uitgebreide evaluatie op GPT-2 (124M/350M) en LLaMA (134M/450M) modellen toont aan dat Pro-KLShampoo consistent KL-Shampoo overtreft in validatieverlies, piekgeheugengebruik per GPU en wandkloktijd om gelijkwaardige verliesniveaus te bereiken over alle geteste rangen van de deelruimte ($r \in \{32, 64, 128\}$).

Experimentele Resultaten

De auteurs evalueerden Pro-KLShampoo tegen AdamW, Muon en KL-Shampoo op vier vooropleidingsconfiguraties.

• Validatieverlies: Pro-KLShampoo behaalde het laagste validatieverlies bij elke geteste rang over alle modelgroottes.
• Geheugenefficiëntie: Het piekgeheugengebruik per GPU was lager dan bij KL-Shampoo en vergelijkbaar met Muon.
• Wandkloktijd: Pro-KLShampoo bereikte specifieke validatieverliesniveaus sneller dan KL-Shampoo.
  • Bij GPT-2 werd de snelheidswinst voornamelijk gedreven door een verminderde tijd per stap (vervanging van $O(n^3)$ QR door $O(r^3)$).
  • Bij LLaMA werd de snelheidswinst gedreven door snellere convergentie (lagere trainingsverliescurves), naast efficiëntie per stap.
• Ablatiestudies:
  • Het verwijderen van orthogonalisatie (Smok-Hop-variant) resulteerde in prestatieverlies, vooral bij grotere schalen (LLaMA 450M), wat de noodzaak bevestigt van orthogonalisatie om whitening per richting te herstellen.
  • Noch de alleen-deelruimte-varianten, noch de alleen-complement-varianten kwamen in de buurt van de prestaties van de volledige Pro-KLShampoo, wat aangeeft dat beide componenten essentieel zijn.

Betekenis en Claims

Het artikel beweert dat Pro-KLShampoo twee eerder geïsoleerde optimalisatiefronten succesvol verenigt. Door gebruik te maken van de intrinsieke "spike-and-flat" structuur van gradiëntcovarianties, bereikt het de statistische efficiëntie van KL-Shampoo terwijl het orthogonalisatie gebruikt om rekenkundige efficiëntie te behouden.

De auteurs benadrukken dat de vervanging van de update met een scalaire staart door orthogonalisatie geen heuristiek is, maar een algebraïsch herstel van de volledige KL-Shampoo-vorm. Dit stelt de methode in staat om de voordelen van KL-gebaseerde preconditionering (zoals robuustheid van de stapgrootte) te behouden, terwijl de rekenkundige overhead die gepaard gaat met het handhaven van volledige Kronecker-factoren aanzienlijk wordt verminderd. Het werk suggereert dat laag-dimensionale spectrale structuur alomtegenwoordig is in de optimalisatie van diepe netwerken en effectief kan worden benut om efficiëntere optimalisatiealgoritmen te ontwerpen voor de vooropleiding van LLM's.

Het artikel blijft bescheiden wat betreft beperkingen, met de opmerking dat de rang $r$ empirisch wordt ingesteld, en dat toekomstig werk nodig is voor per-laag adaptieve kalibratie van het menggewicht $\alpha_{kl}$ en een convergentieanalyse van de praktische variant die Newton-Schulz-iteraties en Nesterov-momentum omvat.