Technical Summary: Sample-Adaptive Latent Rewards for Uncertainty-Guided Diffusion Post-Training

1. Problem Statement

Post-training visual diffusion models to align with human preferences typically relies on evaluating decoded RGB outputs, which necessitates repeated and costly VAE decoding within the feedback loop. While "diffusion-native" latent reward models (LRMs) have emerged to score noisy latents directly—avoiding decoding costs and enabling intermediate-state supervision—existing approaches suffer from two critical limitations:

1. Scalar Outputs Lack Uncertainty: Current latent reward models output only a scalar score. They fail to estimate the uncertainty of each prediction. Consequently, the generator cannot distinguish between reliable and unreliable feedback.
2. Static Uncertainty Assumptions: Prior work like DiNa-LRM assigns uncertainty based solely on the noise level (timestep) of the latent. It assumes all samples at the same timestep share the same reliability. However, the paper argues that reward difficulty is also content-dependent; a "noisy" sample with a clear subject may be more predictable than a "cleaner" sample with ambiguous content.
3. Optimization Risks: Without uncertainty estimation, the generator may optimize based on unreliable local feedback, leading to "reward hacking" (where the model exploits the reward function rather than improving quality) and unstable training dynamics.

2. Methodology

The authors propose SURE (Sample-Adaptive Latent Reward for Uncertainty-Guided Diffusion Post-Training), a unified framework consisting of two components: SURE-LRM (the reward model) and SURE-REFL (the post-training algorithm).

2.1 SURE-LRM: Sample-Adaptive Latent Reward Model

SURE-LRM is designed to predict not just a reward score, but a full Gaussian utility distribution for each noisy latent, conditioned on the prompt and timestep.

• Architecture: It reuses the Diffusion Transformer (DiT) backbone of the target generator (for both images and videos) to process noisy latents, prompts, and timesteps without VAE decoding.
• Prediction Targets: For a noisy latent $z_\sigma$, the model predicts:
  • A mean reward $\mu_\theta$.
  • A sample-adaptive residual variance parameter $a_\theta$.
• Variance Modeling: The total predictive variance $v_\eta(\sigma)$ is the sum of a noise-dependent prior $v_0(\sigma)$ (capturing average information loss as noise increases) and a learned, sample-specific residual variance $v_{\eta, res} = \text{softplus}(a_\theta) + \epsilon_v$.
• Training Objective: The model is trained on standard pairwise preferences ($z^+ \succ z^-$) using a heteroscedastic preference objective:
  • Likelihood Loss ($L_{pref}$): Based on the Thurstone random-utility model, it optimizes the probability that the preferred sample has higher utility. The gradients naturally encourage the model to increase variance (uncertainty) when the ranking is incorrect or ambiguous, and decrease variance when the ranking is confident and correct.
  • Regularization:
    • Variance Regularization ($L_{var}$): Prevents the model from explaining all ranking errors by simply inflating uncertainty (variance inflation).
    • Margin Regularization ($L_{margin}$): Ensures a minimum separation between preferred and rejected scores.
    • Anchor Loss ($L_{anchor}$): Stabilizes the overall score scale.

2.2 SURE-REFL: Uncertainty-Guided Dense Reward Feedback Learning

Once SURE-LRM is trained, its parameters are frozen. SURE-REFL utilizes this model to guide the generator's post-training in latent space.

• Local Feedback Mechanism: Instead of backpropagating through the entire denoising graph, SURE-REFL queries the frozen reward model at selected transitions ($j$) along the denoising trajectory.
• Uncertainty-Guided Weighting:
  1. The model predicts the mean reward $\mu$ and the residual variance for the state at transition $j$.
  2. The total variance is converted into a reliability weight $\rho = V^{-1/2}$ (inverse standard deviation).
  3. This weight is detached (stop-gradient) from the reward model, meaning the generator learns from the weight but does not update the reward model's uncertainty estimates during this step.
  4. Weights are normalized across samples evaluated at the same transition to ensure fair comparison.
• Optimization: The generator is updated using a weighted reward objective where samples with lower predicted variance (higher reliability) contribute more strongly to the gradient. This dense feedback is applied locally to each transition, avoiding the need for full graph retention or pixel-space decoding.

3. Key Contributions

1. Formulation of Sample-Adaptive Uncertainty: The paper formulates reward uncertainty in latent space as a learnable distribution derived from standard pairwise preferences, without requiring additional human confidence annotations.
2. SURE-LRM: A latent reward model that learns sample-dependent uncertainty, demonstrating that samples at the same timestep can have different reliability estimates based on their content.
3. SURE-REFL: A unified post-training method for both image and video that leverages detached uncertainty to weight dense local feedback, eliminating the need for VAE decoding and full denoising graph retention.
4. Theoretical Insight: The authors derive variance gradients to demonstrate how the preference loss provides an intrinsic signal for learning uncertainty.

4. Experimental Results

4.1 Reward Model Accuracy (SURE-LRM)

• Image Benchmarks: SURE-LRM achieved state-of-the-art (SOTA) pairwise preference accuracy across four benchmarks (ImageReward, HPD v2, HPD v3, GenAI-Bench) on three different backbones (SD3.5, Z-Image-Turbo, FLUX.1-dev), outperforming strong baselines including DiNa-LRM and VLM-based pixel rewards.
• Video Benchmarks: On VisionRewardDB and T2V Ranking, SURE-LRM achieved the highest accuracy, surpassing decoded-video evaluators and video-specific reward models.
• Uncertainty Diagnostics: Experiments confirmed that the learned residual variance correlates with prediction errors. Samples with incorrect rankings had significantly higher residual variance (24.26% higher on average). Filtering samples by low predicted variance improved pairwise accuracy from 79.40% to 90.10%.

4.2 Post-Training Performance (SURE-REFL)

• Image Generation: SURE-REFL achieved the best performance across HPSv2.1, ImageReward, and PickScore metrics on SD3.5, Z-Image-Turbo, and FLUX.1-dev, outperforming baselines like Flow-GRPO, ViPO, and DiNa-LRM.
• Video Generation: On VBench, SURE-REFL achieved the highest scores in Quality, Semantic, and Total metrics for the Wan2.1-14B video generator, outperforming DanceGRPO, ViPO, Video-DPO, and HY-PRFL.
• Stability and Reward Hacking: Optimization dynamics showed that while baselines (e.g., DiNa-LRM-REFL) exhibited reward hacking (reward scores rising while independent quality metrics like HPSv3 fell), SURE-REFL maintained alignment between the optimized reward and independent quality metrics for longer, reaching higher final scores.

4.3 Ablation Studies

Ablation of the reward model's loss components confirmed that all terms are necessary. Removing the variance regularization ($L_{var}$) caused the largest drop in accuracy, validating the need to control variance inflation.

5. Significance and Claims

The paper claims that SURE provides a unified, efficient framework for aligning visual diffusion models by addressing the "reliability gap" in latent reward modeling. Its significance lies in:

• Efficiency: By operating entirely in latent space without VAE decoding, it reduces computational overhead during the feedback loop.
• Robustness: By explicitly modeling and utilizing sample-adaptive uncertainty, it prevents the generator from overfitting to noisy or ambiguous feedback, thereby reducing reward hacking and improving training stability.
• Generality: The framework is modality-agnostic, successfully applying the same uncertainty-weighting logic to both image and video generation tasks.

The authors conclude that learning sample-dependent uncertainty from standard pairwise preferences is a viable and effective strategy for improving the reliability of dense reward feedback in diffusion post-training.