Iterative Exact Discrete Guidance for Energy-Based Sampling

Yuwen Qian; Yidong Ouyang; Zhengyan Wan; Hongyuan Zha

cs.LGSep 29, 2026

Iterative Exact Discrete Guidance for Energy-Based Sampling

Authors: Yuwen Qian, Yidong Ouyang, Zhengyan Wan, Hongyuan Zha

Organizations: School of Science and Engineering, The Chinese University of Hong Kong, Shenzhen · Department of Statistics, University of California, Los Angeles · School of Data Science, The Chinese University of Hong Kong, Shenzhen

Abstract

Sampling from unnormalized distributions over large discrete state spaces becomes difficult when a multimodal target is far from a tractable reference. We introduce Iterative Exact Discrete Guidance (IEDG), a population-exact, trajectory-wise guidance framework for unnormalized discrete targets. Rather than learn the full reference-to-target correction in one step, IEDG introduces a global Boltzmann tilt along an annealing trajectory. Each stage learns a stage-local posterior correction for an incremental Boltzmann tilt of the current source, while the resulting corrections are accumulated relative to a fixed analytic posterior. At the population optimum, exact stage posteriors recover the correct reverse dynamics, whose exact simulation reproduces the target distribution. IEDG chooses stage increments by relative effective sample size (rESS), which controls Rényi-2 displacement and locally adapts the step size to the thermodynamic geometry of the annealing path. Our stagewise total-variation analysis shows that limited overlap amplifies Bregman fitting error by $1/\sqrt{\mathrm{rESS}}$ , while posterior, simulation, and truncation errors enter additively. IEDG improves all distribution-level errors over the neural baselines on ordered, exactly enumerated Ising $4\times4$ , while substantially reducing one-shot errors on Ising/Potts $16\times16$ across thermodynamic regimes and attaining the best neural-sampler result on several reported local-statistic and phase-coverage metrics. On Max-Cut, its best-of-512 and average-sample ratios exceed all the baselines. Code and artifacts are available at https://github.com/StillFantasy123/iterative-exact-discrete-guidance.

Figures & tables

Figure 1: IEDG training. Samples from the current buffer

\mathcal{B}_{k}

estimate rESS and set the next annealing increment

\Delta\alpha_{k}

. The energy weight

w_{k}

and frozen posterior ratio

A_{k-1,t}^{d}

form the regression response

R_{k,t}^{d}

, whose conditional mean is learned by positive Bregman regression as the accumulated guide

G_{\psi_{k}}

. Combined with the analytic reference posterior, this defines

q_{\psi_{k},1\mid t}^{d}

, whose direct-

q

rollout produces the next buffer

\mathcal{B}_{k+1}

Method	$\mathrm{TV}$	$\mathrm{KL}(\widehat{p}\\|\pi)$	$\chi^{2}(\widehat{p}\\|\pi)$
Exact MC	$0.00411{\pm}0.00019$	$0.00426{\pm}0.00006$	$0.0515{\pm}0.0122$
One-shot DGM	0.779	2.55	722
UDNS (reproduced)	0.150	0.165	0.237
MDNS ( $F_{\rm WDCE}$ ) †	0.0418	0.0282	1.66
IEDG	0.0314	0.00993	0.229

Table 1: Distribution-level errors on exactly enumerated Ising

4\times 4

\beta=0.6

(lower is better). Exact MC is the equal-budget i.i.d. reference;

\dagger

marks the MDNS-reported

F_{\rm WDCE}

result. Bold marks the best neural sampler.

Figure 2: Phase coverage on $16\times 16$ lattices: Ising $x_{\uparrow}$ histograms (a–c) and Potts sorted dominant-mode mass residuals relative to SW (d–f) across three regimes. In (a–c), the upper subpanels show the full distributions, while the maps below show signed per-bin residuals on a common symmetric-log scale (red: excess; blue: deficit). Panel (f) omits the off-scale one-shot curve.

	$\beta=0.28$			$\beta=0.4407$			$\beta=0.6$
Method	Mag.	Corr. agg.	$x_{\uparrow}$ JS	Mag.	Corr. agg.	$x_{\uparrow}$ JS	Mag.	Corr. agg.	$x_{\uparrow}$ JS
Analytic preconditioner	0.00615	1.51	0.0863	0.00409	16.5	0.609	0.0169	27.1	0.693
MDNS (official ckpt.)	0.0112	0.161	0.00937	0.0105	0.171	0.00980	0.0173	0.0582	0.00336
MetaDNS (official ckpt.)	0.00839	0.195	0.00889	0.0317	0.124	0.0132	0.00490	0.0892	0.00574
One-shot DGM	0.0112	0.402	0.0159	0.0587	16.3	0.599	0.170	26.8	0.693
IEDG	0.00613	0.242	0.00742	0.0106	0.150	0.0252	0.0157	0.0537	0.0193

Table 2: Sampling

16\times 16

lattice models across thermodynamic regimes (lower is better). Bold marks the best result.

	$n\in[20,32]$		$n\in[40,64]$		$n\in[100,128]$
Method	Max.	Avg.	Max.	Avg.	Max.	Avg.
Uniform	0.888	0.696	0.825	0.687	0.769	0.678
One-shot DGM	0.962	0.854	0.884	0.798	0.814	0.748
PDNS †	0.963	0.911	0.897	0.869	0.876	0.840
IEDG	1.00	0.965	0.987	0.920	0.915	0.871

Table 3: BA Max-Cut approximation ratios on 32 held-out graphs (higher is better), using 512 samples per graph.

\dagger

marks PDNS values ( Guo et al., 2025 ) . Bold marks the best result.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

R_{k,t}^{d}\leftarrow\frac{w_{k}(X_{1})}{\widehat{c}_{k}}\frac{\widehat{p}_{k-1,1\mid t}^{d}(X_{1}^{d}\mid X_{t})}{p_{\mathrm{ref},1\mid t}^{d}(X_{1}^{d}\mid X_{t})}

Appendix

Algorithm 1 Iterative Exact Discrete Guidance training

Benchmark	Size	$\|\mathcal{S}\|$	Target parameter	Reference	Generated	NFE
Ising-exact	$4\times 4$	2	$\beta=.28,.4407,.6$	Enumeration	$2^{20}$	128/256/256
Ising	$16\times 16$	2	$\beta=.28,.4407,.6$	SW	4096	256
Potts	$16\times 16$	3	$\beta=.5,1.005,1.2$	SW	4096	256
BA Max-Cut	$[20,32]$	2	$\lambda_{\rm cut}=5$	Certified optimum	512/graph	128
	$[40,64]$	2	$\lambda_{\rm cut}=5$	Certified optimum	512/graph	128
	$[100,128]$	2	$\lambda_{\rm cut}=5$	Certified optimum	512/graph	128

Appendix

Table 4: Benchmark and terminal-evaluation configurations. “Reference” denotes the distribution used to compute evaluation errors.

Target	States	Mag.	Corr. agg.	$x_{\uparrow}$ JS	Mode $\ell_{1}$	CV JS
Ising $\beta=.4407$	4,096	.00416	.190	.0144	–	–
	65,536	.00201	.0573	.00102	–	–
Potts $\beta=1.005$	4,096	.0635	.143	–	.0244	.0555
	65,536	.0280	.0183	–	.00497	.00379

Appendix

Table 5: Near-critical SW reference stability (seed 0 versus seed 1; zero is ideal). Metrics use the same fixed definitions as model evaluation; dashes denote inapplicable phase statistics.

Benchmark	Main-text metrics	Supplemental evidence
Ising $4\times 4$	TV, KL, $\chi^{2}$	regime sweep; observables
Ising $16\times 16$	Mag., Corr. agg., $x_{\uparrow}$ JS	$J_{E}$ , Corr.-curve; marginal plots
Potts $16\times 16$	Mag., Corr. agg., Mode $\ell_{1}$	$J_{E}$ , Corr.-curve, CV JS; marginal/CV plots
BA Max-Cut	$R_{\max}$ , $R_{\rm avg}$	stage-count and path-schedule ablations

Appendix

Table 6: Metric map. Main-text columns are the primary evidence for each benchmark; Appendix G.4 reports the listed complementary distributional diagnostics.

Reported target	Backbone	W/D/H	Batch	Updates	Optimizer	LR	EMA	Source buffer
Ising $4\times 4$ , $\beta=.28$	tiny MLP	–	512	90000	Adam	$10^{-4}$	.999	1000000
Ising $4\times 4$ , $\beta=.4407$	tiny MLP	–	1024	65000	Adam	$10^{-4}$	.995	1000000
Ising $4\times 4$ , $\beta=.6$	tiny MLP	–	1024	70000	Adam	$10^{-4}$	.995	1000000
Ising $16\times 16$ , $\beta=.28$	RoPE DeiT2D	64/4/4	1024	35000	Adam	$4\!\times\!10^{-5}$	.999	131072
Ising $16\times 16$ , $\beta=.4407$	RoPE DeiT2D	64/4/4	512	65000	Adam	$4\!\times\!10^{-5}$	.999	100000
Ising $16\times 16$ , $\beta=.6$	RoPE DeiT2D	64/4/4	512	65000	Adam	$4\!\times\!10^{-5}$	.999	100000

Appendix

Table 7: Training configurations for the formal IEDG results. Each lattice row is the path segment ending at the reported

\beta

; “updates” is its realized segment budget or global cap. Width/depth/heads (W/D/H) describes the principal guide; LR and EMA denote learning rate and exponential-moving-average decay, and source buffer is the number of initially generated source states (per graph for Max-Cut).

Reported target	Frozen segment endpoints	Stages $K$
Ising $4\times 4$ , $\beta=.28$	$0,.1,.19,.28$	3
Ising $4\times 4$ , $\beta=.4407$	$.28,.36,.4407$	2
Ising $4\times 4$ , $\beta=.6$	$.4407,.50232,.6$	2
Ising $16\times 16$ , $\beta=.28$	$0,.0769,.1502,.2186,.28$	4
Ising $16\times 16$ , $\beta=.4407$	$.28,.3195,.3555,.3865,.413,.4407$	5
Ising $16\times 16$ , $\beta=.6$	$.4407,.472,.51,.552,.6$	4

Appendix

Table 8: Frozen ESS-controlled path segments used for the formal IEDG results. Lattice endpoints are inverse temperatures

\beta

; Max-Cut endpoints are effective inverse temperatures

\lambda

in Equation F.1 . The adaptive Ising-

16\times 16

\beta=.28

interior endpoints are rounded to four decimals.

Comparison	Parameters or source	Terminal generation	Evidence
Uniform (Max-Cut)	none	i.i.d. base samples	local
Analytic preconditioner	fixed BP field, $\Delta_{\psi}\equiv 0$	direct- $q$ , matched samples/NFE	local
One-shot DGM	trained by us	direct- $q$	local
UDNS	DASBS-based reproduction	factorized uniform CTMC	local
MDNS (Ising $4\times 4$ )	value from MDNS	reported native setup	published †
MDNS ( $16\times 16$ lattices)	released checkpoints	native sampler	local

Appendix

Table 9: Baseline provenance and terminal generation protocol.

Setting	$\{\theta_{k}\}_{k=0}^{K}$	$K$	Mag.	Corr. agg.	Phase error
Ising $0.28\!\to\!0.4407$	$\{.280,.329,.373,.409,.441\}$	4	$7.26{\times}10^{-3}$	$5.38{\times}10^{-1}$	$6.37{\times}10^{-2}$
	$\{.280,.320,.356,.387,.413,.441\}$	5 (canonical)	$1.06{\times}10^{-2}$	$\mathbf{1.50{\times}10^{-1}}$	$\mathbf{2.52{\times}10^{-2}}$
	$\{.280,.313,.344,.373,.398,.420,.441\}$	6	$\mathbf{5.42{\times}10^{-3}}$	$2.10{\times}10^{-1}$	$3.75{\times}10^{-2}$
Potts $0.5\!\to\!1.005$	$\{.500,.637,.759,.865,.946,1.005\}$	5	$1.30{\times}10^{-1}$	$6.32{\times}10^{-1}$	$3.91{\times}10^{-2}$
	$\{.500,.625,.738,.839,.920,.970,1.005\}$	6 (canonical)	$\mathbf{1.08{\times}10^{-1}}$	$\mathbf{1.23{\times}10^{-1}}$	$\mathbf{1.09{\times}10^{-2}}$
	$\{.500,.598,.692,.775,.851,.914,.963,1.005\}$	7	$3.93{\times}10^{-1}$	$4.91{\times}10^{-1}$	$9.62{\times}10^{-2}$

Appendix

Table 10: Sensitivity to the number of annealing stages. Paths are shown in the result rows (endpoints rounded to three decimals); lower lattice errors and higher Max-Cut ratios are better. The phase metric is

x_{\uparrow}

JS for Ising and Mode

\ell_{1}

for Potts.

Setting	Schedule	$\{\theta_{k}\}_{k=0}^{K}$	$K$	Mag.	Corr. agg.	$x_{\uparrow}$ JS
Ising $0.28\!\to\!0.4407$	Matched- $K$ uniform	$\{.280,.312,.344,.376,.409,.441\}$	5	$\mathbf{5.58{\times}10^{-3}}$	$3.17{\times}10^{-1}$	$2.61{\times}10^{-2}$
	ESS-controlled (canonical)	$\{.280,.320,.356,.387,.413,.441\}$	5	$1.06{\times}10^{-2}$	$\mathbf{1.50{\times}10^{-1}}$	$\mathbf{2.52{\times}10^{-2}}$

Appendix

Table 11: ESS-controlled and uniform schedules at matched stage count and total training/evaluation budgets. Paths are rounded to three decimals.

Setting	$\lambda_{\mathrm{CR}}$	Mag.	Corr. agg.	Phase error
Ising $0.28\!\to\!0.4407$	0.5	$2.37{\times}10^{-2}$	$3.14{\times}10^{-1}$	$4.01{\times}10^{-2}$
	1.0 (canonical)	$1.06{\times}10^{-2}$	$\mathbf{1.50{\times}10^{-1}}$	$2.52{\times}10^{-2}$
	2.0	$\mathbf{5.73{\times}10^{-3}}$	$9.59{\times}10^{-1}$	$\mathbf{1.89{\times}10^{-2}}$
Potts $0.5\!\to\!1.005$	0.5	$1.99{\times}10^{-1}$	$2.81{\times}10^{-1}$	$1.24{\times}10^{-2}$
	1.0 (canonical)	$\mathbf{1.08{\times}10^{-1}}$	$\mathbf{1.23{\times}10^{-1}}$	$\mathbf{1.09{\times}10^{-2}}$
	2.0	$3.31{\times}10^{-1}$	$1.23$	$1.53{\times}10^{-2}$

Appendix

Table 12: Sensitivity to conditional-reweight strength. All other components, including analytic preconditioning and stagewise response normalization, remain enabled at their canonical settings.

Setting	$\rho_{\mathrm{pre}}$	Mag.	Corr. agg.	Phase error
Ising $0.28\!\to\!0.4407$	0.5 (canonical)	$\mathbf{1.06{\times}10^{-2}}$	$\mathbf{1.50{\times}10^{-1}}$	$\mathbf{2.52{\times}10^{-2}}$
	1.0	$1.73{\times}10^{-2}$	$1.43$	$8.92{\times}10^{-2}$
Potts $0.5\!\to\!1.005$	0.5 (canonical)	$\mathbf{1.08{\times}10^{-1}}$	$\mathbf{1.23{\times}10^{-1}}$	$\mathbf{1.09{\times}10^{-2}}$
	1.0	$3.38{\times}10^{-1}$	$2.84$	$4.26{\times}10^{-2}$

Appendix

Table 13: Sensitivity to the analytic-preconditioner ratio scale in the near-critical regime; preconditioning remains enabled in every row.

$\beta$	Method	$\mathrm{TV}$	$\mathrm{KL}(\widehat{p}\\|\pi)$	$\chi^{2}(\widehat{p}\\|\pi)$
0.28	Exact MC	$0.0664{\pm}0.00018$	$0.0323{\pm}0.00021$	$0.0626{\pm}0.00064$
	One-shot DGM	0.193	0.149	0.341
	UDNS (reproduced)	0.0909	0.0412	0.0931
	MDNS ( $F_{\rm WDCE}$ ) †	0.0799	0.0382	0.0868
	IEDG	0.0860	0.0380	0.0823
0.4407	Exact MC	$0.0221{\pm}0.00024$	$0.0193{\pm}0.00010$	$0.0627{\pm}0.00285$

Appendix

Table 14: Full-distribution recovery in the remaining Ising

4\times 4

regimes (lower is better). Exact MC is the equal-budget i.i.d. reference; UDNS is our DASBS-based reproduction and

\dagger

denotes values reported by MDNS. Bold marks the best neural sampler.

$\beta$	Method	$\|\Delta E\|$	$\|\Delta C_{\rm nn}\|$
0.28	Exact MC	$0.00523{\pm}0.00609$	$0.000150{\pm}0.000188$
	One-shot DGM	3.53	0.108
	UDNS (reproduced)	0.740	0.0224
	IEDG	0.661	0.0203
0.4407	Exact MC	$0.00940{\pm}0.00562$	$0.000284{\pm}0.000177$
	One-shot DGM	10.7	0.319

Appendix

Table 15: Observable errors on Ising

4\times 4

(lower is better). Exact MC is the equal-budget i.i.d. reference; bold marks the best neural sampler.

$\beta$	Metric	Analytic precond.	MDNS	MetaDNS	One-shot DGM	IEDG
0.28	$J_{E}$	0.620	0.00230	0.00377	0.0123	0.00433
	Corr.-curve	0.0462	0.00235	0.00125	0.0121	0.00682
0.4407	$J_{E}$	0.693	0.00189	0.00548	0.693	0.0223
	Corr.-curve	0.516	0.00467	0.00360	0.507	0.00275
0.6	$J_{E}$	0.693	0.00117	0.00314	0.693	0.00445
	Corr.-curve	0.848	0.00100	0.000908	0.803	0.000623

Appendix

Table 16: Two complementary residual diagnostics for the

16\times 16

lattices: energy-marginal JS (

J_{E}

) and translational correlation-curve MAE (Corr.-curve; lower is better). Bold marks the best result.

$\beta$	Analytic precond.	MDNS	MetaDNS	One-shot DGM	IEDG
0.5	0.0143	0.0104	0.0119	0.00897	0.00868
1.005	0.646	0.0347	0.0548	0.658	0.0818
1.2	0.693	0.00332	0.0111	0.693	0.0329

Appendix

Table 17: Fixed-grid Potts CV JS across the disordered, near-critical, and ordered regimes (lower is better). Bold marks the best result.

Figure 3: Regime-wide one-dimensional marginals from the frozen 4,096-sample payloads. Columns traverse the disordered, near-critical, and ordered regimes; rows show Ising energy per site and

x_{\uparrow}

, followed by Potts energy per site and

m_{\mathrm{P}}

. All methods share histogram edges fixed from the analytic support, with SW (black) as the test reference.

Figure 4: Near-critical signed-displacement correlation profiles for Ising at

\beta=0.4407

and Potts at

\beta=1.005

. Unlike Corr. agg., these profiles resolve spatial dependence at each separation. Upper panels compare each empirical curve with SW; lower panels show

|c_{\rm method}(r)-c_{\rm SW}(r)|

on a logarithmic scale. The deterministic self-correlation at

r=0

is omitted. Relative to one-shot DGM, the IEDG residual is reduced across separations rather than at a single displacement.

Figure 5: Potts phase geometry at $\beta=1.005$ . Unlike Mode $\ell_{1}$ , which compares only ranked sector masses, the two-dimensional CV plane exposes within-sector geometry. IEDG recovers the three symmetry-related directions absent from one-shot DGM, while excess inter-sector probability explains its remaining CV JS gap. All methods use the same fixed $50\times 50$ grid and logarithmic probability scale.

Figure 6: Representative ordered-endpoint configurations. Columns denote methods; rows show Ising at $\beta=0.6$ and Potts at $\beta=1.2$ . Indices $\{973,2172,2598,4081\}$ are drawn once with seed 2701 and reused across each method’s independently generated 4,096-state payload. This qualitative check is not used for checkpoint selection; Potts colors are shared across methods.

Explore similar work

Mar 18, 2026

One-Step Sampler for Boltzmann Distributions via Drifting

We present a drifting-based framework for amortized sampling of Boltzmann distributions defined by energy functions. The method trains a one-step neural generator by projecting samples along a Gaussian-smoothed score field from the current model distribution toward the target Boltzmann distribution. For targets specified only up to an unknown normalization constant, we derive a practical target-side drift from a smoothed energy and use two estimators: a local importance-sampling mean-shift estimator and a second-order curvature-corrected approximation. Combined with a mini-batch Gaussian mean-shift estimate of the sampler-side smoothed score, this yields a simple stop-gradient objective for stable one-step training. On a four-mode Gaussian-mixture Boltzmann target, our sampler achieves mean error

0.0754

, covariance error

0.0425

, and RBF MMD

0.0020

. Additional double-well and banana targets show that the same formulation also handles nonconvex and curved low-energy geometries. Overall, the results support drifting as an effective way to amortize iterative sampling from Boltzmann distributions into a single forward pass at test time.

Wenhan Cao, Keyu Yan, Lin Zhao

Department of Electrical and Computer Engineering, National University of Singapore

One-Step Generative Modeling

Ask PDF

Oct 1, 2026

Fixed-point neural samplers on discrete spaces

Sampling from discrete, unnormalized distributions without access to data is a challenging problem. Neural samplers offer a promising approach by training generative models from density evaluations directly. Despite recent progress, existing discrete neural samplers are prone to mode collapse, come without convergence guarantees when trained via fixed-point iterations, and are often tied to a specific reference process such as masked or uniform diffusion. In this work, we introduce Discrete Gibbs Iterative Neural Sampler, a fixed-point neural sampler that addresses these limitations, enabling efficient, scalable learning, substantially reducing mode collapse in practice. Our framework builds on masked diffusion and also extends to transport between pairs of distributions. We demonstrate that the resulting method scales effectively to high-dimensional systems, supports amortized sampling across different conditions, and enables accurate estimation of alloy phase diagrams.

Jiajun He, Denis Blessing, Mouyang Cheng +2

MIT · Microsoft Research New England

Diffusion Models

Ask PDF

Mar 10, 2026

A Generative Sampler for distributions with possible discrete parameter based on Reversibility

Learning to sample from complex unnormalized distributions is a fundamental challenge in computational physics and machine learning. While score-based and variational methods have achieved success in continuous domains, extending them to discrete or mixed-variable systems remains difficult due to ill-defined gradients or high variance in estimators. We propose a unified, target-gradient-free generative sampling framework applicable across diverse state spaces. Building on the fact that detailed balance implies the time-reversibility of the equilibrium stochastic process, we enforce this symmetry as a statistical constraint. Specifically, using a prescribed physical transition kernel (such as Metropolis-Hastings), we minimize the Maximum Mean Discrepancy (MMD) between the joint distributions of forward and backward Markov trajectories. Crucially, this training procedure relies solely on energy evaluations via acceptance ratios, circumventing the need for target score functions or continuous relaxations. We demonstrate the versatility of our method on three distinct benchmarks: (1) a continuous multi-modal Gaussian mixture, (2) the discrete high-dimensional Ising model, and (3) a challenging hybrid system coupling discrete indices with continuous dynamics. Experiments show that our framework accurately reproduces thermodynamic observables and captures mode-switching behavior across all regimes, offering a physically grounded and universally applicable alternative for equilibrium sampling.

Lei Li, Zhen Wang, Lishuo Zhang

School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai, 200240, P.R.China. · Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University, Shanghai, 200240, P.R.China.

Generative Modeling Markov Chain Monte Carlo

Ask PDF