🤖 AI

Beyond VLM-Based Rewards: Diffusion-Native Latent Reward Modeling

This paper introduces DiNa-LRM, a computationally efficient, diffusion-native latent reward model that formulates preference learning directly on noisy diffusion states to overcome the domain mismatch and high costs of Vision-Language Model-based rewards, achieving superior alignment performance with significantly lower resource requirements.

Original authors: Gongye Liu, Bo Yang, Yida Zhi, Zhizhou Zhong, Lei Ke, Didan Deng, Han Gao, Yongxiang Huang, Kaihao Zhang, Hongbo Fu, Wenhan Luo

Published 2026-05-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Gongye Liu, Bo Yang, Yida Zhi, Zhizhou Zhong, Lei Ke, Didan Deng, Han Gao, Yongxiang Huang, Kaihao Zhang, Hongbo Fu, Wenhan Luo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot artist to paint pictures based on your descriptions. The robot uses a complex internal system (called a Diffusion Model) that starts with a blurry, noisy static and slowly cleans it up until a clear image appears.

To teach this robot to paint what you actually like, you need a "Teacher" (a Reward Model) to look at the robot's drafts and say, "Good job!" or "Try again."

The Problem: The Old Teacher Was Too Heavy and Out of Place

Currently, the best "Teachers" are massive Vision-Language Models (VLMs). Think of these as super-intelligent, giant libraries that have read every book and seen every picture in the world.

The Issue: These giant libraries are heavy. They take a lot of computer power to run, and they are slow.
The Mismatch: The robot artist works in a "compressed, abstract world" (Latent Space), but the giant library teacher looks at the final, high-definition photo (Pixel Space). It's like asking a chef to judge a soup by tasting the raw ingredients after they've been cooked, rather than tasting the soup while it's simmering. This mismatch makes teaching inefficient and confusing.

The Solution: DiNa-LRM (The Native Teacher)

The authors of this paper created a new kind of teacher called DiNa-LRM. Instead of hiring a giant, external library, they turned the robot artist's own internal brain into the teacher.

Here is how it works, using simple analogies:

1. The "Noise-Calibrated" Eye

The robot artist creates images by starting with noise (static) and cleaning it up over time.

Old Way: The teacher only looks at the final, clean image.
DiNa-LRM Way: This new teacher is comfortable looking at the image while it is still noisy and blurry.
The Analogy: Imagine a music producer teaching a band. A traditional teacher waits until the song is fully mixed and mastered to critique it. DiNa-LRM is like a producer who sits in the studio during the recording, listening to the raw, noisy tracks. Because the teacher understands the "noise" naturally, it knows exactly how much uncertainty to expect at each stage. If the image is very blurry, the teacher says, "I'm not 100% sure yet, so I'll be a bit more cautious in my grading." If the image is clear, the teacher grades it with high confidence.

2. The "Time-Travel" Ensembling

When the teacher needs to give a final score to a finished image, it doesn't just look at it once.

The Analogy: Imagine you are trying to guess the plot of a movie. You could watch just the final scene, but that might be misleading. Or, you could watch the movie at 10%, 50%, and 90% completion and combine those views to get a better understanding.
How DiNa-LRM does it: It looks at the image at several different stages of "cleanliness" (different noise levels) and combines all those opinions into one final, super-accurate score. This is called Noise Ensembling. It makes the teacher much more robust and reliable without needing a bigger brain.

3. The Results: Faster, Cheaper, and Smarter

The paper tested this new teacher against the old giant libraries and other attempts.

Performance: DiNa-LRM is almost as good as the best giant libraries at judging which pictures humans prefer.
Efficiency: Because it lives inside the robot's own "compressed world," it doesn't need to decode the image into a full photo to judge it.
- It uses 51% less memory (like needing half the RAM on your computer).
- It uses 71% less computing power for the actual judging.
Training: When used to train the robot artist, the robot learns faster and doesn't get confused by the mismatch between the teacher's view and the artist's view.

Summary

The paper introduces DiNa-LRM, a reward model that speaks the same "language" as the image generator it is teaching. Instead of forcing a giant, external expert to translate its thoughts into the generator's language, DiNa-LRM uses the generator's own internal structure to grade its work. It is smarter about handling uncertainty (noise), can look at a work-in-progress to make better final judgments, and does all of this while using significantly less computer power.

Technical Summary: DiNa-LRM – Diffusion-Native Latent Reward Modeling

1. Problem Statement

The alignment of diffusion and flow-matching models with human preferences relies heavily on reward functions. Current state-of-the-art approaches predominantly utilize Vision-Language Models (VLMs) as reward backbones. While VLMs offer robust discriminative capabilities, they present two critical limitations for diffusion model alignment:

Computational and Memory Overhead: VLM-based rewards are expensive to query repeatedly during optimization, creating a bottleneck for training efficiency.
Domain Mismatch: VLMs typically operate in pixel space, whereas latent diffusion generators are trained and optimized in latent space. This pixel-to-latent mismatch complicates alignment, particularly for reward-gradient methods, and introduces unnecessary system overhead.

Conversely, while diffusion models have shown promise as discriminative backbones, prior attempts to use them as rewards have often been tied to specific step-level optimization algorithms or focused on noisy intermediate states without establishing a general-purpose framework comparable to VLMs for scoring clean samples.

2. Methodology: DiNa-LRM

The authors propose DiNa-LRM (Diffusion-Native Latent Reward Model), a framework that formulates preference learning directly on noisy diffusion states within the latent space. The method consists of three core components:

A. Diffusion-Native Preference Formulation

Instead of evaluating clean samples ( $x_0$ ) directly, DiNa-LRM extends the Thurstone model to noisy diffusion states ( $x_t$ ).

Noise-Calibrated Uncertainty: The authors recognize that preference judgments on noisy states are inherently less certain as the noise level increases. They introduce a noise-calibrated comparison uncertainty $\sigma_u^2(t)$ that scales with the diffusion noise level $\sigma^2(t)$ :
$\sigma_u^2(t) = k\sigma^2(t) + \sigma_u^2$
This formulation ensures the reward objective remains consistent with the noisy manifold encountered during diffusion pre-training. It acts as an uncertainty-aware regularizer, preventing high-variance gradients from destabilizing training in high-noise regimes.
Training Objective: The model is optimized using a fidelity loss over preference pairs $(x_t^+, x_t^-)$ , where the comparison probability is modulated by the timestep-dependent uncertainty.

B. Latent Reward Architecture

DiNa-LRM is instantiated on top of a pretrained latent diffusion backbone (e.g., SD3.5-Medium) with a lightweight, query-based scoring head.

Multi-layer Feature Extraction: Visual and textual features are extracted from multiple layers of the backbone.
Timestep-Conditioned Adaptation: Features are modulated by timestep embeddings using FiLM (Feature-wise Linear Modulation) to make the reward head explicitly aware of the noise level.
Gated Q-Former Scoring: A gated Q-Former aggregates these multi-layer, timestep-conditioned features to produce a scalar reward score. This architecture operates entirely in the VAE latent space, avoiding pixel-space decoding.

C. Inference-Time Scaling via Noise Ensembling

To improve robustness and reduce sensitivity to a specific evaluation timestep, DiNa-LRM supports noise ensembling.

Instead of scoring a clean sample at a single fixed noise level, the model evaluates the sample at multiple timesteps $\{t_k\}$ .
Features extracted at these different timesteps are concatenated and processed by the same Q-Former head to produce an aggregated score. This serves as a "diffusion-native" test-time scaling mechanism.

3. Key Contributions

Diffusion-Native Preference Formulation: The paper extends the Thurstone preference model to noisy diffusion states by introducing a noise-calibrated uncertainty that scales with the diffusion noise level, bridging the distributional gap between pretraining and reward learning.
Inference-Time Scaling: The authors propose a noise ensembling strategy that aggregates multi-timestep features, providing a mechanism for test-time scaling that leverages the global diffusion trajectory rather than a single point estimate.
Empirical Validation: The work demonstrates that diffusion-native rewards can achieve competitive performance with strong VLM-based judges while operating natively in latent space.

4. Experimental Results

The authors evaluated DiNa-LRM across several benchmarks and alignment scenarios:

Reward Modeling Accuracy:
- DiNa-LRM substantially outperforms existing diffusion-based reward baselines (e.g., LRM-SD1.5, LRM-SDXL) across ImageReward, HPDv2, HPDv3, and GenAI Bench.
- It achieves performance competitive with state-of-the-art VLM-based rewards (e.g., HPSv3, UnifiedReward) on average, with a notable gap narrowing on specific datasets like HPDv2 and HPDv3.
- Noise Ensembling: Applying noise ensembling further improves average accuracy (e.g., from 71.49% to 72.48% on average), validating the benefit of aggregating evidence across noise levels.
Preference Alignment (ReFL):
- In post-training alignment experiments using ReFL (Reward Feedback Learning) on SD3.5-M, DiNa-LRM demonstrated faster convergence in proxy scores and improved "golden" metrics (held-out PickScore) compared to HPSv3.
- Efficiency: DiNa-LRM significantly reduces resource costs compared to VLM-based rewards. It reduces peak VRAM usage by 51.4%, reward-calculation FLOPs by 71.1%, and optimization-phase FLOPs by 46.4%.
Generalization:
- The formulation generalizes across different diffusion backbones (SDXL, FLUX.1-Dev, Z-Image-Turbo), consistently outperforming prior diffusion-based baselines.
- Transferability experiments show a reward model trained on SD3.5-M can effectively align SD3.5-L within the same latent space.

5. Significance and Claims

The paper claims that DiNa-LRM offers a practical and efficient alternative to pixel-space VLM rewards for preference alignment. Its primary significance lies in:

Structural Advantage: By operating natively in the latent space, it eliminates the latent-to-pixel mismatch and the computational overhead of decoding, making it highly compatible with reward-gradient alignment methods.
Optimization-Friendly: It provides a differentiable, efficient signal that improves training dynamics and reduces resource requirements without sacrificing alignment quality.
Scalability: The noise ensembling mechanism provides a novel, diffusion-native way to scale inference-time performance.

The authors acknowledge limitations, noting that the reward is tied to specific latent spaces (limiting cross-backbone transfer without retraining) and that latent-space modeling may under-emphasize certain pixel-level artifacts, potentially leading to reward hacking behaviors (e.g., spurious human injection or stylistic drift) during long-horizon optimization. They suggest future work should focus on stronger unified backbones and incorporating lightweight pixel-space regularization.