MuonMD: новый метод нормализации весов для оптимизации нейросетей

🔬 Метод Авторы замечают, что обновления весов в радиальном и угловом направлениях не являются независимыми. Для оптимизаторов типа Адама/Мюона величина шага не зависит от нормы весов, поэтому изменение угла обратно пропорционально норме весов. Аналогично, при движении вдоль касательной (по радиальному градиенту) растет норма весов. И чтобы она не улетела в космос, приходится накладывать 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% у Мюона). 💡 Выводы Выглядит как многообещающий и интересный результат. Интересно, насколько скоро данная техника будет внедрена в обучение больших моделей и станет ли она такой же стандартной практикой, как Мюон.