Résumé Technique : CacheMuon

Énoncé du Problème

Muon est un optimiseur sensible à la géométrie, conçu pour les paramètres de couches cachées sous forme de matrices dans le deep learning. Contrairement aux optimiseurs par coordonnées (ex: Adam), Muon remplace les mises à jour de momentum standard par une direction orthogonalisée dérivée du facteur polaire de la matrice de momentum. Bien que Muon ait démontré de solides performances empiriques et une grande scalabilité pour les grands modèles de langage, son efficacité pratique est limitée par l'étape d'orthogonalisation.

Les implémentations actuelles approximent le facteur polaire via des itérations de Newton–Schulz (plus précisément Gram Newton–Schulz) pour éviter le coût élevé d'une SVD (décomposition en valeurs singulières) exacte. Cependant, cette approximation nécessite toujours plusieurs multiplications matricielles à chaque étape d'optimisation. Par conséquent, une partie significative du coût de calcul de Muon est consacrée à la résolution répétée de problèmes d'orthogonalisation de matrices très proches. Bien que des travaux récents aient optimisé l'efficacité d'un appel d'orthogonalisation unique (ex: PolarExpress, Gram Newton–Schulz), ces méthodes recomputent systématiquement la transformée d'orthogonalisation à partir de zéro à chaque étape, ignorant la corrélation temporelle inhérente au processus d'entraînement.

Méthodologie : CacheMuon

Les auteurs proposent CacheMuon, une méthode de préconditionnement temporel qui exploite l'évolution lisse de la matrice de momentum durant l'entraînement. L'hypothèse centrale est que si la matrice de momentum change graduellement, la transformée d'orthogonalisation calculée à une étape précédente reste informative pour l'étape actuelle.

Mécanisme Central

CacheMuon modifie la routine d'orthogonalisation de Muon en introduisant un mécanisme de mise en cache :

1. Résolution Fraîche (Étape de Rafraîchissement) : À des étapes de « rafraîchissement » spécifiques (déterminées par un déclencheur), la méthode exécute un solveur complet Fresh Gram Newton–Schulz (FreshGNS). Cela produit le facteur polaire approximatif $\tilde{D}_t$ et, de manière cruciale, une matrice de transformée gauche mise en cache $Q_t$.
2. Réutilisation en Cache (Étapes Intermédiaires) : Lors des étapes de non-rafraîchissement, au lieu de lancer un solveur complet, la méthode réutilise la transformée mise en cache $Q_{a_t}$ (où $a_t$ est l'indice du dernier rafraîchissement) pour approximer le facteur polaire : $\hat{D}_t = Q_{a_t} X_t$, où $X_t$ est la matrice de momentum normalisée.
3. Déclencheur Basé sur le Résiduel : Pour déterminer quand la transformée en cache n'est plus précise, les auteurs introduisent un déclencheur de réutilisation léger basé sur le résiduel. Avant de réutiliser un cache, l'algorithme calcule le résiduel d'orthogonalité de la mise à jour candidate :
   $$\text{err}_t = \frac{\| \hat{D}^{\text{cand}}_t (\hat{D}^{\text{cand}}_t)^\top - I \|_F}{\|I\|_F}$$
   Si $\text{err}_t \leq \gamma$ (un seuil prescrit), la transformée en cache est acceptée. Sinon, une résolution fraîche est déclenchée pour mettre à jour le cache.

Analyse Théorique

Le papier analyse CacheMuon comme une mise à jour Muon inexacte. L'écart entre la direction mise en cache $\hat{D}_t$ et la direction Muon exacte $D_t$ est borné par un terme d'erreur additif $\delta_t$ :
$$\| \hat{D}_t - D_t \|_{\text{sp}} \leq \delta_t = L_Q \tau_t + \varepsilon_{\text{ref}}$$
où :

• $\varepsilon_{\text{ref}}$ est l'erreur du solveur frais (due au nombre fini d'itérations de Newton–Schulz).
• $\tau_t$ est la « vétusté » (staleness) de la matrice de Gram en cache par rapport à la matrice actuelle.
• $L_Q$ est une constante de Lipschitz dérivée de la application polynomiale induite par le solveur.

Sous des hypothèses standards (lissage, variance bornée, vétusté bornée), les auteurs prouvent une garantie de convergence stochastique pour l'optimisation non convexe. Le taux de convergence dépend du paramètre d'inexactitude agrégé $\bar{\delta} = L_Q \bar{\tau} + \varepsilon_{\text{ref}}$. Tant que $\bar{\delta} < 1$, la méthode converge, avec un taux d'échelle de $O((1-\bar{\delta})^{-1} N^{-1/4})$.

Principales Contributions

1. Algorithme CacheMuon : Une variante de Muon qui amortit le coût du calcul du facteur polaire en réutilisant les transformées d'orthogonalisation mises en cache à travers les étapes d'optimisation, en tirant parti de la lissité temporelle.
2. Théorie de la Convergence : Un cadre théorique traitant CacheMuon comme un optimiseur inexact, fournissant des garanties de convergence stochastique contrôlées par l'erreur du solveur frais et la vétusté du cache.
3. Déclencheur Basé sur le Résiduel : Un heuristique pratique et peu coûteux qui décide dynamiquement quand rafraîchir le cache en fonction de l'orthogonalité de la mise à jour candidate, servant de proxy pour la fidélité de la mise à jour.
4. Validation Empirique : Démonstration d'une frontière qualité–efficacité contrôlable sur des charges de travail de modèles de langage (GPT-2) et de vision (ResNet-18).

Résultats Expérimentaux

Les auteurs ont évalué CacheMuon sur GPT-2 Large/Small (OpenWebText) et ResNet-18 (CIFAR-10).

• Compromis Qualité–Efficacité : En faisant varier le seuil de résiduel $\gamma$, les auteurs démontrent une frontière fluide.
  • Réutilisation Conservatrice ($\gamma=2$) : Suit de près les performances de Muon frais (GramMuon et PolarExpMuon) tout en réduisant les FLOPs d'orthogonalisation de 13,6 % (vs GramMuon) et 30,9 % (vs PolarExpMuou) sur GPT-2 Large.
  • Réutilisation Agressive ($\gamma=15$) : Atteint jusqu'à 65,0 % de réduction des FLOPs d'orthogonalisation (vs GramMuon) avec seulement une dégradation modeste de la perte de validation (ex: 3,096 vs 3,037 pour GPT-2 Large).
• Aperçus Diagnostiques : Les expériences montrent que le déclencheur de résiduel est efficacement corrélé à l'alignement entre les mises à jour en cache et les mises à jour fraîches. Au début de l'entraînement, les résiduels sont élevés, forçant des rafraîchissements fréquents ; à mesure que l'entraînement se stabilise, les résiduels chutent, permettant des taux de succès de cache plus élevés et une plus grande efficacité.
• Scalabilité : La méthode maintient des performances compétitives à travers différentes échelles de modèles (GPT-2 Small vs Large) et modalités (Vision vs Langage).

Signification et Revendications

Le papier affirme que CacheMuon offre une approche pratique pour réduire la surcharge de calcul de Muon sans altérer sa règle de mise à jour géométrique sous-jacente. En traitant l'orthogonalisation comme un problème temporel plutôt que purement spatial, la méthode réduit les calculs redondants.

Les auteurs soulignent que bien que la méthode réduise considérablement les FLOPs d'orthogonalisation, le gain de vitesse réel (wall-clock speedup) dans leur implémentation de haut niveau actuelle est modeste en raison des surcharges du système (ex: synchronisation hôte-périphérique pour la vérification du résiduel). Ils postulent qu'une implémentation plus optimisée de bas niveau (ex: fusion des sondes dans les noyaux de calcul/kernels) pourrait traduire une plus grande fraction de ces économies arithmétiques en gains de temps d'exécution réels.

En fin de compte, ce travail suggère que la réutilisation temporelle est une direction prometteuse pour optimiser les optimiseurs sensibles à la géométrie, offrant un levier ajustable pour équilibrer l'efficacité computationnelle et la qualité de l'entraînement.