Flash-DMD: к высококачественной генерации изображений за несколько шагов с помощью эффективного дистилляции и совместного обучения с подкреплением

Flash-DMD: Towards High-Fidelity Few-Step Image Generation with Efficient Distillation and Joint Reinforcement Learning

November 25, 2025

Авторы: Guanjie Chen, Shirui Huang, Kai Liu, Jianchen Zhu, Xiaoye Qu, Peng Chen, Yu Cheng, Yifu Sun

cs.AI

Аннотация

Диффузионные модели стали ведущим классом генеративных моделей, однако их итерационный процесс сэмплирования остается вычислительно дорогим. Дистилляция по временным шагам является перспективной техникой для ускорения генерации, но часто требует длительного обучения и приводит к деградации качества изображений. Более того, тонкая настройка таких дистиллированных моделей для конкретных задач, таких как эстетическая привлекательность или пользовательские предпочтения, с использованием обучения с подкреплением (RL) печально известна своей нестабильностью и легко приводит к взлому функции вознаграждения. В данной работе мы представляем Flash-DMD — новую архитектуру, которая обеспечивает быструю сходимость при дистилляции и совместную RL-доработку. В частности, мы сначала предлагаем эффективную стратегию дистилляции с учетом временных шагов, которая значительно снижает стоимость обучения при повышении реалистичности, превосходя DMD2 при использовании всего 2.1% от ее вычислительных затрат на обучение. Во-вторых, мы вводим схему совместного обучения, в которой модель тонко настраивается по RL-цели, в то время как обучение дистилляции по временным шагам продолжается параллельно. Мы показываем, что стабильная, хорошо определенная функция потерь от продолжающейся дистилляции действует как мощный регуляризатор, эффективно стабилизируя процесс RL-обучения и предотвращая коллапс политики. Многочисленные эксперименты на моделях, основанных на оценке и потоковом соответствии, демонстрируют, что предложенный метод Flash-DMD не только сходится значительно быстрее, но и достигает передового качества генерации в режиме сэмплирования с малым числом шагов, превосходя существующие методы по визуальному качеству, пользовательским предпочтениям и метрикам соответствия текста и изображения. Наша работа представляет эффективную парадигму для обучения эффективным, высокоточным и стабильным генеративным моделям. Код будет опубликован в ближайшее время.

English

Diffusion Models have emerged as a leading class of generative models, yet their iterative sampling process remains computationally expensive. Timestep distillation is a promising technique to accelerate generation, but it often requires extensive training and leads to image quality degradation. Furthermore, fine-tuning these distilled models for specific objectives, such as aesthetic appeal or user preference, using Reinforcement Learning (RL) is notoriously unstable and easily falls into reward hacking. In this work, we introduce Flash-DMD, a novel framework that enables fast convergence with distillation and joint RL-based refinement. Specifically, we first propose an efficient timestep-aware distillation strategy that significantly reduces training cost with enhanced realism, outperforming DMD2 with only 2.1% its training cost. Second, we introduce a joint training scheme where the model is fine-tuned with an RL objective while the timestep distillation training continues simultaneously. We demonstrate that the stable, well-defined loss from the ongoing distillation acts as a powerful regularizer, effectively stabilizing the RL training process and preventing policy collapse. Extensive experiments on score-based and flow matching models show that our proposed Flash-DMD not only converges significantly faster but also achieves state-of-the-art generation quality in the few-step sampling regime, outperforming existing methods in visual quality, human preference, and text-image alignment metrics. Our work presents an effective paradigm for training efficient, high-fidelity, and stable generative models. Codes are coming soon.