扩散模型偏好对齐新范式：基于潜在奖励寄存器的细粒度时序信用分配

在扩散模型生成高质量图像的过程中，如何有效地将人类偏好融入生成过程是一个极具挑战性的课题。传统的对齐方法通常依赖于在生成过程结束时对最终样本进行评估，从而获得一个稀疏的终端奖励信号。然而，扩散模型的生成是一个多步骤的去噪过程，这种仅关注终点的评估方式导致了严重的时序信用分配问题，即模型难以判断去噪过程中的哪一个具体步骤对最终满足人类偏好起到了关键作用。为了解决这一核心痛点，本研究提出了一种名为"潜在奖励寄存器"（Latent Reward Registers）的全新机制。

其核心贡献在于打破了传统终端奖励的局限，提出了一种能够从中间噪声潜变量中直接估计终端偏好的方法。通过这种方式，研究团队成功地将稀疏的终端奖励转化为贯穿整个去噪过程的稠密、可微分的奖励信号，为后续的优化和采样策略提供了坚实的理论基础和数据支持，从而在根本上缓解了时序信用分配的难题。在技术实现层面，该机制巧妙地利用了冻结的扩散Transformer（DiT）架构。具体而言，研究者通过在DiT的输入序列前prepend（前置）一组可学习的、位置无关的寄存器令牌（register tokens），构建了一个独立的读取机制。

这些寄存器令牌能够与输入序列中的潜变量进行交互，从而提取出关于终端偏好的潜在证据。这一设计的精妙之处在于，它完全独立于生成器的主干网络，既没有改变生成器的隐藏状态，也没有干扰其速度场的演化，确保了生成过程的稳定性。基于这种稠密的奖励信号，研究提出了两种对齐策略：在训练阶段，采用奖励梯度在线策略蒸馏（RG-OPD），该方法通过蒸馏奖励引导的更新沿着在线策略轨迹进行，巧妙地避开了标准策略梯度中计算昂贵的滚动（rollouts）过程；在推理阶段，则采用奖励引导采样（RGS），通过引入与奖励梯度幅度匹配的引导信号，在不更新模型参数的情况下直接 steer（引导）生成轨迹，实现了推理阶段的偏好对齐。实验部分在多个基准上验证了所提方法的有效性。

首先，在评估潜在奖励模型的性能时，研究发现该寄存器机制在高噪声水平（u = 0.8）下达到了所有被评估的潜在奖励模型中最高的成对准确率，证明了其捕捉偏好信号的准确性。其次，在训练策略方面，RG-OPD不仅优于各种在线强化学习基线方法，更关键的是，它将所需的GPU计算时长大幅减少了高达33倍，极大地降低了训练成本。在推理阶段，RGS作为一种免训练（training-free）的方法，建立了新的最先进性能（state-of-the-art），在严格增强图像对齐度的同时，也显著提升了感知指标。这些结果有力地证明了该机制在提升对齐效果和计算效率方面的双重优势，消融实验进一步揭示了寄存器令牌在不同去噪步骤中对奖励信号提取的贡献，验证了其在处理时序依赖关系上的有效性。

这项研究对扩散模型领域具有重要的行业意义和潜在影响。首先，它为解决扩散模型偏好对齐中的时序信用分配问题提供了全新的视角和工具，证明了在中间潜变量层面进行奖励估计的可行性。对于开源社区而言，代码和权重的公开将促进相关技术的复现和改进。在工业落地方面，RG-OPD大幅降低的训练成本使得大规模偏好对齐训练变得更加经济可行，而RGS则在推理阶段无需额外参数更新即可提升生成质量，非常适合资源受限的应用场景。此外，这种独立的奖励提取机制为后续研究提供了新的思路，例如可以探索更复杂的奖励模型结构或将其应用于其他生成式任务中，如视频生成或多模态生成，进一步推动生成式人工智能在满足人类复杂偏好方面的能力发展。