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

AI summary

Overview

Research area: Computer Vision — reward modeling and preference alignment for generative image models (diffusion and flow-matching).

Technical level: Intermediate to Advanced. The paper assumes familiarity with diffusion/flow-matching training objectives, latent-space (VAE) generation, preference optimization algorithms (ReFL, DPO, GRPO), and reward-model training.

Scope: The paper introduces DiNa-LRM, a latent-space reward model built on a pretrained diffusion backbone that scores noisy diffusion states rather than decoded pixels, and evaluates it on preference benchmarks and reward-gradient alignment.

What This Paper Is About

Preference alignment for diffusion and flow-matching generators depends on a reward model that must be both accurate and cheap to query repeatedly during optimization. Current practice leans on vision-language models (VLMs), which are accurate but expensive and operate in pixel space, while the generator being optimized lives in VAE latent space — creating a domain mismatch. This paper asks whether a pretrained latent diffusion backbone can itself serve as a general-purpose reward model that scores clean samples in latent space, remaining competitive with VLM judges while being far more optimization-friendly.

Key Contributions

1. Diffusion-native preference formulation. The authors extend the Thurstone preference model from clean samples to noisy diffusion states by introducing a comparison uncertainty that grows with the diffusion noise level, calibrating the confidence of preference judgments at each timestep.

2. Inference-time scaling via noise ensembling. They build a timestep-conditioned latent reward model on a pretrained latent diffusion backbone and propose aggregating features across multiple noise levels within the reward head, giving a diffusion-native test-time scaling knob.

3. A latent reward architecture. A query-based (Q-Former-style) scoring head with FiLM timestep modulation operates entirely in frozen VAE latent space, fusing multi-layer visual and text features into a single scalar reward.

4. Empirical validation of alignment behavior. They demonstrate that diffusion-native rewards substantially beat prior diffusion-based reward baselines, narrow the gap to strong VLM rewards on benchmarks, and improve optimization dynamics while cutting memory and compute.

Main Findings

• Diffusion rewards beat prior diffusion rewards by a wide margin. DiNa-LRM reaches 71.49% average pairwise accuracy (72.48% with ensembling) versus 61.62% for LRM-SD1.5 and 61.73% for LRM-SDXL, while staying slightly behind the strongest VLM judge, HPSv3 (74.84%).

• Noise ensembling gives consistent but modest gains. Aggregating rewards at three noise levels (t ∈ {0.2, 0.5, 0.7}) improves average accuracy from 71.49% to 72.48%, with the largest benefits on out-of-domain benchmarks, supporting the idea that different timesteps expose complementary discriminative evidence.

• Distributional timestep schedules matter for generalization. Fixed-single-noise training performs reasonably in-domain (HPDv3) but generalizes poorly elsewhere (avg 64.93% at t=0, up to 68.75% at t=0.7). Uniform sampling over t ∈ (0,1) yields the strongest and most stable default, outperforming logit-normal schedules in their setting.

• Noise-calibrated variance improves both accuracy and robustness. Replacing a constant-variance Thurstone term with noise-calibrated variance lifts HPDv2 accuracy from 78.72% to 82.13% (single inference) and from 78.16% to 84.31% (with ensembling), and raises the ensembled average from 70.41% to 72.48%.

• Intermediate inference noise levels are optimal. Accuracy peaks for t ∈ [0.3, 0.7], with t = 0.4 chosen as the default single-noise setting; near-clean (t = 0) and high-noise (t = 0.8) evaluations both degrade accuracy.

• Lightweight backbone adaptation helps. LoRA fine-tuning of the diffusion backbone outperforms freezing it (71.49% versus 70.27% average), while the frozen variant remains competitive, indicating that pretrained diffusion representations already carry useful preference signal.

• Alignment becomes faster and cheaper. In ReFL optimization on SD3.5-Medium, DiNa-LRM improves the proxy score faster than HPSv3 while the held-out PickScore golden metric rises in tandem, showing no early reward hacking. It reduces peak VRAM by 51.4%, reward-calculation FLOPs by 71.1%, and optimization-phase FLOPs by 46.4% relative to HPSv3 at 1024×1024 resolution.

• The formulation generalizes across backbones. Additional results on SDXL, FLUX.1-Dev, and Z-Image-Turbo show the approach is not specific to SD3.5-Medium, and transfer from an SD3.5-M reward model to SD3.5-L alignment works within a shared latent space.

Methodology in Plain English

The authors start from the observation that diffusion backbones are pretrained to handle noisy inputs, yet reward models are typically asked to judge clean images. Instead of fighting that mismatch, they deliberately train the reward model on noisy latents — versions of the image that have had varying amounts of Gaussian noise added.

They formalize preference as a Thurstone comparison: each sample's perceived quality is a deterministic score plus random judgment noise, and the probability that sample A beats sample B depends on the score difference divided by the judgment uncertainty. Their key twist is making that uncertainty grow with the diffusion noise level, on the intuition that heavily noised images carry less semantic information, so preference judgments on them should be treated as less confident. This prevents uninformative, high-variance gradients early in training.

The reward model itself reuses a pretrained latent diffusion backbone (SD3.5-Medium in most experiments). Multi-layer visual and text features are extracted, modulated by the timestep embedding via FiLM, compressed, and passed to a Q-Former-style head with learnable query tokens. The queries attend to visual and text tokens, refine through a visual-only attention block, and a small MLP produces a scalar reward — all inside the frozen VAE latent space.

Training uses a fidelity loss over preference pairs, with the chosen and rejected samples perturbed by the same noise so the comparison is fair. At inference, a clean sample is scored either at a single low-noise timestep (t = 0.4 by default) or by extracting features at multiple timesteps and concatenating them into a longer token sequence for the same head — a form of test-time scaling. For alignment, the authors plug the reward into ReFL, a reward-gradient method, backpropagating through the reward directly into the generator.

Why This Matters

Impact on research. The work reframes reward modeling as a discriminative use of diffusion pretraining rather than a task owned by VLMs. It shows that latent diffusion backbones — already the generator being aligned — can supply the supervision signal, eliminating the latent-to-pixel round trip that complicates reward-gradient methods. This opens a direction where reward modeling and generation share representations, and where noise level becomes a tunable inference-time compute knob analogous to test-time scaling in language models.

Real-world applications:

• Post-training pipelines for image and video generators. Studios and labs aligning diffusion models can reduce GPU memory and compute costs during preference optimization, making large-scale alignment more accessible.
• Automated content evaluation. A fast latent-space scorer can rank or filter large batches of generated images without decoding them, useful for dataset curation and moderation workflows.
• Iterative creative tools. Interactive interfaces that repeatedly score and refine outputs benefit from a reward model that runs natively in the generator's latent space.
• Cost-sensitive deployment. Reducing reward FLOPs by roughly 71% lowers the barrier for smaller teams and for on-premise alignment runs.

Industry relevance. The efficiency numbers (halved peak VRAM, large FLOP reductions) target a real bottleneck in production alignment loops, where reward evaluation is queried thousands of times per run. Because the approach reuses an existing generator backbone, it integrates cleanly into pipelines that already have a diffusion model checkpoint, avoiding a separate large VLM serving stack. The work is a collaboration between HKUST, Huawei's Hong Kong AI lab, Tsinghua, and ANU, with code released publicly.

Future Directions

• Cross-backbone generality. The reward is learned in one backbone's latent space; scoring or transferring across different VAEs is not guaranteed. Training on stronger, more unified backbones is a natural next step.
• Latent-invisible artifacts. The authors note that latent-space rewards may under-emphasize pixel-level artifacts, and that long-horizon optimization can produce spurious object insertion or stylistic drift. Adding lightweight pixel-space regularization or perceptual constraints is a proposed remedy.
• Denser reward signals. Moving from a scalar score toward generative or feedback-rich reward modeling that produces structured, dense supervision.
• Broader alignment algorithms. The paper validates ReFL and mentions Flow-GRPO-Fast online RL; extending the noise-ensembling and calibration ideas to other on-policy and offline preference methods remains open.

Target Audience

Researchers and engineers working on preference alignment for generative vision models, particularly those building reward models or running reward-gradient post-training. The paper is also relevant to practitioners optimizing the cost of alignment pipelines, and to readers interested in how pretrained diffusion representations can be repurposed for discriminative tasks. A background in diffusion training objectives and preference optimization will make the method sections substantially easier to follow.

Authors’ abstract

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.