Оптимизатор Muon показал лучшие результаты при тонкой настройке языковых моделей

Экспериментальное исследование показало, что гибридный подход Muon + AdamW превосходит как чистый Muon, так и его вариацию MuonClip при тонкой настройке языковых моделей. Результаты бросают вызов общепринятым практикам и демонстрируют важность адаптации методов оптимизации к конкретным задачам.

Проблема оптимизатора Muon

Muon представляет собой 2D-оптимизатор, который использует итерации Ньютона-Шульца для ортогонализации обновлений весов. Ключевое ограничение метода — он работает исключительно с двумерными матрицами (линейными слоями) и не может оптимизировать одномерные тензоры.

Для решения этой проблемы исследователи протестировали три различные конфигурации:

• Только Muon: оптимизация 2D-слоев, одномерные слои остаются замороженными
• Muon + AdamW (гибрид): Muon для 2D-слоев, стандартный AdamW для 1D-слоев
• MuonClip (использовался в «Kimi K2»): Muon для 2D, Adam для 1D с добавлением «QK Norm Clipping» для стабилизации больших предварительных тренировок

Экспериментальная установка

Исследование проводилось на аппаратной платформе NVIDIA GH200 с моделью Qwen3 4B. В качестве фреймворка использовался trl (Transformer Reinforcement Learning), а набор данных состоял из более чем 10 тысяч строк высококачественных инструкций. Методология включала полную параметрическую тонкую настройку.

Результаты сравнения

Источник: huggingface.co

Muon и Muon+AdamW продемонстрировали превосходство над MuonClip в скорости сходимости. Основная причина — техника клиппинга, которая оказывается более оптимальной и стабильной для длительных тренировок.

Источник: huggingface.co

Это также объясняет более низкую среднюю точность токенов у MuonClip.

Изначально чистый Muon сходился быстрее, чем Muon+AdamW, но затем гибридный подход догнал и превзошел его. Причина этого явления раскрывается при анализе градиентов.

Источник: huggingface.co

Как видно из графиков, Muon+AdamW демонстрирует значительно более стабильную норму градиента по сравнению с Muon и MuonClip, где наблюдаются резкие скачки.

При сравнении с классическим AdamW результаты оказались еще более впечатляющими.

Источник: huggingface.co

Окончательным победителем стал гибридный подход Muon+AdamW, который не только превзошел Muon и MuonClip, но и показал лучшие результаты по сравнению с классическим AdamW.

Интересно наблюдать, как теоретически элегантный подход Muon на практике требует компромиссов. Гибридная схема оказалась не просто удобным решением, а фундаментально лучшим выбором для тонкой настройки. Это напоминает, что в машинном обучении часто побеждает не самая красивая математика, а наиболее практичная комбинация методов.

Код, использованный в эксперименте, доступен на GitHub.

Выводы и перспективы

Исследование позволило сделать два ключевых вывода:

1. Техника клиппинга из «Kimi K2» (MuonClip) неэффективна для мелкомасштабной тонкой настройки
2. Гибридный подход — оптимальное решение: комбинация Muon + AdamW обеспечивает наилучший баланс производительности и стабильности

Эксперименты проводились на модели объемом 4B параметров и примерно 10 тысячах строк данных. Следующим шагом станет масштабирование исследования на более крупные модели и наборы данных для проверки, становятся ли скачки в чистом Muon катастрофическими при длительных запусках, и сможет ли MuonClip превзойти Muon и Muon+AdamW в других условиях.

По материалам Hugging Face