Fixed-point neural samplers on discrete spaces

Jiajun He; Denis Blessing; Mouyang Cheng; Yuanqi Du; Carles Domingo-Enrich

cs.LGOct 1, 2026

Fixed-point neural samplers on discrete spaces

Authors: Jiajun He, Denis Blessing, Mouyang Cheng, Yuanqi Du, Carles Domingo-Enrich

Organizations: MIT · Microsoft Research New England

Abstract

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.

Figures & tables

Figure 1: DIG-IN iteration. One training iteration consists of data collection, importance weight smoothing using ESS control, and parameter updates using a weighted cross-entropy loss.

Ising, $25\times 25$ , $q=2$
	disordered			critical			ordered
	$\|\Delta\ln Z\|$	$\ESS$	$\TVD_{\mathrm{mode}}$	$\|\Delta\ln Z\|$	$\ESS$	$\TVD_{\mathrm{mode}}$	$\|\Delta\ln Z\|$	$\ESS$	$\TVD_{\mathrm{mode}}$
Method	$\times 10^{3}\,{\downarrow}$	${\uparrow}$	${\downarrow}$	$\times 10^{3}\,{\downarrow}$	${\uparrow}$	${\downarrow}$	$\times 10^{3}\,{\downarrow}$	${\uparrow}$	${\downarrow}$
MDNS	$1.40_{\pm 0.79}$	$.961_{\pm.002}$	$.012_{\pm.004}$	$1.93_{\pm 1.00}$	$.925_{\pm.006}$	$.016_{\pm.006}$	$693_{\pm 0}$	$.983_{\pm.007}$	$.500_{\pm.000}$
PDNS	$2.86_{\pm 1.86}$	$.956_{\pm.001}$	$.014_{\pm.006}$	$1.47_{\pm 0.61}$	$.917_{\pm.008}$	$.016_{\pm.008}$	$693_{\pm 0}$	$.988_{\pm.003}$	$.500_{\pm.000}$
DASBS	–	–	$.007_{\pm.004}$	–	–	$.042_{\pm.028}$	–	–	$.447_{\pm.075}$

Table 1: Ising and Potts lattice spin models. We consider the Ising and Potts model at three different inverse temperatures

\beta

corresponding to disordered, critical and ordered phases. We consider three evaluation criteria: The effective sample size (

\ESS

), the error in normalization constant estimation

|\Delta\ln Z|

and the total variation distance between the empirical distribution over the

q

modes of the reference target density, denoted as

\TVD_{\mathrm{mode}}

. The table reports Mean

{\pm}

standard deviation over 3 seeds. In each column, blue shading is used to indicate methods whose pessimistic value lies within one standard deviation of the best method. Table entries with dashes indicate that the method does not support computing

\ln Z

and

\ESS

Panel of Figure 2 — Figure 2: Amortized training across $(\beta,h)$ of DIG-IN . Left: a low $\TVD_{\mathrm{mode}}$ across all $(\beta,h)$ indicates that our method avoids mode collapse and accurately learns the relative mode weights, which change under an external field $h$ . Right, top: Potts samples of the amortized model at $h=0$ in the disordered phase, near the critical point and in the ordered phase. Right, bottom: average $\ESS$ , $|\Delta\ln Z|$ and $\TVD_{\mathrm{mode}}$ over the $9\times 9$ grid of $(\beta,h)$ values, all computed with the importance weights. The results show mean $\pm$ standard deviation over three seeds.

Panel of Figure 3 — Figure 3: Alloy phase diagrams. Left: Ag–Pd ( $D=125$ , top) and Cu–Au ( $D=128$ , bottom). The grey dashed lines are the reference phase boundaries and the colored lines are the phase boundaries the model predicted. Right, top: example structures of different phases: Ag–Pd is a random fcc alloy, shown with mixed occupancy, and separates into Ag-rich and Pd-rich fcc below its miscibility gap; Cu–Au orders into Cu 3 Au, CuAu and CuAu 3 . Right, bottom: average composition error $|\Delta c|$ , $|\Delta\ln Z|/D$ and $\ESS$ over the $(T,\Delta\mu)$ values, which are the thermodynamic conditions the model was trained on. The results show mean $\pm$ standard deviation over three seeds. The best results are highlighted using blue shading.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Example configurations of the Ising and Potts model at different phases.

Systems	Ising $25\times 25$ ( $q=2$ , $D=625$ ); Potts $16\times 16$ ( $q=3$ , $D=256$ )
Phases per system	3 (disordered, critical, ordered); see Section E.1.1
Seeds	3 per (system, phase, method)
Training steps	150k
Evaluation draws	8192 per seed, independent, no MCMC correction
Backbone	periodic residual CNN, 2-D circular convolutions, 96 channels, 7 residual blocks, kernel 3, dilations covering the torus (all methods but LEAPS)
Optimizer	AdamW ( Loshchilov and Hutter, 2019 ) without weight decay (Adam for LEAPS)

Appendix

Table 2: Settings shared by every method and state point.

	Ising $25\times 25$			Potts $16\times 16$
		energy calls			energy calls
Method	train (h)	local	black box	train (h)	local	black box
DIG-IN (ours)	$6.33_{\pm 0.11}$	$1.5{\cdot}10^{6}$	$4.8{\cdot}10^{8}$	$1.89_{\pm 0.07}$	$1.5{\cdot}10^{6}$	$3.9{\cdot}10^{8}$
MDNS	$6.05_{\pm 0.09}$	$7.7{\cdot}10^{5}$	$7.7{\cdot}10^{5}$	$1.79_{\pm 0.10}$	$7.7{\cdot}10^{5}$	$7.7{\cdot}10^{5}$
PDNS	$6.42_{\pm 0.35}$	$7.7{\cdot}10^{5}$	$7.7{\cdot}10^{5}$	$2.22_{\pm 0.30}$	$7.7{\cdot}10^{5}$	$7.7{\cdot}10^{5}$
DASBS	$11.57_{\pm 0.09}$	$1.1{\cdot}10^{8}$	$6.4{\cdot}10^{10}$	$5.60_{\pm 0.28}$	$1.1{\cdot}10^{8}$	$5.3{\cdot}10^{10}$

Appendix

Table 3: Computational cost. Training time in hours (mean

{\pm}

standard deviation over 3 seeds and the three phases) and the total number of energy calls during training, with the single-site difference table counted as one call (local) or as

D(q-1)+1

calls (black box), the cost it would have if the energy were not local.

	Ising	Potts, $q=3$
Lattice	$50\times 50$ torus, $D=2500$	$24\times 24$ torus, $D=576$
Training box	$\beta\in[0.39,0.65]$ , $\|h\|\leq 0.1$	$\beta\in[0.8,1.3]$ , $\|h\|\leq 0.01$
Zero-field $\beta_{c}$	$\tfrac{1}{2}\ln(1+\sqrt{2})=0.4407$	$\ln(1+\sqrt{3})=1.00505$
Network	residual CNN on the torus: 7 blocks, 96 channels, $3\times 3$ circular convolutions, conditioned by FiLM ( Perez et al., 2018 ) on Fourier features of $(t,\beta,h)$ ; dilations $(1,1,2,4,8,1,1)$ for Ising so that 7 blocks span the $50\times 50$ torus, none for Potts
Optimizer	AdamW ( Loshchilov and Hutter, 2019 ) , learning rate $3\times 10^{-4}$ , no weight decay, gradient clipping at 1.0, EMA 0.995
Batch / buffer	1536 / 49 152	1536 / 196 608

Appendix

Table 4: Settings of the amortized runs.

	Ag–Pd	Cu–Au
Lattice, cell	fcc, $5\times 5\times 5$ primitive	fcc, $8\times 4\times 4$ primitive
Sites $D$	125	128
$\sigma_{i}=+1$	Ag	Au
Phases	one miscibility gap	Cu 3 Au (L1 2 ), CuAu (L1 0 ), CuAu 3 (L1 2 )
Training box	$T\in[200,900]$ K, $\Delta\mu\in[-0.4,0.4]$ eV	$T\in[200,900]$ K, $\Delta\mu\in[-0.4,0.4]$ eV
Evaluation $T$	200–900 K, step 25 K (29)	200–900 K, step 25 K (29)

Appendix

Table 5: The two alloy tasks.

Figure 5: Common-tangent construction on Cu–Au. Mixing free energy against composition at 200 K (left pair) and 700 K (right pair), measured from the chord through the pure compositions. In each pair: exact (left) and from the draws of one DIG-IN run (right, exact curve in grey). Dashed chords are the accepted tie lines, whose ends are the phase boundaries at that temperature.

Interpolant and sampler
Noising process	masked (absorbing state)
Reveal schedule	one site per network evaluation ( $D$ evaluations per draw)
Training labels	heat-bath conditionals from the exact single-site energy-difference table, cached with the buffered endpoint
Loss	soft cross-entropy on masked sites
Network
Backbone	periodic residual CNN, 3-D circular convolutions over the cell

Appendix

Table 6: DIG-IN hyperparameters, identical for Ag–Pd and Cu–Au.

Transport
Rate network	locally equivariant CNN: centre-masked convolution (kernel 3), then a state-dependent convolution (kernel 5); 24 channels; FiLM conditioning
Rates	Metropolis-reparametrized, clamped at 10
Free-energy head	MLP, hidden width 200, input $(t,T,\Delta\mu)$
Path discretization	$\Delta t=0.02$ (50 CTMC steps)
Horizon curriculum	from 1 step to 50, quadratic over the first 2000 optimizer steps
Prior

Appendix

Table 7: FrOGS hyperparameters, identical for Ag–Pd and Cu–Au.

Panel of Figure 6 — Figure 6: Component ablation on Ising $25\times 25$ . $\ESS$ , $|\Delta\ln Z|$ and $\TVD_{\mathrm{mode}}$ against training steps for the ordered (top), critical (middle) and disordered phase (bottom), mean over 3 seeds. Dashed: $\ln 2$ and $\TVD_{\mathrm{mode}}=0.5$ , the values of a sampler that holds only one of the two ordered modes.

	disordered, $\beta=0.30$			critical, $\beta=\beta_{c}=0.4407$			ordered, $\beta=0.60$
	$\|\Delta\ln Z\|$	$\ESS$	$\TVD_{\mathrm{mode}}$	$\|\Delta\ln Z\|$	$\ESS$	$\TVD_{\mathrm{mode}}$	$\|\Delta\ln Z\|$	$\ESS$	$\TVD_{\mathrm{mode}}$
variant	$\times 10^{3}\,{\downarrow}$	${\uparrow}$	${\downarrow}$	$\times 10^{3}\,{\downarrow}$	${\uparrow}$	${\downarrow}$	$\times 10^{3}\,{\downarrow}$	${\uparrow}$	${\downarrow}$
DIG-IN (baseline)	$1.1_{\pm 0.4}$	$.995_{\pm.000}$	$.006_{\pm.004}$	$1.6_{\pm 0.3}$	$.980_{\pm.001}$	$.015_{\pm.004}$	$0.7_{\pm 0.4}$	$.979_{\pm.011}$	$.014_{\pm.007}$
no damping	$1.2_{\pm 0.5}$	$.992_{\pm.001}$	$.010_{\pm.007}$	$1.1_{\pm 0.4}$	$.966_{\pm.007}$	$.017_{\pm.012}$	$462.0_{\pm 326.4}$	$.994_{\pm.005}$	$.336_{\pm.232}$
no Gibbs score	$2.1_{\pm 1.9}$	$.965_{\pm.001}$	$.008_{\pm.004}$	$1.5_{\pm 0.9}$	$.895_{\pm.009}$	$.011_{\pm.008}$	$26479.6_{\pm 456.9}$	$.001_{\pm.000}$	$.064_{\pm.022}$
no Gibbs, no smoothing	$1.2_{\pm 0.8}$	$.974_{\pm.002}$	$.009_{\pm.005}$	$1.0_{\pm 0.2}$	$.943_{\pm.003}$	$.019_{\pm.004}$	$693.1_{\pm 0.8}$	$.993_{\pm.000}$	$.500_{\pm.000}$

Appendix

Table 8: Component ablation, final evaluation. Mean

{\pm}

standard deviation over 3 seeds,

|\Delta\ln Z|

10^{-3}

nats. PDNS and MDNS are repeated from Table 1 .

Panel of Figure 7 — Figure 7: Bridging two Ising temperatures. Left: one row per transport, time running left to right. The map shows $P(m)$ along the path; beside it, the magnetization at $\beta_{0}$ and $\beta_{1}$ (grey: Swendsen–Wang reference, blue: the model’s draws, dashed: reweighted), and the two target modes magnified. Right, top: one draw carried along each transport, at $t=0$ , $\tfrac{1}{2}$ and $1$ (upper row: across $\beta_{c}$ , lower row: ordered). Right, bottom: $\ln Z(\beta_{1})-\ln Z(\beta_{0})$ error, $\ESS$ and reweighted $\TVD_{\mathrm{mode}}$ , mean $\pm$ standard deviation over 3 seeds.

Figure 8: Effective sample size of the amortized samplers across

(\beta,h)

. Fraction of the 5000 draws per cell, mean over seeds.

Figure 9: Effective sample size across the alloy sweep. Fraction of the 5000 draws per state point, mean over 3 seeds, for Ag–Pd (top) and Cu–Au (bottom).

Explore similar work

MetaDNS: Enhancing Exploration in Discrete Neural Samplers via Well-Tempered Metadynamics
May 20, 2026Xiaochen Du, Juno Nam, Jaemoo Choi +7Accelerated Sampling Molecular Dynamics
Massachusetts Institute of Technology · Georgia Institute of Technology
FrOGS: Discrete Neural Sampler for Independent Alloy Configurations Across Chemical Conditions
Sep 1, 2026Kyucheol Min, Elyssa Hofgard, Tess SmidtMarkov Chain Monte Carlo Accelerated Sampling
Massachusetts Institute of Technology Cambridge, MA 02139
Discrete Diffusion Models via Evolving Variational Autoregressive Networks
Sep 23, 2026Kewen Pan, Ying TangScore-Based Diffusion Model Spatial Photonic Ising Machines
Institute of Fundamental and Frontier Sciences, University of Electronic Science and Technology of China, Chengdu 611731, China · School of Physics, University of Electronic Science and Technology of China, Chengdu 611731, China · Key Laboratory of Quantum Physics and Photonic Quantum Information, Ministry of Education, University of Electronic Science and Technology of China, Chengdu 611731, China +1

May 20, 2026, cond-mat.stat-mechJ/K move · Enter open · S save

MetaDNS: Enhancing Exploration in Discrete Neural Samplers via Well-Tempered Metadynamics

Xiaochen Du, Juno Nam, Jaemoo Choi, Wei Guo+6

Massachusetts Institute of Technology · Georgia Institute of Technology

Accelerated Sampling Molecular Dynamics Metastable Samples Phase Transitions Exploration

Open paper AskPDF arXiv

Sampling from discrete distributions with multiple modes and energy barriers is fundamental to machine learning and computational physics. Recent discrete neural samplers like MDNS suffer from mode collapse and fail to sample high-energy barrier regions between modes, which is critical for free energy estimation and understanding phase transitions. We propose Metadynamics Discrete Neural Sampler (MetaDNS), a general framework integrating well-tempered metadynamics into discrete diffusion or autoregressive samplers. By maintaining an adaptive, history-dependent bias potential along selected low-dimensional coordinates, MetaDNS forces exploration of previously inaccessible regions, enabling free energy reconstruction infeasible with standard neural samplers due to a lack of high-energy samples. On challenging low-temperature benchmarks including Ising, Potts, and the copper-gold binary alloy, MetaDNS reproduces the thermodynamic distribution. Compared to MCMC-based metadynamics, MetaDNS also achieves comparable exploration requiring fewer bias deposition steps.

Abstract

Explore similar work

MetaDNS: Enhancing Exploration in Discrete Neural Samplers via Well-Tempered Metadynamics

FrOGS: Discrete Neural Sampler for Independent Alloy Configurations Across Chemical Conditions

Discrete Diffusion Models via Evolving Variational Autoregressive Networks