Исследователи представили CMuon (Chunked Momentum Orthogonalization) — метод оптимизации, предназначенный для повышения эффективности обучения диффузионных трансформеров (DiT). Новый подход решает проблему нестабильности и замедления сходимости, характерную для стандартного оптимизатора Muon на поздних этапах обучения, что позволяет значительно сократить вычислительные затраты при сохранении высокого качества генерации изображений.

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

Метод CMuon разбивает процесс оптимизации на отдельные блоки, что позволяет стабилизировать обновление весов и избежать «затухания» прогресса. Это решение позволяет достичь SOTA-результатов быстрее, чем при использовании классических методов оптимизации, что критически важно для обучения крупномасштабных моделей с миллиардами параметров.

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

  • CMuon устраняет проблему нестабильности обучения DiT, возникающую при использовании стандартного оптимизатора Muon.
  • Метод основан на технике «Chunked Momentum Orthogonalization», которая предотвращает деградацию весов на поздних этапах сходимости.
  • Технология позволяет снизить общие вычислительные затраты на обучение диффузионных моделей без потери качества генерации.
  • Исследование подтверждает эффективность подхода для крупномасштабных архитектур, используемых в современных визуальных генеративных системах.