用中间噪声隐状态直接估算偏好,解决扩散模型对齐的信用分配难题。
Latent Reward Registers for Diffusion Preference Alignment

- 引入可学习的注册令牌,从中间噪声隐状态提取偏好信号。
- 在高噪声水平下(t=0.8)达到最优成对准确率,比基线快33倍。
- 无需参数更新即可生成高质量结果,适合快速推理场景。
将扩散模型与人类偏好对齐通常依赖于仅在最终生成样本上评估的稀疏终端奖励,这在去噪过程中引发了严重的时序信用分配问题。我们提出潜在奖励注册机制(Latent Reward Registers),直接从中间噪声隐状态估计终端偏好。可学习的、位置无关的注册令牌作为辅助读取路径附加到冻结的扩散变换器(DiT)上,提取偏好信号而不改变生成器的隐藏状态或速度场。由此产生的密集、可微奖励场覆盖整个去噪轨迹,并支持两种对齐策略:训练阶段采用奖励梯度在线策略蒸馏(RG-OPD),将密集奖励场转化为当前生成器访问状态的每步目标,取代计算量大的滚动策略梯度;推理阶段采用奖励引导采样(RGS),以幅度匹配的奖励梯度修正引导轨迹,无需参数更新。实验表明,在高噪声水平(t=0.8)下,该方法在评估的潜在奖励模型中达到最高成对准确率;RG-OPD优于在线强化学习基线,同时减少高达33倍的GPU小时数;RGS在不增加训练成本的情况下显著提升奖励表现,实现良好的奖励-质量平衡。代码与权重将于https://github.com/Guanys-dar/latent-reward-register发布。
原文摘要 · Abstract (English)
Aligning diffusion models with human preferences usually relies on a sparse terminal reward evaluated on the final generated samples, which creates a severe temporal credit-assignment problem across the denoising process. We propose Latent Reward Registers, a mechanism that estimates terminal preference directly from intermediate noisy latents. Learnable, position-free register tokens are appended as an auxiliary read path to a frozen Diffusion Transformer (DiT), extracting preference signals without altering the generator's hidden states or velocity field. The resulting dense, differentiable reward field spans the full denoising trajectory and supports two alignment strategies. For training, Reward-Gradient On-Policy Distillation (RG-OPD) converts this dense reward field into per-step targets at states visited by the current generator, replacing rollout-intensive policy gradients with direct on-policy distillation. For inference, Reward-Guided Sampling (RGS) steers trajectories with magnitude-matched reward-gradient corrections and no parameter updates. Empirically, at high noise levels (t=0.8) the registers reach the highest pairwise accuracy among the evaluated latent reward models. RG-OPD outperforms online reinforcement learning baselines while reducing GPU hours by up to 33x. RGS achieves significant reward improvement with a favorable reward-quality balance against training-free baselines. Code and weights are to be available at https://github.com/Guanys-dar/latent-reward-register
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。