Résumé Technique : OrScale – Optimisation Orthogonalisée avec Mise à l'Échelle du Ratio de Confiance par Couche

1. Énoncé du Problème

L'optimiseur Muon s'est imposé comme un outil pratique pour l'entraînement à grande échelle, utilisant des mises à jour orthogonalisées pour fixer la direction des pas à valeur matricielle via des itérations de Newton–Schulz (approximant le facteur polaire). Cependant, Muon souffre d'une limitation critique : bien qu'il optimise la direction des mises à jour, il laisse la magnitude contrôlée par un seul taux d'apprentissage essentiellement global.

Cette mise à l'échelle globale est sous-optimale car différentes couches (par exemple, les projections d'attention, les projections montantes/descendantes des MLP, les embeddings) présentent des distributions de gradients et des courbures effectives radicalement différentes. Bien que la variante Moonlight ait introduit un facteur statique par forme pour égaliser la RMS par élément, elle ne peut pas s'adapter à la dérive dynamique des échelles de couches pendant l'entraînement.

Les tentatives précédentes de portage de la mise à l'échelle adaptative du taux d'apprentissage par couche (LARS) et de la mise à l'échelle adaptative du moment par couche (LAMB) — qui ont réussi à découpler les tailles de pas des taux d'apprentissage globaux dans SGD et Adam — vers Muon ont échoué. L'article identifie que les ports naïfs s'effondrent en raison de trois raisons structurelles spécifiques :

1. Dénominateurs Dégénérés : L'utilisation de la norme de Frobenius de la direction orthogonalisée ($\|Q_\ell\|_F$) comme dénominateur du ratio de confiance échoue car elle est presque constante, déterminée par la forme de la matrice (rang) et non par la magnitude de la mise à jour.
2. Incompatibilité d'Unités et Saturation du Recadrage : L'utilisation de la norme brute du momentum ($\|\tilde{M}_\ell\|_F$) introduit une incompatibilité d'unités (unités de gradient vs unités de poids), ce qui fait que le ratio de confiance atteint presque universellement la borne supérieure de recadrage ($r_{max}$), rendant l'adaptation inefficace.
3. Emballement de la Décroissance des Poids Découplée : La combinaison de dénominateurs calibrés avec une décroissance des poids découplée crée une boucle de rétroaction positive où les normes de poids croissent de manière incontrôlable (10–50$\times$), déstabilisant l'entraînement.

2. Méthodologie

Les auteurs proposent OrScale, un cadre fondé sur un seul principe de conception : Le dénominateur d'un ratio de confiance par couche doit mesurer la norme de Frobenius de la vraie direction de mise à jour que l'optimiseur soustraira des paramètres.

Principe Central

Soit $D_{\ell,t}$ la direction réelle dans l'espace des paramètres soustraite des poids $W_{\ell,t}$ avant l'application du taux d'apprentissage global. Le ratio de confiance est défini par :
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
Crucialement, $D_{\ell,t}$ doit inclure à la fois le terme de décroissance des poids et la direction orthogonalisée.

Deux Variantes

Le cadre produit deux implémentations spécifiques :

1. OrScale (Défaut Général/Vision) :

   • Direction de Mise à Jour : $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$ (où $\lambda$ est la décroissance des poids et $Q$ est la direction orthogonalisée).
   • Facteur de Forme : $s_\ell = 1$ (aucune mise à l'échelle statique).
   • Calibration : Aucune.
   • Décroissance des Poids : Intégrée directement dans l'étape de mise à jour mise à l'échelle par le ratio de confiance.

2. OrScale-LM (Défaut LLM) :

   • Direction de Mise à Jour : $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$.
   • Facteur de Forme : Adopte le facteur statique de Moonlight $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ pour correspondre à la RMS par élément d'AdamW.
   • Calibration : Introduit une constante de calibration paresseuse par couche ($c_{\text{denom}, \ell}$) calculée une seule fois à l'initialisation. Cela ancre le ratio de confiance initial exactement à 1 ($r_{\ell,0}=1$).
   • Signification de la Calibration : Cela permet à OrScale-LM d'hériter des hyperparamètres de taux d'apprentissage et de décroissance des poids ajustés pour Moonlight (et par extension, AdamW) sans nécessiter une nouvelle recherche d'hyperparamètres. Cela assure l'invariance de la largeur à l'initialisation.

Détails Algorithmiques

• Décroissance des Poids Couplée : La décroissance des poids est intégrée dans le calcul du ratio de confiance ($\lambda W$ fait partie du dénominateur) plutôt que d'être appliquée comme un rappel séparé. Cela empêche la boucle de rétroaction d'emballement observée dans les variantes découplées.
• Efficacité : La méthode ajoute une surcharge négligeable (<1% du temps réel), nécessitant uniquement un scalaire supplémentaire par couche matricielle et deux calculs de norme de Frobenius par étape, qui sont dominés par les itérations de Newton–Schulz existantes.

3. Contributions Clés

• Algorithme : Une recette unifiée et légère en paramètres (Algorithme 1) qui intègre avec succès la mise à l'échelle du ratio de confiance par couche dans Muon. Les auteurs identifient le seul "coin fonctionnel" unique dans l'espace de conception (Dénominateur de direction de mise à jour réelle + Décroissance des poids couplée).
• Analyse des Échecs : Une cartographie rigoureuse de l'espace de conception du ratio de confiance de Muon (Tableau 1) qui caractérise analytiquement et empiriquement trois modes d'échec (dénominateur dégénéré, saturation du recadrage et croissance incontrôlée des poids), démontrant pourquoi ils échouent et pourquoi la solution proposée réussit.
• Théorie :
  • Convergence : Prouve une garantie de convergence non convexe en norme nucléaire $O(1/\sqrt{T})$.
  • Gain Adaptatif : Établit une constante de descente adaptative par couche $\kappa_{\text{eff}} > 1$ sous une hétérogénéité de couche mesurable, améliorant strictement la constante homogène de Muon.
  • Séparation : Prouve une séparation stricte par rapport à MuTrust/MuScale dans des conditions de saturation du recadrage.
  • Calibration : Démontre que OrScale-LM préserve le transfert de taux d'apprentissage de style $\mu$P à l'initialisation.
• Empirisme : Évaluation complète sur des tâches de modèles de vision et de langage.

4. Résultats Expérimentaux

Vision (CIFAR-10 / DavidNet)

• Performance : OrScale a atteint une précision top-1 de validation de 94,05 % $\pm$ 0,08, se classant premier parmi cinq bases de référence.
• Amélioration : Cela représente une amélioration de +0,35 point par rapport à Muon standard (93,70 %) et de +0,30 par rapport à Muon+Moonlight (93,75 %).
• Stabilité : OrScale a démontré une fenêtre de taux d'apprentissage stable plus large que Muon, cohérente avec les prédictions théoriques de gains adaptatifs par couche.
• Ablation : Une étude d'ablation des modes d'échec a confirmé que les variantes correspondant aux modes d'échec identifiés (par exemple, utilisant $\|Q_\ell\|_F$ ou le momentum brut) présentaient les symptômes prédits (par exemple, saturation du recadrage à 99,5–100 % ou emballement de la norme des poids) et n'ont pas réussi à égaler les performances d'OrScale.

Grands Modèles de Langage (Pré-entraînement FineWeb-Edu)

• Mise à l'Échelle : Évalué sur quatre échelles de modèles (de 125M à 1,1B de paramètres).
• Comparaison : OrScale-LM a surpassé Muon+Moonlight sur trois échelles sur quatre (125M, 545M, 1,1B) et a battu AdamW à chaque échelle.
• Réduction de la Perte : À l'échelle de 1,1B, OrScale-LM a atteint une entropie croisée de validation finale de 2,6251, contre 2,6360 pour Muon+Moonlight et 2,7304 pour AdamW.
• Efficacité : L'avantage a été préservé sur toute la plage de calcul, OrScale-LM montrant une victoire nette dans le régime où le calcul compte le plus pour les LLM.

5. Signification et Revendications

L'article revendique qu'OrScale comble avec succès le fossé entre la géométrie spectrale de Muon et l'adaptativité par couche de LARS/LAMB. Sa signification principale réside dans :

1. Résolution du Problème de Magnitude : Il fournit le premier contrôle de magnitude par couche et variant dans le temps pour Muon, robuste aux modes d'échec structurels.
2. Transférabilité Pratique : Grâce à la variante OrScale-LM et à la calibration paresseuse, il permet aux praticiens de transférer directement les taux d'apprentissage ajustés pour AdamW ou Moonlight vers OrScale-LM sans réajustement, réduisant ainsi la barrière à l'adoption.
3. Solidité Théorique : Il offre une justification théorique rigoureuse expliquant pourquoi les tentatives hybrides précédentes ont échoué et prouve que la méthode proposée produit des constantes de convergence strictement meilleures sous hétérogénéité.
4. Validation Empirique : Il démontre que ces améliorations théoriques se traduisent par des gains tangibles à la fois dans des tâches de vision à petite échelle et dans le pré-entraînement de LLM à grande échelle, établissant OrScale comme un optimiseur de l'état de l'art pour les couches à valeur matricielle.

Les auteurs maintiennent une modestie concernant les résultats sur les LLM, notant que l'avantage à l'échelle de 1,1B est comparable au bruit d'une seule graine à plus petite échelle, et que les résultats doivent donc être considérés comme une preuve directionnelle de l'efficacité de la méthode plutôt que comme une supériorité statistique définitive dans tous les régimes. Cependant, l'amélioration constante à travers les échelles et l'ablation réussie des modes d'échec soutiennent fortement le principe de conception proposé.