跳转至

Distribution Matching Distillation Meets Reinforcement Learning

会议: ECCV2026
论文: ECCV 原文
作者: Dengyang Jiang 等
领域: 图像生成
关键词: 分布匹配蒸馏、少步生成、强化学习、奖励倾斜、偏好对齐

一句话总结

DMDR 将少步扩散蒸馏与偏好优化放进相互约束的两阶段训练流程,使 4 步学生在多个总体指标上超过多步教师,但偏好收益仍伴随生成多样性的下降。

研究背景与动机

高质量文本到图像模型通常需要多次去噪,每一步都要调用大型网络,因此减少采样步数直接关系到生成成本。 分布匹配蒸馏(Distribution Matching Distillation,DMD)并不要求学生逐步重演教师轨迹,而是让少步学生产生的图像分布接近教师分布。 这解决了速度问题,却不自动解决用户偏好:教师分布中既有更符合提示词、更美观的区域,也有不那么理想的区域。 如果蒸馏只追求匹配原分布,就缺少主动偏向高质量样本的目标。 与此同时,强化学习(reinforcement learning,RL)能够使用奖励模型推动偏好对齐,但奖励升高不一定意味着语义能力和覆盖范围得到保留。

把这两件事简单串起来也存在问题:先蒸馏、再独立进行 RL 时,作为参考的少步模型已经损失了一部分教师分布的模式。 继续围绕这个较窄的参考分布优化,可能进一步集中到少量容易获得高分的图像,产生奖励投机和语义退化。 反过来,先对多步模型进行 RL 再蒸馏,既承担昂贵的采样成本,也可能在蒸馏时丢失已经获得的偏好收益。 直接从蒸馏第一步就反向传播奖励同样不稳,因为冷启动学生生成的图像尚不连贯,奖励反馈很难可靠地区分值得学习的方向。

本文因此把“如何使用奖励”与“何时使用奖励梯度”分开考虑:早期只用奖励调整蒸馏样本的权重,学生稳定后再进行直接奖励优化。 整个过程中,DMD 不只是前置的压缩步骤,也成为持续保留教师生成先验的约束。 这不是取消训练阶段,而是让两个阶段都包含偏好信号,只是信号进入优化的方式不同。 核心 idea:先以静态奖励权重和动态分布重叠稳定少步蒸馏,再联合优化奖励与 DMD,让偏好优化决定学习重点,让教师分布约束模型不要只会刷分。

方法详解

整体框架

训练输入是文本提示词和噪声,核心可训练对象是少步生成器;辅助对象包括真实分数估计器、伪分数估计器和奖励模型。 这里的“分数”是分布的 score,即对输入求导的对数密度梯度,不是奖励模型给出的偏好分数。 真实分数估计器提供多步教师分布的方向,伪分数估计器学习当前学生的生成分布,两者的差异构成 DMD 更新信号。 伪分数估计器随学生改变而更新,不能把它理解成另一个固定教师。 论文沿用 backward simulation:学生先走若干采样步得到中间状态,再从该状态预测干净图像,以训练实际少步采样会遇到的状态。

第一阶段由奖励倾斜蒸馏与动态分布重叠协同完成,后者包含 DynaDG 和 DynaRS,两者不是额外的推理模块。 第二阶段改用联合奖励优化,继续保留 DMD,直接优化偏好而不再仅靠样本重加权。 部署时只运行训练好的少步生成器;奖励模型与两套分数估计器属于训练监督,不随每次生成上线。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
        A["提示词 + 噪声"] --> G["少步生成器<br/>完整轨迹与中间状态"]
        subgraph S1["阶段 1:协同蒸馏"]
                direction TB
                B["奖励倾斜蒸馏<br/>完整轨迹奖励作权重"]
                C["动态分布重叠<br/>DynaDG + DynaRS"]
        end
        G --> B
        G --> C
        R["奖励模型"] -.-> B
        T["真实与伪分数估计器"] -.-> C
        B --> W["加权 DMD 更新<br/>得到稳定学生"]
        C --> W
        W --> D["联合奖励优化<br/>直接奖励 + DMD"]
        R -.-> D
        T -.-> D
        D --> O["推理:仅少步生成器"]

关键设计

1. 奖励倾斜蒸馏:先改变学习重点,不急着对奖励反向传播

普通 DMD 用教师与学生的 score 差异更新生成器,但同样的更新强度并不区分样本是否符合人类偏好。 Reward-Tilted Distribution Matching(RT-DM)为这条蒸馏信号加上与奖励相关的静态权重,使更受偏好的样本对更新产生更大影响。 “静态”指在本次生成器更新中不沿权重反向传播奖励梯度,并不是说所有样本共享一个常数,或奖励在整个训练中固定不变。 这样,奖励负责判断哪些生成结果更值得学,教师与伪分数估计器仍负责提供如何修改生成器的方向。 论文把这种机制解释为向奖励倾斜的教师分布学习,而不是均匀地保留教师已有的生成偏好。

奖励的评估图像来自学生的完整少步轨迹,例如完整走完 4 步,而不必只依赖任意中间状态的一次干净图像估计。 冷启动期间,完整轨迹通常比中间状态的一步预测更连贯,因此反馈更有参考价值。 由于奖励只作权重,无需为这条完整奖励评估轨迹保留反向传播图,避免了把较可靠反馈与昂贵梯度图绑定起来。 这不意味着教师生成了逐样本的正确答案,也不意味着奖励模型已经成为第二个去噪教师。 它本质上是对分布匹配更新进行偏好选择,仍受教师分布及当前生成质量的限制。 论文图 5 中,持续只用 RT-DM 的设置后期出现平台期,因此后续必须允许直接优化奖励。

2. 动态分布重叠:同时改善教师方向和加噪状态的可靠性

RT-DM 解决更新重点,但如果学生输出与教师熟悉的图像区域相距太远,DMD 的 score 差异本身就可能不可靠。 Dynamic Distribution Guidance(DynaDG)在真实分数估计器中加入可训练 LoRA,让早期的目标分布暂时向学生分布靠近。 其作用不是永久降低教师标准,而是先建立可用于优化的分布重叠,避免要求冷启动学生一步跨过很大的分布差距。 训练推进时,教师中的 LoRA 缩放逐渐退火到零,真实分数估计器重新回到原教师所表示的分布。 因此“教师固定”需要区分冻结的原始能力与阶段性启用的适配分支,不能把早期真实分数估计器描述为完全不变。

Dynamic Renoise Sampling(DynaRS)则不移动目标,而是改变 DMD 比较学生与教师时使用的噪声水平。 初期更多采样高噪声状态,因为两种分布加上足够噪声后都会更接近高斯噪声,学生与教师更容易处于相似区域。 这些状态的优化更侧重整体结构,而不是让尚未成形的图像过早接受细节层面的纠正。 随着学生变好,采样偏置退火为均匀采样,再覆盖较低噪声状态以学习纹理等细节。 这里的 renoise 时间用于训练监督,不等同于增加部署时的去噪步数。 两项策略分别从“比较的目标是什么”和“在什么噪声水平比较”入手,合起来提高冷启动期的分布匹配可靠性。

3. 联合奖励优化:把 DMD 从压缩目标变成持续正则化

学生具备稳定生成能力后,第二阶段直接加入奖励优化损失,而不再仅依靠高奖励样本获得更大的蒸馏权重。 ReFL 可以通过冻结奖励模型的输入梯度更新图像生成器;DPO 或 GRPO 则提供不同形式的偏好或策略优化信号。 论文把它们作为可替换的 RL 分支验证,并不是每次训练同时运行三种算法。 生成器持续接受 DMD 信号,伪分数估计器仍需跟踪更新后的学生分布,因而此阶段不是“蒸馏结束后只做 RL”。 与把已蒸馏学生作为唯一参考不同,持续 DMD 仍然借助多步教师较丰富的生成先验约束偏好更新。

作者用“奖励上升方向加上回到教师流形的方向”解释这种协同:奖励推动模型偏向高分区域,分布差异提供纠偏信号。 更审慎的理解是,它是一种经验上有效的软正则化,而不是严格保证每次更新都留在教师支持集中的硬投影。 图 4 的对照显示,独立 RL 可出现训练奖励升高但 DPG 下降,联合优化则较好保留提示词遵循能力。 然而这种约束也不能保证所有模式都被保留,后文多样性实验直接说明了这一点。 默认 ReFL 在少步模型上的额外好处是,各采样状态都被训练为可直接预测干净图像,奖励梯度更容易影响较早状态的预测。 这与多步 ReFL 为节省显存而主要优化末端去噪步骤的困难有所不同,但不应据此声称所有实现都对完整轨迹反向传播。

损失函数 / 训练策略

按正文第 3.3 节对两项损失及系数的定义,第二阶段目标可简记为:

\[ \mathcal{L}_{\mathrm{total}}=\mathcal{L}_{\mathrm{dmd}}+\lambda_{\mathrm{rl}}\mathcal{L}_{\mathrm{rl}}. \]

其中 \(\mathcal{L}_{\mathrm{rl}}\) 按最小化损失的约定实现奖励最大化,\(\lambda_{\mathrm{rl}}\) 控制直接偏好优化的强度。 DMD 提供生成器更新方向,伪分数估计器另用学生样本上的扩散损失训练;去噪模型预测噪声,流模型预测速度。 第一阶段则用 RT-DM 的奖励权重调整 DMD 更新,并配合上述两项动态策略,不把它写成第二阶段的直接奖励损失。 正文主实验使用 t2i-2M 的提示词,覆盖 SDXL-Base、SD3-Medium 和 SD3.5-Large。 默认奖励组合为 DFN-CLIP 与 HPSv2.1,默认 RL 算法为 ReFL,两个阶段各训练 1.5K 步。 第 3 节机制验证中出现的 1000 次迭代切换点,不是第 4 节主实验采用的完整阶段预算。 表 5(c) 中默认权重为 \(\lambda_{\mathrm{rl}}=0.5\),增大它并不保证总体能力继续上升。

原文缓存的式 (1)、(2)、(4) 存在明显排版抽取损坏,上式仅依据正文可明确辨认的加权目标整理。 尤其式 (4) 同时涉及奖励倾斜分布与指数形式的权重,不能据损坏文本确认精确归一化方式,故不补写其完整推导。 缓存包含主文和参考文献,但未包含文中指向的附录,因此 LoRA 具体配置、退火日程、完整伪代码与用户研究细节不在本笔记中猜补。

实验关键数据

主实验

下表摘自原文表 2,第 11 页;评估提示词采样自 ShareGPT-4o-Image,与训练提示词不重叠,四项分数均越高越好。 NFE 表示模型函数评估次数,不应直接当作端到端时延的等比例缩放。

基座 方法 NFE CLIP Score Aesthetic Score Pick Score HP Score
SDXL-Base Base-Model 50 34.7588 5.6480 22.1085 27.1477
SDXL-Base DMDR 4 35.4940 6.0324 22.7122 32.9832
SD3-Medium Base-Model 50 34.9025 5.5942 22.1801 28.4021
SD3-Medium DMD2 4 33.9421 5.6137 21.7688 27.3675
SD3-Medium DMDR 4 35.0142 5.8876 22.4892 32.1145
SD3.5-Large Base-Model 50 35.5509 5.7014 22.4856 28.8135
SD3.5-Large DMDR 4 35.9757 6.1541 22.9072 32.7368

SD3-Medium 的 DMDR 相对 DMD2 将 HP Score 从 27.3675 提高到 32.1145,同时维持相同的 4 NFE;表中该 DMD2 为作者复现结果。 表 3 和表 4(第 12 页)进一步报告 SD3-Medium 的 GenEval 总体分数从教师的 0.62 提高到 0.65,DPG 总体分数从 84.08 提高到 85.30。 但 GenEval 计数从 0.63 降至 0.57、位置从 0.34 降至 0.27,说明“超过教师”是总体结果,不是逐项支配。 原文说明未直接针对 GenEval 和 DPG 指标进行 RL 优化,这使它们比训练奖励本身更能检验能力是否被保留。

消融实验

下表为原文表 5(c),第 13 页;均为 SD3-Medium、4 步采样,考察第二阶段奖励权重,所有指标越高越好。

奖励权重 CLIP Score Aesthetic Score Pick Score HP Score DPG overall
0.1 34.7422 5.8508 22.3064 31.6565 85.24
0.5 35.0142 5.8876 22.4892 32.1145 85.30
1.0 35.0087 5.84430 22.4709 32.8982 85.04
2.0 35.0322 5.7806 22.1098 32.7683 82.48

权重从 0.5 增大到 2.0 时,HP Score 仍然较高,DPG 却从 85.30 降至 82.48,不能只凭奖励选择模型。 表 5(b) 的 ReFL、DPO、GRPO 分别得到 85.30、85.00、85.08 的 DPG,支持兼容性,但不是所有指标都由 ReFL 取得最高值。 表 5(a) 中 CLIP + HPS 的 DPG 为 85.30,单独 HPS 为 85.04;单独 HPS 的 HP Score 32.9065 则高于组合奖励的 32.1145。 这说明组合奖励改善的是指标平衡,而非让每个奖励模型自己的分数都最大化。

关键发现

以下为原文表 7,第 15 页;SD3-Medium 在 ShareGPT-4o-Image 上的多样性分析,HPS 与 LPIPS 均按原文标为越高越好。 这里 LPIPS 用作生成差异性的代理,不能把它解释成单张图像的质量分数。

指标 Base-Model DMD2 DMDR
HPS 28.4021 27.3675 32.1145
LPIPS 0.6840 0.5832 0.5548

DMDR 的偏好分数高于两项对照,但 LPIPS 更低,因此 DMD 正则化缓解奖励投机不等于消除多样性代价。 图 5、图 7 分别支持 RT-DM 暖启动和动态策略的作用;缓存没有可可靠转录的曲线点值,本笔记不从曲线猜出精确增益。 表 1(第 7 页)的采样耗时条件是 SD3-M、1024×1024、H100 上本地 batch size 16,少步联合训练为 0.43–1.21 秒,多步 RL 为 9.24–10.57 秒。 这是 rollout 耗时比较,不是完整训练成本或部署服务延迟的证明,也不能机械选取区间端点宣称固定加速倍数。

亮点与洞察

  • 奖励不必总以梯度形式进入训练。RT-DM 在学生还不稳定时用可靠程度更高的完整轨迹反馈筛选更新重点,提供了偏好冷启动的可复用思路。
  • DynaDG 与 DynaRS 处理的是两个不同的分布差距问题。一个暂时移动目标,一个改变观察分布的噪声尺度,组合比笼统地降低学习率更有针对性。
  • 蒸馏目标可以在完成加速后继续发挥作用。把教师分布作为持续约束,有助于区分“偏好改善”与“刷高奖励但丢失提示词语义”。

局限与展望

  • 作者明确承认质量与多样性的权衡;偏好优化把概率质量推向高奖励区域,DMD 并未完全阻止模式覆盖收缩。
  • “流形投影”应视作机制解释而非严格保证,增大奖励权重后的 DPG 退化也表明约束强度需要调节。
  • 原文表 8 在权重 0.1、1 下的 HPS 为 30.4024、32.4786,与表 5(c) 的 31.6565、32.8982 不同;正文未明确解释这些设置差异,不能合并成一条可直接复现的曲线。
  • Flash-DMD 与 Diff-Instruct++ 因作者无法获得满足公平复现要求的权重或训练代码而未进入主实验,最接近工作的定量优势仍有比较缺口。
  • 附录未包含在缓存中,无法核验用户研究统计及所有训练超参数;可进一步研究自适应奖励权重和保多样性的目标,但这些属于展望而非本文已经验证的方案。

相关工作与启发

  • DMD / DMD2:核心是将少步分布匹配到教师;DMDR 不只提高模仿精度,而是引入对教师分布中不同偏好区域的选择。
  • Hyper-SD / PSO:代表蒸馏后进行偏好优化的路线;DMDR 的差异在于更早引入奖励信息,并在直接优化奖励时继续保留分布匹配。
  • Flash-DMD / Diff-Instruct++:同样已有联合蒸馏与 RL 的思想,本文的新意不宜概括为首次相加两个损失,而是 RT-DM 冷启动及配套动态策略。
  • ReFL / DPO / GRPO:是本文可替换的偏好优化分支;读者可以借鉴其训练阶段安排,而不是把 DMDR 误当作新的通用策略梯度算法。

评分

  • 新颖性: 4/5。主要增量在奖励倾斜冷启动及动态重叠策略,联合损失本身已有近邻工作。
  • 实验充分度: 4/5。覆盖多基座、算法和多样性代价,但最近邻比较与缓存外附录仍限制核验。
  • 写作质量: 3/5。两阶段动机清晰,部分理论措辞偏强,表间数值差异和公式抽取问题需谨慎处理。
  • 价值: 4/5。对少步生成的偏好对齐有实际指导意义,同时明确提醒高奖励不等于全面能力提升。