Calibrated Test-Time Guidance for Bayesian Inference

Daniel Geyfman, Felix Draxler, Jan Niklas Groeneveld, Hyunsoo Lee, Theofanis Karaletsos, Stephan Mandt
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:34639-34666, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-geyfman26a, title = {Calibrated Test-Time Guidance for {B}ayesian Inference}, author = {Geyfman, Daniel and Draxler, Felix and Groeneveld, Jan Niklas and Lee, Hyunsoo and Karaletsos, Theofanis and Mandt, Stephan}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {34639--34666}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/geyfman26a/geyfman26a.pdf}, url = {https://proceedings.mlr.press/v306/geyfman26a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Calibrated Test-Time Guidance for Bayesian Inference %A Daniel Geyfman %A Felix Draxler %A Jan Niklas Groeneveld %A Hyunsoo Lee %A Theofanis Karaletsos %A Stephan Mandt %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-geyfman26a %I PMLR %P 34639--34666 %U https://proceedings.mlr.press/v306/geyfman26a.html %V 306 %X 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.
APA
Geyfman, D., Draxler, F., Groeneveld, J.N., Lee, H., Karaletsos, T. & Mandt, S.. (2026). Calibrated Test-Time Guidance for Bayesian Inference. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:34639-34666 Available from https://proceedings.mlr.press/v306/geyfman26a.html.

Related Material