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

Aligning visual diffusion models with human preferences typically involves querying reward models at various stages of the denoising process. Existing latent-space reward models often provide only scalar scores, failing to account for the reliability of these predictions. This lack of uncertainty estimation can lead to unstable optimization and reward hacking, where the generator exploits unreliable feedback to artificially inflate scores. This paper asks whether explicitly modeling reward uncertainty can improve the stability and effectiveness of diffusion post-training.

To address this, the authors introduce SURE, a two-part framework consisting of the Sample-Adaptive Latent Reward Model (SURE-LRM) and Uncertainty-Guided Reward Feedback Learning (SURE-REFL). SURE-LRM processes noisy latents to predict both a mean reward score and a sample-specific residual variance, which serves as a proxy for prediction uncertainty. This is achieved using a heteroscedastic preference objective that learns from standard pairwise comparisons without requiring explicit confidence labels.

Building on this, SURE-REFL uses the frozen SURE-LRM to guide the generator. During the denoising process, it queries the reward model at selected transitions and converts the predicted variance into reliability weights. These weights ensure that high-uncertainty (unreliable) feedback has a smaller influence on the generator's updates, while high-confidence feedback is prioritized. The entire process remains in the latent space, avoiding the computational overhead of decoding intermediate states into pixel space.

Experiments demonstrate that SURE-LRM consistently outperforms existing diffusion-native reward models in preference prediction accuracy across multiple backbones, including SD3.5-Medium and FLUX.1-dev. When applied to post-training, SURE-REFL achieves state-of-the-art performance on various image and video quality metrics. Crucially, the inclusion of uncertainty-guided weighting leads to more stable training dynamics and significantly reduces instances of reward hacking compared to standard dense feedback methods.