Résumé Technique : MuonSSM – Orthogonalisation des Modèles d'Espace d'États pour la Modélisation de Séquences

Énoncé du Problème

Les modèles d'espace d'états (SSM) sont apparus comme des alternatives efficaces, avec une complexité temporelle linéaire, aux mécanismes d'attention pour la modélisation de séquences longues. Cependant, les SSM existants (par exemple, Mamba, DeltaNet) souffrent souvent d'instabilité et de dégradation de la mémoire sur des horizons étendus. Le papier identifie que les SSM actuels reposent sur des mises à jour de premier ordre, de rang un de leurs états de mémoire. Ces mises à jour entraînent trois problèmes principaux :

1. Anisotropie Spectrale : Les valeurs singulières de la matrice de mémoire deviennent hautement non uniformes, provoquant des paysages d'optimisation mal conditionnés.
2. Dégradation du Gradient : La contraction répétée à travers des matrices de transition dépendantes de l'entrée entraîne une disparition des gradients sur de longues séquences.
3. Interférence de Mémoire : Les nouvelles mises à jour peuvent écraser les informations précédentes en raison d'un manque d'inertie temporelle et d'une géométrie de mise à jour déséquilibrée.

Bien que des variantes récentes tentent d'atténuer ces problèmes via le fenêtrage (gating) ou la normalisation, elles ne conditionnent pas explicitement la géométrie des mises à jour accumulées ni n'introduisent d'inertie temporelle pour améliorer la propagation du signal à longue portée.

Méthodologie : MuonSSM

Les auteurs proposent MuonSSM, un cadre qui stabilise l'entraînement des SSM en conditionnant explicitement la géométrie des mises à jour de la mémoire plutôt que la matrice de transition récurrente. L'approche augmente la dynamique standard des SSM avec deux composantes clés tout en préservant la structure affine requise pour des scans parallèles efficaces :

1. Dynamique Augmentée par le Momentum

MuonSSM introduit une matrice de momentum auxiliaire $M_t$ qui accumule les directions de mise à jour à travers les pas de temps. Cela imite le momentum utilisé dans l'optimisation, fournissant une « inertie temporelle » aux mises à jour d'état.

• L'état de momentum est mis à jour comme suit : $M_t = \gamma M_{t-1} + \text{NS}(\tau \beta_t v_t k_t^\top)$.
• L'état de mémoire $S_t$ est mis à jour en combinant l'état précédent avec le terme de momentum : $S_t = S_{t-1}(\alpha_t(I - \beta_t \eta k_t k_t^\top)) + M_t$.
• Cela crée un canal supplémentaire pour le flux de gradient, atténuant la décroissance exponentielle dans l'attribution de crédit à longue portée.

2. Normalisation de Newton–Schulz (NS) à Itération Unique

Pour remédier au déséquilibre spectral, l'injection d'entrée (une matrice de rang 1) est normalisée à l'aide d'une itération de Newton–Schulz légère à une seule étape avant d'être ajoutée à l'état de momentum.

• Opérateur : $\text{NS}(X) = (a + b \tilde{X}\tilde{X}^\top + c (\tilde{X}\tilde{X}^\top)^2)\tilde{X}$, où $\tilde{X}$ est l'entrée mise à l'échelle par sa norme de Frobenius.
• Fonction : Cet opérateur limite les valeurs singulières (les maintenant proches de l'unité) et préserve le rang de l'entrée sans nécessiter de SVD explicite.
• Effet Géométrique : Contra contrairement à la normalisation de Frobenius standard, le Jacobien du passage arrière de l'opérateur NS amplifie les directions orthogonales à l'écriture de rang un actuelle. Cela encourage l'« enrichissement de rang », empêissant l'état de mémoire de s'effondrer dans un sous-espace de faible dimension.

Parallélisabilité

Crucialement, la dynamique couplée de la mémoire et du momentum peut être formulée comme une récurrence affine par blocs. Cela permet au modèle d'utiliser des scans associatifs parallèles (réduisant la profondeur récurrente de $O(L)$ à $O(\log L)$), garantissant que la complexité computationnelle reste linéaire par rapport à la longueur de la séquence $L$.

Principales Contributions

1. Conception du Cadre : Introduction de MuonSSM, qui intègre une dynamique de second ordre basée sur le momentum et un conditionnement spectral implicite dans la récurrence basée sur le scan.
2. Analyse Théorique :
   • Preuve que les dynamiques augmentées préservent la parallélisabilité.
   • Démonstration que la voie de momentum introduit une composante de gradient non nulle, stabilisant la propagation à longue portée.
   • Analyse montrant que l'opérateur NS limite les valeurs singulières et favorise l'enrichissement de rang en biaisant les signaux de rétropropagation vers des directions orthogonales.
3. Validation Empirique : Expériences exhaustives sur des tâches de langage, de vision et de séries temporelles montrant des améliorations constantes de la précision, de la robustesse et de la généralisation de la longueur.

Résultats Expérimentaux

Les auteurs ont évalué MuonSSM par rapport aux modèles de référence (Mamba, LongHorn, Gated DeltaNet) à travers trois modalités :

• Modélisation de Langage : Sur le jeu de données FineWeb-Edu, les variantes de MuonSSM ont atteint une perplexité plus faible et des performances supérieures sur les tests de raisonnement de sens commun. Notamment, dans les tâches de recherche « Needle-in-a-Haystack », MuonSSM a maintenu une précision élevée jusqu'à des longueurs de contexte de 8K, surpassant nettement les modèles de référence qui se dégradent rapidement au-delà de leur horizon d'entraînement (2K tokens).
• Vision : Le remplacement des mélangeurs Mamba par MuonSSM dans le cadre MambaVision a conduit à des gains constants dans la classification ImageNet-1K, la détection d'objets (COCO) et la segmentation sémantique (ADE20K). Crucialement, MuonSSM a montré une robustesse nettement améliorée contre les corruptions (IN-C), les renditions (IN-R) et les attaques adverses (IN-A), suggérant un passage d'un biais de texture vers un apprentissage de caractures invariantes de forme.
• Séries Temporelles : Dans les tâches de reconnaissance d'activité humaine (Human Activity Recognition - MuWiGes, UESTC-MMEA-CL, MMAct), les variantes de MuonSSM ont surpassé les modèles de référence, l'écart de performance s'élargissant à mesure que la complexité des jeux de données augmentait.

Ablation & Analyse :

• Conditionnement Spectral : Les visualisations ont confirmé que les SSM standards souffrent d'un effondrement spectral (nombres de condition élevés), tandis que MuonSSM maintient un spectre quasi-isométrique.
• Capacité vs Géométrie : Doubler la dimension de l'état des modèles de référence n'a produit que des gains marginaux comparés aux améliorations substantielles de MuonSSM, confirmant que les bénéfices proviennent du conditionnement géométrique plutôt que de la simple expansion de capacité.
• Efficacité : Bien qu'ajoutant un facteur de surcharge constant par étape, MuonSSM accélère la convergence (environ 1,3$\times$ plus vite), résultant en moins d'étapes d'entraînement totales pour atteindre la performance cible.

Signification et Revendications

Le papier pose que le conditionnement géométrique des mises à jour est une voie fondamentale pour une modélisation de séquences stable et scalable. En contrôlant explicitement la géométrie des mises à jour de la mémoire via le momentum et la normalisation de Newton–Schulz, MuonSSM traite l'instabilité fondamentale des SSM de premier ordre sans sacrifier l'efficacité temporelle linéaire des scans parallèles.

Les auteurs affirment que cette approche :

• Améliore la propagation du gradient et atténue l'amplification spectrale.
• Enrichit les représentations de la mémoire sur de longs horizons en empêchant l'effondrement de rang.
• Fournit une alternative robuste à l'attention pour les tâches à long contexte, démontant que la stabilité peut être obtenue par la dynamique de mise à jour plutôt que uniquement par des paramétrisations complexes de la matrice de transition.

Ce travail suggère que la modélisation de séquences future peut bénéficier du découplage de la stabilité du mécanisme de mise à jour de la structure de récurrence, offrant un mécanisme simple mais efficace pour un entraînement stable à travers diverses modalités.