Muon-OGD: Muon-based Spectral Orthogonal Gradient Projection for LLM Continual Learning

Résumé Technique : Muon-OGD

Énoncé du Problème
Le défi central abordé dans ce travail est l'oubli catastrophique dans les modèles de langage de grande taille (LLM) lors de l'apprentissage continu. Lorsque les LLM sont affinés séquentiellement sur de nouvelles tâches, leurs performances sur les tâches précédemment apprises se dégradent souvent rapidement. Les méthodes existantes d'apprentissage continu basées sur la projection (par exemple, OGD, GEM, O-LoRA) atténuent ce problème en restreignant les mises à jour des paramètres à des sous-espaces orthogonaux aux directions associées aux tâches passées. Cependant, ces méthodes sont généralement formulées sous une géométrie des paramètres euclidienne, où les magnitudes des mises à jour et l'orthogonalité sont régies par la norme de Frobenius.

Les auteurs identifient un décalage fondamental : les optimiseurs modernes pour les paramètres matriciels, tels que l'optimiseur Muon récemment proposé, fonctionnent sous une géométrie de norme spectrale (norme d'opérateur). Muon applique des mises à jour matricielles orthogonalisées (approchées via des itérations de Newton–Schulz) et peut être interprété comme effectuant une descente la plus raide sous une contrainte de norme spectrale. L'article soutient que s'en remettre à la géométrie de norme de Frobenius pour les contraintes d'apprentissage continu est sous-optimal pour les paramètres matriciels des LLM, car cela ne correspond pas à la géométrie des dynamiques d'optimisation sous-jacentes.

Méthodologie : Muon-OGD
L'article propose Muon-OGD (Projection Spectrale de Gradient Orthogonal basée sur Muon), un cadre qui intègre une optimisation consciente de la norme spectrale avec des contraintes d'apprentissage continu basées sur la projection.

1. Formulation : Au lieu du problème contraint standard de norme de Frobenius, Muon-OGD formule la mise à jour $\Delta$ comme un problème d'optimisation contraint par la norme spectrale :
   $$\min_{\Delta \in \mathbb{R}^{m \times n}} \langle G, \Delta \rangle \quad \text{s.t.} \quad \|\Delta\|_2 \leq \eta, \quad \langle C_i, \Delta \rangle = 0, \forall i=1,\dots,k$$
   Ici, $G$ est la matrice de gradient (ou de momentum), $\|\cdot\|_2$ est la norme spectrale, $\eta$ est la taille de l'étape, et $\{C_i\}$ sont les directions de contrainte (par exemple, les gradients des tâches passées) qui doivent rester orthogonales à la mise à jour.

2. Optimisation Duale : Puisque le problème ne possède pas de solution sous forme close, les auteurs dérivent une formulation duale. L'objectif dual implique la minimisation de la norme nucléaire d'une matrice décalée $H = G + \sum \lambda_i C_i$. La mise à jour optimale est récupérée via la fonction signe matricielle : $\Delta^* = -\eta \cdot \text{msgn}(H)$.

3. Algorithme : La méthode résout le problème dual de manière itérative en utilisant une descente de sous-gradient sur les variables duales $\lambda$.

   • Boucle Intérieure : Met à jour itérativement $\lambda$ pour faire respecter les contraintes d'orthogonalité.
   • Approximation du Signe Matriciel : La fonction signe matricielle $\text{msgn}(H)$ est efficacement approchée en utilisant un petit nombre d'itérations de Newton–Schulz (spécifiquement 5 itérations, notées $NS_5$), évitant ainsi le coût computationnel d'une SVD explicite.
   • Instanciation de Rang Faible : Pour plus d'efficacité, les auteurs proposent une instanciation spécifique où les contraintes sont dérivées des $k$ vecteurs singuliers principaux des tâches passées, permettant d'organiser les variables duales en une matrice de rang faible $\Lambda$, réduisant ainsi la dimensionnalité du problème dual.

Contributions Clés

1. Géométrie de Norme Spectrale pour l'AC : L'article formule l'apprentissage continu basé sur la projection sous une géométrie de norme spectrale, alignant les contraintes de stabilité de l'apprentissage continu avec la géométrie de norme d'opérateur des optimiseurs modernes comme Muon.
2. Règle de Mise à Jour Pratique : Il dérive un algorithme pratique (Muon-OGD) qui combine les contraintes d'orthogonalité des tâches passées avec des mises à jour de signe matriciel de style Muon, résolvant le problème contraint efficacement via une formulation duale et des approximations de Newton–Schulz.
3. Validation Empirique : Le cadre est instancié pour l'affinage continu de LLM et évalué sur divers benchmarks, démontrant des améliorations constantes par rapport aux contreparties basées sur la norme de Frobenius.

Résultats Expérimentaux
Les auteurs évaluent Muon-OGD sur :

• Benchmarks Standards : Des benchmarks d'apprentissage continu à 5 et 15 tâches utilisant T5-Large.
• Benchmark TRACE : Un benchmark pour les LLM affinés par instruction (utilisant LLaMA-2-7B-Chat) couvrant des compétences diverses telles que le codage, les mathématiques et la compréhension multilingue.
• Curricula Spécifiques au Domaine : Apprentissage séquentiel sur le Codage (BigCodeBench) $\to$ Mathématiques (GSM8K) $\to$ Médical (HuatuoGPT-o1) en utilisant les modèles Qwen2.5 et LLaMA3.2.

Constats :

• Performance : Muon-OGD surpasse systématiquement l'affinage séquentiel (SeqFT) et les bases compétitives de gradients orthogonaux (par exemple, O-LoRA, Sculpting Subspaces, OSFT) en termes de Précision Moyenne (AA) et de Transfert Arrière (BT).
  • Sur le benchmark T5-Large à 15 tâches, Muon-OGD a atteint 72,0 % de précision moyenne, surpassant O-LoRA (69,6 %) et OSFT (71,3 %).
  • Sur le benchmark TRACE, Muon-OGD a atteint la plus haute AA (49,4 %) et le taux d'oubli le plus faible (-4,7 %) parmi les méthodes comparées.
• Évolutivité : La méthode reste évolutive sur le plan computationnel, tirant parti de l'efficacité des itérations de Newton–Schulz et des représentations de contraintes de rang faible.
• Agnosticisme Architectural : L'approche est efficace à la fois sur les architectures encodeur-décodeur (T5) et décodeur uniquement (LLaMA, Qwen).

Signification et Revendications
L'article revendique que la géométrie de mise à jour consciente de la norme spectrale fournit une alternative pratique et efficace à la projection de norme de Frobenius pour l'apprentissage continu dans les LLM. En comblant l'écart entre la géométrie des optimiseurs modernes (comme Muon) et les exigences de stabilité de l'apprentissage continu, Muon-OGD améliore le compromis stabilité-plasticité. Les auteurs suggèrent que les méthodes existantes basées sur la projection peuvent être systématiquement mises à niveau pour fonctionner sous une géométrie de norme spectrale afin de mieux préserver les avantages de conditionnement et de stabilité observés dans Muon tout en empêchant l'oubli catastrophique. Le travail ne revendique pas résoudre tous les défis de l'apprentissage continu, mais offre une alternative géométrique fondée sur des principes qui produit des gains empiriques en rétention et en adaptation à travers divers domaines et échelles de modèles.