Авторы замечают, что обновления весов в радиальном и угловом направлениях не являются независимыми. Для оптимизаторов типа Адама/Мюона величина шага не зависит от нормы весов, поэтому изменение угла обратно пропорционально норме весов. Аналогично, при движении вдоль касательной (по радиальному градиенту) растет норма весов. И чтобы она не улетела в космос, приходится накладывать weight decay.
🧮 Дабы обеспечить контроль над величиной веса, его можно представить в следующем виде:
W = diag(γ_row) W_hat diag(γ_col)Где W_hat — нормализованный вес, а γ_row и γ_col — построчные и поколоночные мультипликаторы, которые могут быть как скалярами, так и векторами. Нормализовать можно как по норме Фробениуса всей матрицы, так и по строкам или столбцам.
📊 Эксперименты
🧪 Авторы аблируют разные варианты нормализации и параметризации gains.
Вариант с нормализацией всей матрицы выглядит не хуже остальных, будучи при этом самым простым, поэтому в итоге используют именно его.
✅ γ_row и γ_col полезно иметь в виде векторов, причем оба должны быть обучаемыми. Использование row и col gains более высокого ранга не помогает. Чтобы gains не обращались в ноль, их можно параметризовать через строго положительную функцию активации, но и без такой параметризации метод работает стабильно.
📈 Хорошим свойством нормализации весов является переносимость learning rate между моделями разной ширины и глубины. Оптимальная величина шага обучения оказывается почти постоянной.
Кроме того, больше не нужен warmup в начале обучения, так как оптимизатор не взрывается по ходу обучения. Weight decay теперь тоже становится лишним, поскольку параметризация сама обеспечивает нахождение нормы весов в разумных пределах.
📉 Лучше всего работает линейное затухание learning rate вместо популярного нынче WSD.
🏆 В итоге лучший вариант MuonMD (Magnitude-Direction) сходится быстрее по числу итераций по сравнению с базовым Мюоном. Причем работает он хорошо как для Dense-, так и для MoE-моделей.
⚙️ На практике веса моделей хранятся так же, как и раньше: разбиение на величину и угол выполняется внутри оптимизатора. Замедление шага обучения при этом небольшое (около 5% у Адама и 1–2% у Мюона).
💡 Выводы
Выглядит как многообещающий и интересный результат. Интересно, насколько скоро данная техника будет внедрена в обучение больших моделей и станет ли она такой же стандартной практикой, как Мюон.