Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models

Zuokai Wen; Louis Grenioux; Weinan E; Jiequn Han

stat.MLOct 1, 2026

Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models

Authors: Zuokai Wen, Louis Grenioux, Weinan E, Jiequn Han

Organizations: School of Artificial Intelligence, Shanghai Jiao Tong University · Center for Computational Mathematics, Flatiron Institute · School of Mathematical Sciences, Peking University

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.

Figures & tables

Figure 1: SMC trajectory lineages for the four combinations of reweighting scheme (FKC or EBFKC) and prescribed path (MIT or EAT). The target

p^{\star}

is a 64-dimensional mixture of six isotropic Gaussians in two groups (right edge). The figure displays one-dimensional trajectories for four lineages from a population of 2048 particles. The reward

r

favors the weak mode, with the desired terminal distribution shown in dashed purple. FKC ( Skreta et al., 2025 ) on MIT suffers from both errors and ends in the wrong modes. EBFKC on MIT tracks its path but inherits the endpoint mismatch, while FKC on EAT has the correct endpoint but its path-tracking error makes the weights degenerate, so that only one lineage survives. Only EBFKC on EAT, which removes both errors, reaches the desired target.

Panel of Figure 2 — Figure 2: GMM results under controlled pretrained model error. (a–b) Sample distributions and (c–d) SWD under increasing pretrained model error for annealing and reward tilting.

Task	Method	$\|\Delta\mathrm{NLL}\|$ $\downarrow$	MMD $\downarrow$	SWD $\downarrow$	Mean $L_{2}$ $\downarrow$	Cov $F$ $\downarrow$
Annealing	FKC on EAT	$1.043\pm 0.196$	$0.076\pm 0.007$	$4.105\pm 0.375$	$21.218\pm 1.462$	$2437.100\pm 358.582$
	EBFKC w/o VCG on MIT	$59.355\pm 9.014$	$0.239\pm 0.035$	$11.913\pm 1.374$	$63.396\pm 13.044$	$6770.604\pm 1682.698$
	EBFKC on MIT	$59.204\pm 1.041$	$0.092\pm 0.002$	$3.359\pm 0.694$	$9.890\pm 0.385$	$1879.349\pm 129.679$
	EBFKC on EAT (Ours)	$\bm{0.133\pm 0.095}$	$\bm{0.022\pm 0.001}$	$\bm{1.110\pm 0.331}$	$\bm{5.493\pm 0.774}$	$\bm{639.193\pm 27.982}$
Reward tilting	FKC on EAT	$0.928\pm 0.374$	$0.061\pm 0.019$	$1.313\pm 1.191$	$8.661\pm 7.779$	$507.355\pm 680.777$
	EBFKC w/o VCG on MIT	$21.575\pm 10.209$	$0.435\pm 0.040$	$10.136\pm 6.044$	$64.321\pm 29.173$	$604.058\pm 878.509$

Table 1: Quantitative comparison on the 30-dimensional GMM under the controlled-error setting

\sigma_{\mu}=10

and

\rho=0.1

. Results are mean

\pm

standard deviation over 5 runs.

Panel of Figure 3 — Figure 3: LJ results comparing pairwise-distance and energy distributions with the reference. (a) Samples from the pretrained model. (b) Reward-tilting samples from the inference-time methods.

Figure 4: Ramachandran plots of the backbone dihedral angles $(\phi,\psi)$ for annealing ALDP from $800\,\mathrm{K}$ to $300\,\mathrm{K}$ . Reference samples at $300\,\mathrm{K}$ are compared with samples generated by EBFKC on MIT with/without VCG and EBFKC on EAT using the same pretrained diffusion model at $800\,\mathrm{K}$ .

Methods	Temperature (K)	$(\phi,\psi)$ TVD $\downarrow$	Basin-pop. error $\downarrow$	Energy TVD $\downarrow$	Distance TVD $\downarrow$
Pretrained reverse diffusion	$800$	$0.3005\pm 0.0022$	$0.1699\pm 0.0016$	$0.6730\pm 0.0013$	$0.0253\pm 0.0001$
EBFKC w/o VCG on MIT	$300$	$0.5089\pm 0.0163$	$0.1643\pm 0.0037$	$0.4460\pm 0.0130$	$0.0366\pm 0.0008$
EBFKC on MIT	$300$	$0.5063\pm 0.0035$	$0.1657\pm 0.0026$	$0.4659\pm 0.0095$	$0.0381\pm 0.0006$
EBFKC on EAT (Ours)	$300$	$\mathbf{0.0855\pm 0.0031}$	$\mathbf{0.0047\pm 0.0041}$	$\mathbf{0.0240\pm 0.0028}$	$\mathbf{0.0030\pm 0.0002}$

Table 2: Quantitative ALDP annealing results from

800\,\mathrm{K}

300\,\mathrm{K}

; the first row evaluates the pretrained model at

800\,\mathrm{K}

against the

800\,\mathrm{K}

reference. Results are mean

\pm

std over 5 runs.

Panel of Figure 5 — Figure 5: One-dimensional molecular free-energy profiles at $300\,\mathrm{K}$ . (a) ALDP along $\phi$ . (b) Ala-4 along $\phi_{3}$ . The references are umbrella sampling for ALDP and OPES for Ala-4. Diffusion-based results show the mean and standard deviation (shaded) over five independent runs.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

System	Layers	Hidden dim.	EBM training epochs
DW-4	4	32	Stage 1: 125 epochs
LJ-13	4	32	Stage 1: 300 epochs
ALDP ( $800\,\mathrm{K}$ )	4	32	Stage 1: 200 epochs; Stage 2: 200 epochs
ALDP ( $300\,\mathrm{K}$ )	4	32	Stage 1: 200 epochs; Stage 2: 200 epochs
Ala-4	6	128	Stage 1: 500 epochs

Appendix

Table 3: EGNN architectures and EBM training details.

Experiment	$(s_{\kappa},q_{\kappa})$	$(s_{\alpha},q_{\alpha})$
DW-4	$(0.999,3)$	(0.999,3)
LJ-13	$(0.999,4)$	$(0.999,2)$
ALDP annealing	$(0.999,3)$	-
ALDP umbrella diffusion sampling	$(0.999,3)$	$(0.9988,1)$
Ala-4 umbrella diffusion sampling	$(0.999,3)$	Reward-first: $(0.9988,1)$ Target-first: $(0.9994,1)$

Appendix

Table 4: Reward and target-correction schedule parameters.

Panel of Figure 6 — Figure 6: Sample distribution of FKC w/o VCG on EAT in the GMM task under (a) annealing and (b) reward tilting.

Panel of Figure 7 — Figure 7: SWD of FKC w/o VCG on EAT across pretrained-model error levels on the GMM under (a) annealing and (b) reward tilting.

Method	Annealing ( $\gamma=2.0$ )	Reward tilting ( $\sigma=100$ )
FKC on EAT	$5.357\pm 0.406$ (s)	$5.996\pm 0.736$ (s)
EBFKC w/o VCG on MIT	$3.098\pm 0.537$ (s)	$3.283\pm 0.445$ (s)
EBFKC on MIT	$3.626\pm 0.376$ (s)	$4.061\pm 0.963$ (s)
EBFKC on EAT (Ours)	$5.267\pm 0.082$ (s)	$5.503\pm 0.149$ (s)

Appendix

Table 5: Runtime comparison on the 30-dimensional GMM under

\sigma_{\mu}=10

and

\rho=0.1

. Results are mean

\pm

std over 5 runs.

Method	Annealing ( $\gamma=2.0$ )	Reward tilting ( $\lambda^{\prime}=0.5$ )
FKC on EAT	$412.594\pm 16.254$ (s)	$416.153\pm 17.053$ (s)
EBFKC w/o VCG on MIT	$342.789\pm 16.330$ (s)	$339.989\pm 13.836$ (s)
EBFKC on MIT	$367.378\pm 14.568$ (s)	$365.100\pm 18.214$ (s)
EBFKC on EAT (Ours)	$406.058\pm 16.066$ (s)	$403.458\pm 14.271$ (s)

Appendix

Table 6: Runtime comparison on the DW-4 tasks. Results are mean

\pm

std over 5 runs.

Figure 8: Comparison of the reference distribution and reverse diffusion samples from the pretrained energy-based diffusion model on the double-well system.

Figure 9: Comparison of the reference distribution, FKC on EAT, EBFKC on MIT, and EBFKC on EAT for the double-well system under annealing.

Figure 10: Comparison of the reference distribution, FKC on EAT, EBFKC on MIT, and EBFKC on EAT for the double-well system under reward tilting.

Figure 11: Comparison of the reference distribution, Tempering SMC, FKC on EAT, EBFKC on MIT and EBFKC on EAT for the Lennard-Jones system under annealing.

System	Task	Method	$\|\Delta\mathrm{NLL}\|$	MMD	SWD	$W_{1}^{\mathrm{RDF}}$	$W_{1}^{E}$
DW	Annealing	FKC on EAT	$23.9565\pm 0.9045$	$0.2659\pm 0.0059$	$0.6263\pm 0.0511$	$1.6010\pm 0.0090$	$23.9565\pm 0.9045$
		EBFKC w/o VCG on MIT	$18.1923\pm 1.0910$	$0.4315\pm 0.0645$	$0.9894\pm 0.1340$	$0.7005\pm 0.0307$	$18.1923\pm 1.0910$
		EBFKC on MIT	$19.7381\pm 0.2039$	$0.0996\pm 0.0029$	$0.2341\pm 0.0217$	$0.7582\pm 0.0069$	$19.7381\pm 0.2039$
		EBFKC on EAT (Ours)	$\bm{0.3389\pm 0.2107}$	$\bm{0.0816\pm 0.0193}$	$\bm{0.2032\pm 0.0441}$	$\bm{0.3943\pm 0.1437}$	$\bm{0.3754\pm 0.1556}$
	Reward tilting	FKC on EAT	$7.2620\pm 0.3554$	$0.1506\pm 0.0037$	$0.3644\pm 0.0313$	$0.9265\pm 0.0152$	$7.2620\pm 0.3554$
		EBFKC w/o VCG on MIT	$8.1395\pm 1.0000$	$0.5666\pm 0.2091$	$1.1923\pm 0.4556$	$0.7240\pm 0.1778$	$8.1395\pm 1.0000$

Appendix

Table 7: Performance comparison on particle systems (DW-4 and LJ-13). Results are reported as mean

\pm

std over 5 runs.

Method	Annealing ( $\gamma=2.0$ )	Reward tilting ( $\lambda^{\prime}=0.5$ )
FKC on EAT	$1484.691\pm 1.749$ (s)	$2074.467\pm 2.337$ (s)
EBFKC w/o VCG on MIT	$670.970\pm 0.693$ (s)	$1199.705\pm 1.273$ (s)
EBFKC on MIT	$671.882\pm 0.804$ (s)	$1201.382\pm 1.070$ (s)
EBFKC on EAT (Ours)	$1339.710\pm 0.454$ (s)	$1820.829\pm 1.568$ (s)

Appendix

Table 8: Runtime comparison on the LJ-13 tasks. Results are mean

\pm

std over 5 runs.

Figure 12: ALDP backbone dihedral marginals at

800\,\mathrm{K}

from the reference distribution and pretrained reverse diffusion:

\phi

(left) and

\psi

(right). The reverse diffusion marginals are estimated from 100,000 generated samples.

Figure 13: ALDP annealing marginals comparison for

\phi

(left) and

\psi

(right). EBFKC on EAT closely matches the

300\,\mathrm{K}

reference, whereas both MIT variants deviate substantially. The marginals are estimated from all 250,000 samples generated by each method and all reference samples at

300\,\mathrm{K}

Method	Sampling time (s)
EBFKC w/o VCG on MIT	$142.202\pm 0.636$
EBFKC on MIT	$142.932\pm 0.119$
EBFKC on EAT (Ours)	$508.433\pm 1.109$

Appendix

Table 9: Runtime comparison for ALDP annealing from

800\,\mathrm{K}

300\,\mathrm{K}

. Results are reported as mean

\pm

std over 5 runs.

Figure 14: ALDP backbone dihedral marginals at

300\,\mathrm{K}

from the reference distribution and pretrained reverse diffusion:

\phi

(left) and

\psi

(right). The reverse diffusion marginals are estimated from 100,000 generated samples.

Figure 15: Ala-4 backbone dihedral marginals at

300\,\mathrm{K}

from the reference distribution and pretrained reverse diffusion. The reverse diffusion marginals are estimated from 100,000 generated samples.

Panel of Figure 16 — Figure 16: EBFKC w/o VCG on MIT free-energy estimates for the molecular systems. (a) ALDP along $\phi$ , compared with traditional umbrella sampling reference. (b) Ala-4 along $\phi_{3}$ , compared with the OPES reference and umbrella sampling results. Results are based on 5 independent runs.

Method	Sampling time (s)
EBFKC w/o VCG on MIT	$119.30\pm 10.43$
EBFKC on MIT	$106.17\pm 3.24$
EBFKC on EAT (Ours)	$269.19\pm 9.11$

Appendix

Table 10: Runtime comparison for ALDP free-energy estimation in window 13 with 200 particles. Results are reported as mean

\pm

std over 5 runs.

Method	Sampling time (s)
EBFKC w/o VCG on MIT	$389.05\pm 14.20$
EBFKC on MIT	$376.89\pm 17.25$
EBFKC on EAT (Ours)	$504.87\pm 12.79$

Appendix

Table 11: Runtime comparison for Ala-4 free-energy estimation in window 13 with 12 particles. Results are reported as mean

\pm

std over 5 runs.

Explore similar work

Inference-Time Scaling in Diffusion Models through Iterative Partial Refinement
May 19, 2026Taegu Kang, Jaesik Yoon, Sungjin AhnDiffusion Models Inference-Time
KAIST · KAIST & SAP · KAIST & NYU
Simple Approximation and Derivative Free Inference-Time Scaling for Diffusion Models via Sequential Monte Carlo on Path Measures
May 18, 2026Chenyang Wang, Weizhong Wang, Yinuo Ren +2Diffusion Models Monte Carlo
School of Mathematical Sciences, Peking University, Beijing, China · School of Mathematical Sciences, Fudan University, Shanghai, China · Institute for Computational & Mathematical Engineering, Stanford University, Stanford, CA, United States +2
Metropolis-Adjusted Diffusion Models
May 10, 2026Kevin H. Lam, Tyler Farghly, Christopher Williams +3Score-Based Diffusion Model Diffusion Models
Department of Statistics, University of Oxford · Inria, CNRS, École Normale Supérieure, PSL Research University · Department of Mathematical Sciences, University of Copenhagen

May 19, 2026, cs.LGJ/K move · Enter open · S save

Inference-Time Scaling in Diffusion Models through Iterative Partial Refinement

Taegu Kang, Jaesik Yoon, Sungjin Ahn

KAIST · KAIST & SAP · KAIST & NYU

Diffusion Models Inference-Time Iterative Refinement Scaling Iterative Constraint Satisfaction

Open paper AskPDF arXiv

Inference-time scaling has emerged as a major approach for improving reasoning capabilities, and has been increasingly applied to diffusion models. However, existing inference-time scaling methods for diffusion models typically rely on external verifiers or reward models to rank and select samples, limiting their scalability to settings where such evaluators are available and reliable. Moreover, while recent diffusion models perform sequential inference with region-wise, mixed-noise conditioning, inference-time scaling tailored to this setting remains relatively underexplored. We propose Iterative Partial Refinement (IPR), an inference-time scaling method for sequential diffusion that requires no external verifier. Starting from an already-generated sample, IPR re-noises a subset of regions and regenerates them conditioned on the remaining regions, enabling the model to revise earlier decisions under a richer context than was available during the initial generation. This iterative partial refinement produces more globally consistent samples without external verification. On reasoning tasks requiring global constraint satisfaction, IPR consistently improves performance: on MNIST Sudoku, the valid solution rate increases from 55.8% to 75.0%. These results show that iterative partial refinement alone can serve as an effective inference-time scaling strategy for diffusion models in sequential, mixed-noise settings. Code is available at: https://github.com/ahn-ml/IPR

Abstract

Explore similar work

Inference-Time Scaling in Diffusion Models through Iterative Partial Refinement

Simple Approximation and Derivative Free Inference-Time Scaling for Diffusion Models via Sequential Monte Carlo on Path Measures

Metropolis-Adjusted Diffusion Models