diffusion fine-tuning with rewarded moment matching distillation¶

会议: NeurIPS 2026
arXiv: 2606.30414
代码: https://github.com/google-deepmind/graphcast
领域: 模型压缩
关键词: 扩散模型蒸馏, 强化学习微调, 矩匹配蒸馏, 在策略对齐, 气象预报加速

一句话总结¶

提出 Rewarded Moment Matching Distillation (RMMD)，通过将中间去噪分布的矩匹配损失转化为对齐正则项并在策略生成轨迹上做单步反传，在保留多步蒸馏生成质量的同时高效最大化任意可微下游奖励，将气象预报大模型 GenCast 加速 7.5 倍并在 93% 变量上超越教师。

研究背景与动机¶

扩散模型因其训练目标可转化为加噪数据上的稳定回归问题，已成为高保真图像生成与复杂科学动力学建模的核心基石。然而在推理阶段，扩散模型必须沿着离散时间网格串行迭代去噪数十至数百步，极高的计算开销严重制约了高吞吐部署。为了突破这一瓶颈，后训练阶段通常演化出两条平行的技术路线：一是利用少步蒸馏技术（如一致性模型、矩匹配蒸馏 MMD 等）将采样压缩至数步；二是通过强化学习或可微奖励微调，引导生成分布贴合人类审美偏好、物理约束或任务特定下游指标。

然而，将蒸馏加速与奖励微调深度结合时面临根本性的两难困境。若在统一训练阶段联合执行蒸馏与奖励最大化（如 DI++ 等单步方法），奖励信号会诱导生成样本迅速漂移出教师模型的输入分布，破坏蒸馏回归目标的有效性；若先做多步蒸馏再进行微调，沿多步计算链做反向传播显存开销极大，而截断反传步骤（如 DRaFT-1 仅反传最后一步、ReFL 在中间带噪隐空间评估奖励）又会带来严重梯度偏差或对抗伪影，噪声扰动类方案（如 HyperNoise）则在结构上被局限于低频扰动。更重要的是，在少步多步蒸馏的生成器上，缺乏一种在数学上能够自然充当“分布保真度约束”的正则化手段来抵御奖励攻击（reward hacking）。

本文的切入角度是：多步矩匹配蒸馏（MMD）在预蒸馏阶段能够显式拟合中间去噪边缘分布；只要我们在微调阶段将生成策略切换为在策略（on-policy），并将 MMD 的矩匹配目标作为中间去噪边缘分布的散度正则项，就能在不展开多步计算链的前提下精确评估单步梯度。核心 idea：提出两阶段 RMMD 框架，先利用 8 步 MMD 构建高保真分布基线，再在当前学生策略生成的加噪轨迹上执行单步奖励梯度反传，并复用冻结的矩匹配损失作为积分散度正则项，兼顾极端推理加速与精准奖励对齐。

方法详解¶

整体框架¶

RMMD 采用系统性的两阶段架构。第一阶段为无奖励的纯多步蒸馏：利用矩匹配蒸馏（MMD）将多步扩散教师模型压缩至固定步数（例如 8 步或 2 步），使得学生模型在离散时间步上的边缘去噪分布与教师对齐，训练完成后将该学生模型权重 $\Phi_{\theta_0}$ 及其辅助矩预测网络 $m_{\text{student}}$ 冻结作为稳定的参考分布源。第二阶段为在策略奖励微调：模型从自身当前策略采样完整轨迹生成无噪样本，加噪至随机时间步后执行单步去噪预测，联合计算可微下游奖励、在策略 MMD 正则项以及参考模型输出的 L2 约束，通过一步显存开销捕获跨频段的精细调整。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["起始噪声 x1 与条件输入"] --> B["阶段一：冻结基准蒸馏<br/>8步 MMD 预训练获得参考网络"]
    B --> C["阶段二：在策略轨迹采样<br/>学生网络 K 步自回归 roll-out 生成样本"]
    C --> D["单步重加噪与去噪预测<br/>截取时间步 t 注入高斯噪声并预测 x0"]
    D --> E["在策略矩匹配正则化<br/>匹配参考辅助网络与当前学生条件矩"]
    D --> F["单步可微奖励反向传播<br/>计算任务奖励梯度并联合 L2 锚定约束"]
    E --> G["梯度聚合与参数更新<br/>优化 RMMD 目标函数更新学生参数"]
    F --> G

关键设计¶

1. 在策略轨迹重加噪与单步奖励梯度：消除离线数据偏差与多步反传开销

直接在 MMD 梯度中加入奖励项的天真做法依赖于离线真实数据集的加噪样本，不仅使得微调上限被数据集内部的最佳奖励卡死，且随着策略演化，静态数据边缘分布 $p_t$ 迅速与学生当前中间分布脱节。RMMD 改变了输入采样的生成来源：在每轮微调迭代中，首先使用当前学生参数 $\theta$ 运行完整的 $K_{\text{student}}$ 步采样过程，生成一批在策略终态样本 $\tilde{x}_0 \sim p_{\theta}$。对这些样本应用停止梯度算子 $\text{sg}[\tilde{x}_0]$ 后，随机选取连续时间步 $t \in (0, 1]$ 并按正向扩散过程注入噪声： $$x^{\text{pol}}_t = \alpha_t \text{sg}[\tilde{x}_0] + \sigma_t \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I_d)$$ 随后，学生网络仅对 $x^{\text{pol}}_t$ 评估一次去噪前向传播 $\hat{x}_0 = \Phi_{\theta}(x^{\text{pol}}_t, t, \xi)$，并直接在预测结果 $\hat{x}_0$ 上计算下游可微奖励函数 $R(\hat{x}_0)$。由于反向传播仅穿透最后这一次单步前向推理，梯度计算为 $\nabla_{\theta}\mathcal{L}_{\text{Reward}}(\theta) = \mathbb{E}_{\hat{x}_0}[(\frac{\partial \hat{x}_0}{\partial \theta})^\top \nabla_{\hat{x}_0}R(\hat{x}_0)]$，完全规避了跨步展开多步采样链的显存爆炸问题，同时使学生网络能在自己真实的推理输入分布上实现全频段内容微调。

2. 在策略 MMD 散度正则化：将分布蒸馏转化为抗奖励退化的核心锚点

在强化学习微调中，生成模型极易利用奖励模型的漏洞（例如 CLIP 的高频纹理偏好）发生“奖励攻击”导致图像结构彻底坍塌，而传统的全局 KL 散度或单步得分匹配（Score Distillation）无法有效约束多步生成路径的动力学稳定性。RMMD 创新性地将 MMD 蒸馏目标改造为在策略正则项。对于当前预测 $\hat{x}_0$ 施加停止梯度后，在子区间 $s \in [t - \delta_{\text{student}}, t)$ 上抽取中间条件加噪点 $x'_s \sim p_{\text{cond}}(x'_s | x^{\text{pol}}_t, \text{sg}[\hat{x}_0])$，并通过辅助网络与冻结教师网络计算条件期望矩差： $$\nabla_{\theta}\mathcal{L}_{\text{MMD}}(\theta) = \mathbb{E}_{t,s,x^{\text{pol}}_t,\hat{x}_0,x'_s}\left[\left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \big(m_{\text{student}}(s,x'_s) - m_{\text{teacher}}(s,x'_s)\big)\right]$$ 在微调阶段，$m_{\text{teacher}}$ 可直接由第一阶段训练完毕并冻结的辅助网络替代。该损失本质上构成了对学生中间边缘分布漂移的严密惩罚，充当了多步离散轨迹上的积分散度约束，迫使学生网络在提升奖励的同时维持与预训练高保真分布的一致性。

3. 多样本排序评分规则扩展：面向系综预报离散度校准的连续分级概率评分优化

在复杂高维科学预测（如全球气象模拟）中，扩散模型的核心任务是捕获未来状态的概率分布，但传统少步蒸馏往往面临“欠离散度”（under-dispersion）缺陷，即生成的多个系综成员方差过窄。针对多样本依赖的目标函数，RMMD 进一步拓展为多样本对齐机制，直接将气象学标准评估准则——连续分级概率评分（CRPS）作为可微奖励函数： $$\text{CRPS}_r(\{\hat{x}^{\tau,i}\}_{i=1}^M, x^\tau) = \frac{1}{M}\sum_{i=1}^M |\hat{x}^{\tau,i}_{(r)} - x^\tau_{(r)}| - \frac{1}{2M(M-1)}\sum_{i=1}^M\sum_{j=1}^M |\hat{x}^{\tau,i}_{(r)} - \hat{x}^{\tau,j}_{(r)}|$$ 其中第一项惩罚系综预测与单一样本真实观测 $x^\tau$ 的一范数偏差，第二项惩罚系综内部不同采样成员间的差异以显式鼓励物理多样性。RMMD 在微调阶段同时生成 $M=2$ 个在策略样本分支，各自反传 MMD 正则化梯度的同时，通过反传联合负 CRPS 梯度，引导蒸馏学生不仅提升单变量预测精度，且直接修复系综的发散度标定。

损失函数 / 训练策略¶

RMMD 微调阶段优化的总体加权标量目标函数定义为： $$\mathcal{L}(\theta) = \mathcal{L}_{\text{MMD}}^{\text{online}}(\theta) - \lambda R(\hat{x}_0) + \lambda_{\text{reg}}\mathcal{L}_{\text{L2}}(\theta)$$ 其中 $\mathcal{L}_{\text{L2}}(\theta) = \mathbb{E}_{t,x^{\text{pol}}_t,\xi}[\|\Phi_\theta(x^{\text{pol}}_t,t,\xi) - \Phi_{\theta_0}(x^{\text{pol}}_t,t)\|^2]$ 惩罚当前参数与冻结 MMD 基线参数 $\Phi_{\theta_0}$ 的无噪声输出差距。在图像实验（ImageNet 64 与 512）中，固定 $\lambda_{\text{reg}} = \lambda / 2$ 并扫描奖励权重 $\lambda$；在 GenCast 气象模型微调中设置 $\lambda = 0.3, \lambda_{\text{reg}} = 1.0$。采用 Adam 优化器结合余弦衰减学习率，并在微调阶段使用停止梯度降低显存开销，计算复杂度仅比离线微调增加约 2 倍，换取了彻底的在策略稳定收益。

实验关键数据¶

主实验¶

论文在 ImageNet（分类条件图像生成）与 ERA5 全球气象数据集（GenCast 蒸馏，状态维度高达 5,313,600）上全面评测了 RMMD。针对气象模型，评价指标为相对于 59 步 NFE GenCast 教师模型的相对 CRPS 提升百分比（Rel. CRPS improv.，跨所有变量和预报时效平均）以及优于教师模型的物理变量比例（Win Rate）。

模型 / 配置	推理步数 (NFE)	加速比	相对 CRPS 提升 (↑)	胜率 Win Rate (↑)
GenCast 原始教师 (Price et al., 2023)	59	1.0×	0.00% (基准)	N/A
原始 MMD 直接蒸馏 (Plain MMD)	8	7.5×	-1.32%	4.9%
MMD 最优基线 (η=0.5, ρ=100)	8	7.5×	+0.82%	75.0%
RMMD w/ CRPS (离线采样版本)	8	7.5×	+1.11%	89.2%
RMMD w/ CRPS (在策略版本，本文)	8	7.5×	+1.51%	93.0%

在 ImageNet 64×64 与 512×512 视觉实验中，8 步 MMD 初始模型取得了与教师极度接近的初始 FID（ImageNet-64 下学生 1.26 vs 教师 1.19，远超 1 步蒸馏最佳基线 DI++ 的 2.65）。在多尺度和多种奖励（黑白转化、拉普拉斯平滑、Inception Score、CLIP 语义对齐）的 Pareto 前沿对比中，RMMD 在相同奖励提升下均保持了显著更低的 FID，消除了 DRaFT-2 的全局颜色崩塌和 DRaFT-1 的高频对抗噪声。

消融实验¶

论文在预报物理一致性超参、第一阶段蒸馏策略以及微调正则化组件上进行了系统的消融分析。

表 1：GenCast MMD 蒸馏超参数与随机扰动（Churn）消融

实验配置	采样策略调整	相对 CRPS 提升	胜率 Win Rate	关键结论说明
Plain MMD	默认 DDPM (η=1.0, ρ=7)	-1.32%	4.9%	原始蒸馏参数在气象动力学上显著落后于教师
MMD w/ η=0.5	降低去噪随机性 (η=0.5)	+0.22%	50.7%	适度降低随机扰动方差大幅提升了预报精度
MMD (best)	引入对数信噪比均匀调度 (ρ=100)	+0.82%	75.0%	调整离散网格时间步密度进一步压低了多尺度误差
MMD w/ Churn	强加随机抖动 (S_noise=1.05)	+0.11%	50.2%	EDM 随机抖动仅对教师有效，在蒸馏模型上导致指标恶化
教师模型 w/o Churn	移除教师随机抖动	-0.57%	14.2%	证明教师模型高度依赖 churn 弥补离散度缺陷

表 2：ImageNet 预蒸馏阶段网络随机性（Dropout）与调度平移消融（FID 指标，越低越好）

模型版本	I64 (8 steps)	I64 (2 steps)	I512 (2 steps)
MMD 原文基线 (Salimans et al., 2024)	1.24	3.86	-
MMD 基础复现	1.35	2.00	9.70
MMD + Dropout (0.1)	1.37	1.66	5.40
MMD + Dropout + Shifted Schedule ($b=\log 2$)	1.26	1.40	-

关键发现¶

在策略采样的绝对优势：在 GenCast 微调中，在策略版本相较于离线采样版本在相对 CRPS 上进一步获得 +0.4% 的净增益，胜率从 89.2% 飙升至 93.0%。由于在策略网络是在自身生成的含噪多步轨迹上微调，有效缓解了自回归 roll-out 过程中的复合误差累积。
单步优化泛化至长时程预测：微调时仅在 12 小时（单步状态转移）上计算并反传 CRPS 奖励，但在自回归推演 15 步（7.5 天全期预报）的整个评估窗口内，相对于教师模型的性能优势随预报天数增加而持续扩大，证明 RMMD 学习到了真实的转移概率分布，而非在单步时间截面上投机取巧。
彻底解决蒸馏模型的欠离散度问题：通过引入 Spread-Skill Ratio 分析，单纯 MMD 蒸馏在短期预报内呈现严重的欠离散（Ratio 远小于 1），而 RMMD 引入 CRPS 对齐后使各层大气物理变量的离散度显著恢复，与全精度系综分布完美吻合。

亮点与洞察¶

将生成蒸馏损失巧妙重构为在策略正则化算子：不同于常规 RL 依赖外挂且难以标定的 KL 惩罚，RMMD 发现冻结的 MMD 矩匹配损失在数学上直接约束了中间特征分布的动态边缘，天然充当了防范奖励攻击的数学屏障。
成功打通扩散模型在超高维物理科学中的蒸馏闭环：GenCast 单一状态向量具有超 530 万自由度，RMMD 证明了基于分级概率评分的在策略微调不仅在理论图像生成上可行，更能直接迁移并赋能超大规模真实物理世界模拟。

局限与展望¶

生成保真度上限受制于初始蒸馏阶段：RMMD 的微调阶段主要负责将现有生成空间引导向高奖励方向，并不能从根本上突破第 1 阶段 MMD 模型的初始 FID 水平，因此前序必须配置强悍的多步蒸馏底模。
强依赖奖励函数的可微性：由于单步梯度必须直接求导穿透 $\nabla_{\hat{x}_0} R(\hat{x}_0)$，目前算法无法原生支持黑盒不可微评估指标，在离散文本或复杂非可微物理仿真环境应用时必须先训练可微代理模型。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [将矩匹配损失升华为在策略微调的积分分布正则项，理论视角极佳]
实验充分度: ⭐⭐⭐⭐⭐ [涵盖 ImageNet 跨分辨率图像实验与全球尺度 530 万维气象动力学真实应用]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严密自洽，从动机、公式到实验剖析层层递进]
价值: ⭐⭐⭐⭐⭐ [攻克了扩散模型蒸馏与对齐无法兼得的痛点，兼具学术价值与工业加速落地潜力]