技术摘要：带扩散奖励的 Diff-Instruct（Didr）

问题陈述

近期单步文本到图像（T2I）生成器的进展，通过将潜在噪声直接映射到图像的单次前向传递，实现了实时合成。然而，通过强化学习（RL）使这些生成器与人类偏好对齐，存在一个根本性的结构挑战。

以往的单步 RLHF 方法（如 Diff-Instruct*、Diff-Instruct++）将图像空间的奖励优化与潜在空间的分布匹配相结合。这些方法通常仅在干净图像终点（$x_0$）施加偏好信号（奖励），而在整个含噪扩散轨迹上针对参考分布施加 Kullback-Leibler（KL）正则化。作者指出了该范式中的一个关键结构失配：奖励仅作用于终端状态，而正则化项则跨越所有噪声水平。

这种失配导致了一种被称为终端奖励主导的失效模式。由于前向噪声自然掩盖了图像细节，KL 正则化对模式发散的惩罚在高噪声水平下显著减弱。因此，优化器利用这一轨迹脆弱性，放弃参考分布并坍缩至高奖励区域，而无需承担足够的 KL 惩罚。这虽然提高了奖励分数，但往往以牺牲图像保真度和分布稳定性为代价。

方法论：带扩散奖励的 Diff-Instruct（Didr）

为了解决这种不对齐问题，作者提出了Didr，这是一种源自积分 KL（IKL）最小化的无数据、轨迹级对齐框架。其核心洞察是系统地将 RLHF 最优的奖励倾斜干净图像分布传播至扩散轨迹的所有噪声水平。

1. 原则性轨迹目标

该方法首先确立了在 KL 正则化 RLHF 下干净图像的正确对齐目标。最优干净图像分布 $q^*(x_0|c)$ 是奖励倾斜的密度：
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
其中 $q_0$ 是参考分布，$r$ 是奖励函数，$\tau$ 是温度参数。

Didr 不是仅在 $x_0$ 处应用奖励，而是通过参考前向过程将 $q^*$ 扩散，从而将此最优目标“提升”至完整轨迹。这产生了一个定义在每个噪声水平上的奖励倾斜边缘分布族 $\{q^*_t\}_{t \in [0, T]}$：
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
作者证明，最小化生成器边缘分布与这些奖励倾斜边缘分布之间的积分 KL 散度，与标准 KL 正则化 RLHF 目标拥有相同的最小化器。这确保了奖励和正则化在每个噪声水平上保持平衡，从而防止终端奖励主导。

2. 扩散奖励分数（DRS）

通过对 IKL 目标求导，作者推导出了目标分数函数。最优分数分解为两个分量：
$$s_{target} = s_{ref} + s_r$$
其中 $s_{ref}$ 是参考分数函数，$s_r$ 是扩散奖励分数（DRS）。DRS 作为一个定义在每个噪声水平上的奖励驱动修正项：
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
与分类器引导（一种一阶近似）不同，DRS 保留了完整的对数指数倾斜，仅在干净图像可从含噪观测中识别的位置传播偏好信息。

3. 扩散奖励代理（DRP）

由于 DRS 需要难以处理的边缘期望 $q(x_0|x_t, c)$，作者引入了**扩散奖励代理（DRP）**以使该方法具有实用性。DRP 利用来自冻结参考模型的可微短步去噪链来近似后验期望。

• 机制：从含噪状态 $x_t$ 开始，该方法使用冻结参考模型运行 $K$ 个独立的 $S$ 步去噪链，以生成近似干净样本 $\{\hat{x}^{(k)}_0\}$。
• 梯度估计：DRS 被近似为路径奖励梯度的 softmax 加权平均：
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  其中权重 $\omega^{(k)}$ 与 $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$ 成正比。
• 效率：这种方法允许奖励梯度稳定地反向传播通过轨迹，而无需存储参考模型的激活值或需要图像训练数据。

4. 训练框架

训练过程在两个阶段之间交替：

1. 助教（TA）更新：通过去噪分数匹配（DSM）训练 TA 分数模型（$s_\psi$），以跟踪生成器不断演变的边缘分布，替代隐式的生成器分数。
2. 生成器更新：使用 IKL 目标的梯度更新生成器（$g_\theta$），该梯度比较 TA 分数与 CFG 校正后的参考分数及 DRP 估计值之和。

主要贡献

1. 识别终端奖励主导：本文正式识别并分析了先前单步 RLHF 方法中的一种结构失效模式，即优化器利用噪声水平坍缩至高奖励区域，违反了预期的 RLHF 平衡。
2. 原则性轨迹级对齐：作者提出了 Didr，这是一个通过向前传播 RLHF 最优干净图像密度来推导数学上合理的轨迹级目标的框架。这确保了奖励和正则化在每个噪声水平上保持平衡。
3. DRS 的解析推导：本文解析推导了扩散奖励分数，表明其将轨迹目标精确分解为参考分数加上奖励诱导的修正项。
4. 实用估计器（DRP）：作者引入了扩散奖励代理，这是一种零数据估计器，通过可微短步去噪计算 DRS，从而在无需图像数据的情况下实现稳定训练。

实证结果

大量实验表明，Didr 始终优于现有的单步基线：

• 帕累托优势：在 SDXL 骨干网络（1024×1024）上，标准 Didr 变体在 PickScore–FID 权衡中帕累托优于所有现有的单步对齐基线，同时在两个指标上超越了多步模型（如 FLUX-dev 和 SDXL-DPO）。
• 偏好指标：扩展训练变体（Didr_longer）在评估的单步模型中实现了最高的 PickScore（23.9），并比最强的先前基线（Diff-Instruct*）将 ImageReward 提高了 8.9%。它在 HPSv2.1 和美学评分方面也处于领先地位。
• 泛化性：Didr 成功迁移至 60 亿参数 Z-Image 骨干网络。在单次生成步骤中，它在偏好指标上超越了其 50 步教师模型，展示了跨架构的可扩展性。
• 定性改进：视觉比较显示，与先前的单步方法相比，Didr 产生了更清晰的细节和更自然的色彩，避免了奖励优化基线中常见的过饱和和绘画风格伪影。

意义与主张

本文主张，原则性轨迹级奖励传播是一种有效的对齐策略，解决了仅针对端点的奖励方法固有的不对齐问题。通过严格地将 RLHF 最优分布通过扩散过程传播，Didr 在单步生成器中实现了偏好对齐与图像保真度之间的最佳权衡。

作者强调，他们的方法不需要图像训练数据，仅依赖文本提示和冻结的参考模型。他们将 Didr 定位为一种解决方案，使单步生成器能够匹配甚至超越多步教师的偏好对齐能力，同时保持实时推理速度。这项工作表明，先前方法中识别出的结构失配是一个根本性限制，只有通过对齐整个轨迹（而不仅仅是终端状态）才能解决。