Diffusion Fine-tuning with Rewarded Moment Matching Distillation

📄 arXiv: 2606.30414v1 📥 PDF

作者: Alexis Jacq, Guillaume Couairon, Valentin De Bortoli, Quentin Berthet, Arnaud Doucet, Romuald Elie

分类: cs.LG

发布日期: 2026-06-29

💡 一句话要点

提出奖励时刻匹配蒸馏以优化扩散模型训练

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 蒸馏训练 强化学习 扩散模型 天气预报 生成模型 高维数据 模型优化

📋 核心要点

  1. 现有方法在蒸馏和强化学习微调的相互作用上研究不足，特别是微调对生成质量的影响尚不明确。
  2. 本文提出奖励时刻匹配蒸馏（RMMD），通过同时进行模型蒸馏和奖励函数最大化来解决上述问题。
  3. 实验结果表明，RMMD在ImageNet上实现了优于单步和多步基线的性能，并在天气预报模型中显著提升了效率和准确性。

📝 摘要（中文）

蒸馏和强化学习微调是扩散后训练的主要支柱。尽管这两个阶段的相互作用尚不明确，尤其是微调如何影响蒸馏模型的生成质量。本文提出了一种新颖的框架——奖励时刻匹配蒸馏（RMMD），该框架同时进行扩散模型的蒸馏和奖励函数的最大化。RMMD通过调整采样循环以进行在政策训练，并将蒸馏损失重新用于积分KL正则化的代理，从而保持了高级蒸馏（如8步时刻匹配）的高保真“自然性”。通过在ImageNet上评估FID-Reward Pareto前沿，我们证明RMMD在与单步基线（DI++）和多步竞争者（DRaFT，HyperNoise）相比时，取得了更优的权衡。最后，我们将RMMD应用于GenCast，一个最先进的天气预报模型，在蒸馏的同时优化连续排名概率评分（CRPS）指标。结果显示，蒸馏模型实现了7.5倍的加速，并在93%的目标天气变量上超越了教师模型，且校准效果更佳。这证明了RMMD在复杂的高维科学领域的可扩展性。

🔬 方法详解

问题定义：本文旨在解决蒸馏和强化学习微调之间相互作用不明确的问题，现有方法未能有效评估微调对生成模型质量的影响。

核心思路：RMMD框架通过同时进行模型蒸馏和奖励函数最大化，优化了扩散模型的训练过程，确保了生成质量的提升。

技术框架：RMMD的整体架构包括调整采样循环以适应在政策训练，并将蒸馏损失用作积分KL正则化的代理，形成一个闭环的优化过程。

关键创新：RMMD的主要创新在于结合了蒸馏和强化学习的优势，通过奖励机制引导模型学习，从而提升生成质量，与传统方法相比具有显著的优势。

关键设计：在技术细节上，RMMD采用了8步时刻匹配的高级蒸馏技术，并在损失函数设计上进行了优化，以确保在训练过程中保持高保真度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，RMMD在ImageNet上实现了优于单步基线（DI++）和多步竞争者（DRaFT，HyperNoise）的性能，尤其是在93%的目标天气变量上超越了教师模型，并实现了7.5倍的速度提升，证明了其在复杂领域的有效性。

🎯 应用场景

该研究的潜在应用领域包括天气预报、图像生成等高维科学领域。通过优化模型的生成质量和训练效率，RMMD能够在实际应用中提供更快速和准确的预测，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.