Исследователи изучили эффективность оптимизатора Muon при обучении архитектур с пространством состояний (SSM), в частности модели Mamba-2. Ранее этот метод, использующий итерации Ньютона-Шульца для ортогонализации обновлений весовых матриц, демонстрировал отличные результаты преимущественно в трансформерах. Новое исследование подтверждает применимость Muon для SSM, предлагая альтернативу стандартному оптимизатору AdamW для повышения качества обучения.

Алгоритм Muon выполняет градиентный спуск в условиях спектральной нормы, что позволяет более эффективно обновлять веса нейронных сетей. В отличие от AdamW, который опирается на оценки моментов первого и второго порядка, Muon фокусируется на сохранении ортогональности обновлений. Это критически важно для глубоких архитектур, где накопление ошибок при обновлении весов может приводить к нестабильности обучения.

Эксперименты на модели Mamba-2 объемом 130 миллионов параметров показывают, что использование Muon позволяет достичь более высокой точности при сопоставимых вычислительных затратах. Авторы работы подчеркивают, что спектральная оптимизация становится перспективным направлением для масштабирования не только трансформеров, но и альтернативных архитектур, которые претендуют на более эффективную обработку длинных последовательностей данных.

Ключевые факты

  • Muon использует итерацию Ньютона-Шульца для ортогонализации обновлений каждой матрицы весов.
  • Исследование проведено на архитектуре Mamba-2 с параметризацией 130M для прямого сравнения с AdamW.
  • Метод обеспечивает градиентный спуск под спектральной нормой, что улучшает сходимость по сравнению с классическими методами.
  • Работа расширяет область применения Muon с трансформеров на модели с пространством состояний (SSM).