Search papers, labs, and topics across Lattice.
This paper introduces Latent Reward Registers, a novel mechanism for aligning diffusion models with human preferences by estimating terminal rewards directly from intermediate noisy latents using learnable register tokens. This approach addresses the temporal credit-assignment challenge by providing a dense, differentiable reward signal throughout the denoising process, enabling two alignment strategies: Reward-Gradient On-Policy Distillation (RG-OPD) and Reward-Guided Sampling (RGS). Empirical results show that RG-OPD significantly outperforms traditional reinforcement learning baselines while reducing computational costs, and RGS achieves state-of-the-art performance in training-free methods for preference alignment.
Latent Reward Registers enable real-time preference alignment in diffusion models, achieving up to 33x reduction in computational costs while enhancing accuracy and perceptual quality.
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
