El Muon supera a Adam en el aprendizaje de memoria asociativa en el extremo de la cola.

Muon Outperforms Adam in Tail-End Associative Memory Learning

September 30, 2025

Autores: Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Cunxiao Du, Chao Du, Tianyu Pang, Zhuoran Yang, Mingyi Hong, Vincent Y. F. Tan

cs.AI

Resumen

El optimizador Muon es consistentemente más rápido que Adam en el entrenamiento de Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés), aunque el mecanismo subyacente a su éxito sigue sin estar claro. Este artículo desmitifica dicho mecanismo a través de la lente de la memoria asociativa. Al eliminar selectivamente los componentes del transformador optimizados por Muon, revelamos que los parámetros de memoria asociativa de los LLMs, específicamente los pesos de atención de Valor y Salida (VO) y las Redes de Alimentación Directa (FFNs, por sus siglas en inglés), son los principales contribuyentes a la superioridad de Muon. Motivados por esta perspectiva de memoria asociativa, explicamos luego la superioridad de Muon en corpus del mundo real, que son intrínsecamente de cola pesada: unas pocas clases (clases de cola) aparecen con mucha menos frecuencia que otras. La superioridad se explica mediante dos propiedades clave: (i) su regla de actualización produce consistentemente un espectro singular más isotrópico que Adam; y, como resultado, (ii) en datos de cola pesada, optimiza las clases de cola de manera más efectiva que Adam. Más allá de la evidencia empírica, confirmamos teóricamente estos hallazgos analizando un modelo de memoria asociativa de una capa bajo datos con desequilibrio de clases. Demostramos que Muon logra consistentemente un aprendizaje equilibrado entre clases, independientemente de los embeddings de características, mientras que Adam puede inducir grandes disparidades en los errores de aprendizaje dependiendo de las propiedades de los embeddings. En resumen, nuestras observaciones empíricas y análisis teóricos revelan la ventaja central de Muon: su regla de actualización se alinea con la estructura de producto externo de las memorias asociativas lineales, permitiendo un aprendizaje más equilibrado y efectivo de las clases de cola en distribuciones de cola pesada en comparación con 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.