Diffusion Fine-tuning with Rewarded Moment Matching Distillation

• URL: http://arxiv.org/abs/2606.30414v1
• Date: Mon, 29 Jun 2026 15:00:56 GMT
• Title: Diffusion Fine-tuning with Rewarded Moment Matching Distillation
• Abstract summary: Rewarded Moment Matching Distillation (RMMD) is a novel framework that simultaneously distills diffusion models and maximizes a reward function.<n>We apply RMMD to GenCast, a state-of-the-art weather forecasting model, to distill it while optimizing the Continuous Ranked Probability Score (CRPS) metric.<n>The resulting distilled model achieves a 7.5x speedup while outperforming the teacher model on 93% of target weather variables, and being better calibrated.
• Score: 38.86104469934903
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Distillation and Reinforcement Learning (RL) fine-tuning are the primary pillars of diffusion post-training. While traditionally studied in isolation, the interaction between these phases remains poorly understood, and in particular how fine-tuning impacts the generative quality of distilled models. We introduce Rewarded Moment Matching Distillation (RMMD), a novel framework that simultaneously distills diffusion models and maximizes a reward function. RMMD preserves the high-fidelity ``naturalness'' characteristic of advanced distillation (such as 8-step Moment Matching) by adapting the sampling loop for on-policy training and repurposing the distillation loss as a proxy for integral KL regularization. By evaluating the FID-Reward Pareto fronts on ImageNet, we demonstrate that RMMD achieves superior trade-offs compared to single-step baselines (DI++) and multi-step competitors (DRaFT, HyperNoise). Finally, we apply RMMD to GenCast, a state-of-the-art weather forecasting model, to distill it while optimizing the Continuous Ranked Probability Score (CRPS) metric. The resulting distilled model achieves a 7.5x speedup while outperforming the teacher model on 93% of target weather variables, and being better calibrated. This proves that RMMD scales to complex, high-dimensional scientific domains.

Related papers

• SGMD: Score Gradient Matching Distillation for Few-Step Video Diffusion Distillation [57.297118390628384]
  Distribution Matching Distillation (DMD) is a widely used paradigm for accelerating inference in few-step video diffusion models.<n>We propose textbfScore Gradient Matching Distillation (SGMD).<n>It adopts a fake-score perspective by directly optimizing the fake score toward the teacher, while using teacher stop-gradient Fisher as a stable distribution-matching objective.<n>It substantially improves motion dynamics for 4-step distilled models while preserving temporal consistency.
  arXiv  Detail & Related papers  (2026-05-28T15:50:55Z)
• Guiding Distribution Matching Distillation with Gradient-Based Reinforcement Learning [41.982957134224904]
  Diffusion distillation has shown great promise in few-step generation but often sacrifices quality for sampling speed.<n>We propose GDMD, a novel framework that redefines the reward mechanism by prioritizing distillation gradients over raw pixel outputs.<n>Our models outperform the quality of their multi-step teacher and substantially exceed previous DMDR results in GenEval and human-preference metrics.
  arXiv  Detail & Related papers  (2026-04-21T02:57:13Z)
• Distribution Matching Distillation Meets Reinforcement Learning [30.960105413888943]
  We propose DMDR, a novel framework that combines Reinforcement Learning (RL) techniques into the distillation process.<n>We show that for the RL of the few-step generator, the DMD loss itself is a more effective regularization compared to the traditional ones.<n>Experiments demonstrate that DMDR can achieve leading visual quality, prompt coherence among few-step methods, and even exhibit performance that exceeds the multi-step teacher.
  arXiv  Detail & Related papers  (2025-11-17T17:59:54Z)
• Phased DMD: Few-step Distribution Matching Distillation via Score Matching within Subintervals [48.14879329270912]
  Phased DMD is a multi-step distillation framework that bridges the idea of phase-wise distillation with Mixture-of-Experts.<n>Phased DMD is built upon two key ideas: progressive distribution matching and score matching within subintervals.<n> Experimental results demonstrate that Phased DMD preserves output diversity better than DMD while retaining key generative capabilities.
  arXiv  Detail & Related papers  (2025-10-31T17:55:10Z)
• Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency [60.74505433956616]
  continuous-time consistency model (sCM) is theoretically principled and empirically powerful for accelerating academic-scale diffusion.<n>We first develop a parallelism-compatible FlashAttention-2 JVP kernel, enabling sCM training on models with over 10 billion parameters and high-dimensional video tasks.<n>We propose the score-regularized continuous-time consistency model (rCM), which incorporates score distillation as a long-skip regularizer.
  arXiv  Detail & Related papers  (2025-10-09T16:45:30Z)
• Adversarial Distribution Matching for Diffusion Distillation Towards Efficient Image and Video Synthesis [65.77083310980896]
  We propose Adrial Distribution Matching (ADM) to align latent predictions between real and fake score estimators for score distillation.<n>Our proposed method achieves superior one-step performance on SDXL compared to DMD2 while consuming less GPU time.<n>Additional experiments that apply multi-step ADM distillation on SD3-Medium, SD3.5-Large, and CogVideoX set a new benchmark towards efficient image and video synthesis.
  arXiv  Detail & Related papers  (2025-07-24T16:45:05Z)
• Restoration Score Distillation: From Corrupted Diffusion Pretraining to One-Step High-Quality Generation [82.39763984380625]
  We propose textitRestoration Score Distillation (RSD), a principled generalization of Denoising Score Distillation (DSD)<n>RSD accommodates a broader range of corruption types, such as blurred, incomplete, or low-resolution images.<n>It consistently surpasses its teacher model across diverse restoration tasks on both natural and scientific datasets.
  arXiv  Detail & Related papers  (2025-05-19T17:21:03Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.