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

Preference optimization for diffusion and flow-matching models relies on reward functions that are both discriminatively robust and computationally efficient. Vision-Language Models (VLMs) have emerged as the primary reward provider, leveraging their rich multimodal priors to guide alignment. 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 performance competitive with state-of-the-art VLMs at a fraction of the computational cost. In preference optimization, we demonstrate that DiNa-LRM improves preference optimization dynamics, enabling faster and more resource-efficient model alignment.

💡 Research Summary

The paper tackles the growing need to align diffusion‑based image and video generators with human preferences while avoiding the high computational and memory overhead of Vision‑Language Model (VLM) rewards. Existing approaches typically use large VLMs (e.g., CLIP, BLIP, or newer multimodal transformers) as reward functions. Although VLMs provide strong semantic understanding, they operate in pixel space, incur substantial inference cost, and create a domain mismatch when the generator works in a latent diffusion space. To address these issues, the authors propose DiNa‑LRM, a diffusion‑native latent reward model that learns preferences directly on the noisy latent states produced by diffusion models.

Core technical contribution
The authors extend the classic Thurstone preference model, which treats human judgments as noisy observations of an underlying scalar reward, by calibrating the observation noise to the diffusion timestep. In the clean‑image case the model assumes u(x)=rθ(x)+η with η∼N(0,σu²). Preference probability is Φ((rθ(x⁺)−rθ(x⁻))/√(2σu²)). For diffusion, the input is a noisy latent x_t = α(t)x₀ + σ(t)ε. The authors define a timestep‑dependent observation variance σu²(t)=k·σ²(t)+σu², where σ(t) is the diffusion schedule’s standard deviation. This “noise‑calibrated Thurstone” formulation acknowledges that as t grows, semantic information degrades and human judgments become more uncertain. The resulting likelihood P(x_t⁺≻x_t⁻|t)=Φ((rθ(x_t⁺,t)−rθ(x_t⁻,t))/√(2σu²(t))) is used as the training objective.

Training uses the fidelity loss (binary cross‑entropy on the calibrated preference probability) and samples timesteps from a distribution q(t). The authors experiment with fixed, uniform, and logit‑normal sampling; uniform sampling provides the most robust performance.

Architecture
DiNa‑LRM builds on a pretrained latent diffusion backbone (SD3.5‑Medium in most experiments) with a frozen VAE encoder. From a chosen subset of backbone layers, both visual and textual token embeddings are extracted. Each layer’s features are modulated by the timestep embedding via FiLM (Feature‑wise Linear Modulation), then projected to a lower dimension. The resulting visual tokens V_t and textual tokens T_t are concatenated and fed into a query‑transformer (Q‑Former). N_q learnable query tokens attend to V_t and T_t through value‑gated cross‑attention, followed by a visual‑only cross‑attention refinement. The pooled query representation is passed through a lightweight MLP to produce a scalar reward rθ(x_t, t, c). Because the Q‑Former processes a sequence of tokens, features from multiple timesteps can be concatenated, enabling “noise‑ensemble” inference where predictions from several diffusion steps are aggregated for a more robust score.

Experiments
The authors evaluate DiNa‑LRM on standard image‑text alignment benchmarks (COCO Captions, Flickr30k) and on human‑preference datasets. Compared to diffusion‑based baselines such as PickScore‑Diffusion, DiNa‑LRM achieves substantially higher Recall@1, R‑Precision, and preference‑agreement scores. When pitted against strong VLM rewards (CLIP‑Score, BLIP‑Score), DiNa‑LRM matches or slightly exceeds their performance while using roughly 30‑50 % fewer FLOPs and far less GPU memory. In preference‑optimization tasks, DiNa‑LRM is used as the reward for Direct Preference Optimization (DPO) and for the RL‑based GRPO algorithm. Both methods converge faster and reach higher human‑preference percentages than when using VLM rewards, demonstrating that the noise‑aware formulation yields smoother gradients, especially at high‑noise timesteps.

Ablation studies explore (i) the impact of different timestep sampling strategies, (ii) the effect of the scaling factor k in σu²(t), and (iii) the number of query tokens and selected backbone layers. Uniform timestep sampling and k = 2 consistently provide the best trade‑off between stability and accuracy. The authors also show that the Q‑Former’s query count can be reduced with minimal performance loss, further decreasing inference cost.

Implications
DiNa‑LRM proves that the discriminative power of large diffusion models, learned during their generative pre‑training, can be repurposed as a high‑quality, low‑cost reward function. By staying within the latent diffusion space and explicitly modeling uncertainty as a function of diffusion noise, the method eliminates the pixel‑space/domain mismatch that plagues VLM‑based rewards. The noise‑ensemble inference offers a simple, scalable knob for test‑time robustness, and the overall framework can be plugged into any diffusion or flow‑matching generator without architectural changes.

Conclusion
The paper introduces a principled, diffusion‑native reward modeling approach that bridges the gap between generative diffusion models and preference alignment. Through a noise‑calibrated Thurstone likelihood, a timestep‑conditioned Q‑Former architecture, and efficient inference‑time ensembling, DiNa‑LRM delivers VLM‑level alignment quality at a fraction of the computational expense. This work opens the door to more resource‑efficient preference‑guided generation and suggests that future alignment pipelines may increasingly rely on the latent representations of the generators themselves rather than external, heavyweight multimodal models.