Flash-DMD: Verso una Generazione di Immagini Fedele in Pochi Passi con Distillazione Efficiente e Apprendimento per Rinforzo Congiunto

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

November 25, 2025

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

cs.AI

Abstract

I modelli di diffusione sono emersi come una classe leader di modelli generativi, ma il loro processo di campionamento iterativo rimane computazionalmente costoso. La distillazione dei timestep è una tecnica promettente per accelerare la generazione, ma richiede spesso un addestramento esteso e porta a un degrado della qualità dell'immagine. Inoltre, l'addestramento di precisione (fine-tuning) di questi modelli distillati per obiettivi specifici, come l'attrattiva estetica o le preferenze dell'utente, utilizzando l'Apprendimento per Rinforzo (RL), è notoriamente instabile e facilmente soggetto a reward hacking. In questo lavoro, introduciamo Flash-DMD, un nuovo framework che consente una convergenza rapida con la distillazione e un affinamento congiunto basato su RL. Nello specifico, proponiamo innanzitutto una strategia di distillazione efficiente e consapevole del timestep che riduce significativamente il costo di addestramento migliorando il realismo, superando DMD2 con solo il 2.1% del suo costo di addestramento. In secondo luogo, introduciamo uno schema di addestramento congiunto in cui il modello viene raffinato con un obiettivo RL mentre l'addestramento della distillazione dei timestep continua simultaneamente. Dimostriamo che la perdita stabile e ben definita derivante dalla distillazione in corso agisce come un potente regolarizzatore, stabilizzando efficacemente il processo di addestramento RL e prevenendo il collasso della policy. Esperimenti estesi su modelli basati su score e su flow matching mostrano che il nostro Flash-DMD proposto non solo converge significativamente più velocemente, ma raggiunge anche una qualità di generazione allo stato dell'arte nel regime di campionamento a pochi passi, superando i metodi esistenti in termini di qualità visiva, preferenza umana e metriche di allineamento testo-immagine. Il nostro lavoro presenta un paradigma efficace per l'addestramento di modelli generativi efficienti, ad alta fedeltà e stabili. I codici saranno disponibili a 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.