Мюон, выстреливший некоторое время назад, стал основным оптимизатором в претрейнах больших языковых моделей, сместив с пьедестала, казалось бы, навечно засевший там Адам. Суть алгоритма в том, чтобы сделать полярное разложение матрицы моментов:
polar(X) = U V^T для X = U S V^Tи подставить его вместо момента.
Однако, есть у него существенный недостаток - значительно более дорогой шаг оптимизации в сравнении с Адамом. Адам производит сравнительно простые поточечные операции над параметрами, а Мюон - сложную машинерию с алгоритмом Ньютона-Шульца для ортогонализации матрицы производных.
Команда авторов не предлагает улучшения асимптотической сложности, но замечают, что некоторые операции можно делать более эффективно.
В выражения, возникающие в процессе подсчета многочленов в Ньютоне-Шульце, несколько раз возникают матричные произведения с матрицей X, которая может быть сильно вытянутой вдоль одной из размерностей. Гораздо дешевле тогда работать с матрицей Грама X^T X.
В качестве альтернативы оригинальному алгоритму предлагают считать полярное разложение в виде:
X (X^T X)^{-1/2}. Матричную степень находят тоже через Ньютон-Шульц, но для матрицы X^T X меньшего размера теперь уже.Для типичных конфигураций больших трансформеров это дает 55% уменьшений FLOPs против оптимизированных реализаций.
Далее замечают, что при перемножении симметричных матриц достаточно считать половину элементов, а остальные отражать по диагонали.
И все - успех, победа 🍾?
Ан нет. Наивная реализация, описанная выше численно нестабильна и приводит к спайкам в лоссе, инфам и нанам.
Анализируют причины произошедшего и, оказывается, что все дело в численной стабильности. Из-за конечной точности матрица X^T X может оказаться не положительно определенной. Да и собственные значения разъезжаются.
Чтобы это побороть делают алгоритм с рестартами. Это позволяет добиться контролями над собственными значениями. И для подстраховки делят матрицы на 1.02 - 1.05, дабы надежнее попасть в интервал собственных значений [0, 1], чтобы сойтись к оптимуму. Это чуть увеличивает время алгоритма, но стабилизирует его.
Между float16 и bfloat16 большой разницы нет, но авторы рекомендуют первый, как более точный в окрестности нуля.
Под используемые в работе матричные операции с симметричным матрицами пишут и выкладывают свои кернелы в Quack.
🧪 Эксперименты
Учат семейства моделей разной архитектуры - Llama-430M, Qwen-600M, Gemma-1B и кастомный MoE на 2х Шиншилла оптимальном бюджете. Вроде учится стабильно. Кривые лосса сходятся с исходным мюоном.
Ускорение оптимизатора в 1.5-2 раза в сравнении с ванильной торчовой реализацией, и несколько меньше против Мюона со специальными кернелами. Правда, в плане производительности (времени на шаг) все еще далеко до Адама.
💡 Выводы
Солидная инженерная работа с очевидной пользой для сообщества. В текущих воркфлоу шаг оптимизатора обычно не очень дорогой - 5-15% от общей стоимости обучения. Но с переходом на обучение в FP8/NVFP4 - доля времени работы алгоритма повысится. Тем самым ускорение станет более актуальным.