TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward¶

日期: 2026-03-08
arXiv: 2603.07700
代码: GitHub
领域: 图像生成
关键词: few-step diffusion, reinforcement learning, non-differentiable reward, trajectory distribution matching, RLHF

一句话总结¶

提出 TDM-R1，首个支持非可微奖励的少步扩散模型强化学习范式——利用 TDM 确定性轨迹为中间步骤提供无偏奖励估计，通过代理奖励学习 + 生成器优化的解耦机制，仅 4 NFE 即在 GenEval 上从 61% 提升至 92%，超越 GPT-4o（84%）。

研究背景与动机¶

领域现状: 少步扩散模型（如蒸馏方法）已能高效生成高质量图像，但在精确指令遵循、文字渲染等方面仍有明显不足。强化学习是提升这些能力的有效手段。
现有痛点: 现有少步扩散 RL 方法严格依赖可微奖励——必须能反向传播穿过奖励函数，这排除了大量重要的非可微奖励信号（人类二元偏好、目标计数、OCR 正确率等）。
核心矛盾: 标准扩散 RL 方法本质上等价于加权去噪损失，在少步设置下会产生模糊结果；直接将标准扩散 RL 应用于少步模型效果很差。
核心 idea: 利用 TDM 的确定性采样轨迹获得无偏中间奖励估计，解耦为代理奖励学习（处理非可微性）和生成器学习（最大化奖励+KL正则化）。

方法详解¶

整体框架¶

预训练 TDM 少步生成器 → 采样生成组 → 外部非可微奖励打分 → 代理奖励模型学习（组偏好优化）→ 代理奖励指导生成器更新（reverse KL 正则化）→ 在线迭代。

关键设计¶

确定性轨迹的中间奖励估计:
- 做什么：为少步采样的每个中间步骤分配准确的奖励信号
- 核心思路：若 \(p(\mathbf{x}|\mathbf{x}_t)\) 是 Dirac 分布（确定性 ODE），则 \(r(\mathbf{x}_t) = \mathbb{E}_{p(\mathbf{x}|\mathbf{x}_t)} r(\mathbf{x})\) 可通过端点的单样本无偏估计，方差为零
- 设计动机：随机轨迹的中间奖励估计有高方差，确定性轨迹消除了这一问题，收敛更快
代理奖励学习（Surrogate Reward）:
- 做什么：学习可微的代理奖励模型来替代非可微奖励
- 核心思路：用扩散模型参数化代理奖励 \(\tilde{r}_\phi(\mathbf{x}_{t_k}) \approx \beta \cdot \mathbb{E} \log \frac{p_\phi}{p_{ref}}\)，通过组偏好 BT 模型优化
- 创新点：组级偏好学习（正/负组由 advantage 分区），样本权重为归一化 advantage 的绝对值
- EMA 动态参考模型防止过拟合
生成器学习:
- 目标：\(\max -\tilde{r}_{sg(\phi)}(\mathbf{x}_{t_k}) + \beta_g \cdot KL(p_{\theta,k} \| p_\psi)\)
- 用分布级（marginal-level）reverse KL 正则化而非实例级约束，比标准扩散 RL 更宽松
- 代理奖励和 KL 正则化协同工作：奖励引导质量提升，KL 防止偏离过远

训练策略¶

代理奖励学习和生成器学习交替进行，形成类 GAN 的对抗循环——生成器产出更好样本，代理奖励提供更精确梯度。

实验关键数据¶

主实验¶

模型	NFE	GenEval	OCR Acc.
SD3.5-M (base)	80	0.63	0.59
TDM-SD3.5-M	4	0.61	—
Flow-GRPO	80	0.95	—
TDM-R1 (Ours)	4	0.92	—
GPT-4o	—	0.84	—

消融实验¶

配置	效果
直接组合 TDM + RL loss	图像模糊，性能差
TDM-R1 (确定性轨迹)	收敛快，质量高
随机轨迹 vs 确定性轨迹	确定性显著更优（更低方差）
w/o EMA 参考模型	训练不稳定

关键发现¶

4 NFE 的 TDM-R1 大幅超越 80 NFE 的基础模型（GenEval 0.92 vs 0.63）
在组合生成（物体数量、属性绑定、空间位置）上提升最大
成功扩展到 6B 参数的 Z-Image 模型，普适性强
非可微奖励（OCR、GenEval 计数）首次在少步模型上取得 SOTA

亮点与洞察¶

确定性轨迹+中间奖励：优雅解决了少步扩散 RL 的核心困难——中间步骤的奖励分配
代理奖励解耦设计：将非可微奖励问题转化为可微优化，思路新颖且通用
4 NFE 超越 GPT-4o：展示了少步扩散 + RL 后训练的巨大潜力，对产业部署有直接意义

局限性 / 可改进方向¶

需要预训练 TDM 基础模型，适用范围限于 TDM 系列少步模型
代理奖励模型增加了训练成本和复杂度
视频生成方向的扩展未验证

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首创少步扩散非可微 RL 范式
实验充分度: ⭐⭐⭐⭐⭐ GenEval + OCR + 美学 + 偏好分数全面评估
写作质量: ⭐⭐⭐⭐ 理论推导完整，动机清晰
价值: ⭐⭐⭐⭐⭐ 对少步生成模型的后训练有重大指导意义