Sample-Adaptive Latent Rewards for Uncertainty-Guided Diffusion Post-Training

2026-08-06

AI Summary

This work addresses the limitations of existing latent reward models, which output only scalar scores without uncertainty estimates, often leading to unreliable feedback and reward hacking in diffusion model post-training. To overcome this, the authors propose SURE, a novel framework that introduces uncertainty estimation into latent reward modeling for the first time. SURE constructs a latent reward distribution via sample-adaptive Gaussian utility prediction and incorporates an uncertainty-guided feedback mechanism (SURE-REFL) to direct dense post-training. The entire process operates in the latent space, eliminating the need for pixel-space decoding or full denoising. Experiments demonstrate that SURE-LRM outperforms strong baselines in preference prediction, while SURE-REFL significantly enhances training stability, achieving state-of-the-art performance on VBench in terms of quality, semantics, and overall score.

Abstract

Latent reward models can supervise visual diffusion models without decoding intermediate states into pixel space. This makes alignment with human preferences more efficient. However, existing latent reward models output only scalar scores. They do not estimate the uncertainty of each prediction. The generator therefore cannot determine which feedback is reliable. This can drive optimization in the wrong direction and lead to reward hacking. We propose \textsc{SURE}, a unified latent-space framework for image and video diffusion models. It learns reward distributions and directly uses their reliability to guide dense post-training. First, we propose sample-adaptive latent reward model (\textsc{SURE-LRM}). It predicts a Gaussian utility for each noisy latent. Its mean predicts the reward score. Its variance reflect the uncertainty of prediction without human annotation. The learned distribution then guides post-training through uncertainty-guided reward feedback learning (\textsc{SURE-REFL}). This method provides uncertainty-guided dense feedback along the denoising trajectory. At selected transitions, \textsc{SURE-REFL} queries the frozen \textsc{SURE-LRM}. It converts detached variance into reliability weights for samples at the same transition. Each weighted reward is backpropagated only through its local transition. The entire process remains in latent space and requires neither pixel-space decoding nor the full denoising graph. Experiments show that \textsc{SURE-LRM} improves preference prediction over strong baselines. \textsc{SURE-REFL} achieves the sota performance among various metrics and further improves optimization stability. It also achieves the highest VBench quality, semantic, and total scores among the evaluated methods.

Problem

• latent reward models
• uncertainty estimation
• reward hacking
• diffusion post-training
• reliability of feedback

Innovation

• latent reward model
• uncertainty estimation
• diffusion post-training
• sample-adaptive feedback
• dense reward learning