Rdm: Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillation¶

会议: ECCV 2026
论文: ECCV 2026 Official
代码: https://github.com/Flq2002/Rdm
领域: 图像生成
关键词: 扩散模型, 扩散蒸馏, 分布匹配, 强化学习, GRPO

一句话总结¶

论文将扩散模型分布匹配蒸馏（DMD）的目标重新概念化为一种内生奖励（\(R_{\mathrm{dm}}\)），并在统一的群相对策略优化（GRPO）框架下引入群归一化分布匹配（GNDM）与自适应幅度平衡，解决了多步扩散蒸馏中大时间步方差发散与外生奖励融合崩塌的难题。

研究背景与动机¶

扩散模型凭借逐步去噪机制在高质量图像合成领域取得了显著成功，但数十至数百步的串行迭代采样带来了极高的计算延迟与推理开销，难以满足实时生成需求。为了压缩采样步数，以分布匹配蒸馏（DMD/DMD2）为代表的少步蒸馏技术被广泛采用，其核心是通过最小化真实教师分布与生成器假分布间的逆向 KL 散度，引导紧凑生成器在 1 至 4 步内生成图像。然而，单纯锚定教师模型输出的传统蒸馏目标存在天然的能力上限，学生生成器不仅继承了教师模型的固有缺陷，且难以进一步针对人类审美或任务对齐进行超越性优化。

近期工作尝试引入强化学习（如 DDPO、ReFL、DPO）来突破教师模型的性能天花板，例如 DMDR 将蒸馏损失与外生强化学习奖励函数做简单的线性加权组合。然而，这种将蒸馏与 RL 割裂为两个独立目标并生硬叠加的做法面临深刻的优化冲突：一方面，扩散去噪过程中随着加噪时间步变大，样本模糊度骤升导致分布匹配的分数估计方差急剧发散，使蒸馏梯度极不稳定；另一方面，外生奖励模型（如 HPS、CLIP Score）属于无界标量优化，容易诱发模式崩溃（Reward Hacking），而静态固定的损失加权系数无法适应蒸馏过程中动态变化的非分数幅度系数，导致外生奖励与蒸馏项失衡。

本文的切入角度是打破蒸馏与强化学习各自为政的二元范式：既然 DMD 的梯度形式与策略梯度在数学结构上高度同构，那么完全可以将教师与学生的分数差异显式重塑为一个内生奖励项。核心 idea：将分布匹配重新形式化为策略奖励 \(R_{\mathrm{dm}}\)，构建统一的群相对策略优化框架（GNDMR），通过群均值归一化滤除高噪声下的梯度方差，并借助动态幅度系数实现蒸馏与多外生奖励的自适应平衡。

方法详解¶

整体框架¶

Rdm 的核心思想是将传统的分布匹配梯度映射为基于马尔可夫决策过程（MDP）的策略梯度，从而将多步少步蒸馏完全纳入强化学习的算法体系。在整个管线中，学生生成器 \(G_\theta\) 针对给定提示词批量采样一组去噪轨迹；对于采样得到的中间潜变量，利用冻结的预训练真实模型 \(\mu_{\mathrm{real}}\) 与同步更新的假样本分数估计器 \(\mu_{\mathrm{fake}}\) 计算分数差分并构建分布匹配奖励 \(R_{\mathrm{dm}}\)，同时针对终态图像计算外生美学或对齐奖励；随后，在群组（Group）维度分别执行群归一化得到优势函数，并通过自适应权重系数动态对齐两类优势的梯度尺度；最后，利用重要性采样修正（Importance Sampling）更新生成器，实现轨迹的高效复用。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入提示词与初始高斯噪声"] --> B["群轨迹采样<br/>生成 M 组少步去噪轨迹"]
    B --> C["假分布分数估计器更新<br/>追踪学生生成分布"]
    C --> D["分布匹配奖励形式化 Rdm<br/>构建内生密集奖励"]
    D --> E["群归一化分布匹配 GNDM<br/>共享时间步减均值降方差"]
    E --> F["自适应多奖励动态对齐<br/>自适应系数平衡尺度"]
    F --> G["重要性采样策略更新<br/>GRPO 多次复用更新生成器"]

关键设计¶

1. 分布匹配奖励形式化 Rdm：打通蒸馏与策略梯度的数学同构

传统 DMD2 在多步去噪推断中，通过最小化分布散度使学生生成器 \(G_\theta\) 的梯度对齐真实与伪分布分数差 \((s_{\mathrm{real}} - s_{\mathrm{fake}})\)。论文揭示了扩散转移概率的高斯对数导数 \(\nabla_\theta \log p_\theta(x_{t-1} \mid x_t) = \frac{x_{t-1} - \mu_\theta(x_t)}{\sigma_{t-1}^2} \nabla_\theta \mu_\theta(x_t)\) 与 DMD 梯度的内在关联。将扩散去噪映射为一个 MDP，单步转移对应动作，即可将分数差除以残差噪声项，显式定义分布匹配奖励 \(R_{\mathrm{dm}}\)：

\[R_{\mathrm{dm}}(x_t, x_{t-1}, t') = \frac{s_{\mathrm{real}}(x_{t'}) - s_{\mathrm{fake}}(x_{t'})}{x_{t-1} - \mu_\theta(x_t)} \cdot \frac{\sigma_{t-1}^2}{\alpha_{t-1}}\]

在此定义下，策略梯度 \(\nabla_\theta \mathcal{J}_{\mathrm{DDPO}_{\mathrm{DM}}} = \mathbb{E}[R_{\mathrm{dm}} \nabla_\theta \log p_\theta(x_{t-1} \mid x_t)]\) 与传统 DMD 损失梯度 \(-\nabla_\theta \mathcal{L}_{\mathrm{DMD}}\) 严格等价。不同于 HPS 等无界贪婪奖励，\(R_{\mathrm{dm}}\) 本质是趋向于 0 的散度度量，为强化学习探索提供了天然的正则化边界。针对实际训练中采样残差 \(x_{t-1} - \mu_\theta(x_t)\) 带来的数值爆炸与正负符号扰动，工程上采用符号归一化 \(\operatorname{sign}(x_{t-1} - \mu_\theta(x_t))\) 并在分母结合真实分数估计与预测起点的 \(L_1\) 偏差权重，保证优化梯度的数值稳定性。

2. 群归一化分布匹配 GNDM：抑制大时间步高方差与精确方向估计

当扩散加噪时间步 \(t'\) 较大时，图像严重退化为各向同性噪声，真实分数与假分数差向量 \(R_s(t') = s_{\mathrm{real}}(x_{t'}) - s_{\mathrm{fake}}(x_{t'})\) 产生剧烈波动和高方差，使原始 DMD 的更新方向在深层去噪区间受到严重误导。为解决这一痛点，论文将群相对策略优化（GRPO）的无 Critic 思想引入分布匹配，提出了群归一化分布匹配（GNDM）。

生成器对同个提示词生成 \(G\) 个独立样本的去噪轨迹，在计算优势时强制组内样本共享生成时间步 \(t\) 与扩散加噪时间步 \(t'\)，进而通过组内均值消除共有漂移：

\[A_{\mathrm{dm}, t}^{i, t'} = \frac{R_{\mathrm{dm}}(x_t^i, x_{t-1}^i, t') - \operatorname{mean}(\{R_{\mathrm{dm}}(x_t^k, x_{t-1}^k, t')\}_{k=1}^G)}{\operatorname{std}(\{R_{\mathrm{dm}}(x_t^k, x_{t-1}^k, t')\}_{k=1}^G)}\]

减去群均值的操作极大地过滤掉了特定高噪声步共有的方差偏置，为学生模型在不同生成样本之间建立了清晰的相对优劣判定面，彻底稳定了大加噪步长下的蒸馏方向。

3. 自适应多奖励动态对齐：消除静态超参敏感与梯度尺度失衡

当引入外生下游奖励模型 \(R_o\)（如 HPSv2.1 审美奖励、CLIP 文本对齐奖励）时，外生奖励作为终态图像的稀疏奖励经过群归一化得到优势 \(A_{o, t}^i\)。如果采用固定静态超参直接相加，会因蒸馏内在幅度随训练迭代与去噪时间步剧烈变化而导致模型崩塌。论文从 \(R_{\mathrm{dm}}\) 的物理量纲出发，提取出底层动态缩放标量 \(w_{\mathrm{dm}, t}\)，并定义了全局自适应匹配系数 \(\beta_{\mathrm{dm}, t}\)：

\[\beta_{\mathrm{dm}, t} = \frac{\|w_{\mathrm{dm}, t}\|_1}{C \cdot S}\]

其中 \(C\) 和 \(S\) 分别为通道数和空间分辨率。将多奖励的总优势构造为 \(A_{\mathrm{sum}, t}^i = w_{\mathrm{dm}, t} A_{\mathrm{dm}, t}^{i, t'} + \beta_{\mathrm{dm}, t} \sum_j w_j A_{o_j, t}^i\)，使外生奖励梯度的动态量纲在任意训练阶段均与分布匹配项保持严格自适应同步。该机制彻底免除了极其脆弱的人工超参微调，既鼓励了生成器在美学特征空间的积极探索，又依靠分布匹配的稳固约束抑制了奖励黑客现象。

4. 重要性采样策略更新：打破单次采样瓶颈的高效轨迹复用

强化学习在扩散模型蒸馏中最大的实际障碍是采样开销极大：收集数步去噪轨迹后传统策略梯度要求严格 On-Policy，往往每个采样轨迹仅能支撑一次梯度反向传播。得益于将蒸馏完全表征为高斯转移下的策略分布，两步之间的似然比 \(r_t^i(\theta) = \frac{p_\theta(x_t^i \mid x_{t-1}^i)}{p_{\theta_{\mathrm{old}}}(x_t^i \mid x_{t-1}^i)}\) 具备解析的闭式解。论文将近端策略优化的截断重要性采样机制集成进 GNDMR，允许在单次轨迹采样后对学生生成器连续执行多次梯度更新（如 2 至 5 次迭代），截断系数取 \(\eta = 0.5\)。此举在保持高保真度生成与稳步美学提升的同时，直接将最耗时的环境交互与轨迹采样开销降低了一半以上。

实验关键数据¶

主实验¶

实验基于 LAION-AeS-6.5+ 提示词进行蒸馏微调，在 COCO2014 验证集（Karpathy 30K 分割中的 10K 提示词）上进行综合评测，对比了包含基础教师模型（50步）、无数据蒸馏及结合 RL 的主流 SOTA 方法（4步生成）。评测指标涵盖人类偏好分（HPS v2.1）、审美分（PickScore / PS）、多维偏好分（MPS）、图文相关性（CLIP Score / CS）、图像生成真实度（FID）以及与教师输出的一致性散度（FID-SD）。

模型	步数 (NFE)	分辨率	免真实图像 (Img-Free)	HPS ↑	PS ↑	MPS ↑	CS ↑	FID ↓	FID-SD ↓	采样计算开销
Stable Diffusion 3 Medium
Base Model (CFG=7)	50	1024	-	29.00	22.72	12.10	38.86	24.48	-	-
Flow-GRPO (CFG=7)	50	1024	-	30.35	22.90	12.56	38.52	26.63	12.39	-
Hyper-SD (CFG=5)	8	1024	✗	27.20	21.90	11.22	37.76	26.94	10.09	-
LCM	4	1024	✗	27.76	22.31	11.61	36.97	27.71	15.90	-
DMD2	4	1024	✗	26.64	22.36	11.37	38.00	27.28	16.51	-
Flash-SD3	4	1024	✗	27.47	22.65	11.98	38.07	26.01	12.21	-
DMDR	4	1024	✓	29.50	22.77	11.98	38.10	29.10	14.11	-
GNDMR-IS (本文)	4	1024	✓	30.00	22.89	12.48	38.15	28.84	12.47	128 × 4k
GNDMR (本文)	4	1024	✓	30.37	22.88	12.53	38.20	28.02	12.21	128 × 8k
Stable Diffusion 3.5 Medium
Base Model (CFG=3.5)	50	512	-	27.78	22.59	11.91	38.46	20.69	-	-
Flow-GRPO (CFG=3.5)	50	512	-	31.81	23.24	12.93	39.21	29.27	9.39	-
DMD2	4	512	✗	30.44	22.92	12.73	38.59	26.64	14.63	-
DMDR	4	512	✓	30.83	23.07	12.80	38.22	26.05	16.73	-
GNDMR-IS (本文)	4	512	✓	30.88	22.94	12.86	38.39	25.60	16.68	128 × 3k
GNDMR (本文)	4	512	✓	31.25	23.15	12.93	38.59	24.44	13.93	128 × 6k

消融实验¶

消融实验使用 SD3-Medium（\(512 \times 512\) 分辨率）作为基底模型，先执行 500 步纯蒸馏，再分别引入不同奖励或采样配置，在 COCO30K 上评估。

表1：群归一化（GN）对分布匹配奖励的消融

方法配置	仅 Rdm (500步) FID ↓	+HPS: FID ↓	+HPS: HPS ↑	+PS: FID ↓	+PS: PS ↑
GNDMR (含 GN)	23.07	24.47	30.22	22.32	22.76
w/o GN (退化为普通均值)	24.94	25.40	30.39	24.61	22.71

表2：群内时间步共享策略与加噪区间消融

时间步采样策略	共享加噪步 \(t'\)	共享生成步 \(t\)	FID ↓	施加 GN 的加噪时间步区间	FID ↓
Vanilla DMD 基线	✗	✗	24.94	\([0, 300]\)	23.71
仅共享加噪步	✓	✗	24.34	\([300, 600]\)	23.46
全共享 (本文推荐)	✓	✓	23.07	\([600, 1000]\) (高方差区间)	23.43

关键发现¶

群归一化大幅降低真实性误差：在纯蒸馏前 500 轮训练中，GNDM 相比无 GN 方案直接将 FID 从 24.94 降至 23.07（下降 1.87 点），表明在蒸馏早期滤除高频伪梯度方差能建立更干净的生成流。
高噪声时间步是方差主要源头：将 GN 限制在不同区间测试发现，加噪步长位于 \([600, 1000]\) 的高模糊区间时施加群归一化带来的收益最显著，证明了数学分析中高加噪步长分数不可靠的核心论断。
自适应权重 \(\beta_{\mathrm{dm}, t}\) 杜绝崩溃：若固定静态权重 \(w=10\) 或 \(100\)，训练中后期 HPS 提升缓慢或骤然崩溃；而引入 \(\beta_{\mathrm{dm}, t}\) 后，HPS 能够伴随步数持续平滑上升，在 DMDR 与 GNDMR 上均表现出极强的鲁棒性。
重要性采样实现近两倍加速：引入重要性采样（GNDMR-IS）后，生成器在单次轨迹采样下进行 2 次迭代更新，总训练采样成本直接从 \(128 \times 8\mathrm{k}\) 砍半至 \(128 \times 4\mathrm{k}\)，而最终 HPS 与 FID-SD 指标与标准 GNDMR 几乎持平，大幅缓解了强化学习蒸馏的训练开销。

亮点与洞察¶

将逆向散度蒸馏内生化为策略奖励：打破了过去将蒸馏视为回归约束、将 RL 视为外挂微调的生硬拼接，首次在数学上证明了多步 DMD 与强化学习策略梯度在形式上的精确等价，将整个蒸馏领域顺理成章地接入了成熟的 RL 生态。
以统计均值滤除高维分数估计噪音：敏锐地指出了高加噪时间步下真实与假分布分数差异向量极易发散的机理，借用 GRPO 的群均值中心化技巧消除了样本共有的估计漂移，极其优雅且无须引入庞大的判别器或 Critic 网络。
动态幅度解耦跨模态奖励优化：推导出的自适应权重 \(\beta_{\mathrm{dm}, t}\) 能够敏锐反映去噪残差的实时能量级，保证了密集物理去噪与稀疏宏观美学评价处于同一梯度量级，彻底解决了多奖励强化学习中常见的奖励挤占与破坏性捷径问题。

局限与展望¶

作者承认的局限：方法依然依赖一个在线持续更新的假分布分数估计器 \(\mu_{\mathrm{fake}}\)，其去噪更新与生成器策略更新之间存在一定的训练延迟；此外重要性采样的更新次数受限于轨迹似然比截断范围，通常无法扩展至大于 5 次更新。
拓展思考：群内强制共享时间步 \(t\) 和 \(t'\) 虽然极大地增强了比较的有效性，但减少了单批次内时间步的多样性覆盖，可能需要更大的批大小或更长的热身轮次；此外，当前框架主要在流匹配与扩散图像生成上验证，尚未探索在视频扩散模型大步长时空蒸馏中的表现。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 概念创新极强，将传统分布匹配蒸馏与策略梯度在数学上打通，为扩散蒸馏与 RL 协同开辟了全新视角。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 SD3 与 SD3.5 两个主流大模型，指标兼顾美学偏好（HPS/PS/MPS）与图像真实度（FID/FID-SD），消融分析层层递进。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严谨，从数学等价性到工程稳定化设计的逻辑链条环环相扣。
价值: ⭐⭐⭐⭐⭐ 为少步高质量图像生成提供了稳定、高效且兼具高审美的工业级蒸馏范式，代码开源且易于复用。