AI 中文总结
针对扩散模型偏好对齐的时间信用分配挑战，提出Latent Reward Registers机制，结合RG-OPD和RGS策略，在高噪声下实现最优性能且大幅降低计算成本。
AI 中文摘要
将扩散模型与人类偏好对齐通常依赖于对最终生成样本评估的稀疏终端奖励，这在多步去噪过程中会带来严重的时间信用分配挑战。我们提出了Latent Reward Registers（潜在奖励寄存器），该机制通过将可学习的、无位置的寄存器令牌添加到冻结的Diffusion Transformer（DiT）的输入序列中，直接从中间噪声潜在变量估计终端偏好。这种独立的读出机制可提取潜在奖励证据，而不会改变生成器的隐藏状态或速度场。由此产生的在整个去噪过程中密集、可微分的奖励信号促进了两种对齐策略：训练时，Reward-Gradient On-Policy Distillation（RG-OPD，奖励梯度在线策略蒸馏）沿在线策略轨迹蒸馏奖励引导的更新，绕过了标准策略梯度计算成本高昂的rollout（展开）；推理时，Reward-Guided Sampling（RGS，奖励引导采样）通过幅度匹配的奖励梯度引导轨迹，无需参数更新。经验证，在高噪声水平（u = 0.8）下，该寄存器在评估的潜在奖励模型中达到最高的成对准确率；此外，RG-OPD的性能优于在线强化学习基线，同时GPU运行时间最多减少33倍，而RGS在无训练方法中建立了新的SOTA（当前最优），严格提升了对齐和感知指标。代码和权重可在https URL获取。
英文摘要
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
