Flash-DMD: Hacia la Generación de Imágenes de Alta Fidelidad en Pocos Pasos mediante Destilación Eficiente y Aprendizaje por Refuerzo Conjunto

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

November 25, 2025

Autores: Guanjie Chen, Shirui Huang, Kai Liu, Jianchen Zhu, Xiaoye Qu, Peng Chen, Yu Cheng, Yifu Sun

cs.AI

Resumen

Los Modelos de Difusión han surgido como una clase líder de modelos generativos, pero su proceso de muestreo iterativo sigue siendo computacionalmente costoso. La destilación de pasos temporales es una técnica prometedora para acelerar la generación, pero a menudo requiere un entrenamiento extenso y conduce a una degradación de la calidad de la imagen. Además, el ajuste fino de estos modelos destilados para objetivos específicos, como el atractivo estético o la preferencia del usuario, utilizando Aprendizaje por Refuerzo (RL) es notoriamente inestable y cae fácilmente en la manipulación de recompensas. En este trabajo, presentamos Flash-DMD, un marco novedoso que permite una convergencia rápida con destilación y un refinamiento conjunto basado en RL. Específicamente, primero proponemos una estrategia de destilación eficiente y consciente del paso temporal que reduce significativamente el costo de entrenamiento con un realismo mejorado, superando a DMD2 con solo el 2.1% de su costo de entrenamiento. En segundo lugar, introducimos un esquema de entrenamiento conjunto donde el modelo se ajusta con un objetivo de RL mientras el entrenamiento de destilación de pasos temporales continúa simultáneamente. Demostramos que la pérdida estable y bien definida de la destilación en curso actúa como un poderoso regularizador, estabilizando efectivamente el proceso de entrenamiento de RL y evitando el colapso de la política. Experimentos exhaustivos con modelos basados en puntuación y de emparejamiento de flujos muestran que nuestro Flash-DMD propuesto no solo converge significativamente más rápido, sino que también logra una calidad de generación de vanguardia en el régimen de muestreo de pocos pasos, superando a los métodos existentes en calidad visual, preferencia humana y métricas de alineación texto-imagen. Nuestro trabajo presenta un paradigma efectivo para entrenar modelos generativos eficientes, de alta fidelidad y estables. Los códigos estarán disponibles pronto.

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.