Latent Reward Registers for Diffusion Preference Alignment

2026-08-04

AI Summary

This work addresses the challenge of aligning diffusion models with human preferences, which is hindered by sparse rewards derived only from final samples and the difficulty of temporal credit assignment across multi-step denoising. To overcome this, the authors propose inserting learnable, position-agnostic register tokens at the input sequence of a frozen Diffusion Transformer, enabling a dense, differentiable reward signal throughout the denoising trajectory. An independent readout mechanism is introduced to prevent interference with the generator. Furthermore, they develop Reward-Gradient On-Policy Distillation (RG-OPD), a replay-free training method, and Reward-Guided Sampling (RGS), a training-free inference strategy. Experiments demonstrate that the approach achieves state-of-the-art pairwise accuracy under high noise levels, with RG-OPD reducing GPU hours by up to 33× compared to reinforcement learning baselines and RGS outperforming existing training-free methods in both alignment and perceptual quality.

Abstract

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

Problem

• diffusion models
• preference alignment
• temporal credit assignment
• latent reward
• denoising process

Innovation

• Latent Reward Registers
• Diffusion Preference Alignment
• Reward-Guided Sampling
• On-Policy Distillation
• Frozen DiT