PaperlayerView exact revision on arXiv

Source-linked AI summary

$R_\text{dm}$: Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillation

Linqian Fan, Peiqin Sun, Tiancheng Wen, Shun Lu, Chengru Song

arXiv:2603.28460v2cs.CVcs.LG

TL;DR

Diffusion sampling is computationally costly, and teacher-derived distillation objectives can limit student performance. This paper reframes distribution matching as a reward, then introduces group normalization and a unified multi-reward framework. GNDM outperforms vanilla DMD, while GNDMR balances aesthetic quality and fidelity.

  • Problem

    Diffusion models incur computational overhead from iterative sampling, while teacher-derived distillation objectives can bottleneck student performance.

    Source: §1 · p. 2

  • Method

    The paper treats distribution matching as a reward, applies group normalization to stabilize it, and combines it with other rewards in a unified framework.

    Source: §1 · p. 2

  • Results

    GNDM achieves lower FID than vanilla DMD, and GNDMR achieves an optimal balance between aesthetic quality and fidelity.

    Source: §5.3 · p. 10, §6 · p. 10

  • Takeaways & Limitations

    The reward-centered framework offers a flexible, stable approach to real-time, high-fidelity image synthesis within the paper’s reported scope.

    Source: §6 · p. 10

    At larger diffused timesteps, higher noise increases the variance of the distribution-matching guidance, which can mislead optimization and make distillation more difficult.

    Source: §4.2 · p. 5

Abstract

from arXiv · show

Diffusion models achieve state-of-the-art generative performance but are fundamentally bottlenecked by their slow, iterative sampling process. While diffusion distillation techniques enable high-fidelity, few-step generation, traditional objectives often restrict the student's performance by anchoring it solely to the teacher. Recent approaches have attempted to break this ceiling by integrating Reinforcement Learning (RL), typically through a simple summation of distillation and RL objectives. In this work, we propose a novel paradigm by re-conceptualizing distribution matching as a reward, denoted as $R_\text{dm}$. This unified perspective bridges the algorithmic gap between Diffusion Matching Distillation (DMD) and RL, providing several primary benefits. (1) Enhanced Optimization Stability: We introduce Group Normalized Distribution Matching (GNDM), which adapts standard RL group normalization to stabilize $R_\text{dm}$ estimation. By leveraging group-mean statistics, GNDM establishes a more robust and effective optimization direction. (2) Seamless Reward Integration: Our reward-centric formulation inherently supports adaptive weighting mechanisms, allowing for the fluid combination of DMD with external reward models. (3) Improved Sampling Efficiency: By aligning with RL principles, the framework readily incorporates Importance Sampling (IS), leading to a significant boost in sampling efficiency. Extensive experiments demonstrate that GNDM outperforms vanilla DMD, reducing the FID by 1.87. Furthermore, our multi-reward variant, GNDMR, surpasses existing baselines by striking an optimal balance between aesthetic quality and fidelity, achieving a peak HPS of 30.37 and a low FID-SD of 12.21. Ultimately, $R_\text{dm}$ provides a flexible, stable, and efficient framework for real-time, high-fidelity synthesis. Codes are coming soon.

1 Introduction

Few-step diffusion distillation reduces sampling overhead, but teacher-anchored objectives can limit student performance. The paper reframes distribution matching as a reward and introduces GNDM and GNDMR to combine distillation with RL and other rewards.

1 Introduction

  • Traditional distillation targets derived exclusively from a pretrained teacher can bottleneck student performance, while prior RL-distillation approaches typically linearly combine their objectives.The paper positions reward-based distribution matching as an alternative unified optimization framework.Source: §1 · p. 2
  • GNDM applies group normalization to stabilize the distribution-matching reward’s optimization direction.The authors identify increasing reward variance at higher image noise as a source of unstable optimization.Source: §1 · p. 2
  • GNDM outperforms vanilla DMD, while GNDMR surpasses existing baselines with a balance between visual aesthetics and distillation efficiency.The framework treats distribution matching as a reward that can be jointly optimized with task-specific rewards.Source: §1 · p. 2
  • The unified reward framework GNDMR performs GRPO simultaneously with the distribution-matching reward and other rewards.This framework is presented as enabling adaptive weighting and importance sampling.Source: §1 · p. 2, Figure 2 · p. 3

2 Related Work

Related work spans methods that improve distribution matching distillation and methods that apply RL to diffusion models. The paper positions its approach within prior work combining DMD and RL while distinguishing its reward-centered formulation.

Distribution Matching Distillation

  • DMD2 adds a GAN loss, Flash-DMD combines timestep-aware strategy with pixel-GAN, TDM combines trajectory distillation with distribution matching, and Decoupled DMD separates the objective.DMDR combines DMD with RL, using distribution matching as a regularizer to permit student exploration beyond the teacher.Source: §2 · p. 3

Reinforcement Learning for Diffusion Models

  • RL methods for diffusion include preference-based approaches and scalar-reward policy optimization, while DMDR first jointly optimizes DMD and RL objectives.The proposed method is also built upon the DMDR* framework.Source: §2 · p. 3

3 Preliminaries

The preliminaries describe DMD as matching a few-step student generator to a multi-step teacher and introduce policy optimization for diffusion denoising. They also motivate treating the difference between real and fake scores as a reward.

3.1 Distribution Matching Distillation

  • DMD distills a multi-step diffusion teacher into a high-fidelity few-step student by minimizing reverse-KL divergence between their distributions.Its gradient uses real and fake score functions evaluated at a diffused sample.Source: §3.1 · p. 3
  • For few-step generation, the student can be trained through backward simulation that alternates denoising and noising across sampling steps.The text describes extending the one-step DML gradient to multi-step DML.Source: §3.1 · p. 4

3.2 Denoising Diffusion Policy Optimization

  • DDPO maps diffusion denoising to a multi-step Markov decision process and applies policy gradients to collected trajectories.Importance sampling reweights trajectories from previous parameters to facilitate multi-step optimization.Source: §3.2 · p. 4
  • The real-minus-fake score term in the policy-gradient formulation motivates treating that term as a reward.The formulas omit conditioning variable c for ease of observation.Source: §3.2 · p. 4

4 Methodology

The method recasts distribution matching as a reward that can be optimized with policy gradients, then uses group normalization and adaptive weighting to address reward-estimation and optimization challenges. It also combines distribution-matching and auxiliary rewards and incorporates importance sampling into training.

4.1 Rdm: Distribution Matching as a Reward

  • Using one trajectory sample makes the distribution-matching policy-gradient objective strictly equivalent to the negative DMD loss gradient.The formulation connects the policy-gradient and DMD objectives while preserving consistency with traditional distribution matching.Source: §4.1 · p. 5

4.2 Revisiting the Distribution Matching Reward

  • At larger diffused timesteps, score variance rises and can mislead optimization; Rdm is also described as a divergence optimized toward zero that guards against over-optimization.The paper frames Rdm as a distribution-alignment reward rather than an unbounded preference score.Source: §4.2 · p. 5, Figure 3 · p. 5

4.3 Group Normalized Distribution Matching

  • Group normalization uses mean subtraction to stabilize Rdm’s reward-direction estimate, while shared timesteps within each group support consistency and importance sampling permits multiple generator updates per sampling pass.The method shares generator timestep t and diffused timestep t′ within a group and says importance sampling increases sampling efficiency while maintaining performance.Source: §4.3 · pp. 5–6

4.4 GNDM with Other Rewards

  • GNDMR combines dense latent-trajectory distribution-matching rewards with sparse final-image auxiliary rewards, weighted into a total advantage and multi-reward objective.The auxiliary reward is group-normalized, and its contribution is controlled by a weighting function.Source: §4.4 · p. 6

4.5 Adaptive Weight Design in Practice Implement

  • Sign-based denominator normalization stabilizes Rdm, and adaptive β_dm,t weighting keeps auxiliary-reward amplitudes consistent with the dynamic distribution-matching weight.The paper attributes a stable optimization signal to reduced variance in raw sampling residuals and says synchronized weighting can improve the target score without collapse.Source: §4.5 · pp. 6–7

5 Experiments

Experiments compare GNDMR with state-of-the-art models and ablate its reward normalization, timestep sampling, and weighting choices. Results report stronger distillation and aesthetic alignment, while importance sampling reduces the sample budget needed for comparable performance.

5 Experiments

  • Experiments use LAION-AeS-6.5+ prompts, with DFN-CLIP and HPSv2.1 as default rewards, and evaluate 4-step models on 10K COCO2014 prompts.Evaluation includes HPS, PickScore, MPS, CLIP Score, FID, and FID-SD; GNDMR-IS updates the generator twice per sampling.Source: §5 · pp. 7–8, §5.1 · p. 8

5.1 Comparison with State-Of-The-Art (SOTA)

  • GNDMR outperforms existing 4-step models on aesthetics on SD3 and SD3.5 and achieves strictly lower FID than data-free DMDR on both; GNDMR-IS halves training expenses while maintaining highly competitive alignment.Source: §5.1 · p. 8
  • GNDMR shows richer detail, more vibrant color, stronger prompt alignment, and fewer artifacts than DMDR in qualitative comparisons.The reported comparison describes an aesthetic-quality and clean-synthesis balance.Source: §5.1 · p. 8, Figure 4 · p. 9

5.2 Importance Sampling Correction

  • With five training iterations per sample and η = 0.5, 32×16 with IS matches standard 32×16 performance while using fewer samples than even 16×8.Larger batches improve reward optimization per training step but increase sampling cost; the reported IS comparison uses SD3-Medium and HPSv2.1.Source: §5.2 · pp. 8–9, Figure 5 · p. 10

5.3 Ablation Study

  • GNDMR has lower FID after the first 500 distillation iterations and during later HPS- or PS-reward training, with comparable target-reward performance.Table 2 compares group normalization during an initial distillation phase followed by training with different rewards.Source: Table 2 · p. 9, §5.3 · p. 10
  • Sharing both generation timestep t and diffused timestep t′ within each group yields the best performance among the tested sampling strategies.The reported rationale is that Rdm depends on both timesteps.Source: Table 3 · p. 9, §5.3 · p. 10
  • Applying group normalization to larger diffused-timestep intervals stabilizes Rdm estimation and leads to lower FID.The passage associates larger intervals with higher noise and increased estimation variance.Source: Table 3 · p. 9, §5.3 · p. 10
  • Scaling the external reward by βdm,t synchronizes its weighting with the dynamic distillation coefficient, supporting steady reward improvement without collapse.The passage contrasts this with static weighting, which can produce unstable optimization or suboptimal reward gains.Source: §5.3 · p. 10

6 Conclusion

The conclusion presents distribution matching as a reward-based framework for diffusion distillation and RL. It reports improved DMD performance with GNDM and a balance of aesthetic quality and fidelity with GNDMR.

6 Conclusion

  • GNDM outperforms vanilla DMD, achieving a notable reduction in FID scores.The conclusion does not specify the size of the reduction.Source: §6 · p. 10
  • GNDMR integrates additional rewards and achieves a balance between aesthetic quality and fidelity.Source: §6 · p. 10
  • In Figure 6, incorporating βdm,t yields more stable and superior reward optimization than static weighting under both DMDR and GNDMR.The figure evaluates HPS on HPDv2 test prompts.Source: Figure 6 · p. 11
  • The framework points toward applying recent RL techniques to diffusion-model distillation.Source: §6 · p. 11

A.1 Experiment Details

The experiments use AdamW for both the generator and fake score network, with model-specific training settings for SD3-Medium and SD3.5-Medium.

A.1 Experiment Details

  • Both the generator and fake score network use AdamW, and the fake score is updated once per generator update.The experiments run on 8 NVIDIA H800 GPUs; β1 and β2 are 0.9 and 0.999.Source: §A.1 · p. 15
  • For SD3-Medium, training uses 1024 × 1024 resolution, batch size 128, CFG 7.0, and 8k total iterations.HPS and CLIP weights rise from 5 for the first 500 iterations to 10 for the next 7.5k.Source: §A.1 · p. 15
  • For SD3.5-Medium, training uses 512 × 512 resolution, batch size 128, CFG 3.5, and 8k iterations.HPS and CLIP weights are both 10.Source: §A.1 · p. 15

A.2 Training Algorithm Details

The training procedure samples a next state, applies forward diffusion, and compares real and fake denoiser outputs to construct distribution-matching rewards and weighted advantages.

A.2 Training Algorithm Details

  • The procedure samples noise to form x_next from the generator prediction, then forward-diffuses that prediction for denoising.The next sample uses the predicted mean and sigma_next-scaled noise.Source: §A.2 · p. 15
  • Real and fake denoiser predictions are combined with the sample residual to calculate R_dm and the DMD weighting factor w_dm.Source: §A.2 · p. 15
  • The algorithm group-normalizes R_dm and each external reward before aggregating the distribution-matching advantage with w_dm.Algorithm 2 details the weighted advantage aggregation used in training.Source: §A.2 · p. 15

B Additional Experiments

The additional experiments examine noise initialization and the granularity of the adaptive distribution-matching weight βdm,t. They report better FID with shared noise and better HPS with sample-wise weighting.

B Additional Experiments

  • The additional ablations use SD3-Medium at 512 × 512, first training GNDM for 500 iterations before the full GNDMR process.The rapid evaluation considers only the HPS reward.Source: B Additional Experiments · p. 16

B.1 Effect of Noise Initialization

  • Shared initial noise across candidates yields more robust distillation and a lower FID than independent initialization.Table 4 compares the noise initialization strategies.Source: §B.1 · p. 16

B.2 Design of the Adaptive Weight βdm,t

  • βdm,t calibrates the pixel-wise distillation weight wdm,t, addressing its dimensionality mismatch with sample-wise RL rewards.Source: §B.2 · p. 16
  • The sample-wise βdm,t, computed by averaging wdm,t across non-batch dimensions, improves HPS more than pixel-wise or baseline configurations.The passage attributes the improvement to a more stable reward signal across generated samples.Source: §B.2 · p. 16

B.3 Sensitivity Analysis of the Reward Weight whps

The sensitivity analyses show how HPS reward weighting and the importance-sampling clip range affect optimization: adaptive weighting supports consistent HPS gains, while a larger clip range benefits convergence.

B.4 Sensitivity of Clip Range η in Importance Sampling

  • Unlike conservative GRPO fine-tuning, this distillation framework benefits from a larger clip range η for rapid, efficient convergence.A larger η permits more aggressive policy updates, whereas an excessively low η can restrict learning progress and fail to correct distribution shifts.Source: §B.4 · p. 16, Figure 8 · p. 17

B.3 Sensitivity Analysis of the Reward Weight whps

  • Increasing w_hps improves overall HPS, while removing β_dm,t makes improvement inconsistent.The tested w_hps values are 1, 10, 15, and 20.Source: §B.3 · p. 16, Figure 9 · p. 17

C Additional Qualitative Results

Additional qualitative comparisons show GNDMR results across complex prompts selected from DMDR and Flash-SD. The authors report more pleasing aesthetics, texture, and structural integrity than state-of-the-art distillation baselines.

C Additional Qualitative Results

  • Across complex prompts, GNDMR yields more aesthetically pleasing results with enhanced textural detail and structural integrity than state-of-the-art distillation baselines.Figure 10 uses prompts from DMDR in its top three rows and Flash-SD in its bottom three rows.Source: C Additional Qualitative Results · p. 17, Figure 10 · p. 18

Related papers

Loading 2603.28460v2…