Latent Reward Registers for Diffusion Preference Alignment

Yuanshen Guan; Zipeng Feng; Chengru Song; Zhiwei Xiong; Peiqin Sun

cs.LGAug 4, 2026

Latent Reward Registers for Diffusion Preference Alignment

Authors: Yuanshen Guan, Zipeng Feng, Chengru Song, Zhiwei Xiong, Peiqin Sun

Organizations: 1Kling Team · University of Science and Technology of China

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

Explore similar work

May 25, 2026cs.CV

DRM: Diffusion-based Reward Model With Step-wise Guidance

Current mainstream methods of aligning diffusion models with human preferences typically employ VLM-based reward models. However, these reward models, pre-trained for semantic alignment, struggle to capture the essential perceptual qualities-such as aesthetics, composition, and visual harmony. In this work, we argue that a model capable of high-fidelity generation must possess a profound understanding of these visual attributes. Based on this insight, we introduce the Diffusion-based Reward Model (DRM), a novel paradigm that use the pre-trained diffusion model as a powerful evaluative backbone. A key advantage of the DRM is its unique ability to assess not only the final image but also the noisy intermediate latents at any stage of the generative process. We leverage this step-wise evaluative capacity in two ways. First, we propose Step-wise GRPO, a reinforcement learning algorithm that provides dense, per-step rewards to resolve the imprecise credit assignment problem in GRPO algorithm, leading to more stable and effective alignment. Second, we introduce Step-wise Sampling, a novel inference strategy that employs the DRM as a dynamic guide to evaluate multiple generation paths at each step, steering the process towards higher-quality outcomes. Extensive experiments confirm that our approach significantly enhances the final quality of generated images. Code: https://github.com/jjaxonx/DRM.

Jaxon Zhang, Binxin Yang, Hubery Yin +2

Few-Step

Ask PDF

May 26, 2026cs.LG

Explicit Critic Guidance for Aligning Diffusion Models

Online reinforcement learning is becoming increasingly important for aligning diffusion models with non-differentiable objectives. However, existing methods still face limitations in assigning fine-grained credit along denoising trajectories and in realizing stable value-based optimization. We propose a state-aligned latent actor-critic framework for diffusion post-training, in which the diffusion model serves as its own timestep-conditioned value function and predicts values directly on noisy latent states. This enables trajectory-level PPO training, supports stable actor-critic optimization with simple conditioning and value pretraining strategies, and naturally allows the learned critic to be reused for inference-time steering. We further extend the framework to multi-reward optimization, where joint training with complementary rewards helps alleviate reward hacking. Across both UNet- and DiT-based backbones, our method consistently outperforms prior group-relative RL and actor-critic baselines on single-reward and multi-reward benchmarks, while test-time steering provides additional gains in generation quality.

Zhengyang Liang, Qihang Zhang, Ceyuan Yang

Critic-Free Reinforcement Learning Soft Actor-Critic Critic

Ask PDF

May 12, 2026cs.LG

The tractability landscape of diffusion alignment: regularization, rewards, and computational primitives

Inference-time reward alignment asks how to turn a pre-trained diffusion model with base law

p

into a sampler that favors a reward

r

while remaining close to

p

. Since there is no canonical distributional distance for this closeness constraint, different choices lead to different "reward-aligned" laws and, just as importantly, different algorithmic problems. We develop a primitive-based approach to reward alignment: rather than assuming arbitrary reward-aligned laws can be sampled, we ask which simple algorithmic primitives suffice to implement alignment for non-trivial reward classes. If closeness is measured in KL distance, the target law is

q(x) \propto p(x) \exp(λ^{-1}r(x))

. For this setting, we show that linear exponential tilts of the form

q(x)\propto p(x)\exp(\langle θ, x \rangle)

-- which according to recent work [MRR26] can be efficiently sampled from -- are a sufficient primitive for aligning to a very broad class of convex low-dimensional rewards. If closeness is measured in Wasserstein distance, the corresponding primitive is a proximal transport oracle: given

x

, solve

\mbox{argmax}_y \{r(y)- λc(x,y)\}

. This oracle can be efficiently implemented for concave or low-dimensional Lipschitz rewards

r(x)=f(Ax)

. Together, these results illustrate that the choice of distribution distance for alignment affects the computational primitive and the tractable reward class.

Ankur Moitra, Andrej Risteski, Dhruv Rohatgi

Regularization Techniques Computable

Ask PDF