Skip to content

Diffusion Model as a Noise-Aware Latent Reward Model for Step-Level Preference Optimization

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

This paper proposes repurposing a pre-trained diffusion model itself as a step-level reward model (LRM) that operates directly in the noisy latent space, avoiding the distribution shift and complex pixel-space transformations that plague VLM-based reward models. Built on LRM, Latent Preference Optimization (LPO) performs DPO-style alignment entirely within the denoising trajectory, achieving state-of-the-art image quality with 2.5–28× training speedup over prior methods.

Problem

Existing preference optimization for diffusion models relies on VLMs (CLIP, BLIP) as pixel-level reward models (PRMs). When applied step-by-step during denoising, PRMs suffer from three compounding issues: (1) they require full denoising and VAE decoding at each timestep to produce a clean image for evaluation; (2) highly noisy latent images at large timesteps create severe distribution shift from VLM training data (clean images), making reward predictions unreliable; and (3) PRMs lack explicit timestep conditioning, making them insensitive to the varying role of noise level during generation.

Method

Latent Reward Model (LRM): The frozen U-Net backbone and text encoder of the diffusion model being aligned are repurposed for reward prediction. Given a noisy latent \(x_t\) and prompt \(p\), U-Net features from down-blocks and mid-block are average-pooled into multi-scale visual representations. A Visual Feature Enhancement (VFE) module improves text-image alignment sensitivity by injecting classifier-free guidance contrast: \(V_{enh} = V_{mid} + (g_s - 1)(V_{mid} - V_{mid\_ucond})\), amplifying text-conditional signal. A dot-product score \(S(p, x_t) = \tau \cdot \ell_2(V) \cdot \ell_2(T)\) is used with Bradley-Terry loss for training on preference pairs.

Multi-Preference Consistent Filtering (MPCF): Pick-a-Pic v1 pairs where the "winner" underperforms in CLIP score or VQAScore (≈40% of pairs) are filtered out, retaining only pairs where winners dominate on text-image alignment while allowing a small aesthetic gap (\(G_A \geq -0.5\)). This keeps 168k of 511k win-lose pairs.

Latent Preference Optimization (LPO): At each denoising step \(t\), \(K=4\) candidates \(\{x_t^i\}\) are sampled from \(x_{t+1}\). LRM scores them directly in latent space—no VAE decode needed—and the highest/lowest-scoring pair becomes the \((x_t^w, x_t^l)\) training signal for the SPO-style DPO loss. A dynamic threshold \(th_t = \frac{\sigma_t - \sigma_{min}}{\sigma_{max} - \sigma_{min}}(th_{max} - th_{min}) + th_{min}\) adapts the selection gap to noise level, enabling reliable sampling across the full range \(t \in [0, 950]\).

Key Contributions

  • LRM: First use of the diffusion model itself as a step-level reward model, natively handling noisy latents at arbitrary timesteps without pixel-space transformation.
  • VFE module: Classifier-free guidance–inspired enhancement that injects text-conditional signal into visual features for better text-image alignment reward.
  • MPCF: Automated multi-criterion filtering of preference datasets to remove inconsistent pairs that would corrupt latent-space reward training.
  • LPO: End-to-end step-level preference optimization in noisy latent space with dynamic thresholding, covering the full denoising range.
  • Generalization demonstrations: LRM applied to GRPO variant, DiT-based SD3, and cross-architecture (heterogeneous) fine-tuning where LRM-1.5 trains SD2.1.

Results

  • SD1.5 general preference (PickScore): LPO 21.69 vs. SPO 21.22, Diffusion-DPO 20.99, original 20.56.
  • SD1.5 ImageReward: LPO 0.6588 vs. SPO 0.1678, SePPO 0.5077.
  • SD1.5 HPSv2.1: LPO 27.86 vs. SPO 25.83, SePPO 27.34.
  • SD1.5 GenEval (20-step): LPO 48.39 vs. SPO 48.39 (tied), Diffusion-DPO 44.97; LPO matches SPO while training 3.5× faster.
  • SDXL: LPO (PickScore 22.86, HPSv2.1 31.89, GenEval 59.85) slightly surpasses InterComp (22.63, 31.52, 59.24), which uses an internal SD3/FLUX dataset.
  • T2I-CompBench++ (SD1.5, Color): LPO 0.5042 vs. SPO 0.4112, SePPO 0.4265.
  • Training speed (SD1.5): LPO 23 A100-hours vs. Diffusion-DPO 240h (10×), SPO 80h (3.5×).
  • Training speed (SDXL): LPO 92h vs. Diffusion-DPO 2,560h (28×), SPO 234h (2.5×).
  • Sampling step time: LRM 0.039s vs. SPM 0.243s per step (6.2× faster at K=4).
  • Ablation—timestep range: \(t \in [750, 950]\) alone achieves GenEval 48.28, nearly matching full-range 48.39; low-noise-only \([0, 200]\) achieves only 40.11.
  • Heterogeneous: LRM-1.5 fine-tuning SD2.1 improves GenEval from 48.59→56.01; fails on SDXL (different VAE).

Limitations

  • Training data (Pick-a-Pic v1) consists of SD1.5/SDXL-generated images with relatively low quality, limiting LRM generalization to higher-fidelity domains.
  • MPCF relies on three automatic reward models (Aesthetic Score, CLIP Score, VQAScore) whose shared biases are inherited by the filtered dataset.
  • LRM and the diffusion model being optimized (DMO) must share the same VAE encoder; cross-family fine-tuning (e.g., LRM-1.5 → SDXL) fails due to latent space mismatch.
  • No formal analysis of reward hacking risk when LRM and DMO are initialized from the same checkpoint (homogeneous optimization).

Relevance to Vision-Language Models

VLMs (CLIP, BLIP) have been the dominant reward signal for aligning text-to-image diffusion models, but this paper directly challenges their suitability for step-level optimization by exposing the distribution shift they suffer at high noise levels. The work reframes reward modeling away from VLMs toward using the generative model itself as an implicit multimodal scorer, which is a conceptually significant departure. For VLM researchers, it raises concrete questions about where CLIP-style models break down and whether VLM-derived reward signals should be supplanted or hybridized for noisy-latent objectives. The MPCF filtering strategy is also broadly applicable to any multi-dimensional preference dataset used to train VLM-based reward models.