Latent Reward Registers for Diffusion Preference Alignment

📄 arXiv: 2608.03929v1 📥 PDF

作者: Yuanshen Guan, Zipeng Feng, Zhiwei Xiong, Peiqin Sun

分类: cs.LG, cs.CV

发布日期: 2026-08-04

🔗 代码/项目: GITHUB

💡 一句话要点

提出潜在奖励寄存器以解决扩散模型偏好对齐问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 奖励机制 深度学习 生成模型 偏好对齐 强化学习 去噪过程

📋 核心要点

  1. 现有方法依赖于稀疏终端奖励，导致在多步去噪过程中面临时间信用分配的挑战。
  2. 提出潜在奖励寄存器，通过在输入序列中添加寄存器标记，直接从中间潜变量中提取奖励信号。
  3. 实验结果显示，RG-OPD在减少GPU计算时间的同时，超越了在线强化学习基线，RGS在无训练方法中建立了新的最优状态。

📝 摘要（中文）

对扩散模型与人类偏好的对齐通常依赖于在最终生成样本上评估的稀疏终端奖励，这在多步去噪过程中带来了严重的时间信用分配挑战。本文提出了一种潜在奖励寄存器机制，通过在冻结的扩散变换器输入序列前添加可学习的、位置无关的寄存器标记，直接从中间噪声潜变量中估计终端偏好。这种独立的读取机制在不改变生成器的隐藏状态或速度场的情况下提取潜在奖励证据。最终的稠密、可微分的奖励信号在整个去噪过程中促进了两种对齐策略。实验表明，在高噪声水平下，寄存器在评估的潜在奖励模型中达到了最高的成对准确率。

🔬 方法详解

问题定义：本文旨在解决扩散模型与人类偏好对齐中的时间信用分配问题。现有方法依赖于稀疏的终端奖励，导致在多步去噪过程中难以有效评估奖励。

核心思路：提出潜在奖励寄存器，通过在输入序列中添加可学习的寄存器标记，直接从中间噪声潜变量中提取奖励信号，从而实现对齐过程中的奖励信号密集化和可微分化。

技术框架：整体架构包括两个主要模块：潜在奖励寄存器和奖励引导策略。潜在奖励寄存器负责从中间潜变量中提取奖励信号，而奖励引导策略则通过RG-OPD和RGS两种方式进行训练和推理。

关键创新：最重要的技术创新在于引入了潜在奖励寄存器机制，使得奖励信号在整个去噪过程中变得稠密和可微分，显著提升了对齐效果。与现有方法相比，该机制不需要改变生成器的隐藏状态或速度场。

关键设计：在设计中，寄存器标记是可学习的且位置无关，确保了在不同输入条件下的灵活性。损失函数设计上，RG-OPD采用了奖励引导的更新策略，避免了传统策略梯度方法中计算开销大的回滚过程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在高噪声水平下，潜在奖励寄存器在评估的潜在奖励模型中达到了最高的成对准确率。RG-OPD在减少GPU计算时间方面表现优异，最高可减少33倍，同时在对比基线中超越了在线强化学习方法。RGS在无训练方法中建立了新的最优状态，显著提升了对齐和感知指标。

🎯 应用场景

该研究的潜在应用领域包括生成模型的优化、图像生成、自然语言处理等。通过提高扩散模型的对齐能力，可以在艺术创作、内容生成等多个领域实现更高质量的输出，具有重要的实际价值和广泛的未来影响。

📄 摘要（原文）

Aligning diffusion models with human preferences usually relies on a sparse terminal reward evaluated on the final generated samples, presenting a severe temporal credit-assignment challenge across the multi-step denoising process. We propose Latent Reward Registers, a mechanism that estimates terminal preference directly from intermediate noisy latents by prepending learnable, position-free register tokens to the input sequence of a frozen Diffusion Transformer (DiT). This independent readout mechanism extracts latent reward evidence without altering the generator's hidden states or velocity field. The resulting dense, differentiable reward signal throughout the full denoising process facilitates two alignment strategies. For training, Reward-Gradient On-Policy Distillation (RG-OPD) distills reward-guided updates along on-policy trajectories, bypassing the computationally expensive rollouts of standard policy gradients. For inference, Reward-Guided Sampling (RGS) steers trajectories via magnitude-matched reward gradients without parameter updates. Empirically, at high noise levels (u = 0.8), the registers reach the highest pairwise accuracy among the evaluated latent reward models. Furthermore, RG-OPD outperforms online reinforcement learning baselines while reducing GPU hours by up to 33x, and RGS establishes a new state-of-the-art among training-free methods, strictly enhancing both alignment and perceptual metrics. Code and weights are available at https://github.com/Guanys-dar/latent-reward-register