Beyond VLM-Based Rewards: Diffusion-Native Latent Reward Modeling
Abstract
Preference optimization for diffusion models relies on reward functions that are both discriminative and computationally efficient. Vision-Language Models (VLMs) have emerged as powerful reward providers. However, their computation and memory cost can be substantial, and optimizing a latent diffusion generator through a pixel-space reward introduces a domain mismatch that complicates alignment. In this paper, we propose DiNa-LRM, a diffusion-native latent reward model that formulates preference learning directly on noisy diffusion states. Our method introduces a noise-calibrated Thurstone likelihood with diffusion-noise-dependent uncertainty. DiNa-LRM leverages a pretrained latent diffusion backbone with a timestep-conditioned reward head, and supports inference-time noise ensembling, providing a diffusion-native mechanism for test-time scaling and robust rewarding. Across image alignment benchmarks, DiNa-LRM substantially outperforms existing diffusion-based reward baselines and achieves competitive performance compared to state-of-the-art VLMs while maintaining a substantially lower computational cost. In preference optimization, we demonstrate that DiNa-LRM improves preference optimization dynamics, enabling faster and more resource-efficient model alignment.
Lay Summary
Modern image generation models can create impressive pictures, but improving them to better match human preferences still requires reliable feedback. Today, this feedback often comes from large vision-language models that judge the final image, which can be slow, memory-intensive, and misaligned with how diffusion generators internally work. In this paper, we ask whether a diffusion model itself can be turned into a better feedback provider. We propose DiNa-LRM, a reward model that evaluates generated images in the same latent space used by diffusion models. Instead of only judging clean final images, it learns from images with different levels of noise, much like the states a diffusion model sees during generation. The model also accounts for the fact that noisier images are harder to judge reliably, and it can combine evidence from multiple noise levels at test time. Our experiments show that this diffusion-native reward model is much stronger than previous diffusion-based rewards and approaches strong vision-language rewards. More importantly, it makes preference-based training faster and cheaper, suggesting a practical route toward more efficient alignment of visual generation models.