TDM-R1：用不可微奖励信号强化少步扩散模型 TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward

TDM-R1 的完整训练流程如 Algorithm 1 所述，包含以下步骤：第一步，从条件数据集 $\mathcal{D}_c$ 中采样条件 $c \sim \mathcal{D}_c$。第二步，用当前生成器 $p_{\theta^-}$ 在条件 $c$ 下通过 $K$ 步采样生成一组样本 $G = \{x_0^{(1)}, \ldots, x_0^{(G)}\}$，其中每个样本对应一条确定性去噪轨迹。第三步，使用不可微奖励函数 $r(\cdot, \cdot)$ 对生成的样本打分，得到 $r_i \leftarrow r(c, x_0^{(i)})$。第四步，更新代理奖励模型 $\tilde{r}_\phi$：将样本按优势值 $A(x_{t_k}^i) = \frac{r_i - \text{mean}(\{r_j\}_{j=1}^G)}{\text{std}(\{r_j\}_{j=1}^G)}$ 划分为正组 $G_k^+$ 和负组 $G_k^-$，然后通过组偏好优化目标 $\mathcal{L} = -\log p(G_k^+ \succ G_k^- | c)$ 训练，其中组级奖励 $R(G_k) = \sum_{x \in G_k} w(x_{t_k}) \tilde{r}_\phi(x_{t_k}, c)$，权重 $w(x_{t_k}^i) = |A(x_{t_k}^i)|$。第五步，通过去噪分数匹配更新在线假分数 $s^{\text{fake}}$。第六步，通过最大化代理奖励并最小化反向 KL 散度更新少步生成器 $p_\theta$，梯度为：$\nabla_\theta \mathcal{L}(\theta) = -\mathbb{E} \left[ \beta \alpha_{t|t_k}(T - t_k) \nabla_{x_t} \log \frac{p_\phi(x_{t-1}|x_t)}{p_{\text{ref}}(x_{t-1}|x_t)} + \beta_g \lambda_t (s^{\text{fake}}(x_t) - s_\psi(x_t)) \frac{\partial x_t}{\partial \theta} \right]$。以上步骤交替执行 $N$ 轮迭代直至收敛。

TDM-R1 的技术新颖性体现在多个层面。在算法设计层面，这是首个针对少步文本到图像扩散模型、使用不可微奖励进行大规模强化学习后训练的工作，填补了该领域的空白。在理论贡献层面，论文推导了基于确定性轨迹的无偏中间步奖励估计，形式化为 $r(x_t, c) = \mathbb{E}_{p(x|x_t)} r(x)$，并证明当 $p(x|x_t)$ 为 Dirac 分布时估计方差最小。在奖励建模层面，扩散参数化的动态代理奖励模型是一个新颖的设计，它将扩散模型既作为生成器又作为奖励函数的参数化工具，实现了 GAN 式对抗框架的效果。在工程实践层面，TDM-R1 采用停止 $\phi$ 的梯度以节省内存，且发现这不影响性能；使用 EMA 动态参考模型替代冻结参考模型，有效提升了训练稳定性和最终性能。特别值得注意的是，论文发现 4 NFE 的 TDM-R1 在 GenEval（0.92）和 OCR 准确率（0.59 到 0.95）上均超越了 80 NFE 的代理奖励模型 $p_\phi$（GenEval 0.89，OCR 0.92），这与 LLM 文献中 DPO 类方法的发现一致：提升模型作为奖励信号的性能并不必然提升其生成性能。

TDM-R1 存在以下几个值得深入探讨的弱点。首先，在方法通用性方面，TDM-R1 强依赖于 TDM 框架的确定性轨迹特性，这限制了其在其他少步蒸馏方法（如 DMD、SDXL-Turbo 等）上的直接应用。改进方向可以是设计一个通用的确定性采样包装器，使得任意少步模型都能利用逐步奖励估计。其次，在奖励信号选择上，论文主要使用 GenEval 分数和 OCR 准确率等可验证指标作为不可微奖励的代表，但这些指标相对「简单」——它们有明确的正确答案。对于更复杂的不可微奖励（如真实的人类二元偏好），论文仅展示了 ImageReward 和 HPS 等已有奖励模型的结果，缺乏直接的人类评估实验。第三，在动态代理奖励的稳定性方面，虽然 EMA 参考模型缓解了参考模型过拟合问题，但论文未深入分析代理奖励模型和生成器之间训练速度不匹配可能导致的不稳定性，特别是在奖励信号稀疏或噪声较大的场景下。第四，论文的消融实验虽然覆盖了关键设计选择，但缺乏对组大小 $G$、学习率、KL 正则化系数 $\beta$ 和 $\beta_g$ 等超参数的敏感性分析。

基于 TDM-R1 的成果，未来研究可以从多个方向展开。首先，作者提到 TDM-R1 在视频生成等模态上的扩展是一个自然的延伸方向，特别是随着 Sora 等视频扩散模型的兴起，少步视频生成的 RL 后训练将有重要应用价值。其次，论文发现训练在一个可验证指标上可以协同提升其他指标，这暗示可能存在一种「代理任务优化」的范式——通过精心选择的代理任务来全面提升少步模型的指令跟随能力，这值得深入探索。第三，将 TDM-R1 与更先进的 RL 算法（如 REINFORCE++、PPO 的扩散版本）结合，探索更高效的训练策略。第四，研究如何将 TDM-R1 扩展到多目标奖励优化场景，同时优化多个不可微奖励信号（如同时优化人类偏好、文字渲染准确性和物体计数）。第五，探索 TDM-R1 在少步视频生成模型上的应用，利用视频的时序结构设计更精细的逐步奖励分配机制。最后，将代理奖励模型的知识蒸馏回生成器本身，使得经过 RL 后训练的生成器能够自我评估，从而减少推理时对外部奖励模型的依赖。

在复现方面，TDM-R1 提供了相对良好的条件。论文开源了代码仓库（https://github.com/Luo-Yihong/TDM-R1），这大大降低了复现门槛。在数据方面，论文使用的条件数据集来自标准的文本到图像数据集，GenEval 基准和 DrawBench 均为公开可用的评估集。在算力方面，论文基于 SD3.5-M 模型进行实验，该模型参数量约为 2B，加上代理奖励模型（同样基于 SD3.5-M 参数化），训练时需要同时维护生成器、代理奖励模型、参考模型和假分数网络，内存开销较大。不过作者提到停止 $\phi$ 的梯度可以节省内存，且不影响性能。对于 Z-Image 实验，模型规模达到 6B 参数，算力需求更高。在难度方面，TDM-R1 的算法实现涉及多个组件的交替训练（代理奖励、假分数、生成器），且需要正确实现组偏好优化和 EMA 参考模型更新，对工程实现有一定挑战。总体而言，有扩散模型 RL 训练经验的研究者应该能够在 1-2 周内完成复现。