Technische Samenvatting: LoRA-Muon: Spectrale Steilste Afdaling op de Laag-Rang Manifold

1. Probleemstelling

Low-Rank Adaptation (LoRA) wordt breed toegepast voor het fine-tunen van grote deep learning modellen vanwege de efficiëntie in het verminderen van trainings-FLOPs en geheugenvereisten. Echter, standaard optimizers zoals AdamW, wanneer deze onafhankelijk worden toegepast op de LoRA-factoren ($A$ en $B$), introduceren aanzienlijke afstemmingsproblemen vergeleken met dense training. Deze problemen omvatten:

• Gevoeligheid voor Initialisatie: Prestaties zijn sterk afhankelijk van de keuzes bij de initialisatie.
• Slechte Hyperparameter-overdraagbaarheid: Optimale leersnelheden (learning rates) laten zich niet goed overdragen over verschillende rangen, breedtes, dieptes of factor-schaalconventies.
• Divergerende Dynamiek: Zelfs wanneer ze starten vanuit identieke samengestelde gewichten ($W = AB^T$), divergeren de trainingsdynamieken vaak als de factor-schalen uit elkaar gaan lopen.
• Suboptimale Prestaties: Natieve laag-rang pretraining heeft vaak moeite om dense baselines te evenaren zonder extra stabilisatie.

De kernuitdaging is het ontwerpen van een optimizer voor de laag-rang manifold die de geometrische eigenschappen van full-rank optimizers (specifiek Muon en de Shampoo-familie) behoudt, terwijl het zorgt voor optimale overdraagbaarheid van leersnelheden over rang, breedte, diepte en factor-rescaling.

2. Methodologie

2.1 Geometrische Fundering

Het artikel herformuleert optimizer-ontwerp als een geometrisch probleem: het identificeren van de norm waaronder een optimizer de steilste afdaling (steepest descent) uitvoert.

• Muon Optimizer: Voert de steilste afdaling uit onder de spectrale norm ($\|\cdot\|_{2\to2}$) in de ambient matrixruimte $\mathbb{R}^{m \times n}$. De update-richting wordt gegeven door de matrix-signfunctie, $-\text{msign}(G)$.
• Laag-Rang Manifold: LoRA beperkt gewichten tot een manifold $\mathcal{M}_r = \{W \in \mathbb{R}^{m \times n} : \text{rank}(W) = r\}$. De auteurs stellen voor om het probleem van de steilste afdaling direct op de raakruimte (tangent space) van deze manifold op te lossen, $T_W \mathcal{M}_r$, in plaats van de factoren $A$ en $B$ als onafhankelijke coördinaten te optimaliseren.

2.2 Afleiding van LoRA-Muon

De auteurs leiden LoRA-Muon af door het geconstreerde lineaire minimalisatieprobleem op de laag-rang raakruimte onder de spectrale norm op te lossen.

1. Raakruimte Decompositie: Updates op $W$ worden gedecomposeerd in $\Delta W = \Delta A B^T + A \Delta B^T$.
2. Trust-Region Splitsing: Om het probleem hanteerbaar te maken, wordt het trust-region budget $\eta$ even verdeeld tussen de twee raakcomponenten. Dit wordt gerechtvaardigd door de empirische observatie dat de twee componenten tijdens de training snel uitlijnen, waardoor de driehoeksongelijkheid bijna nauw aansluit.
3. Projector Formulering: Gebruikmakend van unitaire invariantie, wordt het probleem gereduceerd tot het vinden van Linear Minimization Oracles (LMOs) op geprojecteerde gradiënten. De ideale update is:
   $$\Delta W = \frac{\eta}{2} \text{LMO}(\nabla W f \cdot P_B) + \frac{\eta}{2} \text{LMO}(P_A \cdot \nabla W f)$$
   waarbij $P_A$ en $P_B$ projectoren zijn op de kolomruimtes van $A$ en $B$.
4. Spectrale Specialisatie: Voor de spectrale norm is de LMO $-\text{msign}(\cdot)$. De resulterende factor-updates (Vergelijkingen 14–15) omvatten het "whitening" van de factor-gradiënten door de huidige Gram-matrices ($S_A = A^TA, S_B = B^TB$) en het toepassen van de spectrale signfunctie:
   $$\Delta A^* = -\frac{\eta}{2} \text{msign}(\nabla_A f \cdot S_B^{-1/2}) S_B^{-1/2}$$
   $$\Delta B^* = -\frac{\eta}{2} \text{msign}(\nabla_B f \cdot S_A^{-1/2}) S_A^{-1/2}$$

2.3 Split Weight Decay

Standaard factor-gewijs weight decay leidt tot een mismatch tussen de full-rank en laag-rang update-dynamiek. De auteurs leiden een split weight-decay regel af waarbij de decay-factor $s = \sqrt{1 - \lambda \eta}$ asymmetrisch wordt toegepast:
$$A_{t+1} = s A_t + \frac{1}{s} \Delta A^*_t, \quad B_{t+1} = s B_t + \frac{1}{s} \Delta B^*_t$$
Dit zorgt ervoor dat de samengestelde gewicht-update overeenkomt met de gewenste full-rank dynamiek, wat de overdraagbaarheid van hyperparameters waarborgt.

2.4 Gauge Invariantie

Een cruciale eigenschap van LoRA-Muon is gauge invariantie. De update op het samengestelde gewicht $W$ blijft invariant onder willekeurige inverteerbare transformaties $(A, B) \to (AR, BR^{-T})$. Dit wordt bereikt omdat de update alleen afhangt van de projectoren $P_A$ en $P_B$, die zelf gauge-invariant zijn.

3. Belangrijkste Bijdragen

1. Afleiding van LoRA-Muon: Een nieuwe familie van optimizers afgeleid door het probleem van de steilste afdaling op de laag-rang manifold op te lossen onder unitaire-invariante normen, met specialisatie naar de spectrale norm.
2. Split Weight-Decay Regel: Een specifieke update-regel die de dynamiek van laag-rang training afstemt op die van full-rank training, wat de overdraagbaarheid van hyperparameters mogelijk maakt.
3. Bewijs van Gauge Invariantie: Theoretisch bewijs dat de LoRA-Muon update invariant is onder de volledige gauge-groep actie, wat garandeert dat de optimizer afhangt van de geïnduceerde gewicht-update in plaats van willekeurige factor-coördinaten.
4. Vergelijking met Spectron: Demonstratie dat Spectron niet gauge-invariant is en gevoelig is voor willekeurige factor-rescaling, in tegen tegenstelling tot LoRA-Muon.
5. Relatie met LoRA-RITE: Bewijs dat de vereenvoudigde QR-coördinaatkern van LoRA-RITE algebraïsch identiek is aan de Loora-Muon spectrale update, hoewel LoRA-Muon de computationele overhead van QR-decomposities en second-moment opslag vermijdt.
6. Empirische Validatie: Bewijs dat LoRA-Muon fungeert als een effectieve proxy voor de dense Muon leersnelheid-zoektocht.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op character-level taalmodellering met behulp van de TinyShakespeare dataset met rekenkracht-gematchte dense en laag-rang baselines.

• Leersnelheid Overdraagbaarheid: De optimale leersnelheid voor dense Muon ($lr = 0.1$) draagt perfect over naar LoRA-Muon over sweeps van rang (2, 4, 8, 16, 32), breedte, diepte en scalar factor rescaling.
  • Bij Rang 2 herstelt LoRA-Muon de beste geteste dense leersnelheid.
  • Bij Rang 32 bereikt LoRA-Muon een lagere gemiddelde validatieverlies ($1.776 \pm 0.002$) dan de dense baseline ($1.789 \pm 0.002$) in de seed-gemiddelde sweep.
• Gauge Rebalancing: Wanneer factoren worden geïnitialiseerd met een grote scalaire onbalans (schaal 9), heeft het toepassen van een rebalancing stap een verwaarloosbaar effect op LoRA-Muon (verandering in loss $\approx -5.1 \times 10^{-5}$), wat de gauge invariantie bevestigt. In tegenstelling hiermee verandert de prestatie van Spectron aanzienlijk ($2.0 \times 10^{-2}$), wat de gevoeligheid voor factor-schaling bevestigt.
• Efficiëntie: Een rang-2 proxy gebruikt slechts ~4,3% van de trainbare parameters van het dense model, terwijl het een ~9,2x groter optimizer-stap budget ondersteunt onder dezelfde rekenkrachtbeperkingen.

5. Betekenis en Claims

Het artikel claimt dat LoRA-Muon een theoretisch gefundeerde oplossing biedt voor de afstemmingsproblemen van LoRA. Door LoRA te behandelen als een optimalisatieprobleem op de laag-rang manifold in plaats van een gefactoriseerd coördinatensysteem, bereiken de auteurs:

• Robuustheid: De optimizer is invariant voor de willekeurige keuze van factorisatie (gauge symmetrie).
• Overdraagbaarheid: Optimale leersnelheden (en daarmee andere hyperparameters via schaalwetten) dragen naadloos over tussen rang, breedte, diepte en factor-schalen.
• Efficiëntie: Het berekent de spectrale update zonder dure QR-decomposities of second-moment opslag, wat het accelerator-vriendelijk maakt.

De auteurs positioneren LoRA-Muon als een "goede laag-rang proxy" voor full-rank Muon en Shampoo-familie optimizers. Ze benadrukken dat de laag-rang parametrisatie de geometrie van de optimizer niet mag dicteren; in plaats daarvan moet de update gedefinieerd worden door de geometrie van de laag-rang matrix manifold zelf.

Beperkingen: De auteurs merken op dat de empirische claims momenteel zijn vastgesteld op TinyShakespeare-schaal modellen en nog niet zijn gevalideerd op grootschalige pretraining of downstream fine-tuning benchmarks. De studie richt zich op leersnelheid-overdraagbaarheid in plaats van een volledige hyperparameter sweep.