Новое исследование выявило критическую проблему при использовании оптимизатора Muon в трансформерах: модели демонстрируют эффект «пост-грокинга» (post-grokking collapse). Несмотря на ускоренное обучение модулярному сложению, системы теряют способность к обобщению после достижения пиковых показателей. В отличие от стандартного AdamW, решения, найденные с помощью Muon, оказываются нестабильными и быстро деградируют в процессе дальнейшего обучения.
Эксперименты проводились на задаче $(a+b) \bmod 113$. В ходе тестов все девять конфигураций с использованием Muon сначала демонстрировали резкий рост качества, а затем закономерный спад точности. Для сравнения, контрольные модели на AdamW показали более устойчивые результаты, хотя и они в ряде случаев не смогли удержать высокую точность, опустившись до 27,59% на определенных этапах.
Авторы работы проанализировали интерфейс между скрытыми матрицами и выходными слоями. Выяснилось, что нестабильность сохраняется при изменении модулей и ширины нейронных сетей. Это ставит под сомнение надежность Muon как полноценной замены AdamW в задачах, требующих долгосрочного сохранения выученных закономерностей, несмотря на его высокую эффективность на начальных этапах оптимизации.
Ключевые факты
- Исследование сфокусировано на сравнении оптимизатора Muon и стандартного алгоритма AdamW в архитектурах трансформеров.
- В задаче $(a+b) \bmod 113$ все протестированные конфигурации на Muon теряли способность к обобщению после фазы «грокинга».
- Контрольные модели на AdamW в 4 из 5 случайных запусков также показали снижение точности до 27,59%.
- Нестабильность подтверждена при варьировании параметров модулей, ширины моделей и различных случайных инициализациях (seeds).