Technische Samenvatting: CacheMuon

Probleemstelling

Muon is een geometrie-bewuste optimizer ontworpen voor matrix-waardige verborgen laagparameters in deep learning. In tegenstelling tot coördinaat-gewijze optimizers (zoals Adam), vervangt Muon standaard momentum-updates door een georthogonaliseerde richting die afgeleid is van de polaire factor van de momentummatrix. Hoewel Muon sterke empirische prestaties en schaalbaarheid naar grote taalmodellen heeft getoond, wordt de praktische efficiëntie beperkt door de orthogonalisatie stap.

Huidige implementaties benaderen de polaire factor met behulp van Newton–Schulz iteraties (specifiek Gram Newton–Schulz) om de hoge kosten van exacte Singular Value Decomposition (SVD) te vermijden. Deze benadering vereist echter nog steeds meerdere matrixvermenigvuldigingen bij elke optimalisatiestap. Bijgevolg gaat een aanzienlijk deel van de computationele kosten van Muon verloren aan het herhaaldelijk oplossen van nauw verwante matrix-orthogonalisatieproblemen. Hoewel recente werken de efficiëntie van een enkele orthogonalisatie-oproep hebben geoptimaliseerd (bijv. PolarExpress, Gram Newton–Schulz), vereisen deze methoden nog steeds het vanaf nul opnieuw berekenen van de orthogonaliserende transformatie bij elke stap, waarbij de temporele correlatie die inherent is aan het trainingsproces wordt genegeerd.

Methodologie: CacheMuon

De auteurs stellen CacheMuon voor, een methode voor temporele preconditioning die de geleidelijke evolutie van de momentummatrix tijdens de training exploiteert. De kernhypothese is dat als de momentummatrix geleidelijk verandert, de georthogonaliseerde transformatie die berekend is bij een vorige stap nog steeds informatief blijft voor de huidige stap.

Kernmechanisme

CacheMuon wijzigt de orthogonalisatie-routine van Muon door een caching-mechanisme te introduceren:

1. Fresh Solve (Refresh-stap): Bij specifieke "refresh"-stappen (bepaald door een trigger), voert de methode een volledige Fresh Gram Newton–Schulz (FreshGNS) solver uit. Dit produceert de benaderde polaire factor $\tilde{D}_t$ en, cruciaal, een gecachte linker transformatiematrix $Q_t$.
2. Cached Reuse (Intermediaire stappen): Bij niet-refresh stappen, in plaats van een volledige solver uit te voeren, hergebruikt de methode de gecachte transformatie $Q_{a_t}$ (waarbij $a_t$ de index is van de meest recente refresh) om de polaire factor te benaderen: $\hat{D}_t = Q_{a_t} X_t$, waarbij $X_t$ de genormaliseerde momentummatrix is.
3. Residual-Based Trigger: Om te bepalen wanneer de gecachte transformatie niet langer accuraat is, introduceren de auteurs een lichtgewicht, op residu gebaseerde reuse-trigger. Voordat een cache wordt hergebruikt, berekent het algoritme het orthogonaliteitsresidu van de kandidaat-update:
   $$\text{err}_t = \frac{\| \hat{D}^{\text{cand}}_t (\hat{D}^{\text{cand}}_t)^\top - I \|_F}{\|I\|_F}$$
   Als $\text{err}_t \leq \gamma$ (een voorgeschreven drempelwaarde), wordt de gecachte transformatie geaccepteerd. Zo niet, dan wordt een fresh solve getriggerd om de cache bij te werken.

Theoretische Analyse

Het artikel analyseert CacheMuon als een inexacte Muon-update. De afwijking tussen de gecachte richting $\hat{D}_t$ en de exacte Muon-richting $D_t$ wordt begrensd door een additieve foutterm $\delta_t$:
$$\| \hat{D}_t - D_t \|_{\text{sp}} \leq \delta_t = L_Q \tau_t + \varepsilon_{\text{ref}}$$
waarbij:

• $\varepsilon_{\text{ref}}$ de fout van de fresh solver is (door eindige Newton–Schulz iteraties).
• $\tau_t$ de "staleness" (veroudering) van de gecachte Gram-matrix ten opzichte van de huidige is.
• $L_Q$ een Lipschitz-constante is afgeleid van de door de solver geïnduceerde polynomiale mapping.

Onder standaard aannames (gladheid, begrensde variantie, begrensde staleness), bewijzen de auteurs een stochastische convergentiegarantie voor niet-convexe optimalisatie. De convergentiesnelheid hangt af van de geaggregeerde inexactheidsparameter $\bar{\delta} = L_Q \bar{\tau} + \varepsilon_{\text{ref}}$. Zolang $\bar{\delta} < 1$, convergeert de methode, met een snelheid die schaalt als $O((1-\bar{\delta})^{-1} N^{-1/4})$.

Belangrijkste Bijdragen

1. CacheMuon Algoritme: Een variant van Muon die de kosten van de berekening van de polaire factor amortiseert door gecachte orthogonalisatie-transformaties over optimalisatiestappen te hergebruiken, gebruikmakend van temporele gladheid.
2. Convergentietheorie: Een theoretisch kader dat CacheMuon behandelt als een inexacte optimizer, en stochastische convergentiegaranties biedt die gecontroleerd worden door de fout van de fresh-solver en de staleness van de cache.
3. Residual-Based Trigger: Een praktische, lage-kosten heuristiek die dynamisch bepaalt wanneer de cache moet worden ververst op basis van de orthogonaliteit van de kandidaat-update, wat dient als een proxy voor de getrouwheid van de update.
4. Empirische Validatie: Demonstratie van een controleerbare kwaliteit–efficiëntie-frontier over taalmodel (GPT-2) en visie (ResNet-18) workloads.

Experimentele Resultaten

De auteurs evalueerden CacheMuon op GPT-2 Large/Small (OpenWebText) en ResNet-18 (CIFAR-10).

• Kwaliteit–Efficiëntie Afweging: Door de residu-drempelwaarde $\gamma$ te variëren, demonstreren de auteurs een vloeiende frontier.
  • Conservatieve Hergebruik ($\gamma=2$): Volgt nauwgezet de prestaties van fresh Muon (GramMuon en PolarExpMuon) terwijl de orthogonalisatie FLOPs worden verminderd met 13,6% (vs. GramMuon) en 30,9% (vs. PolarExpMuon) op GPT-2 Large.
  • Agressieve Hergebruik ($\gamma=15$): Bereikt tot wel 65,0% reductie in orthogonalisatie FLOPs (vs. GramMuon) met slechts een bescheiden degradatie in validatieverlies (bijv. 3.096 vs. 3.037 voor GPT-2 Large).
• Diagnostische Inzichten: Experimenten tonen aan dat de residu-trigger effectief correleert met de uitlijning tussen gecachte en fresh updates. In het begin van de training zijn de residuen hoog, wat frequente refreshes afdwingt; naarmate de training stabiliseert, dalen de residuen, wat een hogere cache hit rate en grotere efficiëntie mogelijk maakt.
• Schaalbaarheid: De methode behoudt concurrerende prestaties over verschillende modelmaten (GPT-2 Small vs. Large) en modaliteiten (Visie vs. Taal).

Betekenis en Claims

Het artikel claimt dat CacheMuon een praktische aanpak biedt om de computationele overhead van Muon te verminderen zonder de onderliggende geometrische update-regel te wijzigen. Door orthogonalisatie te behandelen als een temporeel probleem in plaats van een puur ruimtelijk probleem, vermindert de methode redundante berekeningen.

De auteurs benadrukken dat hoewel de methode de orthogonalisatie FLOPs aanzienlijk vermindert, de gerealiseerde wall-clock versnelling in hun huidige high-level implementatie bescheiden is vanwege systeem-niveau overheads (bijv. host-device synchronisatie voor de residu-check). Zij stellen dat een meer geoptimaliseerde low-level implementatie (bijv. het fuseren van probes in device kernels) een groter deel van deze rekenkundige besparingen kan vertalen in werkelijke runtime-winsten.

Uiteindelijk suggereert het werk dat temporele hergebruik een veelbelovende richting is voor het optimaliseren van geometrie-bewuste optimizers, waarbij het een regelbare knop biedt om computationele efficiëntie af te wegen tegen trainingskwaliteit.