Flash-DMD : Vers une génération d'images haute fidélité en peu d'étapes grâce à une distillation efficace et un apprentissage par renforcement conjoint

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

November 25, 2025

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

cs.AI

papers.abstract

Les modèles de diffusion sont devenus une classe prédominante de modèles génératifs, mais leur processus d'échantillonnage itératif reste coûteux en calcul. La distillation des pas de temps est une technique prometteuse pour accélérer la génération, mais elle nécessite souvent un entraînement intensif et entraîne une dégradation de la qualité de l'image. De plus, le réglage fin de ces modèles distillés pour des objectifs spécifiques, tels que l'attrait esthétique ou les préférences des utilisateurs, à l'aide de l'apprentissage par renforcement (RL), est notoirement instable et tombe facilement dans le détournement de récompense. Dans ce travail, nous présentons Flash-DMD, un nouveau cadre qui permet une convergence rapide avec distillation et un raffinement conjoint basé sur le RL. Plus précisément, nous proposons d'abord une stratégie de distillation efficace prenant en compte les pas de temps, qui réduit significativement le coût d'entraînement tout en améliorant le réalisme, surpassant DMD2 avec seulement 2,1 % de son coût d'entraînement. Deuxièmement, nous introduisons un schéma d'entraînement conjoint où le modèle est affiné avec un objectif de RL tandis que l'entraînement de distillation des pas de temps se poursuit simultanément. Nous démontrons que la perte stable et bien définie provenant de la distillation en cours agit comme un régularisateur puissant, stabilisant efficacement le processus d'entraînement RL et empêchant l'effondrement de la politique. Des expériences approfondies sur les modèles basés sur le score et l'appariement de flux montrent que notre Flash-DMD proposé converge non seulement beaucoup plus rapidement, mais atteint également une qualité de génération de pointe dans le régime d'échantillonnage à faible nombre d'étapes, surpassant les méthodes existantes en termes de qualité visuelle, de préférence humaine et de métriques d'alignement texte-image. Notre travail présente un paradigme efficace pour l'entraînement de modèles génératifs efficaces, fidèles et stables. Les codes seront bientôt disponibles.

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.