Le Muon surpasse Adam dans l'apprentissage de la mémoire associative en fin de distribution

papers.abstract

L'optimiseur Muon est systématiquement plus rapide qu'Adam pour l'entraînement des grands modèles de langage (LLMs), mais le mécanisme sous-jacent à son succès reste mal compris. Cet article démystifie ce mécanisme à travers le prisme de la mémoire associative. En éliminant sélectivement les composants du transformateur optimisés par Muon, nous révélons que les paramètres de mémoire associative des LLMs, à savoir les poids d'attention Valeur et Sortie (VO) et les réseaux feed-forward (FFNs), sont les principaux contributeurs à la supériorité de Muon. Motivés par cette perspective de mémoire associative, nous expliquons ensuite la supériorité de Muon sur les corpus réels, qui sont intrinsèquement à queue lourde : quelques classes (classes de queue) apparaissent bien moins fréquemment que les autres. Cette supériorité s'explique par deux propriétés clés : (i) sa règle de mise à jour produit systématiquement un spectre singulier plus isotrope que celui d'Adam ; et par conséquent, (ii) sur des données à queue lourde, il optimise les classes de queue plus efficacement qu'Adam. Au-delà des preuves empiriques, nous confirmons théoriquement ces résultats en analysant un modèle de mémoire associative à une couche sous des données déséquilibrées en classes. Nous prouvons que Muon atteint systématiquement un apprentissage équilibré entre les classes, indépendamment des embeddings de caractéristiques, tandis qu'Adam peut induire de grandes disparités dans les erreurs d'apprentissage en fonction des propriétés des embeddings. En résumé, nos observations empiriques et analyses théoriques révèlent l'avantage fondamental de Muon : sa règle de mise à jour s'aligne sur la structure de produit extérieur des mémoires associatives linéaires, permettant un apprentissage plus équilibré et efficace des classes de queue dans les distributions à queue lourde par rapport à Adam.

English

The Muon optimizer is consistently faster than Adam in training Large Language Models (LLMs), yet the mechanism underlying its success remains unclear. This paper demystifies this mechanism through the lens of associative memory. By ablating the transformer components optimized by Muon, we reveal that the associative memory parameters of LLMs, namely the Value and Output (VO) attention weights and Feed-Forward Networks (FFNs), are the primary contributors to Muon's superiority. Motivated by this associative memory view, we then explain Muon's superiority on real-world corpora, which are intrinsically heavy-tailed: a few classes (tail classes) appear far less frequently than others. The superiority is explained through two key properties: (i) its update rule consistently yields a more isotropic singular spectrum than Adam; and as a result, (ii) on heavy-tailed data, it optimizes tail classes more effectively than Adam. Beyond empirical evidence, we theoretically confirm these findings by analyzing a one-layer associative memory model under class-imbalanced data. We prove that Muon consistently achieves balanced learning across classes regardless of feature embeddings, whereas Adam can induce large disparities in learning errors depending on embedding properties. In summary, our empirical observations and theoretical analyses reveal Muon's core advantage: its update rule aligns with the outer-product structure of linear associative memories, enabling more balanced and effective learning of tail classes in heavy-tailed distributions than Adam.

Le Muon surpasse Adam dans l'apprentissage de la mémoire associative en fin de distribution

Muon Outperforms Adam in Tail-End Associative Memory Learning

papers.abstract

Support