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

本文提出了 TDM-R1，一种基于轨迹分布匹配（TDM）的新型强化学习范式，通过解耦奖励学习与生成器训练并引入每步奖励信号，成功解决了少步扩散模型无法利用非可微奖励（如人类偏好）进行优化的难题，在文本渲染、视觉质量及偏好对齐等任务中实现了最先进的性能。

要点

这篇文章介绍了一种名为 TDM-R1 的新方法，它就像是为“极速版”AI 绘画模型装上了一套**“智能教练系统”**，让 AI 在极少的步骤内就能画出高质量、符合人类喜好的图片。

为了让你更容易理解，我们可以把整个过程想象成**“教一个天才但急躁的画家”**。

1. 背景：快枪手画家的烦恼

现在的 AI 绘画（扩散模型）有两种流派：

• 慢工出细活派（传统模型）： 像老画家，画一幅画要慢慢描 100 笔（100 步），画得很稳，但太慢了。
• 快手派（Few-Step 模型）： 像天才速写画家，只画 4 笔（4 步）就能出图，速度极快，适合商业应用。

问题出在哪？
虽然快手派画得快，但经常“翻车”。比如你让它画“一只猫在桌子左边”，它可能把猫画在右边，或者把猫画成狗。
以前，我们想教这些快手画家，只能靠**“可计算的分数”（比如数学公式算出来的误差）。但这就像只教画家做数学题，却没法教他“人类喜不喜欢这幅画”、“图里的字写得对不对”或者“物体数量对不对”。因为这些人类的主观喜好是“不可计算的”**（非可微的），传统的训练方法根本“教不会”快手画家。

2. 核心突破：TDM-R1 的“三步走”策略

TDM-R1 就像一位高明的教练，它发明了一套新的训练法，专门解决“怎么教快手画家听懂人类的主观评价”这个问题。

第一步：把“模糊的直觉”变成“清晰的地图”

• 旧方法： 以前教画家，只能等画完最后一笔（生成完图片）才给打分。如果画错了，画家不知道是哪一步（第 1 笔还是第 3 笔）出了问题，只能瞎猜。
• TDM-R1 的妙招： 它利用快手画家特有的**“确定性路径”**（就像画家走的是固定的直线，而不是乱跳的曲线）。
  • 比喻： 想象画家在画画时，每一步都留下了清晰的脚印。TDM-R1 能顺着脚印往回看，精准地告诉画家：“你第 2 笔的时候，颜色选错了；第 3 笔的时候，位置偏了。”
  • 效果： 即使奖励（比如人类点赞）是在最后才给的，TDM-R1 也能把它公平、准确地分摊到画画的每一个中间步骤上。

第二步：请一位“替身教练”（Surrogate Reward）

• 难题： 人类的点赞（非可微奖励）没法直接用来做数学推导（反向传播），就像你没法直接通过“点赞”来修改画家的肌肉记忆。
• TDM-R1 的妙招： 它训练了一个**“替身教练”**（一个小型的 AI 模型）。
  • 比喻： 这个替身教练专门负责观察画家的每一笔，把人类模糊的“喜欢/不喜欢”翻译成画家能听懂的“数学指令”。
  • 怎么练？ 它采用**“分组 PK"**的方式。比如，让画家画 10 张图，把画得好的 5 张归为“红队”，画得差的 5 张归为“蓝队”。替身教练通过对比这两队，学会如何给出更精准的指导信号。
  • 动态更新： 这个替身教练不是死板的，它会随着画家的进步不断调整自己的标准，就像教练会根据学生的水平动态调整训练难度。

第三步：双管齐下，既求快又求好

• 训练过程：
  1. 画家（生成器） 努力画得更好，以赢得替身教练的“高分”。
  2. 替身教练 不断进化，给出更精准的反馈。
  3. 同时，为了防止画家为了拿高分而“走火入魔”（画成奇怪的抽象画），TDM-R1 还加了一个**“紧箍咒”**（KL 正则化），确保画家画的图依然符合基本的物理规律和审美常识。

3. 成果：小步快跑，超越大师

实验结果非常惊人：

• 速度极快： 只需要 4 步（4 NFE）就能生成高质量图片。
• 质量极高： 在“物体数量”、“文字渲染”、“空间位置”等硬核测试中，TDM-R1 的得分（92%）甚至超过了需要画 80 步的传统大师模型（63%），也超过了目前最强的商业模型 GPT-4o（84%）。
• 通用性强： 无论是 SD3.5 还是最新的 Z-Image 大模型，只要用了这套方法，都能瞬间变强。

总结

TDM-R1 就像是给“极速赛车”装上了“顶级导航和教练系统”。
以前，极速赛车（Few-Step 模型）虽然快，但容易跑偏，且无法理解复杂的赛道规则（人类主观奖励）。TDM-R1 通过精准的路径分析和智能的替身教练，让赛车手在保持极速的同时，也能完美理解并执行复杂的指令，最终实现了**“又快又好”**的终极目标。

这项技术意味着，未来我们在手机上使用 AI 绘画，不仅能秒出图，而且画出来的东西能精准听懂你的话，甚至能画出完美的文字和复杂的场景。

技术摘要

1. 研究背景与问题 (Problem)

• 背景：少步生成模型（Few-step generative models，如基于扩散蒸馏的模型）在图像和视频生成领域取得了巨大成功，能够以极低的计算成本（如 4 步采样）实现高质量生成，已成为工业界的标准。
• 核心痛点：现有的强化学习（RL）方法在应用于少步扩散模型时，存在严重的局限性：
  1. 依赖可微分奖励：现有方法通常要求奖励信号（Reward Signal）必须是可微分的，以便通过反向传播将奖励梯度传递给生成模型。这排除了大量现实世界中至关重要的不可微分奖励（Non-differentiable Rewards），例如：人类二元偏好（喜欢/不喜欢）、离散的对象计数、OCR 文本识别的正确性等。
  2. 中间步骤奖励分配困难：少步模型通过确定性轨迹（Deterministic Trajectories）从噪声生成图像。传统的 RL 方法通常将最终图像的奖励直接分配给整个生成路径，导致中间去噪步骤的奖励估计存在偏差（Bias）和高方差。
  3. 生成质量下降：直接将标准扩散模型的 RL 方法（通常基于去噪损失）应用于少步模型，往往会导致生成图像模糊，无法兼顾少步生成的高效性与高质量。

2. 方法论 (Methodology)

作者提出了 TDM-R1，一种基于轨迹分布匹配（Trajectory Distribution Matching, TDM）的新型强化学习范式。其核心思想是将学习过程解耦为**代理奖励学习（Surrogate Reward Learning）和生成器优化（Generator Optimization）**两个部分。

2.1 核心组件

1. 基于确定性轨迹的精确中间奖励估计 (Accurate Intermediate Reward Estimation)

   • 原理：利用 TDM 的确定性采样轨迹（ODE 采样）。在确定性路径下，从中间噪声状态 $x_t$ 到最终图像 $x_0$ 的路径是唯一的。
   • 优势：这使得我们可以无偏地估计中间步骤的奖励。公式上，中间状态 $x_t$ 的奖励可以定义为最终图像奖励在条件分布 $p(x|x_t)$ 下的期望。由于路径是确定的，该期望的方差极低，从而为 RL 提供了稳定且准确的中间步骤反馈信号。

2. 代理奖励学习 (Surrogate Reward Learning)

   • 动机：为了处理不可微分的奖励（如人类偏好），不能直接对生成器进行梯度回传。
   • 方法：引入一个参数化的代理奖励模型 $\tilde{r}_\phi$（由扩散模型参数化）。
   • 训练机制：采用组级偏好优化（Group-based Preference Optimization）。
     • 将采样生成的样本分为正例组（$G^+$）和负例组（$G^-$）。
     • 利用 Bradley-Terry (BT) 模型学习组间的偏好关系。
     • 通过最大化组内优势（Advantage）的加权奖励，训练代理模型 $\tilde{r}_\phi$ 来拟合不可微分的真实奖励信号。
   • 动态参考模型：使用代理模型的指数移动平均（EMA）版本作为参考模型，防止过拟合噪声信号并提高训练稳定性。

3. 少步生成器优化 (Few-Step Generator Learning)

   • 目标函数：结合代理奖励最大化与反向 KL 散度正则化。
     $$L(\theta) = \mathbb{E}[-\tilde{r}_{sg}(\phi)(x_{t_k}, c) + \beta_g KL(p_{\theta,k}(x_t) || p_\psi(x_t))]$$
   • 机制：
     • 奖励最大化：利用训练好的代理奖励 $\tilde{r}_\phi$ 指导生成器 $\theta$ 更新，使其生成符合奖励偏好的样本。
     • 分布对齐：通过反向 KL 散度将生成器的分布约束在预训练的基础模型（Teacher）附近，防止模式坍塌（Mode Collapse）和分布偏移。
   • 解耦优势：代理奖励负责学习复杂的不可微分信号，生成器负责在保持分布稳定的前提下提升生成质量。

3. 关键贡献 (Key Contributions)

1. 首个针对少步模型的不可微分奖励 RL 范式：TDM-R1 是第一个能够大规模利用非可微分奖励（如人类偏好、OCR 准确率）来强化少步扩散模型的方法，打破了现有方法对可微分奖励的依赖。
2. 确定性轨迹的利用：创新性地利用 TDM 的确定性轨迹特性，实现了对中间去噪步骤的无偏、低方差奖励估计，解决了少步 RL 中奖励分配的关键难题。
3. 解耦的代理奖励机制：提出了基于组偏好优化的动态代理奖励学习机制，成功将不可微分的离散/二元奖励转化为可微分的梯度信号，同时避免了直接 RL 导致的图像模糊问题。
4. 无需额外真值数据：该方法仅需条件提示（Prompts）和奖励函数，无需额外的 Ground-truth 图像数据即可进行强化后训练。

4. 实验结果 (Results)

作者在多个基准测试和模型上进行了广泛实验，结果显著：

• GenEval 基准（组合图像生成）：

  • 在 SD3.5-M (4 NFE) 基础上，TDM-R1 将 GenEval 分数从 61% 提升至 92%。
  • 超越基线：不仅超越了其 80 步（80 NFE）的基线模型（63%），甚至超越了商业 SOTA 模型 GPT-4o (84%)。
  • 在物体计数、空间关系、属性绑定等细分任务上均取得 SOTA 表现。

• 视觉文本渲染 (Visual Text Rendering)：

  • 在 OCR 准确率指标上，TDM-R1 显著优于直接结合 RL 损失的基线方法，且训练过程更稳定（见图 4）。

• 人类偏好对齐 (Human Preference Alignment)：

  • 使用 ImageReward 和 HPSv3 作为奖励信号，TDM-R1 显著提升了少步模型的人类偏好评分，且未牺牲图像质量。

• 大规模模型扩展 (Z-Image)：

  • 在 60 亿参数的 Z-Image 模型上，TDM-R1 (4 NFE) 在域内和域外指标上均超越了其 100 步版本和 Turbo 版本。

• 消融实验：

  • 证明了确定性轨迹比随机轨迹收敛更快、性能更好。
  • 证明了动态代理奖励优于使用冻结的奖励模型。
  • 证明了直接组合 RL Loss 会导致图像模糊和性能下降，验证了解耦设计的必要性。

5. 意义与影响 (Significance)

• 打破少步模型的 RL 瓶颈：TDM-R1 解决了少步扩散模型难以利用丰富、复杂的现实世界奖励信号（如人类反馈）的难题，使得少步模型不仅能“快”，还能“准”和“好”。
• 性能超越多步模型：实验表明，经过 TDM-R1 强化的 4 步模型，其生成质量（包括指令遵循能力和视觉质量）可以超越昂贵的 80 步甚至 100 步的基础模型，极大地提升了推理效率与质量的性价比。
• 通用性：该方法不依赖于特定的奖励模型架构，适用于各种不可微分的评估指标，为未来少步生成模型的对齐（Alignment）和微调提供了通用的解决方案。
• 工业应用价值：鉴于少步模型在工业界的大规模部署需求，TDM-R1 提供了一种低成本、高效率的模型优化路径，有助于提升 AIGC 产品的实际用户体验。

总结：TDM-R1 通过巧妙结合确定性轨迹估计与代理奖励学习，成功将强化学习的优势引入少步扩散模型，实现了在极低采样步数下（4 步）超越传统多步模型及商业 SOTA 的性能，是 AIGC 领域少步生成与强化学习结合的重要突破。