Оптимизатор 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.
Выводы и перспективы
Исследование позволило сделать два ключевых вывода:
- Техника клиппинга из «Kimi K2» (MuonClip) неэффективна для мелкомасштабной тонкой настройки
- Гибридный подход — оптимальное решение: комбинация Muon + AdamW обеспечивает наилучший баланс производительности и стабильности
Эксперименты проводились на модели объемом 4B параметров и примерно 10 тысячах строк данных. Следующим шагом станет масштабирование исследования на более крупные модели и наборы данных для проверки, становятся ли скачки в чистом Muon катастрофическими при длительных запусках, и сможет ли MuonClip превзойти Muon и Muon+AdamW в других условиях.
По материалам Hugging Face
