Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models
Abstract
Inference-time scaling adapts pretrained diffusion models to new sampling tasks without additional training. Existing methods rely primarily on Monte Carlo sampling with more particles, yet are premised on the pretrained model being exact. In practice, data and training limitations make the model imperfect, and these methods inherit its error. More particles reduce Monte Carlo error but cannot remove the mismatch between the endpoint and the desired target or the error in tracking the prescribed probability path. We introduce the Energy-based Feynman-Kac Corrector (EBFKC), a framework for energy-based diffusion models that corrects these errors on the fly given a reference energy. We first derive Feynman-Kac dynamics that track a prescribed path exactly in the continuous-time population limit even when the model is imperfect, and approximate these dynamics using sequential Monte Carlo with variance-controlling guidance. To remove the endpoint mismatch, we use the pretrained energy as a surrogate along the diffusion path and progressively incorporate the discrepancy between the learned and target terminal energies. Experiments on Gaussian mixture models, particle systems, alanine dipeptide, and alanine tetrapeptide show that our method closely matches target distributions and molecular free-energy profiles under annealing and reward tilting, whereas standard inference-time scaling baselines retain substantial sampling errors.
Related
- Imperfect Verifiers, Exact Targets: Inference-Time Scaling for Continuous Generation
- Honest Inference for Synthetic Controls with Imperfect Fit: Time-Split Validity and the Choice of Anchor
- Debiased Inference for AI-Generated Data without Gold-Standard Labels: Identification via Multiple Imperfect Measurements
- Bias-Corrected Empirical-Likelihood Inference for Single-Index Expectile Regression Under Missing-at-Random Responses
- On the convergence of Lattice Boltzmann method for advection-diffusion equation with intermediate scaling
BibTeX
@misc{wen2026errorcorrected,
title = {{Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models}},
author = {Zuokai Wen and Louis Grenioux and Weinan E and Jiequn Han},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.01933},
eprint = {2610.01933},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.01933},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models
AU - Zuokai Wen
AU - Louis Grenioux
AU - Weinan E
AU - Jiequn Han
PY - 2026
DA - 2026/10/01
PB - arXiv
DO - 10.48550/arxiv.2610.01933
UR - https://arxiv.org/abs/2610.01933
AB - Inference-time scaling adapts pretrained diffusion models to new sampling tasks without additional training. Existing methods rely primarily on Monte Carlo sampling with more particles, yet are premised on the pretrained model being exact. In practice, data and training limitations make the model imperfect, and these methods inherit its error. More particles reduce Monte Carlo error but cannot remove the mismatch between the endpoint and the desired target or the error in tracking the prescribed probability path. We introduce the Energy-based Feynman-Kac Corrector (EBFKC), a framework for energy-based diffusion models that corrects these errors on the fly given a reference energy. We first derive Feynman-Kac dynamics that track a prescribed path exactly in the continuous-time population limit even when the model is imperfect, and approximate these dynamics using sequential Monte Carlo with variance-controlling guidance. To remove the endpoint mismatch, we use the pretrained energy as a surrogate along the diffusion path and progressively incorporate the discrepancy between the learned and target terminal energies. Experiments on Gaussian mixture models, particle systems, alanine dipeptide, and alanine tetrapeptide show that our method closely matches target distributions and molecular free-energy profiles under annealing and reward tilting, whereas standard inference-time scaling baselines retain substantial sampling errors.
ER -