Latent Reward Registers for Diffusion Preference Alignment

arXiv:2608.03929 · cs.LG, cs.CV · Submitted 2026-08-14 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Latent Reward Registers for Diffusion Preference Alignment".

Jane: The paper was written by Yuanshen Guan, Zipeng Feng, Chengru Song, Zhiwei Xiong and Peiqin Sun from Kling Team and University of Science and Technology of China.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back to the show, everyone. Today we’re digging into a fresh arXiv paper called “Latent Reward Registers for Diffusion Preference Alignment.” Jane, I gotta say, that title is a mouthful, but the idea behind it is genuinely clever.

Jane: It really is, Tom. And the title actually tells you exactly what they did. They took a diffusion model — you know, the kind that generates images from text — and they found a way to read out how good the final image will be, way before the image is finished. Those “reward registers” are like little sticky notes they attach to the model’s internal processing.

Tom: Sticky notes, I like that. So instead of waiting until the very end to check if the image matches what you asked for, these registers peek at the noisy intermediate steps and predict the final quality. That’s a huge shift, because normally you’d have to generate the whole image, then score it, then go back and fix things.

Jane: Exactly. And the authors are from Kling Team and the University of Science and Technology of China. They’ve got this really practical angle — they’re not just proposing a theory, they’re showing it works on real models like SD3-Medium and FLUX.one-dev.

Tom: Right, and that matters because those are the models people actually use. But here’s what excites me, Jane — this could change how we align AI with human preferences. Right now, if you want a model to generate images people like, you have to train it with feedback on the final image. That’s slow and expensive.

Jane: And it’s also kind of blind, Tom. The model makes dozens of small decisions along the way, but you only tell it at the end whether the whole thing worked. It’s like giving someone directions but only telling them if they arrived at the right city, not whether they took the right turns.

Tom: That’s the credit assignment problem they mention in the paper. And their solution is to give feedback at every single turn, not just at the destination. That’s the big idea.

Jane: And the beauty is that they do it without messing with the original model. The registers are added on the side, so the generator’s core behavior stays untouched. That means you can use this for training, but also for guiding the model at inference time, without retraining.

Tom: So you could take a frozen, already-trained model and just steer it toward better outputs on the fly. That’s a big deal for practical use.

Jane: It is. And I’m curious to hear how they actually pull that off technically. That’s where it gets really interesting.

Tom: Yeah, let’s get into that next. We’ve got the high-level picture — now we need to understand the machinery.

Summary: Tom: So, Jane, we’ve established that “Latent Reward Registers for Diffusion Preference Alignment” is about giving feedback during the generation process, not just at the end. But how do they actually build these registers?

Jane: Great question. So inside a diffusion model, there’s this transformer backbone that processes the noisy image latents step by step. The authors add a small set of learnable tokens — they call them registers — into that processing stream. These registers don’t change what the model generates; they just watch what’s happening and accumulate information about how well the generation is going.

Tom: Watching from the sidelines, basically. And they’re trained to predict what the final reward would be — like a human preference score — based on the intermediate state. So at any point in the denoising process, you can ask, “Hey, how is this looking?” and get a prediction.

Jane: Exactly. And they train these registers using pairwise comparisons, which is how a lot of preference models work. You show them two images, one preferred by humans, and they learn to rank them correctly. But here’s the twist — they train on noisy latents, not clean images. So the registers learn to extract reward information even when the image is mostly noise.

Tom: And that’s the part that surprised me. They tested this at a noise level of zero point eight, which means the latent is eighty percent pure noise. And their registers still achieved competitive accuracy compared to models that score the fully clean image. That’s pretty remarkable.

Jane: It is. The internal representations of the diffusion model apparently encode a lot of information about the final output quality very early on. The paper shows that even under heavy corruption, the model’s features contain signals about layout, prompt alignment, and overall viability.

Tom: So they’re essentially saying the model knows where it’s going long before it gets there. And their registers learn to read that knowledge.

Jane: Right. And this dense reward signal — available at every step — enables two things. First, for training, they use it to distill reward-guided updates along the model’s own trajectory, which they call RG-OPD. Second, for inference, they use the gradient of the predicted reward to steer the sampling process, which they call RGS.

Tom: And both of those are big claims. They say RG-OPD beats online reinforcement learning baselines while using up to thirty-three times fewer GPU hours. That’s a massive efficiency gain.

Jane: Yeah, because standard RL methods need to roll out full trajectories and estimate policy gradients, which is noisy and sample-inefficient. Their method gives step-wise, dense gradients, so the model learns faster and more stably.

Tom: And RGS is training-free — you just apply the reward gradient during sampling. That means you can align a model to a new preference without retraining it at all.

Jane: Exactly. And they show it improves both alignment and perceptual quality, which is usually a trade-off. That’s the part I want to dig into — how they manage to get both.

Tom: Let’s bring in Lu and Meng for that. They’ll have strong opinions on whether this actually holds up.

Improvements: Tom: Okay, so we’re back with Lu and Meng. We’ve been talking about “Latent Reward Registers for Diffusion Preference Alignment” and how it gives dense, step-wise feedback. Lu, what do you think is the most significant improvement this paper brings?

Lu: Thanks, Tom. I think the most significant thing is that they’ve turned a sparse, delayed signal into a dense, local one without touching the generator. That’s a conceptual leap. In reinforcement learning, credit assignment is the hardest problem — you have a reward at the end, but you don’t know which actions caused it. Here, they’ve essentially built a learned value function that works in latent space, at every step.

Meng: But Lu, I want to push back a little. A learned value function is only useful if it’s accurate. They show high accuracy at noise level zero point eight, but what about the full range? I saw in the paper that their accuracy dips at very low noise levels, around zero point two, compared to some baselines.

Jane: That’s a fair point, Meng. But they also show that their registers are more stable across the entire noise range than the baselines. DiNa-LRM, for example, performs well near clean data but drops sharply at high noise. The registers maintain reliability throughout, which matters more for early-stage guidance.

Lu: And that early-stage guidance is where you set the global structure. If you get the layout wrong at the beginning, no amount of fine-tuning at the end will fix it. So having a reliable reward signal at high noise is actually more valuable than at low noise.

Meng: Okay, that makes sense. But what about the practical side? They claim RG-OPD is up to thirty-three times faster than Flow-GRPO. How are they getting that speedup?

Tom: That’s the part I want to understand too. Is it just because they’re not doing full rollouts?

Jane: Right. Standard policy gradient methods need to sample complete trajectories, compute rewards, and then estimate gradients with high variance. RG-OPD instead constructs a one-step target at each state along the student’s own trajectory. The reward register gives a local gradient, and they distill that into the student model. No backprop through the whole chain, no rollout-level variance.

Meng: So it’s like a teacher-student setup where the teacher is the frozen reference model plus the reward register. And the student learns to match the teacher’s tilted predictions.

Lu: Exactly. And the “tilt” is controlled by a coefficient that scales the reward gradient relative to the reference displacement. That’s clever because it lets you control how much you want to deviate from the original model.

Meng: And what about RGS? Is that just classifier guidance in disguise?

Jane: It’s similar in spirit, but with a key difference. Classifier guidance typically requires a classifier trained on clean images, and you have to decode the latent to apply it. Here, the register operates directly on the noisy latent, so there’s no decoding overhead. And they magnitude-match the reward gradient to the sampler’s own displacement, which keeps the correction stable.

Lu: And that stability is why they can improve both reward and perceptual quality. The correction is proportional to what the model would have done anyway, so it doesn’t push the trajectory off-manifold.

Meng: That’s the part that impressed me. Usually, optimizing for a reward metric hurts perceptual quality. They show MUSIQ and CLIP-IQA staying high while HPSv3 goes up. That’s not easy.

Tom: So we’ve got efficiency, stability, and quality. Sounds like a win across the board. But I’m wondering — what are the limitations? What’s the catch?

Jane: Well, the registers need to be trained on paired preference data, and they’re specific to the backbone they’re attached to. You can’t just take a register trained on SD3 and slap it on FLUX. They do show it transfers to FLUX with retraining, but that’s still a cost.

Lu: And the reward heads are trained to match specific endpoint models. So if you want to align to a new reward, you need to train a new head. That said, the multi-head setup they show — combining HPS and ImageReward — suggests you can compose objectives without retraining the whole thing.

Meng: Right, and the inference overhead is real but manageable. They report about one point eight times the latency of standard CFG, which is better than the decode-and-score baselines that are four times slower.

Tom: So it’s a practical trade-off, not a free lunch. But it’s a much better trade-off than what we had before.

Jane: Exactly. And that’s why I think this paper could have real impact. Let’s wrap up with that thought.

Conclusion: Tom: We’ve spent the whole episode on “Latent Reward Registers for Diffusion Preference Alignment,” and I think we’ve only scratched the surface. Jane, how would you sum it up for someone who just tuned in?

Jane: Sure, Tom. The paper tackles a fundamental problem in aligning diffusion models with human preferences — you only get feedback at the very end, but the model makes decisions at every step. Their solution is to add lightweight, learnable registers to a frozen diffusion transformer that predict the final reward from intermediate noisy latents. That gives you a dense, step-wise reward signal without changing the generator at all.

Tom: And that signal powers two things — a training method that’s up to thirty-three times faster than online RL, and an inference-time guidance method that improves alignment without retraining. Both of those are big practical wins.

Lu: I’d add that the conceptual contribution is just as important. They’ve shown that diffusion models internally encode reward-relevant information much earlier than we thought. That opens the door for more efficient alignment methods across modalities — video, audio, even multi-step reasoning.

Meng: And from an engineering standpoint, the fact that it works with frozen models is huge. You can deploy this on existing infrastructure without retraining the base generator. The overhead is a small auxiliary network and a few extra milliseconds per step.

Jane: Right. And they’ve made the code and weights available, so people can actually try it. That’s going to accelerate adoption.

Tom: So, big picture — this could make preference alignment cheaper, faster, and more flexible. That means better image generation, but also potentially better alignment in other generative domains.

Lu: And the multi-head composition is a glimpse of what’s possible — aligning to multiple objectives at once, balancing quality and preference, without retraining. That’s a direction I’m excited to see explored.

Meng: I just hope they publish more details on the register training stability. The EMA and the noise-aware ranking loss are interesting, but I’d want to see how sensitive it is to hyperparameters.

Jane: Fair point, Meng. But for now, this is a solid contribution with clear results and practical implications. We’ll be watching for follow-ups.

Tom: And with that, we’ll say goodbye to “Latent Reward Registers for Diffusion Preference Alignment.” Thanks to Lu, Meng, and Jane for a great discussion. Next up on the show, we’ve got another paper that’s been making waves — but that’s for next time. Stay tuned.

Jane: Thanks for listening, everyone. See you on the next episode.

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

Kling Team · University of Science and Technology of China

cs.LG, cs.CV

Submitted: 2026-08-14

Updated: 2026-08-17

Code: https://github.com/Guanys-dar/latent-reward-register

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 69/100

The gist: "Latent Reward Registers estimate step-wise rewards directly from intermediate noisy latents.

Key concepts

Diffusion Model
A type of generative AI model used to create images from text. The process involves gradually transforming pure noise into a coherent image over many steps.
Latent Reward Registers
Small, learnable tokens added to the model's internal processing stream. They predict the final quality (reward) of an image by 'watching' the noisy intermediate steps, providing step-wise feedback.
Credit Assignment Problem
The challenge in training AI where feedback is only given at the end (the destination), making it difficult to determine which specific actions or decisions early on were responsible for the final outcome.
RGS (Reward Guidance Sampling)
An inference-time method that uses the predicted reward gradient to steer a frozen, already-trained model toward better outputs. This allows alignment without needing full retraining.

Terminology

Summary

Summary

The paper introduces Latent Reward Registers, a mechanism designed to address the temporal credit-assignment challenge in aligning diffusion models with human preferences. Traditional alignment methods rely on a sparse terminal reward evaluated on the final generated samples, which presents a severe credit-assignment problem across the multi-step denoising process. The proposed method 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 central insight is that "DiT internal representations already encode sufficient evidence about the final output quality, well before the trajectory reaches a fully denoised state. Even under heavy noise, these features capture emerging global structure, prompt alignment, and overall trajectory viability. The method turns a frozen DiT into an accurate, step-wise reward model by adding a small set of learnable, position-free register tokens as global readouts to aggregate reward-relevant evidence from the intermediate representations of the frozen transformer. A lightweight readout head predicts the expected terminal preference from any intermediate latent, and differentiating this prediction with respect to the latent yields a local reward gradient at every denoising step. The registers do not modify the generator’s hidden states or velocity predictions, so the pretrained generative dynamics remain fully intact."

The dense reward signal supports 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, allowing multiple objectives to be composed at test time.

The paper reports three main contributions. First, Reliable latent reward signals: "Latent Reward Registers estimate step-wise rewards directly from intermediate noisy latents. Even at high noise levels, their pairwise ranking accuracy remains competitive with external reward models that require fully denoised images." Second, Efficient training-time alignment: By distilling register-guided on-policy trajectories, RG-OPD surpasses online RL baselines in alignment quality while reducing GPU hours by up to 33×. Third, Effective inference-time alignment: "RGS steers trajectories using the derived latent reward gradients without parameter updates, improving target rewards while maintaining a superior reward–quality trade-off relative to existing training-free sampling methods."

The method section details the three core components. The Reward Register Mechanism initializes K learnable, position-free register tokens as global readouts. For the initial L transformer blocks, the registers share a single trainable query projection while reusing the frozen key, value, output, normalization, and noise-level-gating modules from each block. The update is computed via attention, with the registers bypassing the block feed-forward networks, ensuring the original denoiser hidden states and velocity predictions stay strictly invariant. The register states are then fused with pooled frozen features from selected blocks through a lightweight attention module. Each reward head is trained with a noise-aware Thurstone pairwise-ranking objective together with the DiNa-LRM variance adjustment, which penalizes disagreement between the signs of the score differences and assumes comparison variance grows with the sampled noise level. The learned score preserves rankings induced by endpoint reward models rather than estimating a calibrated reward value.

RG-OPD generates step-wise supervision at states actively visited by the current student model. From a student rollout, the detached state is evaluated with the frozen reference generator to obtain a reference mean and displacement. The shared reward direction is evaluated at the detached state, and the reward-tilted teacher mean is computed as the reference mean plus a scaled reward correction. The student is optimized with a mean squared error between its one-step mean and the detached teacher target, with the reference generator and reward register remaining frozen.

RGS uses the original sampler as a baseline reference and applies a reward-gradient correction after each active solver step. The trajectory advances according to a formula that adds a magnitude-matched reward correction to the reference mean, with a three-band schedule for the guidance strength: higher strength for early steps, lower for mid steps, and zero for the low-noise tail.

Experiments evaluate the framework on three axes. For latent reward prediction, the reward register with K=32 tokens on a frozen SD3-Medium backbone is tested on four benchmarks totaling 54,170 pairs. At high noise levels (u=0.8), it outperforms all latent baselines on three of the four benchmarks and remains competitive on the fourth. For training-time alignment, RG-OPD attains the highest HPSv3 scores on both architectures while keeping a balanced multi-reward profile, outperforming reward-backpropagation and online RL baselines. For inference-time alignment, RGS reaches the highest HPSv3 score among all evaluated training-free methods, with a favorable reward–quality trade-off under ImageReward optimization.

Ablation studies show that maintaining a persistent register state is essential for accurate step-wise reward estimation, with removing the register tokens causing the most severe performance degradation. Training efficiency analysis shows RG-OPD achieves Flow-GRPO-equivalent HPSv3 scores 14.0×–33.2× faster in GPU hours across both backbones. Multi-head RGS effectively trades off conflicting objectives, preserving both reward gains and perceptual quality. The register mechanism demonstrates stability at high noise levels, maintaining reliable prediction capabilities across the entire timestep range, unlike baselines whose accuracy drops sharply when latent inputs become predominantly noise.

The paper concludes that intermediate DiT representations capture terminal reward information early in the denoising process, and the proposed framework establishes a reliable latent reward model for diffusion RL by resolving the credit assignment problem.

Improvements for AI systems

Based on the paper, here are the specific improvements I can implement in AI systems:

Improvement: Replace sparse terminal-only reward evaluation with a non-intrusive latent reward register mechanism that predicts expected terminal preference from intermediate noisy latents.

Implementation:

  • Append 32 learnable, position-free register tokens to a frozen Diffusion Transformer (DiT) backbone

  • Train a lightweight readout head with pairwise ranking loss (Thurstone model with noise-aware variance adjustment)

  • Use an exponential moving average of parameters for stable evaluation

Capability: The system can now assign reward scores at any denoising step (u=0.2 to u=0.8) with 63-78% pairwise accuracy, matching or exceeding endpoint reward models that require fully denoised images.

Sources

Related papers