Оптимизатор Muon от Microsoft значительно ускоряет гроккинг в трансформерах
Исследователи Microsoft показали, что оптимизатор Muon существенно ускоряет гроккинг в моделях трансформеров, обеспечивая более быстрый переход от запоминания к обобщению по сравнению с AdamW.
Особенности гроккинга в глубоком обучении
Гроккинг — это явление, при котором модели глубокого обучения после длительного периода запоминания данных внезапно начинают хорошо обобщать на новые данные. Это наблюдалось в простых алгоритмических задачах, таких как модульная арифметика. Модель достигает почти идеальной точности на обучении, но долгое время показывает слабые результаты на валидации, а затем резко улучшает обобщение.
Роль оптимизаторов в гроккинге
Исследователи из Microsoft изучили влияние выбора оптимизатора на этот эффект, сравнив популярный оптимизатор AdamW и новый Muon, который использует ограничения спектральной нормы и вторую производную для обновлений. Цель была проверить, может ли Muon ускорить фазу обобщения.
Экспериментальная настройка
Были протестированы семь алгоритмических задач — операции с модульной арифметикой и классификация чётности, реализованные на архитектуре трансформера в PyTorch. Также изучались три варианта softmax — стандартный, stablemax и sparsemax — чтобы понять, влияет ли нормализация выхода на гроккинг, но основное внимание было на оптимизаторах.
Архитектура модели и детали оптимизации
Модель трансформера включает многоголовое самовнимание, позиционные эмбеддинги RoPE, RMS-нормализацию, активации SiLU и регуляризацию с помощью dropout. Входные данные кодируются простыми идентичными эмбеддингами.
AdamW использует адаптивные скорости обучения с раздельным затуханием весов и служит базовой линией.
Оптимизатор Muon применяет ортогонализированные градиенты, накладывает ограничения спектральной нормы для стабильности обучения и приближает вторые производные, чтобы сделать обновления более информативными. Это помогает избежать проблем типа «коллапса softmax» и обеспечивает более эффективное обучение за счёт согласования обновлений с размерностями слоёв.
Результаты и выводы
Эксперименты проводились на графических процессорах NVIDIA H100 с несколькими инициализациями для статистической достоверности. Гроккинг определялся как эпоха, когда точность на валидации впервые превысила 95% после стабилизации на обучении.
Muon достиг этой отметки значительно быстрее — в среднем за 102.89 эпох, тогда как AdamW — за 153.09 эпох. Разница статистически значима (t = 5.0175, p ≈ 6.33e−8). Кроме того, Muon показал более стабильное время появления гроккинга по всем задачам.
Значение для обучения нейросетей
Результаты подчёркивают, что дизайн оптимизатора, особенно с учётом геометрии обновлений и ограничений спектральной нормы, сильно влияет на момент и качество обобщения моделей. Подход Muon помогает моделям избегать затяжных фаз запоминания и быстрее находить истинную структуру данных.
Это исследование показывает, что выбор и архитектура оптимизатора — важнейший фактор в стратегии обучения нейросетей, наряду с данными и регуляризацией.
Для подробностей ознакомьтесь с оригинальной статьёй и следите за обсуждениями в соцсетях.
Switch Language
Read this article in English