Diffusion Fine-tuning with Rewarded Moment Matching Distillation

2026-06-29

AI Summary

This work addresses the disconnect between distillation and reinforcement learning (RL) fine-tuning in existing diffusion models, which often leads to degraded generation quality when optimizing reward objectives. The authors propose a unified framework that, for the first time, integrates high-order moment-matching distillation with reward-driven online policy RL. Within a shared sampling loop, distillation and policy updates are performed concurrently, with the distillation loss serving as a surrogate KL regularization to jointly preserve sample fidelity and enhance task-specific performance. Evaluated on ImageNet, the method achieves a superior trade-off between FID and reward metrics. When applied to the GenCast weather forecasting model, it accelerates inference by 7.5× while outperforming the teacher model on 93% of meteorological variables and demonstrating improved calibration.

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.

Problem

• diffusion distillation
• reinforcement learning fine-tuning
• generative quality
• reward optimization
• model calibration

Innovation

• Rewarded Moment Matching Distillation
• diffusion distillation
• reinforcement learning fine-tuning
• Pareto optimization
• knowledge distillation