Оптимизатор Muon добавили в PyTorch
Тут вышел PyTorch v2.9:
https://github.com/pytorch/pytorch/releases/tag/v2.9.0
Для меня главный хайлайт релиза — это torch.optim.Muon. Если знаете про Muon, то новость на этом всё :)
Если не слышали про Muon, то это такой оптимизатор, который согласно легенде появился в 2024 году в результате дискуссий на сайте икс-точка-ком. Официальной "статьей" про Muon является этот пост: https://kellerjordan.github.io/posts/muon. У Muon уже есть истории успеха на бенчмарках и реальных задачах, в том числе при обучении больших моделей.
Технический TL;DR:
(1) Muon — это как SGD с моментом, но момент градиента ортогонализируется перед его применением (код). То есть ищется ортогональная матрица, наиболее похожая на оригинальный градиент, и вот она уже применяется к параметру с учетом learning rate. Это существенное отличие от "поэлементных" оптимизаторов по типу SGD или Adam, которые смотрят на все параметры модели как на один плоский вектор.
(2) Muon используется только для матриц в скрытых слоях. Все остальные параметры (входные эмбеддинги, финальные слои, bias-ы в скрытых слоях, и т.д.), нужно продолжать оптимизировать AdamW или чем-то подобным. То есть в коде у вас будет два оптимизатора.
(3) Из интуиций, почему Muon работает, я слышал примерно такое: "By orthogonalizing the update matrix, Muon ensures that updates are distributed more evenly across all parameters, preventing some features from being over-represented or neglected" (конкретно эту формулировку я просто нагуглил).
(4) Из практических рекомендаций для Muon знаю про важность использования weight decay. Еще вижу в доках интересную опцию adjust_lr_fn, которая вроде как позволит переиспользовать гиперпараметры, которые вы затюнили для AdamW. Но сам не пробовал.
———
Я сам видел случаи, когда Muon давал качество лучше, чем AdamW, но наверняка тут зависит от задачи. Моё личное ощущение такое, что сообщество пока еще копит интуицию о пользе Muon для скорости обучения и финального качества в разных сетапах, но в целом все уже согласны, что это интересный апдейт для области. Вообще, добавление новых оптимизаторов в PyTorch — это редкость, и вот Muon решили добавить.
Если у вас есть какой-то опыт работы с Muon, будет интересно узнать.
Тут вышел PyTorch v2.9:
https://github.com/pytorch/pytorch/releases/tag/v2.9.0
Для меня главный хайлайт релиза — это torch.optim.Muon. Если знаете про Muon, то новость на этом всё :)
Если не слышали про Muon, то это такой оптимизатор, который согласно легенде появился в 2024 году в результате дискуссий на сайте икс-точка-ком. Официальной "статьей" про Muon является этот пост: https://kellerjordan.github.io/posts/muon. У Muon уже есть истории успеха на бенчмарках и реальных задачах, в том числе при обучении больших моделей.
Технический TL;DR:
(1) Muon — это как SGD с моментом, но момент градиента ортогонализируется перед его применением (код). То есть ищется ортогональная матрица, наиболее похожая на оригинальный градиент, и вот она уже применяется к параметру с учетом learning rate. Это существенное отличие от "поэлементных" оптимизаторов по типу SGD или Adam, которые смотрят на все параметры модели как на один плоский вектор.
(2) Muon используется только для матриц в скрытых слоях. Все остальные параметры (входные эмбеддинги, финальные слои, bias-ы в скрытых слоях, и т.д.), нужно продолжать оптимизировать AdamW или чем-то подобным. То есть в коде у вас будет два оптимизатора.
(3) Из интуиций, почему Muon работает, я слышал примерно такое: "By orthogonalizing the update matrix, Muon ensures that updates are distributed more evenly across all parameters, preventing some features from being over-represented or neglected" (конкретно эту формулировку я просто нагуглил).
(4) Из практических рекомендаций для Muon знаю про важность использования weight decay. Еще вижу в доках интересную опцию adjust_lr_fn, которая вроде как позволит переиспользовать гиперпараметры, которые вы затюнили для AdamW. Но сам не пробовал.
———
Я сам видел случаи, когда Muon давал качество лучше, чем AdamW, но наверняка тут зависит от задачи. Моё личное ощущение такое, что сообщество пока еще копит интуицию о пользе Muon для скорости обучения и финального качества в разных сетапах, но в целом все уже согласны, что это интересный апдейт для области. Вообще, добавление новых оптимизаторов в PyTorch — это редкость, и вот Muon решили добавить.
Если у вас есть какой-то опыт работы с Muon, будет интересно узнать.