Search papers, labs, and topics across Lattice.
This paper introduces Rewarded Moment Matching Distillation (RMMD), a novel framework that integrates distillation and reinforcement learning (RL) fine-tuning to enhance the generative quality of diffusion models. By adapting the sampling loop for on-policy training and utilizing the distillation loss as a proxy for integral KL regularization, RMMD maintains the high-fidelity characteristics of advanced distillation methods. The framework demonstrates superior performance on the FID-Reward Pareto fronts and achieves significant improvements in a weather forecasting model, outperforming the teacher model while providing a 7.5x speedup and better calibration across 93% of target variables.
RMMD not only accelerates model inference by 7.5x but also outperforms its teacher model on nearly all target weather variables, showcasing a breakthrough in distillation techniques.
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.
DeepMind