Guiding Distribution Matching Distillation with Gradient-Based Reinforcement Learning

1. Problem Statement

Diffusion models have revolutionized image synthesis but suffer from high computational costs due to their iterative sampling nature. Distillation methods (e.g., Distribution Matching Distillation, or DMD) compress these models into efficient few-step generators but often sacrifice image quality and alignment with human preferences to achieve speed.

Recent attempts to integrate Reinforcement Learning (RL) into distillation (e.g., DMDR) aim to bridge this quality gap. However, existing approaches suffer from critical limitations:

• Naive Sample-Based Scoring: They evaluate raw generated samples ($x_0$) to compute rewards. In early training stages, these samples are noisy and unstable, leading to unreliable reward signals.
• Optimization Divergence: The RL objective (maximizing reward on $x_0$) often conflicts with the DMD objective (minimizing distribution divergence). This creates gradient conflicts, causing instability, "reward hacking" (artifacts to maximize score), and a requirement for a "cold-start" phase where RL is disabled initially.
• Inefficiency: Traditional RL data collection often requires expensive Stochastic Differential Equation (SDE) sampling to generate diverse rollouts.

2. Methodology: GDMD Framework

The authors propose GDMD, a framework that redefines the reward mechanism by prioritizing distillation gradients over raw pixel outputs. Instead of evaluating the final image $x_0$, GDMD evaluates the direction of the update required to improve the model.

Core Concept: Gradient-Based Scoring

The authors reinterpret the DMD gradient as an implicit target tensor. By transforming the DMD loss objective, they derive a target vector ($x_{tar}$) that encapsulates the quality of the gradient update. This allows existing reward models (trained on images) to evaluate the potential of an update rather than the noisy current state.

The pipeline consists of three key components:

1. Distillation-aware Gradient Collection (DaGC):

   • Instead of SDE sampling, GDMD leverages the inherent randomness of the DMD process to collect grouped gradients.
   • It employs two strategies:
     • $t$-based collection: Estimates scores with varying noise intensities to fit the time-step distribution.
     • $\mu_{fake}$-based collection: Tracks gradients across different versions of the fake score estimator during optimization.
   • This eliminates the need for external SDE rollouts.

2. Implicit Gradient Scoring (IGS):

   • Since standard reward models (e.g., CLIP, HPS, ImageReward) require image inputs, GDMD transforms the gradient information into an implicit target image $x_{tar}$.
   • Mathematically, the DMD gradient is reformulated as a loss term: $x_{tar} = G_{\theta^-}(z) + (s_{real} - s_{fake})$.
   • The reward model $R_\psi$ evaluates $x_{tar}$ (decoded via VAE) against the text prompt. This score implicitly reflects the quality of the gradient update, providing a stable signal even in early training.

3. Negative-aware Preference Optimization (NaPO):

   • Inspired by DiffusionNFT, the framework uses implicit positive and negative policies for preference learning.
   • It constructs positive ($v^+$) and negative ($v^-$) velocity fields based on the raw reward scores.
   • The RL loss ($L_{RL}$) optimizes the student model to move toward the positive policy and away from the negative policy, ensuring alignment with the distillation direction.

Training Objective:
The final loss combines the standard DMD loss and the RL loss:
$$\mathcal{L}_{GDMD} = \lambda \mathcal{L}_{DMD} + \gamma \mathcal{L}_{RL}$$
This joint optimization synchronizes the RL policy with the distillation objective, preventing divergence.

3. Key Contributions

• Paradigm Shift: Pioneers a gradient-based scoring strategy for DMD-RL joint distillation, moving away from unstable sample-based evaluation.
• Stability & Efficiency: Eliminates the need for cold-start phases and expensive SDE sampling. The gradient-based approach provides reliable guidance from the very first training step.
• Specialized Pipeline: Introduces a comprehensive architecture including DaGC (gradient collection), IGS (implicit scoring), and NaPO (preference optimization).
• Performance: Demonstrates that aligning RL with distillation gradients resolves optimization conflicts, allowing the student model to exceed its multi-step teacher.

4. Experimental Results

The authors evaluated GDMD on SDXL-Base and SD3-Medium models, comparing against baselines like DMD, DMD2, DMDR, LCM, and Flash Diffusion.

• Quantitative Performance (SOTA):

  • GenEval: GDMD outperforms all competitors (including DMDR and the multi-step teacher) across all dimensions (Single, Two, Count, Colors, Pos, Attr).
  • Quality Metrics: Trained solely on CLIP Score and Human Preference (HP) Score, GDMD achieved state-of-the-art results on unseen metrics (Aesthetic Score, Pick Score, ImageReward).
  • Comparison: A 4-step GDMD model outperformed its 50-step teacher model and surpassed DMDR significantly (e.g., ImageReward: 1.2702 vs. 1.0120 for SD3).

• Qualitative Improvements:

  • GDMD produces images with fewer artifacts, less oversaturation, and better text alignment compared to DMDR and DMD2.
  • It effectively handles complex prompts involving counting, spatial relationships, and attribute binding.

• Human Evaluation:

  • In a user study, GDMD (4-step) was preferred over the Teacher model (55.1% win rate), DMDR (62.7%), and DMD+DiffusionNFT (61.9%) for both image quality and prompt alignment.

• Ablation Studies:

  • Confirmed that combining $t$-based and $\mu_{fake}$-based gradient collection yields the best generalization.
  • Showed that gradient-based scoring (GDMD) outperforms sample-based RL strategies (ReFL, GRPO, DiffusionNFT) when applied to distillation.

5. Significance

GDMD represents a significant advancement in efficient generative modeling. By decoupling the reward signal from the noisy output of early-stage generation and anchoring it to the optimization gradient, the authors have solved the stability issues plaguing previous DMD-RL attempts.

The significance lies in:

1. Breaking the Quality-Speed Trade-off: It proves that few-step models can not only match but surpass the quality of their multi-step teachers.
2. Scalability: The method removes the dependency on cold-starts and complex SDE sampling, making high-quality distillation more accessible and robust.
3. Generalization: The ability to improve unseen metrics (like Aesthetic Score) while training on specific rewards (CLIP/HP) suggests a more robust alignment with human preferences, reducing the risk of reward hacking.

In summary, GDMD establishes a new state-of-the-art for few-step diffusion models by fundamentally rethinking how reinforcement learning signals are integrated into the distillation process.