Latent Reward Registers for Diffusion Preference Alignment

AI 中文解读

核心亮点：这项研究让AI绘画模型在生成过程中每一步都能“自我纠错”，不再等到画完才打分，既省电又画得更好。通俗解读：过去AI画图就像蒙眼作画，画完才有人告诉你哪里不对，导致中途错误无法修正。现在科学家给AI加了个“实时评分员”，它能在AI画到一半时，就从模糊的草稿里看出最终效果好坏，并悄悄调整方向。这相当于给AI戴上了“透视眼镜”，边画边检查，还能根据反馈即时微调，甚至不用重新训练就能提升画质。实际影响：对普通人来说，这意味着AI绘图工具会更快、更便宜、效果更惊艳。比如你用手机App生成海报，以前可能要等几十秒甚至更久，现在速度能提升几十倍，而且画面细节和审美更贴合你的要求。同时，这项技术也能让AI视频生成更流畅，未来你随手输入一句话，就能在几秒内得到高质量的视频素材，创作门槛大幅降低。

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