Мюон превосходит Adam в обучении ассоциативной памяти на заключительных этапах.

Аннотация

Оптимизатор Muon стабильно превосходит Adam по скорости обучения больших языковых моделей (LLM), однако механизм, лежащий в основе его успеха, остается неясным. В данной статье мы раскрываем этот механизм через призму ассоциативной памяти. Исключая компоненты трансформера, оптимизируемые Muon, мы обнаруживаем, что параметры ассоциативной памяти LLM, а именно веса внимания Value и Output (VO) и сети прямого распространения (FFN), являются основными факторами превосходства Muon. Вдохновленные этим взглядом на ассоциативную память, мы затем объясняем превосходство Muon на реальных корпусах, которые по своей природе имеют тяжелые хвосты: несколько классов (хвостовые классы) встречаются значительно реже, чем другие. Это превосходство объясняется двумя ключевыми свойствами: (i) правило обновления Muon стабильно приводит к более изотропному спектру сингулярных значений по сравнению с Adam; и, как следствие, (ii) на данных с тяжелыми хвостами он оптимизирует хвостовые классы более эффективно, чем Adam. Помимо эмпирических данных, мы теоретически подтверждаем эти выводы, анализируя однослойную модель ассоциативной памяти на данных с дисбалансом классов. Мы доказываем, что Muon стабильно достигает сбалансированного обучения по классам независимо от свойств эмбеддингов, тогда как Adam может вызывать значительные различия в ошибках обучения в зависимости от свойств эмбеддингов. В итоге, наши эмпирические наблюдения и теоретические анализы раскрывают ключевое преимущество Muon: его правило обновления согласуется со структурой внешнего произведения линейных ассоциативных память, что позволяет более сбалансированно и эффективно обучать хвостовые классы в распределениях с тяжелыми хвостами по сравнению с 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.

Мюон превосходит Adam в обучении ассоциативной памяти на заключительных этапах.

Muon Outperforms Adam in Tail-End Associative Memory Learning

Аннотация

Support