PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2511.20549v1cs.CVcs.AI

TL;DR

Diffusion distillation seeks to reduce the costly iterative sampling of diffusion models, but existing approaches can require extensive training and distribution matching can be unstable. Flash-DMD combines timestep-aware distillation with reinforcement learning during the distillation process. It reports improved generation quality and human preference at reduced training cost.

  • Problem

    Diffusion distillation can require thousands of GPU hours, while distribution matching objectives may train unstably and exhibit mode-seeking.

    Source: §1 · p. 2, §3.2 · p. 4

  • Method

    Flash-DMD assigns distribution matching and adversarial losses by timestep, then integrates latent reinforcement learning into the distillation process for few-step models.

    Source: §1 · p. 2

  • Results

    Flash-DMD reports superior image quality, human preference, and text-image alignment; stage 1 reports best performance at 2.1% of DMD2’s training cost.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The authors conclude that Flash-DMD makes diffusion distillation more efficient and accessible for low-step generative modeling.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Flash-DMD combines timestep-aware distillation with reinforcement learning tailored to few-step models. It targets costly, unstable distillation and reports strong generation quality at substantially reduced training cost.

1. Introduction

  • Existing distillation methods often require thousands of GPU hours, limiting accessibility, while DMD objectives can be unstable and mode-seeking.DMD2’s naive combination of adversarial training and distribution matching neglects timestep-aware behavior and can compromise efficiency.Source: §1 · p. 2
  • The paper asks how to coordinate distribution matching with perceptual realism early, then directly refine visual details and perceptual fidelity later.Source: §1 · p. 2
  • At high-noise timesteps, Flash-DMD uses DMD loss to align the student with the teacher; at low-noise timesteps and on real images, Pixel-GAN enhances realism and texture.The method assigns different objectives to the model’s global-structure and fine-detail phases.Source: §1 · p. 2
  • For few-step distilled models, Flash-DMD integrates latent reinforcement learning with ongoing distillation to refine details and mitigate reward hacking.The paper contrasts this with prior few-step RL work that can produce “oil painting” artifacts.Source: §1 · p. 2
  • Stage 1 reports best performance at only 2.1% of DMD2’s training cost, while Flash-DMD surpasses the teacher and baseline on image quality, human preference, and text-image alignment.The reported results also generalize across score-based diffusion and flow matching models.Source: §1 · pp. 2–3

2. Related Work

The related work introduces diffusion’s forward-noising and reverse-generation processes, then reviews methods for reducing inference steps and aligning distilled outputs. It also situates reinforcement learning for text-to-image generation, where few-step models remain relatively understudied.

2. Related Work

  • Diffusion generation perturbs real images with noise in a forward process and reconstructs images through a reverse process; repeated reverse steps make generation time-consuming.The reverse DDPM step models x_(t-1) given x_t with a predicted mean and covariance.Source: §2 · p. 3
  • Progressive distillation halves inference steps iteratively but is costly and inherits teacher-quality limits; consistency distillation can degrade few-step inference.Segmented trajectory distillation is cited as a response to the consistency issue.Source: §2 · p. 3
  • Adversarial distillation aligns few-step student outputs with multi-step teachers, while DMD-family methods use distribution matching; naive combinations can hinder efficiency.DMD2 adds latent adversarial training, ADM uses a GAN framework with Hinge loss, and SenseFlow optimizes scorers and discriminators.Source: §2 · pp. 3–4
  • Text-to-image reinforcement learning includes DPO preference objectives and GRPO methods based on group-sampled advantages, but work on few-step models remains limited.Pairwise Sample Optimization strengthens the relative likelihood margin between training and reference sets.Source: §2 · p. 4

3. Methodology

Flash-DMD distills a teacher into a few-step generator by coordinating distribution matching and adversarial training across timesteps, then continues distillation alongside preference optimization. Its methods address the training overhead of DMD-series distillation and the reward-hacking risks of few-step reinforcement learning.

3.1. Preliminary

  • DMD-series distillation minimizes reverse KL between teacher and student distributions, using score estimators to track each distribution.DMD2 adds adversarial training with real images to the distribution-matching objective.Source: §3.1 · p. 4

3.2. Training Inefficiency of DMD Series

  • DMD2’s combined gradients can conflict, while its score estimator must both track generated samples and discriminate them, contributing to training inefficiency.DMD2 updates the score estimator five times per generator update under TTUR.Source: §3.2 · p. 4

3.3. Faster Convergence at First Stage

  • Flash-DMD assigns distribution matching to high-noise steps and adversarial training against real images to the final low-noise step, reducing objective interference.Its SAM-based PixelGAN provides pixel-level supervision intended to counter the mode-seeking tendency of distribution matching.Source: §3.3 · p. 5
  • The score estimator is trained only to track the generator, with one or two updates per generator update sufficient for stable distribution tracking.An EMA update injects generator parameters into the score estimator after each generator update; the approach reports more stable dynamics and superior sample fidelity than DMD2 at TTUR=5.Source: §3.3 · p. 5

3.4. Reinforcement Learning for Distilled Model

  • Few-step preference optimization uses LRM-ranked win–lose pairs from shared noise at high-noise timesteps and combines its likelihood loss with Flash-DMD’s ongoing training.The authors report that this combination can further improve human preference and text-image alignment; their motivation notes that high-noise samples show better layout and fine-detail diversity than low-noise samples.Source: §3.4 · p. 6

4. Experiments

Experiments evaluate Flash-DMD on SDXL and SD3-Medium, then test reinforcement-learning refinement and ablations. The results cover benchmark quality, training cost, stability, and preference-related metrics.

4.1. Implementation Details

  • The experiments use filtered LAION 5B image-text pairs for stage-one SDXL and SD3-Medium distillation, and the Latent Reward Model for stage-two RL on distilled four-step SDXL.Evaluation uses 10,000 COCO 2014 prompts and metrics including CLIP, HPSv2, ImageReward, PickScore, and MPS.Source: §4.1 · pp. 6–7

4.2. Experiment Analysis

  • At TTUR=5, Flash-DMD surpasses DMD2 on all benchmarks at 37.5% of its training cost; at TTUR=2 it retains superior human preference and comparable text-image consistency at 8.3% cost.With TTUR=1, 1,000 steps, and 2.1% of DMD2’s cost, it still yields a higher human preference score; it also consistently outperforms the teacher across tested settings.Source: §4.2 · p. 7, Table 1 · p. 7
  • On SD3-Medium, Flash-DMD with flow matching achieves significantly better results than the teacher and SD3-Flash using only 4K training steps.The authors report this as evidence that the distillation strategy is effective and efficient in the flow-matching paradigm.Source: §4.2 · p. 7, Table 2 · p. 7
  • Flash-DMD scores highest on PickScore and MPS, while its image fidelity surpasses SDXL and competing reinforcement approaches.Hyper-SDXL has the highest ImageReward and HPSv2 scores but produces overexposed, unnatural images; LPO-SDXL has the highest CLIP score but produces oversmoothed images. The authors say distribution-matching and PixelGAN constraints alleviate reward hacking.Source: §4.2 · pp. 7–8, Table 3 · p. 7, Figure 4 · p. 8

4.3. Ablation Studies

  • With TTUR=2, Flash-DMD improves steadily during training, whereas DMD2’s initial gains quickly degrade.The authors interpret this comparison as showing better training stability and efficiency for Flash-DMD.Source: §4.3 · p. 8, Figure 5 · p. 8
  • EMA yields higher later-stage ImageReward and Pickapic scores, while HPSv2 remains nearly identical to training without EMA.The authors conclude that EMA enhances visual quality and human preference without compromising text alignment.Source: §4.3 · p. 8, Figure 6 · p. 8
  • Among reinforcement-loss to distribution-matching-loss frequency ratios of 1:1, 2:1, 5:1, and 10:1, 5:1 achieves the highest score.The comparison trains for 2,000 iterations on one H20 GPU, using a generator initialized from TTUR1-1k weights.Source: §4.3 · pp. 8–9, Table 4 · p. 8
  • Online training outperforms post-training with LPO alone, high-noise-only training beats all-noise training, and adding PixelGAN gives a marginal metric improvement.The selected 5,000-step model with GAN loss is the final enhanced version and costs 12 H20 GPU hours.Source: Figure 7 · p. 8, §4.3 · p. 9

5. Conclusion

The conclusion presents Flash-DMD as a two-phase approach to efficient, high-quality few-step generation. It reports improved generation quality and human preference at reduced training cost.

5. Conclusion

  • Flash-DMD coordinates timestep-aware objectives to improve distillation efficiency, then applies latent reinforcement learning to refine visual details while preventing mode collapse and artifacts.The early phase coordinates distribution matching with perceptual-realism enhancement; the later phase performs the refinement.Source: §5 · p. 9
  • Experiments report superior generation quality and the highest human preference scores with significantly reduced training costs.Source: §5 · p. 9
  • The authors conclude that the approach makes diffusion distillation more efficient and accessible for low-step generative modeling.Source: §5 · p. 9

Supplementary Material

The supplementary material describes a table comparing mainstream reward models across four aspects.

Supplementary Material

  • The table compares reward-model scope and evaluation dimensions.Source: Table 5 · p. 12
  • It also compares the feature extractors used by the reward models.Source: Table 5 · p. 12
  • Timestep awareness is the fourth comparison aspect; the caption defines Aes, Align, and Fid as aesthetics, text-image alignment, and fidelity.Source: Table 5 · p. 12

6. Latent Reward Model: Selection Rationale

The section explains why prior online preference optimization generally uses rewards at the final denoising stage and motivates the Latent Reward Model. Its stated rationale combines compute constraints with reward models’ lack of timestep sensitivity.

6. Latent Reward Model: Selection Rationale

  • Prior online preference optimization applies reward models in the final denoising stage, a practice the authors associate with reward hacking.Source: §6 · p. 12
  • Evaluating rewards at multiple denoising steps would require VAE decoding at each step, increasing GPU memory and computation requirements.Source: §6 · p. 12
  • The authors select LRM because it evaluates noisy latents at any timestep without switching to pixel space, using a pretrained diffusion model’s latent representations.Source: §6 · p. 12

7. Algorithm of Flash-DMD

Flash-DMD combines timestep-aware distribution matching with a Pixel-GAN constraint, then incorporates a latent reward model in a second stage. The stage-1 comparison table covers SDXL on COCO-10k.

Algorithm of Flash-DMD

  • Flash-DMD’s two stages combine fast-converging distribution matching with timestep awareness and a Pixel-GAN constraint, then add a latent reward model to enhance fidelity and aesthetics.The first-stage strategy is described as alleviating mode-seeking and improving realism.Source: §7 · p. 12
  • Table 6 compares stage-1 Flash-DMD on SDXL with other distillation methods using COCO-10k; the caption does not provide the table values.Source: Table 6 · p. 12

8. 4-steps and 8-steps Flash-DMD on SDXL

On SDXL, Flash-DMD is reported to outperform other distillation and reinforcement-learning methods in the 8-step setting. The section describes 4- and 8-step results and qualitative comparisons.

8-step and qualitative results

  • Flash-DMD outperforms other distillation and reinforcement-learning methods in the 8-step setting, and the paper reports high-quality results with realism and aesthetic appeal at few steps.The authors also state that its image quality surpasses the SDXL teacher in a qualitative comparison.Source: §8 · p. 12

Qualitative comparisons

  • Figure 8 presents qualitative comparisons with other models; the prose says Flash-DMD surpasses other distillation models and the teacher in image quality.Source: §8 · p. 12, Figure 8 · p. 13

Stage-2 evaluation

  • Stage 2 is compared with other reinforcement-learning models on COCO-10k, while the supplied table caption gives no scores or winner.Source: Table 7 · p. 13

9. 4-step Flash-DMD on SD3-Medium

The paper extends Flash-DMD to SD3-Medium after stage 1 and presents visual results. It uses a manually curated, proprietary 100,000-instance training set rather than LAION.

Training data

  • Rather than training on LAION, the authors curated a proprietary, high-quality training set of 100,000 instances spanning diverse subjects and scenes, including portraits, architecture, flora and fauna, and food in a realistic photographic style.Source: §9 · p. 13

SD3-Medium results

  • The SD3-Medium extension is demonstrated after stage 1, with Figure 14 showing visual results.Source: §9 · p. 13

10. Additional Visualizations and Captions

The supplementary material includes the Flash-DMD training algorithm and qualitative visualizations across SDXL and SD3-Medium settings. The algorithm combines distribution-matching and adversarial training, with reward optimization enabled in stage 2.

Algorithm 1

  • The algorithm takes a pretrained teacher, real data, inference-step and timestep sets, a Pixel-level discriminator, VAE decoder layers, and a latent reward model, and outputs a trained few-step generator.Source: §10 · p. 14
  • In the generator update, distribution matching at noisy timesteps is combined with an adversarial loss for realism and detail.Source: §10 · p. 14
  • When stage 2 is active, the algorithm samples latents for reward evaluation, constructs win–lose pairs, and updates the generator with a preference-optimization loss.Source: §10 · p. 14

Additional visualizations

  • SDXL visualizations cover stage 1 and stage 2 at both four and eight steps, with captions specifying training schedules and checkpoint initialization.The four-step stage-1 examples include TTUR settings of 1 for 1,000 steps and 2 for 4,000 steps; the eight-step stage-1 examples use TTUR 2 for 3,000 steps.Source: Figure 9 · p. 15, Figure 10 · p. 16, Figure 11 · p. 17
  • Figure 14 shows stage-1 four-step SD3-Medium visual results, trained with TTUR=2 for 7,000 steps.Source: Figure 14 · p. 20

Related papers

Loading 2511.20549v1…