EasyData

Классический градиентный спуск двигает все веса на один и тот же шаг вдоль антиградиента. Adam умнее: он хранит моментум (сглаженные предыдущие направления) + подбирает индивидуальный размер шага для каждого параметра, исходя из его истории. Но “взгляд” у него поштучный: каждый вес - сам по себе, а градиент по слою рассматривается как длинный вектор независимых чисел.

Например, веса полносвязного слоя записываются матрицей, физически она задаёт линейное преобразование (привет всем, кто спрашивал про практическое применение теории на парах по линалу✌️). Градиент по такой матрице весов тоже можно записать в матричной форме - а не вектором, как делают обычные оптимизаторы.
Так и поступает Muon: он берёт матрицу обновления (сглаженный моментумом градиент), оставляет её в родной матричной форме -> а затем превращает в ортогональную (как именно - вспомним ниже). Это основная идея, и отсюда пошло название: MomentUm Orthogonalized by Newton-Schulz.

Есть проблема, что у матрицы обновления весов обычно пара направлений сильно доминирует, а остальные почти нулевые. Получается, обучение пойдет вдоль этих немногих осей, а вся остальная “ёмкость” слоя простаивает. Как раз с этим и помогает ортогонализация!
Теперь “как это реализуется”: вспомним SVD - он раскладывает любую матрицу на две ортогональные (которые задают повороты) + одну диагональную посередине (она задаёт растяжение/сжатие, числа на диагонали называют “сингулярными”). Ортогонализацией мы превращаем все сингулярные значения в единицу, как раз убирая проблему “перекоса”. За счёт этого шаг делается во все стороны поровну, и слой учится целиком.

Итак, честно ортогонализировать матрицу можно через SVD - но высчитывать его на каждом шаге и для каждого слоя неподъёмно дорого... Поэтому, как обычно, пользуемся эвристиками: Newton-Schulz даёт почти тот же результат за несколько матричных умножений (а их видеокарта уважает). Обычно хватает 5 итераций - так что Muon вполне пригоден для реального обучения.

• Быстрее сходится: примерно вдвое эффективнее AdamW по замерам Moonshot (они посчитали, что один прогон обходится в 2x дешевле).
• Экономит память: у Adam на каждый вес два состояния, у Muon - только моментум. Около трети памяти под оптимизатор долой.

• Применим только к слоям, где вес - это полноценное преобразование: например, полносвязные и attention. Но нормировки, свободные члены, выходной слой и эмбеддинги (обновляются построчно) учатся через AdamW.
• Каждый шаг дороже: Newton-Schulz добавляет матричные умножения, которых у AdamW нет. Шагов, конечно, нужно меньше, но с ростом модели расходы приходится отдельно оптимизировать.
• На больших моделях чаще всплывает нестабильность: значения внимания разгоняются и обучение может разойтись.

Похоже, что да - Muon уже вышел из разряда экспериментов:
• Обучает топовые открытые модели: Kimi K2 (1 трлн параметров) и GLM-5 (744 млрд), DeepSeek-V4.
• Нестабильность научились лечить: например, Moonshot добавили приём MuonClip, который придерживает разгоняющиеся значения внимания.
• Уже адаптируется под инфраструктуру: NVIDIA встроила Muon в свой Megatron + команда DeepSpeed реализует поддержку.
• Оброс модификациями: за год вышли десятки ускоренных версий (Gram Newton-Schulz от Три Дао, NorMuon, Hierarchical Muon и другие) - так что направление живое.

• Оригинальный пост Keller Jordan - с интуицией и картинками
• Референс-реализация
• Вывод "с нуля" для любителей математики
• Newton-Schulz в доках Modula - с визуализациями сходимости