Muon Optimizer Accelerates Grokking
Amund Tveit, Bjørn Remseth, Arve Skogvold
Статья: https://arxiv.org/abs/2504.16041

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

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

Текущая работа определяет гроккинг следующим образом: это первая эпоха, на которой точность на валидации достигает или превосходит 95%, после того как точность на трейне стабилизировалась около 100%.

Если не читали разбор оригинальной статьи про гроккинг (https://t.me/gonzo_ML/831), очень рекомендую.

Вот ведь, было время, когда OpenAI публиковал интересные статьи, и эту, и про double descent, и много всего другого…

Muon (MomentUm Orthogonalized by Newton-Schulz) -- это свежий оптимизатор (https://github.com/KellerJordan/Muon), предназначенный для обучения внутренних слоёв моделей, работает с 2D+ матрицами параметров. Разные векторы и скаляры надо обучать чем-то другим, например, AdamW. Эмпирически нашли также, что и эмбеддинги с классификационными головами лучше обучать не Muon’ом, у них другая динамика оптимизации (по крайней мере у первых).

Алгоритм берёт градиенты из SGD-Nesterov-momentum (можно и без Нестерова, но с ним лучше) и делает их постпроцессинг через итерацию Newton-Schulz, которая приближённо ортогонализует матрицу апдейтов. Больше подробностей есть в описании алгоритма тут: https://kellerjordan.github.io/posts/muon/.

Muon родился осенью прошлого года (https://x.com/kellerjordan0/status/1842300916864844014) и уже показал себя хорошо на скейлинге MoE LLM c 3B/16B параметров (вариант deepseek-v3-small) в обучении на 5.7T токенов, будучи примерно в два раза эффективнее AdamW (https://arxiv.org/abs/2502.16982).

В текущей работе проверили несколько датасетов. Датасет как в оригинальной работе про гроккинг и арифметику по модулю 97, а также датасет про чётность 10-битных бинарных строк. Задачи выбраны потому что на них феномен хорошо демонстрируется.

Поскольку одна из свежих работ про гроккинг (https://arxiv.org/abs/2501.04697) показывала влияние нестабильности софтмакса (Softmax Collapse) на проявление феномена, попробовали несколько альтернатив: стандартный Softmax, кусочнолинейный Stablemax с лучшей стабильностью, и приводящий к разреженным результатам Sparsemax.

Проверяли на классическом трансформере, с identity эмбеддингами (токен “42” маппится в вектор 42), RoPE, RMSNorm, SiLU активации, с дропаутом.

Сравнивали AdamW (β1 = 0.9, β2 = 0.98) и Muon. Делали множественные запуски с разными случайными сидами.

Как результат, у Muon статистически значимо преимущество, он существенно раньше приводит к гроккингу, примерно на 100-й эпохе вместо 150-й. Ещё и распределение grokking times более узкое.

Так что хотите быстрого гроккинга, юзайте мюон!

Авторы немного рассуждают на тему, что именно в мюоне помогает гроккингу. Ограничения спектральных норм и сигналы второго порядка уводят модель в сторону от тупого запоминания и помогают обнаружить настоящий паттерн.

В целом требуется более широкое исследование на более крупных моделях и разнообразных задачах.