Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers

Résumé technique : Effondrement post-grokking à l'interface représentation–lecture dans les Transformers entraînés avec Muon

Énoncé du problème
L'article étudie une instabilité spécifique survenant dans les modèles Transformer entraînés sur l'arithmétique modulaire (plus précisément l'addition et la soustraction modulaires) après avoir atteint le « grokking » — le phénomène où la généralisation émerge longtemps après que la précision d'entraînement a saturé. Alors que des travaux antérieurs ont établi que Muon (un optimiseur d'orthogonalisation) accélère le grokking par rapport à AdamW, cette étude rapporte que les modèles entraînés avec Muon souffrent fréquemment d'un effondrement récurrent et abrupt de la performance de généralisation post-grokking. Crucialement, cet échec n'est pas limité à Muon ; les modèles AdamW présentent également une instabilité. L'article démontre que ni l'un ni l'autre optimiseur n'est strictement stable sur cinq graines (seeds) : la ligne de base AdamW sélectionnée enregistre des évaluations sous le seuil sur quatre graines sur cinq (atteignant un minimum de 27,59 %), tandis que Muon les enregistre sur les cinq. Le problème central est que ces modèles perdent la capacité de généraliser alors que l'ensemble d'entraînement reste résolu, et que les mesures de progrès mécaniste standard (support spectral et distribution de puissance) ne parviennent pas à détecter l'échec, signalant un circuit intact même lorsque la précision chute à des niveaux proches du hasard.

Méthodologie
Les auteurs emploient un régime d'entraînement à optimiseurs divisés où les matrices de poids cachées sont mises à jour par Muon, tandis que les plongements (embeddings) de jetons/positions et le dé-plongement de sortie (readout/unembedding) sont mis à jour par AdamW. L'étude utilise des tâches d'arithmétique modulaire ($a+b \pmod p$ et $a-b \pmod p$) avec $p=113$ (et des variantes avec $p=97$) sur des Transformers de type décodeur seul de profondeurs variables (1, 2 et 4).

Composantes méthodologiques clés :

1. Ramification appariée (Matched Branching) : Les trajectoires sont ramifiées à partir d'états de paramètres bit-identiques pour isoler les facteurs causaux. Une branche continue avec l'optimisation standard, tandis que d'autres figent des groupes de paramètres spécifiques (matrices cachées, plongements ou lecture) pour tester la nécessité.
2. Décomposition de Fourier : Les auteurs décomposent les états résiduels internes du modèle en modes de Fourier 2D sur la grille des opérandes. Ils définissent des « familles » spécifiques de modes (par exemple, la famille de l'addition $(k, k)$, la famille de la soustraction $(k, -k)$) et utilisent des interventions de filtrage pour mesurer la suffisance et la nécessité de ces familles pour la performance de la tâche.
3. Suite d'interventions : L'étude applique des interventions telles que le brouillage de phase (phase scrambling), la relocalisation de fréquence et la substitution de lecture croisée pour déterminer la nature exacte de l'algorithme appris et la source du désalignement.
4. Ablation : Le chemin d'orthogonalisation normalisée de Muon est ablaté pour tester son rôle dans la dispersion spectrale et l'instabilité observées.

Principaux résultats et conclusions

• La nature de l'effondrement : L'échec est un « désalignement caché–lecture ». À mesure que la perte d'entraînement approche de zéro, le gradient disparaît, mais les deux optimiseurs réagissent différemment. Le pas de mise à jour de Muon devient largement indépendant de la magnitude du gradient (élasticité $\approx -0,03$), tandis que celui d'AdamW reste proportionnel (élasticité $\approx +1,5$). Cela provoque une dérive entre la représentation cachée et la lecture à un taux de 8,0 fois par paramètre. L'ensemble d'entraînement reste résolu, mais la lecture ne peut plus décoder la représentation qui a été entraînée à ses côtés.
• Masquage vs Échec de circuit : L'article distingue deux modes de défaillance. Dans un échec de circuit, la famille de Fourier elle-même cesse de calculer la tâche. Dans un masquage de circuit, la famille alignée sur la tâche (par exemple, la famille de l'addition) continue de résoudre la tâche parfaitement de manière isolée (précision de 100 %), mais sa contribution est submergée par un reste adversaire dans le modèle complet. Les auteurs montrent que le simple fait de redimensionner la composante alignée sur la tâche restaure la pleine précision du modèle, prouvant que le circuit était intact mais « voté » par la minorité.
• Différences d'optimiseurs : Muon atteint le grokking significativement plus vite (médiane de 13 011 étapes contre 30 486 pour AdamW) mais est universellement instable sur toutes les configurations testées. AdamW est également instable sur cinq graines, bien que la sévérité de l'effondrement (mesurée par le nombre d'évaluations sous le seuil) soit de deux ordres de grandeur inférieure à celle de Muon. Dans un balayage sur une seule graine, les configurations AdamW ont montré une instabilité conditionnelle où la sévérité augmentait avec le taux d'apprentissage, mais la réplication sur cinq graines a confirmé que la ligne de base « stable » sélectionnée n'était pas strictement stable selon les graines. L'instabilité de Muon est attribuée à sa règle de mise à jour, qui disperse la représentation sur un large spectre (326 paires conjuguées effectives contre environ 5 pour AdamW), laissant la famille alignée sur la tâche avec une part plus faible de la puissance totale.
• Sélection de la famille par la tâche : La tâche, et non l'optimiseur, sélectionne la famille de Fourier. Les modèles entraînés pour l'addition utilisent les modes $(k, k)$ ; ceux entraînés pour la soustraction utilisent $(k, -k)$. La base au sein de la famille est exacte et non partagée ; deux modèles entraînés à partir de la même graine divergent pour utiliser des fréquences et des phases différentes, et leurs lectures deviennent presque orthogonales, échouant à décoder les représentations l'une de l'autre.
• Profondeur et alignement : Dans les modèles plus profonds, le code de Fourier est écrit par un MLP dans le bloc pénultième, mais ne devient « lisible » pour le dé-plongement (unembedding) que dans le dernier bloc. Le désalignement se produit spécifiquement à l'interface entre le résidu du dernier bloc et le dé-plongement.
• Prévention : Figer soit les matrices cachées, soit les paramètres non cachés (plongements et lecture) après la formation du circuit empêche complètement l'effondrement. Un « Muon stable », où les plongements et la lecture sont figés post-grokking, atteint une stabilité stricte sur 451 400 étapes post-grokking et cinq graines, contrairement au Muon standard qui échoue de manière répétée.

Signification et affirmations
L'article prétend identifier une limitation fondamentale dans l'entraînement à optimiseurs divisés et la dynamique post-grokking : le flux résiduel n'a pas de base privilégiée, et lorsque la surface de perte s'aplatit, la perte n'exerce aucune pression pour maintenir l'alignement spécifique entre la représentation et la lecture. Les auteurs soutiennent que :

1. Les mesures standard sont insuffisantes : Les mesures spectrales (support et distribution de puissance) peuvent rester invariantes pendant un effondrement car l'échec est un changement de base (désalignement) plutôt qu'une perte de l'algorithme lui-même.
2. Le grokking est un masquage réversible : La transition de la généralisation vers l'effondrement est l'inverse de la transition de la mémorisation vers la généralisation ; les deux impliquent que la part de puissance représentationnelle de la famille alignée sur la tâche augmente ou diminue par rapport au reste adversaire.
3. L'interface est le lieu de l'échec : L'instabilité n'est pas une propriété de la représentation cachée ou de la lecture seule, mais de leur identification conjointe. Fixer un côté de l'interface stabilise le système.
4. La vitesse de Muon a un coût : Si Muon accélère la découverte du circuit, sa règle de mise à jour crée une représentation plus susceptible de désalignement avec la lecture sous des conditions d'entraînement standard, nécessissant une intervention (gel) pour une stabilité à long terme.

L'étude ne propose pas de nouvelles applications mais offre un compte rendu mécaniste d'un mode de défaillance spécifique en interprétabilité mécaniste, soulignant qu'un modèle peut posséder un circuit correct tout en échouant à l'utiliser en raison d'une dérive représentationnelle.