这篇论文给扩散模型加了个奖励寄存器，不用改生成器就能获得稠密奖励，训练比在线 RL 快 33 倍，做生成对齐的可以看看。
论文提出 Latent Reward Registers，在冻结的 Diffusion Transformer（DiT）输入序列中加入可学习的 register token，直接从中间噪声潜在表示估计最终偏好奖励。相比稀疏终端奖励，该方法在整个去噪过程提供稠密、可微的奖励信号。训练策略 RG-OPD 在在线策略轨迹上蒸馏奖励梯度，GPU 时间最多减少 33 倍，超过在线强化学习基线。推理策略 RGS 无需参数更新即可引导采样，在无训练方法中达到新 SOTA。在高噪声水平 u=0.8 时，寄存器在潜在奖励模型中成对准确率最高。
Latent Reward Registers for Diffusion Preference Alignment
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