Calibrated Test-Time Guidance for Bayesian Inference

Test-time guidance is a widely used mechanism for steering pretrained diffusion models toward outcomes specified by a reward function. Existing approaches, however, focus on maximizing reward rather than sampling from the true Bayesian posterior, leading to miscalibrated inference. In this work, we show that common test-time guidance methods do not recover the correct posterior distribution and identify the structural approximations responsible for this failure. We then propose consistent alternative estimators that enable calibrated sampling from the Bayesian posterior. We significantly outperform previous methods on a set of Bayesian inference tasks, and set a new state-of-the-art PSNR in black hole image reconstruction. We publish our code at https://github.com/mandt-lab/Calibrated-Guidance.

Many modern generative AI systems can be adjusted at test time to satisfy a goal, such as matching a measurement or improving an image. In scientific settings, however, the goal is not just to find one appealing answer but to represent all plausible answers and their uncertainty correctly. We prove mathematically that existing methods make mistakes in capturing this uncertainty, and propose a new method that fixes it: As we use more computational power, we capture the uncertainty better. This makes generative AI more reliable for scientific inverse problems and other settings where uncertainty matters.