Flash-DMD: Rumo à Geração de Imagens de Alta Fidelidade em Poucos Passos com Destilação Eficiente e Aprendizado por Reforço 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

Resumo

Os Modelos de Difusão emergiram como uma classe líder de modelos generativos, mas o seu processo iterativo de amostragem permanece computacionalmente dispendioso. A destilação de timesteps é uma técnica promissora para acelerar a geração, mas frequentemente requer treino extensivo e leva à degradação da qualidade da imagem. Além disso, o ajuste fino destes modelos destilados para objetivos específicos, como apelo estético ou preferência do utilizador, usando Aprendizagem por Reforço (RL), é notoriamente instável e facilmente cai em "reward hacking". Neste trabalho, introduzimos o Flash-DMD, uma nova estrutura que permite uma convergência rápida com destilação e um refinamento conjunto baseado em RL. Especificamente, propomos primeiro uma estratégia eficiente de destilação consciente do timestep que reduz significativamente o custo de treino com realismo aprimorado, superando o DMD2 com apenas 2,1% do seu custo de treino. Em segundo lugar, introduzimos um esquema de treino conjunto onde o modelo é ajustado com um objetivo de RL enquanto o treino de destilação de timesteps continua simultaneamente. Demonstramos que a perda estável e bem definida da destilação em curso atua como um poderoso regularizador, estabilizando efetivamente o processo de treino de RL e prevenindo o colapso da política. Extensivas experiências com modelos baseados em scores e de flow matching mostram que o nosso Flash-DMD proposto não só converge significativamente mais rápido, mas também alcança uma qualidade de geração de última geração no regime de amostragem de poucos passos, superando métodos existentes em qualidade visual, preferência humana e métricas de alinhamento texto-imagem. O nosso trabalho apresenta um paradigma eficaz para treinar modelos generativos eficientes, de alta fidelidade e estáveis. O código será disponibilizado em breve.

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.