AI技术文章

ICML 2026 | 我们让蒸馏模型当了采样器,Diffusion RL采样成本降低一个数量级

作者:机器之心0 次浏览 来源:机器之心
ICML 2026 | 我们让蒸馏模型当了采样器,Diffusion RL采样成本降低一个数量级

过去的 Diffusion RL 多聚焦于奖励设计与优化算法,训练时的采样成本被忽视。DMSampler 指出:在在线 RL 中,限制规模化的不只是奖励信号或优化器,很多时候是 rollout 本身太贵。

扩散模型已经成为图像和视频生成的重要底座,但当研究者希望用在线强化学习进一步优化审美质量、文本一致性和视频综合表现时,训练成本很快成为瓶颈:每次 policy rollout 往往要完成大量去噪步骤,在线采样占据了大量训练时间。

过去的 Diffusion RL 多聚焦于奖励设计与优化算法,训练时的采样成本被忽视。DMSampler 指出:在在线 RL 中,限制规模化的不只是奖励信号或优化器,很多时候是 rollout 本身太贵。

最直接的想法是减少采样步数,例如从 50 步降至 10 步,但这会导致样本质量下降。一旦奖励模型基于劣化样本进行评估,策略更新便可能偏离高质量分布。

文章配图

图 1:传统 Diffusion RL 框架(左)与 DMSampler(右)的概念对比。传统方法每次 RL 更新依赖约 50 步采样;DMSampler 引入共同演化的少步蒸馏采样器,图中给出了 VBench 训练耗时从 900 小时降至 288 小时的对比。

让蒸馏模型跟着 policy 一起进化

DMSampler 的核心判断是:蒸馏模型不应仅是训练结束后的推理加速工具,更可作为低成本采样器,深度参与 Diffusion RL 的每一轮策略更新。

文章配图

  • 论文标题:Distillation Models are Good Samplers for Diffusion Reinforcement Learning
  • 论文地址:https://openreview.net/pdf?id=VkMWujvv2c
  • 项目地址:https://github.com/HiDream-ai/DMSampler

具体来说,该框架维护一个与 policy model 共同演化的 few-step distilled sampler。此采样器仅需 4 至 8 步(图像任务 7 步,视频任务 16 步),并通过周期性重蒸馏来追随更新后的策略。

文章配图

图 2:DMSampler 整体框架

整个框架由两个阶段交替组成。

RL 阶段:蒸馏采样器参数冻结,用混合蒸馏采样快速生成样本;样本经 reward model 打分后,用于更新 policy model。实验以 DiffusionNFT 为底层优化器,但 DMSampler 的改动在采样环节。

蒸馏阶段:更新后的 policy model 充当 teacher,对 distilled sampler 进行重蒸馏,使其重新贴近当前 policy 分布。其中,轨迹分布匹配(TDM)保证基础对齐,而奖励感知蒸馏则确保高奖励轨迹的能力得以保留。

两阶段循环推进后,policy 通过 RL 获得更好的生成能力,sampler 通过蒸馏同步追上新的 policy 分布,下一轮 RL 又能以更低的采样成本继续训练。

蒸馏:从后处理变为 RL 训练的一环

传统流程中,蒸馏往往发生在训练之后。DMSampler 改变了这个顺序,让蒸馏从「后处理」变成训练闭环的一部分。

在此框架中,RL 阶段借助蒸馏模型降低 rollout 成本,蒸馏阶段又利用 RL 优化后的 policy 提升蒸馏模型本身。两者不再是先后关系,而是共同演化关系。

混合采样:保住结构,提速后半程

在采样策略上,DMSampler 并未把整个去噪链路都交给蒸馏模型。扩散采样的前半程更决定结构、语义布局和主体内容,后半程更多影响纹理与细节。因此,让 policy model 先完成早期去噪,再由蒸馏模型接手后期步骤,可以在保持分布对齐的同时减少计算。

文章配图

图 3:四种采样策略对比

论文比较了四种策略:S1 使用原 policy model 完成完整采样并启用 CFG,质量强但成本高;S2 全程使用少步蒸馏模型,速度最快但分布漂移明显;S3 先由 policy model 完成早期去噪,再切换到蒸馏模型处理后期步骤,是最终采用的方案;S4 则反过来先用蒸馏模型,早期偏差会影响后续结构,policy model 难以完全修正。

结果显示,S3 在图像任务中将步数从 40 步减至 7 步,视频任务从 50 步减至 16 步,同时保持了较高的 OCR 表现。

奖励感知蒸馏

只让蒸馏模型追上 policy 的平均分布还不够,Diffusion RL 真正要保留的是高 reward 样本背后的能力。

DMSampler 在 RL 阶段记录生成样本、初始噪声和对应 reward;进入蒸馏阶段后,框架从中筛选高 reward 轨迹。若无样本超过阈值,则至少保留 reward 最高的一条。随后,reward-weighted trajectory loss 会让蒸馏模型更重视这些轨迹,权重随 reward 单调增加。

消融实验显示,去掉奖励感知蒸馏后,图像 OCR 从 0.97 降到 0.96,视频 OCR 从 0.50 降到 0.48。提升幅度有限,但说明高 reward 轨迹有助于稳定迭代。

少步采样带来加速,性能基本不掉

在文本渲染任务上,DMSampler 将图像采样从 40 步降到 7 步、视频采样从 50 步降到 16 步,性能基本保持。与直接减少采样步数(DRS)相比,优势显著:DRS1 将步数减半并保留 CFG;DRS2 在此基础上关闭 CFG;DRS3 进一步降到 10 步。视频任务对样本质量更敏感,DRS3 的视频 OCR 回到基线水平,而 DMSampler 在 7 步图像采样和 16 步视频采样下仍保持接近完整采样的表现。

文章配图

表 1:DMSampler 与直接减步策略的对比。DMSampler⁻ 为去掉 Reward-aware Distillation 的消融版本。

文章配图

图 4:不同加速策略生成的视频效果对比。DRS 导致明显质量退化,DMSampler 保持高质量。

GenEval 综合评测:从 0.63 提升到 0.96

在 GenEval 基准测试上,DMSampler 取得 Overall 0.96,超过 FlowGRPO 和 DiffusionNFT 的 0.95;相比 Base(SD3.5-M)的 0.63,提升更为明显。作为参照,GPT-4o 为 0.84,Qwen-Image 为 0.91。训练效率方面,论文报告 DMSampler 需要 360 GPU hours,低于 FlowGRPO 的 2000 GPU hours 和 DiffusionNFT 的 480 GPU hours。

文章配图

表 2:GenEval 基准测试结果。

1.3B 模型超过 14B,训练时间降至约三分之一

在 VBench 视频评测上,DMSampler 将 Wan2.1-1.3B 的总分从 81.23 提升到 85.00,超过 Wan2.1-14B(83.69)及 CausVid(83.88)、HunyuanVideo(83.43)。相比 DiffusionNFT,DMSampler 将训练时间从约 900 小时降至约 288 小时,总分从 84.74 提升至 85.00。

文章配图

表 3:VBench 基准测试结果(完整 16 个评测维度,分两组展示)。

文章配图

续表 3:其余 8 个评测维度

文章配图

图 5:VBench 生成视频效果对比。

不绑定优化器,可作为通用采样模块

DMSampler 不依赖某一个 RL 优化器。研究者将其分别接入 FlowGRPO 和 DGPO,并在 GenEval 上验证:FlowGRPO 的训练成本从超过 1000 GPU hours 降到 400 GPU hours,GenEval 分数从 0.95 提升到 0.96;DGPO 的训练成本从 240 GPU hours 降到 192 GPU hours,GenEval 分数保持 0.97。

文章配图

表 4:DMSampler 接入不同 RL 优化器的效果。

这表明,DMSampler 更像是一个面向在线 Diffusion RL 的采样加速模块。它不改变优化器本身的目标函数,而是降低 rollout 成本,因此可以与不同 RL 方法组合使用。

副产品:蒸馏模型同步增强

DMSampler 的主要目标是训练更强的 policy model,但共同演化的蒸馏采样器本身也获得了性能提升。在 VBench 上,该蒸馏模型取得 84.21 的总分和 85.60 的 Quality 分数,超过 CausVid(82.88)、Self Forcing(83.80)和 BLADE(83.38)等蒸馏方法。

文章配图

表 5:DMSampler 蒸馏模型与其他蒸馏方法的对比。

传统蒸馏侧重推理加速,而 DMSampler 的蒸馏采样器会在训练过程中不断吸收 RL 优化后的高 reward 能力。因此,它不仅是加速器,也成为经过 reward 对齐的少步生成器。

文章配图

图 6:DMSampler 蒸馏模型的生成示例。

从训练加速到协同进化

整体来看,DMSampler 首次将可训练、共同演化的蒸馏模型作为 Diffusion RL 的采样引擎。通过双阶段交替训练、混合蒸馏采样和奖励感知蒸馏,将 rollout 采样成本降到更适合在线训练的范围。

其意义不仅在于加速,更在于重新审视了采样成本这一瓶颈,并将蒸馏从后处理前移,使之与 RL 后训练在同一闭环中相互增益。实验中,DMSampler 在 GenEval 上取得 0.96,在 VBench 上取得 85.00,均达到当前最优水平。训练效率上,GenEval 任务从 DiffusionNFT 的 480 GPU hours 降到 360 GPU hours,VBench 任务从约 900 小时降到约 288 小时。它不绑定特定 RL 优化器,也可以与 FlowGRPO、DGPO 等方法组合使用。

对后续研究来说,DMSampler 的协同进化思路可以迁移到其他需要大量 rollout 的生成模型 RL 场景,也为扩散模型蒸馏与强化学习的结合提供了新的技术参考。

文章来自于微信公众号 “机器之心”,作者 “机器之心”

相关主题

读完之后

加入 AINRK 交流群

了解加入交流群的方式,与更多 AI 从业者交流。

查看交流群说明(在新窗口打开)