Latent Reward Registers for Diffusion Preference Alignment

2026-08-04 · 17:00 UTC ·Yuanshen Guan, Zipeng Feng, Zhiwei Xiong... ·1 min read

Brief

Latent Reward Registers introduce learnable, position-free register tokens prepended to a frozen Diffusion Transformer to predict terminal human-preference reward from intermediate noisy latents. This yields dense, differentiable signals used for Reward-Gradient On-Policy Distillation (RG-OPD) and Reward-Guided Sampling (RGS). Experiments report highest pairwise accuracy at u=0.8, RG-OPD cutting GPU costs up to 33x, and RGS achieving training-free SOTA on alignment and perceptual metrics.

Why it matters

Latent Reward Registers estimate terminal preference from intermediate noisy latents by prepending learnable, position-free register tokens to the input of a frozen Diffusion Transformer (DiT), extracting reward evidence without altering the generator's hidden states or velocity field and producing dense, differentiable rewards across denoising steps.

Key details

The dense reward enables two alignment strategies: Reward-Gradient On-Policy Distillation (RG-OPD), which distills reward-guided updates along on-policy trajectories and avoids expensive rollouts, and Reward-Guided Sampling (RGS), which steers trajectories at inference via magnitude-matched reward gradients without parameter updates.

Empirical results: at high noise (u = 0.8) the registers achieve the highest pairwise accuracy among evaluated latent reward models; RG-OPD outperforms online reinforcement-learning baselines while reducing GPU hours by up to 33x; RGS sets a new state-of-the-art among training-free methods, improving alignment and perceptual metrics.

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