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

TL;DR - SURE is a latent-space framework that makes diffusion post-training uncertainty-aware: its reward model predicts a Gaussian (mean + variance) instead of a scalar, and that variance is used to weight reward feedback so unreliable signals don't drive reward hacking. It matters because it improves alignment stability for image and video diffusion without pixel-space decoding.

• SURE-LRM: a sample-adaptive latent reward model that outputs a Gaussian utility per noisy latent — mean as the reward score, variance as prediction uncertainty, learned without human uncertainty annotation.
• SURE-REFL: uncertainty-guided reward feedback learning that queries the frozen LRM at selected denoising transitions, converts detached variance into per-sample reliability weights, and backpropagates each weighted reward only through its local transition.
• Efficiency: the whole pipeline stays in latent space, avoiding pixel decoding and backprop through the full denoising graph.
• Results as reported: better preference prediction than strong baselines, SOTA across several metrics with improved optimization stability, and top VBench quality/semantic/total scores among evaluated methods.