MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control
Authors: Yuchen Zhu
Wei Guo
Jaemoo Choi
Guan-Horng Liu
Yongxin Chen
Molei Tao
Venue: NeurIPS 2025 poster ()
Keywords: Masked diffusion model
Continuous-time Markov Chain
Sampling
Markov-chain Monte Carlo
Optimal Control
Primary area: probabilistic_methods

Abstract: We study the problem of learning a neural sampler to generate samples from discrete state spaces where the target probability mass function $\pi\propto\mathrm{e}^{-U}$ is known up to a normalizing constant, which is an important task in fields such as statistical physics, machine learning, combinatorial optimization, etc. To better address this challenging task when the state space has a large cardinality and the distribution is multi-modal, we propose **M**asked **D**iffusion **N**eural **S**ampler (**MDNS**), a novel framework for training discrete neural samplers by aligning two path measures through a family of learning objectives, theoretically grounded in the stochastic optimal control of the continuous-time Markov chains. We validate the efficiency and scalability of MDNS through extensive experiments on various distributions with distinct statistical properties, where MDNS learns to accurately sample from the target distributions despite the extremely high problem dimensions and outperforms other learning-based baselines by a large margin. A comprehensive study of ablations and extensions is also provided to demonstrate the efficacy and potential of the proposed framework. Our code is available at https://github.com/yuchen-zhu-zyc/MDNS.

MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control
Yuchen Zhu Wei Guo Jaemoo Choi', Guan-Horng Liu?, Yongxin Chen', Molei Tao
Georgia Institute of Technology
FFIR at Meta
{yzhu738, wei guo, jchoi843, yongchen, mtao)@gatech..ddy ghliulmeta.com
Abstract
We study the problem learning neural sampler to generate samples from discrete state spaces where the target probability mass function x e is known up to a normalizing constant, which an important task in fields such as statistical physics, machine learning, combinatorial optimization, etc. To better address this challenging task when the state space has large cardinality and the distribution is multi-modal, we propose Masked Diffusion Neural Sampler (MDNS), a novel framework for training discrete neural samplers by aligning path measures through a family of learning objectives, theoretically grounded in the stochastic optimal control of the continuous-time Markov chains. We validate the efficiency and scalability of MDNS through extensive experiments on various distributions with distinct statistical properties, where MDNS learns to accurately sample from the target distributions despite the extremely high problem dimensions and outperforms other learning-baad based baselines by a large margin. A comprehensive study of ablations and extensions is also provided to demonstrate the efficacy and potential of the proposed framework. Our code available at https: // thub. wwwccoom/auooe hucyyuuchhhnhhhce zhu- 2yc/MDNS.
Introduction
Drawing samples from an unnormalized target distribution e some state space Xg is fundamental problem with wide-ranging applications across fields such as statistical physics [LB14], Bayesian inference [Gel+ 13], machine learning [And+03; WJ08), etc. For decades, Markov chain Monte Carlo (MCMC) methods, such as the Langevin Monte Carlo, the gorithm, and the Glauber dynamics, have been cornerstone. These methods simulate Markov chain whose stationary distribution the target distribution, and converge provably fast under certain cireumstances [Che22]. Despite their widespread adoption, MCMC algorithms face significant chal- lenges, especially when the state space is high-dimensional and the target distribution is multimodal, which hinders efficient exploration and convergence (see, e.g., the lower-bound results in [Ran06; GT06; GLRI8 88; HZ25]).
al-
Recently, inspired by the success of diffusion models in generative modeling for continuous data [HJA20; SMÉ21; Son+21], significant progress has been made in leveraging neural networks to learn stochastic differential equation (SDE) generate samples from I,, which collectively refer to as neural samplers (e.g., |ZC22; VGD23; AVE25; Hav+25]). Concurrently, the diffusion paradigm has been effectively extended to discrete state spaces [Aus+21; Cam+22; Sun+23b; LME24), finding broad applications in domains involving sequences of categorical data. However, while neural samplers have been extensively developed for distributions Id and discrete diffusion models
Equua contribution, random order.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

have shown promise in modeling discrete data, the study of leveraging discrete diffusion models for sampling remains relatively unexplored. Addressing this gap is crucial, as numerous real- world sampling problems in areas such as statistical physics [F1.24] and combinatorial optimization [Sun+23a; SHI.24; San+25] necessitate specialized methods navigate the unique challenges they pose. This work aims to develop such specialized approach.
Contributions, Our contributions can be summarized in the following points.
1. We introduce Masked Diffusion Neural Sampler (MDNS), a novel framework for training discrete neural samplers that leverages optimal control of CTMCs and masked discrete diffusion models.
2. To cope with the optimization challenges caused by the discontinuity of CTMC trajectories, we introduce several learning objectives that are free of differentiability requirements.
3. To achieve scalable training in high dimensions, we propose novel training loss named Weighted Denoising Cross-entropy (16) that directly implements score learning on the target distribution using importance sampling.
4. We conduct comprehensive experiments and ablation studies on Ising and Potts models validate the efficacy of our method. MDNS can accurately sample from target distributions even when the state space has cardinality 10 substantially surpassing other learning- based methods.
Related works. We refer readers to App. for a detailed discussion of related works.
2 Preliminaries
2.1 Continuous-time Markov Chains
continuous-time Markov chain (CTMC) X (X.) is stochastic process defined on probability space F,, Pr) and takes value in finite state space X. Its law is completely characterized by the generator Q = (Qi E XxXx ( defined by
lim 1 (Pr(Xr+se= s/x, = x) -1 Ar-0 At
(=(
where LA is the indicator function of the statement A, being one if A holds and zero if otherwise. By definition, the generator satisfies Qx y) 2 for x # and Q2(1, = Eyzn Qz(z. y).
The path of X, i.e., t..tte + Xe(w), is a piecewise constant and càdlàg' function. The path measure of a CTMC X is a probability measure on F2) defined as Px = Pr(X E A), which is the distribution of X on Ω. We define P.X as the marginal distribution of X,, and P. (-Ix() as the distribution of X, conditional from the literature of CTMCs, which shows how to compute the Radon-Nikodým (RN) derivative between two path measures driven by CTMCs with different generators and initial distributions, and crucial for developing our proposed sampling algorithms in (14):
X,
The following lemma is an important result to know
Lemma 1. Given two CTMCs with generators Q7 Q² and initial distributions Pi, on X, let pl 1p2 be the associated path measures. Then for any trajectory & (E.)ve/u.1))
log dp (E) = -log
dp2
djz
(Eu)
+
Σ log
(E) (E)
/' Σ (Q)(())) Q? (2)
dup
1:E #B
A detailed proof is provided in App. C.3. An intuitive interpretation of (2) is to view the RN derivative between path measures as the limit of density ratios between fiite-dimeeniionl joint distributions, and approximate the transition distribution by (1). We remark that for masked diffusion models, (2)
can be precisely calculated without discretization error, as will be seen later (14).
Acronyym of the French phrase standing for "right continuous with left limits".

---

2.2 Discrete Diffusion Models
In discrete diffusion models, one learns the generator of a CTMC that starts from simple distribution Piuit and reaches the target distribution Pdata at the final time. This CTMC is typically chosen as the time-reversal of noising process that converts any distribution to Pinit An especially effective subclass of discrete diffusion is the masked discrete diffusion models [LME24; Ou+25; Sah+24; Shi+24a; Zhe+25], which corresponds to choosing Pinit to Prnakks the Dirac distribution on the fully masked sequence. Adding a mask token M into the original state space Xo == (1, 2,,,... N) containing length-D sequences with vocabulary size N, denote the mask- mask-au state space by X ;= 21, ..., N, N, Throughout the paper,
(x* x!. denotes sequence of tokens, r de represents the sequence constructed by replacing the d-th position of z by n, and -UM = (x"
x." * M) represents the unmasked part of x E X.
Suppose need to model sequential data X E Xy following the distribution Pdata. The intuition behind the masked discrete diffusion model is to define a noising CTMC that independently and randomly masks each token according to certain schedule, then reverse this CTMC to generate data from a sequence of pure mask tokens. It is proved in [Ou+25] that for x # y E X, the generator for the unmasking generative process enjoys special structure and can be written
Qe v) = (t) Pr (x* = = UM = x UM,, if = M and y = x den X--Pdata
(3)
and O if otherwise, for some noise schedule y : T] + 10, oo) such that S (t) = oc. In
practice, practitioners typically leverage neural network se that takes a partially masked sequence x E X' as input and outputs (x E RD whose (d, n) entry approximates Prx~paat (X" n [X UM = - UM if x² M.
A standard way for training a masked discrete diffusion model given i.i.d. samples from Pdata is using the denoising cross-entropy (DCE) loss [Ou+25; Sah+24; Shi+24a]:
min LE parta (x) Ex-UUiit(0.)) wo E (a
Σ d:
- log so so())
where the transition kernel px(-lxr means independently masking each entry of x with probability and the weight w-) can be any positive function. In particular, the loss with w(A) x corresponds to the any-order autoregressive loss [Ou+25, Eq. (3.7)).
Sampling from a masked discrete diffusion model can be achieved through multiple schemes, such as the Euler method, variants of r-leaping [LME24; Ou+25; Ren+25b; Cam+22], uniformization [CY25], and seemi-Jautoregressiv sampler [Ar++25; Nic+25J. In this paper. we mainly use an exact sampling method known the random order autoregressive sampler [Ou+25, App. J.4], imple- mented by choosing a uniformly random permutation of {1...., D} and autoregressively unmasking each position along the permutation conditional on the observed positions.
3
Control-based Learning of Discrete Neural Sampler
3.1 Discrete Sampling with CTMCs
In this paper, we focus on the task of drawing samples from target distribution π(x) - (xx) on the finite state space X, {1, ...,...N Here, we only have access the potential function U, and the normalizing constant 7 " z e Ux() is unknown. Moreover, using CTMCs to sample from this target distribution, in the sense that we hope to find generator
interested in
Q'= (00)) such that it drives CTMC X (Xi)te(o,T) with X, Pinit (a readily sampleable initial distribution) to reach the target distribution at the final time T, i.e.,. Xx ~ π.
While this setup is simple and ideal, the problem is essentially ill-posed and hard to solve due to the non-uniqueness the feasible generator Q", We thus seek to restrict the solution space by enforcing stricter constraints: finding generator Q' such that the associated path measure I* satisfies
1 p ({) = ((Go.))(r)mm( = (E) dx (ET) = (E) 7 er(Er), where == -U - llg Phases
(5)
3

---

for any trajectory& ç = (E))7 Here, Fo is reference path measure with known generator Q", readily sampleable initial distribution P =: Pinit, and known final distribution PP =: Pbase-
Such formulation has natural connection the stochastic optimal control (SOC) of CTMCs. In fact, parameterizing our candidate generator as Q" and assuming that it induces path measure Ps, the learning of Q" through matching |d IP can be understood as controlling a CTMC with generator Qo to reach a new terminal distribution т. We demonstrate this connection in detail by considering the following SOC problem on X::
//'
Qi Qi
- + 20++03) (X,, (x..))dt r(Xx) -(()))
min Ex
Σ X
1g
log
Q2
(6)
s.t. X (X)))e(o.7) is CTMC on X with generator Q", Xo ~ Piait:
It can be shown that the problem (6) is equivalent minimizing the Kullback- (KI.) divergence between two path measures P" and IP", (L[P" (P*) Ep log de= Define the value function as
(v) = inf Ex~P"
Σ XFX.
Q² log
Q
Q. 29 Q
(X., u)d v)ds
r(XT) Xx
(7)
In order to guarantee the existence and uniqueness of the solution to (6) for any r and Pinits we need to choose a reference path measure Ipo that is memoryless [DE+25], i.e., Xo and Xr are independent for X P0_ In this case, the unique optimal solution Q" can be expressed as multiplicative perturbation of the reference generator Q":
(7 v) Q. (z. y) exp(V.(y) Vi(x)), Yz zy
(8)
Moreover, the optimal path measure IP" associated with Q" has the identical form as (5): dF
(E) =
a
z er((rr) Vg, where Z = Epg e". These results can be shown using Lem. and standard SOC theory, and we include the proofs in App. C.2 for self-consistency and better readability.
3.2 Optimal Control Formulation of Discrete Neural Sampler Training
As discussed in the previous section, if we manage learn the generator Q" that produces a path measure 1d matching IP', we can sample from the target distribution by simulating the CTMC with this generator. However, minimizing (L(P")P))
not the only available choice of objective reach the goal. In fact, it has some disadvantages due to the inherent discontinuous nature of the problem. Instead, we propose the following general framework for learning the discrete neural sampler:
Framework: Given a target distribution π X e- on Xg. a choice of reference path measure 560 with generator Q" learn the generator Q" to sample from π through optimizing an efficiently estimable objective F(P", IP*), where 1P* is given in (5) and Q argminq= F(P", P*).
is straightforward to see that the problem (6) a special case of the proposed framework upon choosing ([[P",") KL(P"M|PP) However, although directly relevant to many theoretical results, naively optimizing the discretized, estimated objective (6) is not an appropriate approach for learning the desired controlled generator Q". The estimation of (6) requires the simulation of CTMC using a neural network parameterized generator with parameter 0, yet the objective itself is inherently non-differentiable with respect to 0 since (1) the trajectories of CTMCs are pure jump processes, and (2) the reward function r is non-differentiable, making it difficult to effectively train the generator.
A workaround is proposed in [Wan+25a], where the authors retain the differentiability of the objectives through relaxing the CTMC trajectories from staying on X, to sequences of probability vectors on D using the Gumbel softmax trick. This partly solves the problem at the cost of introducing high approximation error into the learning process, since the # gradients coming from the Gumbel softmax trick are known to be biased and casily cause numerical instability [JGP17; Liu+23]. This leads to failure to converge to the correct target distribution, even low dimensions, as empirically validated in our experiments (see App. D.5). Moreover, backpropagation over the entire trajectory is a memory-intensive operation, and we would like to avoid it as much as possible. the following, we propose several alternative learning objectives that operate without these disadvantages.

---

Relative-entropy with REINFORCE (RERF). One of the key reasons that KL(P"I|" is intractable for direct optimization originates from the fact that the expectation is with respect to P", which tracks the 0 gradient. Alternatively, we can estimate Vo (T[[I""| by introducing gradient-nontracking trajectory simulated using ū = stopgrad (u), also known as the REINFORCE trick [Wil92; RGBI 14; MG14]. We denote the path measure produced with ū to be p82, and express the gradient of the relative-entropy by the following identity: Vp (LPPP" = V&JRERF (IP", PP), where
dIP*
dIp*
FRERF (TP", P*) Rg log
log
dPa + C) VC
dP"
See App. C.3 for detailed derivation. Thus, FREEF can be introduced for the purpose of gradient estimation of relative entropy. We remark that, in contrast to the approach proposed in [Wan+25a], (9) yields an unbiased estimator of the relative-entropy gradient, which is crucial for accurate learning of the target distribution. However, FRERR (PP", IP*) is not a valid "loss function" but only a computational surrogate, in the sense that reduction in the objective value does not guarantee an improvement in learning performance.
Log-variance (LV). Other than KL divergence on path measures, can also consider the variance type of losses, such as Log-variance (LV), Originally proposed in [NR2I] to solve SOC problems for SDEs on Rd, LV considers minimizing the following e-dependent objective
FLV FLv(PP,P) = Varp. log
(10)
dl
where u for now is generic and Pd a chosen sampling path measure driven by the generator Q" that does not require gradient backpropagation, Note that the optimality of the solution (10) is guaranteed under weak regularity assumptions on IP". In practice, we choose v = ū as in [NR21] I] for algorithmic effectiveness. (10) can be efficiently computed for CTMCs leveraging Lem. and (5).
Cross-entropy (CE). While the aforementioned RERF and LV objectives have an optimality guar- antee of the solution, they often do not enjoy optimization guarantees as in general these objectives have a nonconvex landscape. To mitigate it, we consider cross-entropy between IP" and F.
dlp"
dlp dr
FTE P²) = KL((P* IP*) ) = Eg- log dF
= Ep
log
(11)
dP
dp
FCE(Pu, IP*) is conver in pa due to the convexity of + log l, thus enjoying benign optimization landscape. However, since pe not directly tractable for simulation, we need to introduce an auxiliary sampling path measure Ipd and equivalently express the objective based on IP". This can be understood as an importance sampling estimation of the original objective, and in practice we use = u.
To sum up, the RERF, LV, and CE losses do not involve the error for approximating the discrete variables by continuous ones, do not require backpropagation over the entire trajectory of states, and thus are more preferable for efficient and stable optimization.
3.3 Masked Diffusion Neural Sampler
Besides the objective F(Pu, IP)), another major component in the design space of the proposed training framework is the choice of reference path measure po0 and the corresponding generator Q" Recall from Sec. 3.1 that we need the reference path measure to be memoryless guarantee the existence and uniqueness of the SOC problem (6). We now introduce a method for choosing such memoryless reference path measure based on masked discrete diffusion model. This approach, which we term the Masked Diffusion Neural Sampler (MDNS), forms the core of our framework. The corresponding learning algorithms are subsequently developed based on the objectives proposed in Sec. 3.2. Furthermore, we demonstrate that this framework can be extended to incorporate uniform discrete diffusion models, an adaptation referred to as the Uniform Diffusion Neural Sampler (UDNS), with detailed discussion deferred to App. F.
We choose p. to be the generative process of masked discrete diffusion model, starting from Pinit + (I) =: I==(M....M) and terminating at Paata + Punit (z) := x Lee the uniform
"This means we sample a trajectory IPM and detach from the computational graph. After that, use the values of the trajectory to compute the objective and do backpropagation. As long 8d is supported almost everywhere on the space of paths.
5

---

Algorithm Training of Masked Diffusion Neural Sampler (MDNS)
Require: score model sg, batch size B, training iterations K, reward r Xo - IR, learning objective F E {FRERF; FLV, JWDCEJ (num. replicates R, resample frequency l: for FWDCE).
1: for step = to K do
2:
if {JRERF, FLV. } then
3:
{x ( Jisss Sample_ Trajectoeeej)))
See Alg. for details. See (15).
4:
Compute with (x(",, ("(
5:
else = FWDCE then
6:
if step mod k then
b Sample new trajectories every l steps.
7:
{x (" Jssis Sample Trajectories(B), Set replay buffer B + (x(,, W" (x(()))) Isiinn (X(1, W
8:
9:
Resample with Mask(B; R). &. See App. B for details. Compute FWDCE with (x(1, (X(())))sss
10:
11: Update the parameters & based on the gradient VaF.
return trained score model sg.
distribution the unmasked data space No. Based (3), the corresponding generator is
7(()
(e x = (()) Pr X-oPans
(x"= = n UM x: 1
N 1M==M-
(12)
One can prove the following lemma (see App. C.3 for the proof):
Lemma 2. Under the assumption that Jo (t))dt co, the reference path measure TPO with generator Qº defined in (12) and starting from Pmask memoryless and satisfies F4. Punif-
With such a choice for pD, the optimal generator Q" also has a special structure:
Lemma 3. The generator Q' corresponding to the optimal solution of (6) with Q defined as the memoryless reference path measure (12) satisfies
Qi (i, = ((t) Pr =n UM 1 XNT
(13)
See App. C.3 for the proof. This suggests that it suffices to parameterize Q nd+n)= ((t)s)(z)d I 1= Here, 8q X' + R D is parameterized to have non-negative entries and each row sums up to 1, representing the one-dimensional marginals of the target distribution conditioned on unmasked entries of the input Moreover, such parameterization implies that the diagonal entries of Qi are
(x, Σ Σ (x, vfz
=(()
E 1= 1= M)}.
d
where the first equality is due to (3) and the second equality is due to the parameterization of Q" Similarly, one can also verify that Eyze Q y) = 7() zd = -M). This turns out to significantly simplify the expression of the RN derivative, as will be shown later.
With the results above, path measures [P* and [p",, the common term for computing FRERF, Fiv, and FCE:
Lem. I to derive the RN derivative between the optimal and the current
dl log dp (X) = r(x,) log Z
Oo log
(X,, / T Σ (Q - (Xe) v)
Σ
Q
:: *X,
VAX
1/N
=r(Xr)+ +
Σ
log so(x-) (x-))((0,xx(10
log Z = W" (X) - log Z.
1: AX
1/N
where (xXX) r(XT) Σ
log so s0(X,- ()()))())xx.11
(14)

---

Here, we assume that the jump from X₁- to X, is at the d(t)-th position, and the total number of jumps is D. With (14), the affrremeniooeed training objectives (9) to (11) are simplified to
FRERF - E W (X))"(XX,, Fix = Var (X), FCE - E X~20
W (X0 ("X),, (15)
X~ps
X~Pp
where we have removed terms related to Z in FRERF and FLv without modifying the optimization landscape (in particular, C' is chosen as log Z in (9)). For the CE loss, as in practice the normalizing constant Z may be prohibitively large, removing z from the loss may lead to numerical instability.
To avoid this, we propose to estimate Z via the equality Z Ex e (X) implied by (14), which is equivalent to applying softmax to the weights {Wo (W"(X)))) (XSsS in a batch.
Scalable training via weighted denoising cross-entropy. While the learning objectives in (15) successfully avoid backpropagation along all the states X and are thus relatively memory efficient, their implementation still wastes much compute. Note that for computing W (XX), we need to call
the model s-(-) D times, but each time only the (d(t), used in optimization. What's worse, the three losses in (15) require backpropagation through all of the D tradknntttracking score outputs in (14), which may be unaffordable due to the GPU memory constraints. To propose a more scalable loss for high- nigh-dimensional data, we start with a further simplification of FCE by discarding the u-independent terms in (X)):
x >-th Xx element of the D x N output matrix
1 (X) 1 (X)(X( min Ex-pa 7 (X)
min Ex-pa 7
Σ - -log so (Xt-) a(0))
Here, the key idea is to treat i.idd samples from I" as importance weighted samples from P Instead of only learning the conditional distribution sp on a set of positions and values on the generation trajectory, we can actually forget about the trajectory and only focus on the achieved clean sample Xr. By remasking XT and computing the DCE loss in (4), we arrive at the following weighted denoising cross-entropy (WDCE) loss:
FWDCE (P",P") =
E X-Ppe
(x)
((A)
Σ
- log so (z)d,xg
Uninff((.1)
(
M
(16)
where we estimate by Ex-ps (X) in FCE. is straightforward to note that (16) is equivalent to the DCE loss (4) with π being the data distribution pdata. except that X7 are now importance weighted instead i.i.d. samples from π. FWDCE is much more efficient than FCE in that we can use all the output of the score model s-(-) to compute the loss, instead of only one element. Moreover, for each pair of (Xr,ew((X) (x)),
sample multiple (say, R) partially masked J compute the loss, so the expensive O(D) computation of the RN derivative can be amortized. Inspired by recent works for reusing samples [Mid+23; Hav+25], we can also create a replay buffer B that stores pairs of final samples and weights (Xr, (x)( for multiple steps of optimization for further reduction of the computational cost. Our algorithm is summarized in Alg. 1.
3.4 Theoretical Guarantees
We can establish the following theoretical guarantees for our learning algorithms (proof in App. C.3); Propusition 1 (Guarantee for sampling). Let Psamp := Py be the distribution of the samples generated from the learned sampler. Then, to ensure (L [T) S s2 (resp., (L(T||Psamp) S )²),
it suffices train the sampler to reach <L(PP"|[") S E² (resp., (L [F < E²),
Proposition 2 (Guarantee for normalizing constant estimation). Ż c (x) X P is an (1
< < 2 3 4. it suffices to train the sampler to reach min min(KL((e) [KL(P" [P*', KL(PP < 5.t One can boost the probability from 3 to 1 - for any < E (0, 2) by taking the median of O (log] 2 i.i.d. estimates.
unbiased estimation of Z. To guarantee Pr
4 Experiments
In this section, we experimentally validate our proposed frameworks by learning to sample from the Ising model and Potts model on square lattices. At different temperatures, these models exhibit
7

---

Aameat 0.4407
Frogh 0.28
10000
2 point Comelation
0.4
0.300
0.0
0,030
0.4
0.2
0:871
Distance
Distance
Distance
-
MENE (oun)
-M0 LEAPS
Metmpsleelbeenng
-4 Ground Trath (5W])
Figure 1: Average of 2-point correlation Crow (ks k r) of samples from 16 16 Ising model.
distinct behaviors, providing a rich ground for demonstrating the effectiveness of our algorithms. We use probabilistic metries (KL divergence, TV distance, x² divergence, etc.) and observables (magnetization, 2-point correlation) originating from statistical physics to benchmark our learning results. We include both learning-basee approaches such as LEAPS [HAJ25], and learning- Satropolis-Hasting (MH) algorithm benchmarks. For fair comparisons,
approaches such as the we run all baselines with a long enough but comparable compute time. To accurately estimate ground truth values for all the observables when an exact computation is intractable, running the Swendsen-Wang (SW) algorithm [SW86; SW87] for a sufficiently long time. We include experiment details such as metric definitions and training hyperparameters in Apps. D and E.
draw samples by
4.1 Ising model on Square Lattice
We consider learning to sample from the Ising model on a square lattice with L sites per dimension, A {1, ..... 2)2. The state space is Xn = {+1]A, where +1 represent the spin states. Given an interaction parameter J E RR an external magnetic field h E R, and an inverse temperature 3 > 0, the probability distribution of any configuration
Xo is given by
1 -H((), where (z) = =J Σ π {σ) 7
x ل
Σ x'' and Z = Σ e (H(())
(17)
In the above equation, i j means that i.j E A are adjacent the lattice. For simplicity, impose periodic boundary conditions in both the horizontal and vertical directions.
A case study on the choice of F. We first compare the effects of different learning objectives F by learning to sample from a 4 4 high temperature Ising model with J = 1, h 0.1, and Bnigh 0.28. The system size is picked so that the sampling problem is not trivial but exact computation of the partition function remains computationally tractable. The results are reported in Tab. 2, and results at other temperatures are available in App. D. For fair comparison, size of 256 for each loss, and use R = 16 resampling replicates for FWDCE so that the effective batch size is the same across all objectives. We draw sufficiently large number (229) of samples from the learned model to estimate all the reported metrics. We also report the average effective sample size (ESS, defined in (24)) of the last 100 steps showcase the learning speed for each objective. We include a baseline produced by running the MH algorithm for sufficiently long as a reference. We can that all four learning objectives can learn good model that generates ground-truth-like samples. In all the evaluation metries except the absolute error of log 2, the effectiveness ranking is uniformly FRERF > FCE.
train for 1000 steps with a batch
FLV > FWDCE >
Scaling to higher dimensions and lower temperatures. We also on significantly harder settings by increasing the system size to 16 x 16 (so the cardinality of the state space is 2256). We keep J h 0 and vary the inverse temperature, ranging from Snigh 0.28, Peritical log(1 + 22)/2 0.4407, to Blow 0.6. The behaviors of the Ising model are known to be distinct at these temperatures by the phase transition theory [KW41; Ons44]. For the purpose of efficient training, we use FWDCE as learning objectives and train on each problem for 50k steps. In Tab. 1, we report the ESS and the absolute error of magnetization and 2-point correlation (respectively defined in (26) and (28)) to the estimated ground truth. We also plot the estimated average 2-point correlation (defined in (27)) between varying distance r in Fig. 1, and defer further results to
8

---

Table 1: Results for learning 16 16 Ising model different temperatures, best in bold. Mag. and Corr. represent the absolute eπror of the magnetization and 2- 2-point correlation to ground truth values.
Temperature
Blow 0.6
Beritical 0.4407
Shigh 0.28
Metrics
Mag. +
Corr.
ESS
Mag.
Cor 4
ESS f
Mag. J
Cor J
ESS t
MDNS (ours)
9.9c-3 3
2.4e-3
0.981
3.7e-3
2.0e-3 3
0.933
8.5e-3
1.0c-3 3
0.962
LEAPS
2.40-2 2
5,8c-1
0.261
7.4-3 3
1.6c-1 1
0.384
7.4e-3 3
1.6e-3
0.987
Baseline (MH)
1.9p-2 2
7.7e-4 4
4.6e-3 3
2.5e 3
6.
1.1e-3
Table 2: Results for learning 4 x Ising model with J 1, 0.1, and Bnigh 0.28, best in bold.
Method
ESS
(V 7 4. KT(Pemplliz) 4 x² (Poampl (1))
KT(P"IP*)
Abs. err. of log
FREAF
0.9621
0.0799
0,0380
0.0845
0.0188
0.00003
FLy
0.9713
0.0748
0.0348
0.0714
0.0141
0.00046
FCB
0.9513
0.0833
0.0393
0.0903
0.0248
01.00999
FWnch
0.9644
0.0799
0.0382
0.0868
0.0177
0.00030
Baseline (MH)
0.0667
0.0325
0.0628
App. D.3. Our learned distribution accurately matches the theoretically predicted trend of correlation functions, which exponentially decays to the non- spontaneous magnetization at Slow. polynomial decays Beritical and exponential decays to at Phigh- As shown the table and figure, MDNS outperforms LEAPS by a large margin across almost all metrics. MDNS accurately learns the high-dimensional distribution across all temperatures, as evidenced by low observables errors and high ESS values, whereas LEAPS fails at Beritical and 3low despite long training.
Warm-up for lower temperatures. Instead of only training from scratch, our training includes warm-up phase. When targeting hard distributions such as Peritical and 3 to help the model better locate the modes of the target distribution, we start the training by fitting easier ones such as Shigh for short (20k steps in this example) period of time. An ablation study to demonstrate the importance of this proposed technique is presented in App. D.2. A more systematic study on the warm-up is deferred to future work [Guo+25].
Preconditioning. In learning diffusion samplers for continuous distribution e Rd. one typically leverages the score information V log y in the neural network architecture to guide the sampling process for faster convergence and better sampling quality, known as preconditioning. We also explore similar techniques to improve the training of MDNS, see App. D.4 for details.
4.2 Potts model on Square Lattice
We now consider a harder problem known as the standard Potts model, which has q(> 2) spin states on a square lattice with L sites per dimension, A = {1, 2) i.c., the state space is No = 21, / Given an interaction parameter J R and an inverse temperature B > 0, the probability distribution of any configuration x E Mu is given by
π (x) 7 AM(x) where H(x) =J] -J= 1= and 7 injj
Σ (8()))
(18)
zEAo
In the above equation, i ~ j also means that i.j E A are adjacent on the lattice, and we impose the same periodic boundary conditions. As generalization of the Ising model, the Potts model also exhibits diverse behaviors across different temperatures, and the critical inverse temperature is known to be Beritical = log(1 + /q) [BDC12]. In the following experiment, we consider the system size to be L. = 16, fix q 3 and J = 1, and vary the inverse temperature B8 ranging from Snigh 0.5. Beritical log(1 v3) 1.005, Blo 1.2.
Similar to the training for the Ising model, we use FWDCE and train on each temperature for 100k steps, among which 30k steps are in the warm-up phase for training Beritical and Stow- We report quantitative metrics (defined (30) and (32)) in Tab. and similarly plot the estimated average 2- 2-point correlation (defined in (31) between varying distance r in Fig. 2. More results can be

---

Adtial 1.00
Stigh
0.5
2. 2-pint Correlation
Distance
Distance
[oum]
(3N))
Figure 2: Average of 2- 2-point correlation Crow (ke, k + rr of samples from 16 x 16 Potts model.
Table Results for learning 16 x 16 Potts model at different temperatures, best in bold. Mag. and Corr. represent absolute errors of magnetization and 2-point correlation ground truth values.
Temperature
Pro 1.2
Peritiea 1.005
Bhigh 0.5
Metries
Mag. 4
Corr. | ESS 1
Mag. 4
Corr. 4
ESS
Mag. 4
Curr. +
ESS f
MDNS (ours)
1.3e-3 3
8.8e
0.933
4.3e-3
2.9c-3 3
0.875
2.2c-
5.8c 4
0.983
LEAPS
2.9e-1
2.5e-
0.012
2.7e-1
2.0e-1
0.004
2.9e-3
1.2e-3
0.991
Baseline (MH)
7.4e-1
5.6e-1 1
5.2e-1
3.5e-1
3.5e-2
1.6e-2
found in App. E. Again, our approach recovers distribution with correlation function whose decay patterns are well aligned with the ground truth, whereas other baselines fail to do so. It is worth noting that the MH algorithm cannot mix even after continuous simulation of more than 20 hours, while remains successful approach for all Ising model experiments. Such a difference likely originates from the exponential increase in the state space cardinality |.t]] despite the system size remaining the same. For example, a 4 x 4 Ising model has around 60k distinct states, while a 4 x 4 Potts model with & = 3 has 40M distinct states, which clearly highlights the increased difficulty of Potts model sampling. Nevertheless, our approach succeeds with moderate increase in training iterations. Moreover, the evaluation shows that MDNS performs best across almost all metries and temperatures, suggesting that our framework and learning objectives are highly scalable.
5 Conclusion, Limitations and Future Directions
This paper introduces Masked Diffusion Neural Sampler (MDNS), a novel framework for training discrete neural samplers based stochastic optimal control and masked diffusion models. While MDNS has proven effective for distributions arising in statistical physics, its performance on other families of discrete distributions, such those graphs related to combinatorial optimization, unknown. We also propose that the framework can be extended to fine-tume tne-tune pretrained discrete diffusion models given (possibly non-differentiable) reward function, which is left for future work [Zhu+25d]. Theoretically, we conjecture that our interpolation between the masked and the target distribution possesses superior properties compared to the geometric annealing (H, x e e-nU)) ))e used in LEAPS [HAJ25]. A rigorous theoretical analysis of annealing paths in discrete spaces, potentially building upon insights in continuous spaces [GTC25b; GTC25a], remains key area for future investigation.
Acknowledgments and Disclosure of Funding
WG and MT thank Peter Holderrieth and Michael Albergo for insightful discussions on the paper [HAJ25] during ICLR 2025. YZ and MT are grateful for partial support by NSF Grants DMS- 1847802, DMS-2513699 DOE Grants NA0004261, SC0026274, and Richard Duke Fellowship. WG, JC, and YC acknowledge support from NSF Grants ECCS-1942523, DMS-2206576, and CMMI-2450378.
10

---

References
[Ala+23]
Sarah Alamdari. Nitya Thakkar, Rianne van den Berg, Alex Lu, Nicolo Fusi, Ava Amini, and Kevin Yang. "Protein generation with evolutionary diffusion". . In: NeurlPS 2023 Generative Al and Biology (GenBio) Workshop. 2023. URL: https: /ppprevveeeeee.ee/ net hett/formmee I 1itpmmmi.ddp6999pP... Christophe Andrieu, Nando De Freitas, Arnaud Doucet, and Michael Jordan. "An introduction to MCMC for machine learning". In: Machine learning 50 (2003), pp. 5 43. DOI: 1 1023/A: A:
[And+03]
1020281327116.
[Am+255]
Marianne Arriola, Subham Sekhar Sahoo, Aaron Gokaslan, Zhihan Yang. Zhixuan Qi, Jiaqi Han, Justin I Chiu, and Volodymyr Kuleshov. "Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models". In: The Thirteenth International Conference on Learning Representations. 2025. URL: https: / /ttp//ppeeeeeeeeeeeeeeeeeeeeeeee net pottffoo htt//formmmiie Byr
[AS+24]
Tara (Khound-Sadegl Jarrid rector-Brooks, Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Alexander Tong. "Iterated Denoising Energy Matching for Sampling from Boltzmann Densities". In: Proceedings of the 41st International Conference on Machine Learning. Ed. by Ruslan Salakhutdinov. Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett. and Felix Berkenkamp. Vol. 235. Proceedings Machine Learning Research. PMIR, 2024, pp. 760-786. URI.: https: http:::://ppoceedigggeeee mlr press/v233 aathound-sadegee
html. Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. "Struc- tured Denoising Diffusion Models in Discrete State Saate-S In: Advances in Neural Information Processing Systems. Fd. by A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan. 2021. URL.: httpa://openreiier
[Aus+21]
meet/forumiddhh7-XXX XixPAAL.
[AVE25]
Michael Samuel Albergo and Eric handen-Eijnden, "NETS: A Non-equilibrium Transport Sam- pler". In: Forty-second International Conference on Machine Learning. 2025. URL: https //openreview www.nne//eeee http://rrmmmnnddddstttteeee
[Bai+25]
Jinbin Bai, Tian Ye, Wei Chow, Enxin Song, Qing-Guo Chen, Xiangtai Li, Zhen Dong. Lei Zhu, and Shuicheng YAN. "Meissonic: Revitalizing Masked Generative Transformers for Efficient High- Resolution Text-to- tomage Synthesis". In: The Thirteenth International Conference on Learning Representations. 2025. URL: http:://///ppeeviieee// net ffoormmiddGGaaaa Jssuhh..
[BDC12]
Vincent Beffara and Hugo Duminil-Copin. "The self-dual point of the two-dimensional random- cluster model is critical for q 1". In: Probability Theory and Related Fields 153.3 (2012). pp. 311-542, DOE: 10000/0004000 011 0353-8.
[Be166]
Richard Bellman. "Dynamic programming". In: science 153.3731 (1966), pp. 34-37.
[Ben+24]
Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. "From denoising diffusions to denoising Markov models", In: Journal of the Royal Statistical Society Series B: Staristical Methodology 86.2 (2024). pp. 286-301.
[Ble+25]
Denis Blessing, Julius Berner, Lorenz Richter, and Gerhard Neumann. "Underdamped Diffusion Bridges with Applications to Sampling". In: The Thirteenth International Conference on Learning Representations. 2025. URL: ht //atts://ppeee /pppnreeeeeeeeeeee//eeee nnet/ pett/formm' formmm 1ddtmmmpdoq0000
[BRU24]
Julius Bemner, Lorenz Richter, and Karen Ullrich. "An optimal control perspective on diffusion- based generative modeling". In: Transactions Machine Learning Research (2024). ISSN: 2835-888 8835- URL: bittps://opereee
[Cam+22] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligianni- dis, and Arnaud Doucet. "A Continuous Time Framework for Discrete Denoising Models" In: Advances in Neural Information Processing Systems. Ed. by S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh. Vol. 35. Curran Associates, Inc., 2022, pp. 28266- 28279. URL: https: : / prwceedingg..s..ooooe/p/.o...e.p neurips cc sccpaper__. files/ paper / 2022/ / file/ wttt////00033001..eeeeeeeeeeaeee
www.aaa000---------------- Conference. pdf.
[Cha+22] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. "MaskGIT: Masked
generative image transformer". In: 2022 JEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR). 2022, pp. 11315-11325. DOI: 10. 1109/CVPR52668. 2022. 01103.
[Che22]
Sinho Chewi. Log-Concave Sampling. Book draft, in preparation, 2022. chewisinho. github.ii iio.
URL: https: /
[Che+25a]
Jannis Chemseddine, Christian Wald, Richard Duong. and Gabriele Steidl. "Neural Sampling from Boltzmann Densities: Fisher-Rao Curves in the Wasserstein Geometry". In: The Thirteenth International Conference on Learning Representations. 2025. URL: httpp://opeettt //tp://v/////eeeeeeeeeee .ntt/ forum??dd-Uugsšudddd
[Che+25b]
Haoxuan Chen, Yinuo Ren, Martin Renqiang Min. Lexing Ying, and Zachary Izzo. "Solving inverse problems via diffusion-based priors: An approximation free ensemble sampling approach". In: arXiv preprint arXiv:2506.03979 (2025).

---

[Che+25c] Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing. Gerhard Neumann, and Anima Anandkumar. "Sequential Controlled Langevin Diffusiios".. In: The Thirteenth International Conference on Learning Representations. 2025. URL: https: http///pppeeeeeeeeeeeeeeeeeeeeee wtwpiiiwwnn///ooumm/eee
dImD2sgy86.
[Che+25d]
Tong Chen, Yinuo Zhang, Sophia Tang. and Pranam Chatterjee. "Muuti-objective-giided discrete flow matching for controllable biological sequence design". In: arXiv preprini arXiv:2505. 07086 (2025).
[Chi+23] Cheng Chi, Zhenjia Xu, Siyuan Feng, Fric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. "Diffusion policy: Visuomotor policy learning via action diffusion". In: The International Journal Robotics Rexeurch (2023), p. 027836692411233668. DOI: 10. 1177/ 02783649241273668. URL: https: / //to//do.//oo /tt/://..rr/....1////1/////55//0 0117111//0077789992777666
[Cho+25]
Jaemoo Choi, Yongxin Chen, Molei Tao, and Guan-Homg Liu. "Non-equilibrimm Annealed Adjoint Sampler". In: arXiv preprint arXiv:2506. 18165 (2025).
[CY25]
Hongrui Chen and Lexing Ying. "Convergence Analysis of Discrete Diffusion Model: Exact Implementation through Uniformization". In: Journal of Machine Learning (2025). DOE: 10. 4208/jm1. 240812.
[CZH23]
Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. "Analog Bits: Generating Discrete Data using Diftusion Models with Self colf-Conditioning" In: The Eleventh International Conference on Learning Representations. 2023. URL: https: h/ http://ppp/eeeeeeeeeeeeeee net forumiidd watimmdddditjjjooal. 3itjR9Qxxw..
[DE24]
Carles Domingo-EEnrich, "A laxonomy of loss functions for stochastic optimal control". In: arXiv preprint arXiv:24J0. 00345 (2024).
[DE+25]
Carles Domingo-Fane Enrich. Michal Drozdzal, Brian Karrer, and Ricky T. Q. Chen, "Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control". In: The Thirteenth International Conference on Learning Representations. 2025. URL: https:
//openreview net/ cntt/forummiin diruuidddxxBBrrQ@BB...
[DG25]
Justin Deschenaux and Caglar Gulcehre. "Beyond Autoregression: Fast L.LMs via Self-Distillation Through Time". In: The Thirteenth International Conference on Learning Representations. 2025. ew httt//orrmmdddu25Rd...
URL:: https://oppneee
[Dos+21]
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer. Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlshy. "An Image is Worth 16x 16 Words: Transformers for Image Recognition at Scale". In: International Conference on Learning Representations. 2021, URI:: net. http:/ormmmddii.cbaaaaaaa
https://openrerii
[Dou+22]
Amaud Doucet, Will Grathwohl, Alexander G Matthews, and Heiko Strathmann. "Score- BBsed Diffusion meets Annealed Importance Sampling". In: Advances in Neural Information Processing Systems. Fd. by S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh. Vol. 35. URL:: https / prwceedinee..eeee//oo neurips Paper-
Curran Associates, Inc., 2022, pp. 21 21482-21494
cc / _ files/ paper 2022 / file/ 8667:288ea3360ffffff660003030200 Conference.pde pdf.
[EHI17]
Weinan E, Jiequn Han, and Arnulf Jentzen. "Deep tearming-based numerical methods for high- dimensional parabolic partial differential equations and hackward stochastic diffcrential equations". In: Communications in mathematics and staristics 5.4 (2017) 7), pp. 349-380. DOT: 10. 1007 /
940304-017-0117-
[Ess+24]
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, and Robin Rombach. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis". In: Proceedings of the 4 International Conference on Machine Learning. Ed. by Ruslan Salakhutdinov. Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp. Vol. 235. Proceedings of Machine Learning Research. PMLR, 2024, pp. 22606-12633, URL: https: / http///ooeeenns..e//e
press/ wwrresss//////eee v255 v233//sser2233 esser24a. html.
[Fen+25]
Guhao Feng, Yihan Geng, Jian Guan, Wei Wu, Liwei Wang. and Di He. "Theoretical Benefit and Limitation of Diffusion Language Model". In: arXiv preprint arXiv:2502.00622. (2025).
[FL24]
Michael Faulkner and Samuel Livingstone. "Sampling Algorithms in Statistical Physics: Guide for Statistics and Machine Learning". In: Statistical Science 39.1 (2024), pp. 137 -164.
DOI: 10. 2214/23-STS994 URL.: https: //doi. org/100 10. 1211/23--1S829
[Gat+24]
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. "Discrete Flow Matching". : In: The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024. URL: https: /ppnreveeeeeeeeeeeeeeeee net net/formm?dd GTDKo3Sv9p.
12

---

[GD23]
Tomas Geffiner and Justin Donke. "Langevin Diffusion Variational Inference". In: Proceedings of The 26th International Conference Artificial Intelligence and Statistics. Ed. by Francisco Ruiz, Jennifer Dy, and Jan- Mil-ellemm van Meent. Vol. 206. Proceedings of Machine Learning Research. PMLR, 2023, pp. 376-593. URL: https: : petps://pree //ppoceediggg...... mlr. prwssreeevvve v22e// get fne223a. html.
[Gel+13]
Andrew Gelman, John B. Carlin, Hal S. Stem, and Donald B. Rubin. Bayesian data analysis. 3rd ed. Chapman and Hall/CRC, 2013.
[GLRI8]
Rong Ge, Holden Lee, and Andrej Risteski. "Simulated tempering Langevin Monte Carlo II: An improved proof using soft Markov chain decomposition". . In: arXiv preprint arviv / (2018).
[Gon+25]
Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye. Lin Zheng. Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, Hao Peng, and Lingpeng Kong. "Scaling Diffusion Lan- guage Models via Adaptation from Autoregressive Models". In: The Thirteenth International Conference on Learning Representations. 2025, URL: httpe: / //oporpeeeeeeeeee werrriiinne//foooe www//eerormmiie j1tSLYKugB.
[Gru+23]
Nate Gruver, Samuel Don Stanton, Nathan C. Frey. Tim G. J. Rudner, Isidro Hotzel, Julien Lafrance-Vanasse, Arvind Rajpal, Kyunghyun Cho, and Andrew Gordon Wilson "Protein Design with Guided Discrete Diffusion". In: shirty-seventh Conference on Neural Information Processing Systems. 2023. URL: attps://opeet evier. pet/forummideee 1k69Ga6p.
[GT06]
David Galvin and Prasad Tetali. "Slow mixing of Glauber dynamics for the hard-core model on regular bipartite graphs". In: Random Structures & Alggriitmms 28.4 (2006). pp. 427-443. DOI:
10. 1002/rra, 20094.
[GTC25a]
Wei Guo, Molei Tao, and Yongxin Chen. "Complexity Analysis of Normalizing Constant Estima- tion: from Jarzynski Equality Annealed Importance Sampling and beyond". arXiv preprint arXiv:2502.0455 (2025).
[GTC25b]
Wei Guo, Molei Tao, and Yongxin Chen. "Provable Benefit of Annealed Langevin Monte Carlo for Non-log-onn con-log--oncave Sampling". In: The Thirteenth International Conference on Learning Repre- sentations. 2025, URL: https: /tops:://ppp /ep/pppeeevieeeeee www.eeee//toee
[Guo+24]
Wei Guo, Yuchen Zhu, Molei Tao, and Yongxin Chen. "Plug-and-Play Controllable Generation for Discrete Masked Models". In: arXiv preprint arXiv:2410.02145 (2024). Wei Guo, Jaemoo Choi, Yuchen Zhu, Molei Tao, and Yongxin Chen. "Proximal Diffusion Neural
[Guo+25]
Sampler". In: arXiv preprint arXiv:2510.00 038224 (2025).
[HAJ25]
Peter Holderrieth. Michael Samuel Albergo, and Tommi Jaakkola. "LEAPS: A discrete neural sampler via locally equivariant networks". In: Forty-second International Conference on Muchine Learning. 2025. URL: https: /ttp:////eeeeeeeeeeeeeeeee wwtp:/eeeeetttttttteeeeeee Hoqn??dd-HQRRRIAAAEET
[Hav+25]
Aaron J Havens. Benjamin Kurt Miller, Bing Yan, Carles Domingo- Enrich, Anuroop Sriram, Daniel S. Levine, Brandon M Wood, Bin Hu, Brandon Amos, Brian Karrer. Xiang Fu, Guan- Horng Liu, and Ricky Q. Chen. "Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching". . In: Forry-second International Conference on Machine Learning. 2025. URL: https: / //ppn//eeeeeeeee///////e/ee net/ net/forum? ? id id- OrHmm22....
Domingo-Enrich,
[Hay+25]
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q Tran, Jonathan Deaton, Marius Wiggert, et al. "Simulating 500 million years of evolution with a language model". In: Science (2025), pp. 850-858. DOI: 10. 2/2///cience... ads0018.
[Hc+25]
Jiajun Hc, Yuanqi Du, Francisco Vargas, Dinghuai Zhang, Shreyas Padhy, RuiKang Ou YuYangg Carla Gomes, and José Miguel Hernández-Lohato "No Trick, No Treat: Pursuits and Challenges Towards Simulation- (2025).
Training Neural Samplers". In: arXiv preprint arXiv:2502.n 2502.06685
[Heo+25] Bycongho Hco, Song Park, Dongyoon Han, and Sangdoo Yun. "Rotary Position Embedding for Vision Transformer". In: Computer Vision ECCV 2024. Ed. by Aleš Leonardis, Elisa Ricci, Stefan Roth, Olga Russakovsky, Torsten Sattler, and Gül Varol. Cham: Springer Nature Switzerland, 2025, pp. 289-305. ISBN: 978-3-031-72684-2.
[HJA20]
Jonathan Ho, Ajay Jain, and Pieter Abheel. "Denoising Diffusion Probabilistic Models". In: Advances in Neural Information Processing Systems. Ed. by H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin. Vol. 33. Curran Associates. Inc.. 2020, pp. 1851. URL: https / proceedings neurips cc / paper files ppaper/ / 2020 / file/ / @ccbcfeecS844ff008671111.0..00.0
Paper .ddf.
[HJI8]]
Jiequn Han, Arnulf Jentzen, and Weinan E. "Solving high-dimensional partial differential equa- tions using deep learning". . In: Proceedings of the National Academy of Sciences 115.34 (2018). pp. 8505-8510. DOr: 10.1073/pnas. 17 1718942115.
[HLVE24] Mengjian Hua, Matthieu Laurière, and Eric Vanden-Eijade Eijuden. "A Ceauulaatioo-Feeee Deep Learning Approach to Stochastic Optimal Control". In: arXiv preprint arXiv:2410.05163 (2024).
13

---

[HRT224]
Ye He. Kevin Rojas, and Molei Tao. "Zeroth-Order Sampling Methods for Non-Log-Concan Distributions: Alleviating Metastability by Denoising Diffusion". In: The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024. URL: https::/t///// htapp://oppereiee..o nett/
forum?id-X3AM x3A]juss.....
[HRT25]
He, Kevin Rojas, and Molei Tao. "What Exactly Does Guidance Do in Masked Discrete Diffusion Models". In: arXiv preprini arXiv:2506. 10971 (2025).
[HZ25]
Yuchen He and Chihao Zhang. "On the query complexity of sampling from non-log-concave distributions (extended abstract)". . In: Proceedings Thirty Eighth Conference on Learnigg Theory. Ed. by Nika Haghtalab and Ankur Moitra. Vol. 291. Proceedings of Machine I leamnng Research. PMLR, 2025. pp. 2786-2787. URL: atttps://proceeee
mmr. press/v291/hrr
html.
[Jar97]
Christopher Jarzynski. "Nonequilibrium Equality for Free Energy Differences". In: Phys. Rev. Lett. 78 (14 1997), pp. 2690-2693. DOI: 10. 1103/ PhysRevLett. 78 2690. URL.: https
//link. a wpw. wrg//ogg/loooo /otg//oo/1....// Po0//PhyRewwett.... -78 . 2690.
[JGP17]
Eric Jang. Shixiang Gu, and Ben Poole. "Categorical Reparameterization with Cumbel-SSottas In: International Conference on Learning Representations. 2017. URL: https::
[Jin+25]
Yang Jin, Zhicheng Sun, Ningyuan Li, Kun Xu, Kun Xu, Hao Jiang, Nan Zhuang, Quzhe Huang. Yang Song. Yadong MU, and Zhouchen Lin. "Pyramidal Flow Matching for Efficient Video Generative Modeling". In: The Thirteenth International Conference Learning Representations. 2025. URL: https: pestps://opereeee /tppr//pppe/eeeeeee/eeee//e wwwevvwwweeet.eeeeeeeeeeee wwt.neet//ruuu 66NzzcRuQQ
[Kim+25]
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham M. Kakade, and Sitan Chen. "Train for the Worst, Plan for Best: Understanding Token Ordering in Masked Diff'usions.. In: Forty-second International Conference on Machine Learning. 2025. URL: https : // openrevies, net forum? ? id-Dj Jmre5IkP.
[KW41]
H. A. Kramers and G. H. Wannier. "Statistics of the Two-Dimensional Ferromagnet. Part I". In: Phys. Rev. 60 (3 1941), pp. 252-262. DOI: 10. 1103 PhhysRev 60. 252. URL: https:
//link. aps oogg//i//// /11/100110/ 1103/PhysRev....m 60. 60..5500.....2 .52..
[LB14]
David P. Landau and Kurt Binder. A Guide Monte Carlo Simulations in Statistical Physics.
4th ed. Cambridge University Press, 2014.
[LH19]
Ilya Loshchilov and Frank Hutter. "Decoupled Weight Decay Regularization". In: International Conference on Learning Representations. 2019. URL:: httpe/////peeeeeeeeeeeee Bkg6Ricqy7.
[Li+25]
Mufei Li, Viraj Shitole, Eli Chien, Changhai Man, Zhaodong Wang. Srinivas, Ying Zhang, Tushar Krishna, and Pan Li. "LayerDAG: A Layerwise Autoregressive Diffusion Model for Directed Acyclic Graph Generation". . In: The Thirteenth International Conference on Learning Representations. 2025. URL:: https: http:///pppeeeeeeeeeeeeeee
[Liu+23]
Liyuan Liu, Chengyu Dong, Xiaodong Liu, Bin Yu, and Jianfeng Gao. "Bridging Discrete and Backpropagation: Straight- Strought-Through Through and Beyond".. In: Thirty- seventh Conference Neural Infor- mation Processing Systems. 2023. URL: https /tppprpeeeeeeeeeeeeeeeeeeeeeeeee wwwpaareieennnn//oruuueee http/forrnniddmmyyllll
[Liu+25a]
Anji Liu, Oliver Broadrick, Mathias Niepert, and Guy Van den Broeck. "Discrete Copula Dif- fusion". In: The Thirteenth International Conference on Learning Representations. 2025. URl: https: //per//epeeeeeee net/ / Srrrum/dddFFOOOkoccdl
[Liu+25b]
Guan-Horgg Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, and Ricky TQ Chen. "Adjoint Schrödinger Bridge Sampler". In: arXiv preprint arXiv:2506.22555 (2025). Sulin Liu, Juno Nam, Andrew Camphell, Hannes Stark, Yilun Xu, Tommi Jaakkola, and Rafael Bombz--lo "Think while You Generate: Discrete Diffusion with Planned Denoising". In: The Thirteenth International Conference Learning Representations. 2025. URL: https: //oenreview http///ummiii///laaaesee
[Liu+25c]
[LME24]
Aaron Lou, Chenlin Meng. and Stefano Ermon. "Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution". In: Proceedings of the 41st International Conference on Machine Learning. Ed. by Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp. Vol. 235. Proceedings of Machine Learning Research. PMLR, 2024, pp. 28819-32448. 1ou24a. html.
URL: ptcttss:/prrcceedeee
mlr press/v236/
[MF23]
Bálint Máté and François Fleuret. "Learning Interpolations between Boltzmann Densities". In: Transactions on Machine Learning Research (2023). ISSN: 2835-8856. URL: https : // openrevi .ne./ee//ee www.//eeemmnne/nee oorunidd-T6VrRebbhh
[MG14]
Andriy Mnih and Karol Gregor. "Neural Variational Inference and Learning in Belief Networks". In: Proceedings of the 31st International Conference on Machine Learning. Fd. by Fric P. Xing and Tony Jebara. Vol. 32. Proceedings of Machine Learning Research 2. Bejing, China: PMLR, 2014, pp. 1 1791-1799, URL: https::/ /ttpee/////cceeeee//e mlr. ttess//32/mmm 222/mn1144..1 htm.html. html.
14

---

[Mid+23]
Laurence Illing Midgley, Vincent Stimper. Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernandezz-obato. "Flow Annealed Importance Sampling Bootstrap". In: The Eleventh International Conference on Learning Representations. 2023. URL: https: haap:://openreeee http///ppeeeeeeeeeeeeeeeeeeeeeee net/ forum? Add-CCVVFJWS9L..
[Nea01]
Radford M. Neal. "Annealed importance sampling". In: Statisticy and Computing 11.2 (2001), pp. 125-139, 139. ISSN; 1575-1375. DOI: 10. 1023/A A 1008923215028 URL: https: //doi. .rrg/
10. 1023/022/AA0 A: 1008923215028.
[Nie+25]
Shen Nie, Fengqi Zhu. Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. "Large language diffusion models". In: arXiv preprint arXiv:2502.0999 (2025).
[Nis+25]
Hunter Nisonoff. Junhao Xiong. Stephan Allenspach, and Jennifer Listgarten. "Unlocking Guid- ance for Discrete State-Space Diffusion and Flow Models". In: The Thirteenth International Conference on Learning Representations. 2025. URL: https: pettpp://openreviiiet pepperrevie... wwwriiiw.neet//eee wwt./eet/fuuutee formmid XagH154y07.
[No6+19]
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. "Boltzmann generators: Sampling equilih- rium states of many-body systems with deep learning". In: Science 365.6457 (2019), caaw 147.
[NR21] Nikolas Nüsken and Lorenz Richter. "Solving high-dimensional Hamilton Jacobi Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space". In: Partial differential equations and applications 2.4 (2021), p. 48. DOt: 10. 1007/
42985-021-00111
[Ons44]
Lars Onsager. "Crystal Statistics. I. A Two-Dimensional Model with an
Transition". In: Phys. Rev. 65 (3-4 1944), pp. 117-149. DOI: 10. 1103/ PhysRev. 65. 117. URL: https:
//link. aww.aaa......................a wwt.//og///oo/ooooo /t////////1//// 1101 /0..///y/eee..... 65. 117.
[Ou+25]
Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan
Li. "Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data". In: The Thirteenth International Conference on Learning Representations. 2025. URL: net htt//formmmdddd sy
[OZL25]
Zijing Ou, Ruixiang Zhang, and Yingzhen Li. "Discrete Neural Flow Samplers with Locally Equivariant Transformer". In: arXiv preprint arXiv:2505. (00225055.7771 1774I (2025).
[Par+25]
Yong-Hyun Park, Chich-Hssin Lai, Satoshi Hayakawa, Yuhta Takida, and Yuki Mitsufuji. "Jump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models". In: The Thirteenth International Conference on Learning Representations. 2025, URL: https: pattps://openee //tp////ee///eeee//////ee/e nnt/ forum?'d-pDDTiipypa
[Ran06]
Dana Randall. "Slow mixing of Glauber dynamics via topological obstructions", In: Proceedings of the Seventeenth Annual ACM-SIAM Symposium on Discrete Algorithm. SODA '06. Miami, Florida: Society for Industrial and Applicd Mathematics, 2006, 370-879, ISBN: 089871603
[Ren+25a]
Yinuo Ren, Haoxuan Chen, Grant M. Rotskoff, and Lexing Ying. "How Discrete and Continuous Diffusion Meet: Comprehensive Analysis of Discrete Diffusion Models via Stochastic Integral Framework". In: The Thirteenth International Conference on Learning Representations. 2025. URL: https: //ppsr////peeeiiii ewreeviieeee.eee nwt id 6awxxxQQ182.
[Ren+25b] Yinuo Ren, Haoxuan Chen, Yuchen Zhu Wei Guo, Yongxin Chen, Grant M. Rotskoff, Molei Tao, and Lexing Ying. "Fast Solvers for Discrete Diffusion Models: Theory and Applications of High- Order Algorithms". In: The Thirty-ninth Annual Conference Neural Information Processing Systems. 2025. URL: https: / http://pperevveeee eviet ew meet/formm? d=Dux11663a0..
[Ren+25cc]
Yinuo Ren, Wenhao Gao, Lexing Ying. Grant M Rotskoff, and Jiequn Han. "Driftl Lite: Lightweight drift control for inference-time scaling of diffusion models". In: arXiv preprint arXiv:2509.3 1655 (2025).
[RGB14] Rajesh Ranganath, Sean Gerrish, and David Blei. "Black Box Variational Inference". In: Pro- ceedings of the Seventeenth International Conference on Artificial Intelligence and Statistics. Ed. by Samuel Kaski and Jukka Corander. Vol. 33. Proceedings of Machine learning Research. Reykjavik, Iceland: PMLR, 2014, pp. 223/ranganath14. html.
URL: https: / proceediggg mlr. press/ /
[Ris+25] Severi Rissanen, RuiKang Ou Yang, Jiajun He, Wenlin Chen, Markus Heinonen, Arno Solin, and José Miguel Hernández-Lobal "Progressive Tempering Sampler with Diffusson".. In: Forty- second International Conference on Machine Learning. 2025. URL: https : htt:///pppeeevieeee net/forum SatrrummidduBuabbBB....
[Roj+25a] Kevin Rojas, Ye He. Chieh-Hsie Lai, Yuta Takida, Yuki Mitsufuji. and Molei Tao, "Theory- Informed Improvements to Blassifier-Free Guidance for Discrete Diffussin Models". In: arXiv preprint arXiv:2507.08965 (2025).
[Roj+25b] Kevin Rojas, Yuchen Zhu, Sichen Zhu, Felix X-F. Ye, and Molei Tao. "Diffuse Everything: Multi- modal Diffusion Models on Arbitrary State Spaces". In: Forty-second International Conference on Machine Learning. 2025. URL: https: / opwrrevieeeeeeeeeeeeeee wtwpleewwweeeeeeeeeeeeeeeeeeeeee Attp:/oooonnnddddiiiieee
15

---

[RRY25]
Yinuo Ren, Grant M Rotskoff. and Lexing Ying. "A Unified Approach to Analysis and Design of Denoising Markov Models". In: arXiv preprint arXiv:2504. 01938 (2025).
[Sah+24]
Subham Sekhar Sahoo, Marianne Arriola, Aaron Gokaslan, Edgar Mariano Marroquin, Alexander M Rush, Yair Schiff, Justin T Chiu, and Volodymyr Kuleshov. "Simple and Effective Masked Diffusion Language Models". In: The eiihty Annual Conference on Neural Information Processing Systems. 2024. URL: https: /ttp::////peeeeew/e/ee wwwp:/rwwwwwwwee.eeeeeeeeeeeeeee htte///oomiidd--ulaaa
[San+25]
Schastian Sanokowski, Wilhelm Franz Berghammer, Haoyu Peter Wang, Martin Ennemoser, Sepp Hochreiter, and Sebastian I Lehner. "Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics". In: The Thirteenth International Conference Learning Representations. 2025. URL: ttttps://pperse
net. http//ourmmidddiiiippppiieeee
[Sch+25]
Yair Schiff, Subham Sekhar Sahoo. Hao Phung, Guanghan Wang. Sam Boshar, Hugo Dalla- torre, Bernardo P de Almeida, Alexander M Rush, Thomas PIERROT, and Volodymyr Kuleshov. "Simple Guidance Mechanisms for Discrete Diffusion Models". In: The Thirteenth International Conference on Learning Representations. 2025. URL: httpa: // /wwpppeeeieeeeeeeee www.ee...eeeeeeeeeeeeeeeeeeeeeee net/forummid- 15Mr J6g5G1.
[SD+15]
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. "Deep Unsu- pervised Learning using Nonequilibrium Thermodynamies". In: Proceedings of the 32nd In- teratiional Conference on Machine Learning. Ed. by Francis Bach and David Blei. Vol. 37. Proceedings of Machine Learning Research. Lille, France: PMLR, 2015, pp. 1256-2265. URL: https://proceeddngg www.eeiiiinnmm.......mmmmmammeee prwss /v37 /t////33//////// /shh/solllaaaa di ckstein15..1.. hhmml.
[Sha+25]
Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Karrer, Yaron Lipman, and Ricky T. Q. Chen. "Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspective". In: The Thirteenth International Conference on Learning Represen- tations. 2025. URL: https: // //tpp//eeeeeeeee////e/eeee net/ wwwooooouuunntttiiiiii
[Shi+24a]
Jiaxin Shi, Kehang Han, Zhe Wang. Arnad Doucet, and Michalis Titsias. "Simplified and Generalized Masked Diffusion for Discrete Data". In: The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024. URL:: https:/ http:///pppeevveeee www.//eeu//eeececem XcqSDfft4g.
[Shi+24b]
Zhekun Shi, Longlin Yu, Tianyu Xie, and Cheng Zhang. "Diffusion Penbisiion-PINNN Sampler". In: arXiv preprint arXiv:2410. 15336 (2024).
[Shi+25]
Qingyu Shi, Jinbin Bai, Zhuoran Zhao, Wenhao Chai, Kaidong Yu, Jianzong Wu, Shuangyong Song, Yunhai Tong, Xiangtai Li, Xuclong Li, et al. "Muddit: Liberating generation beyond text-to- image with unified discrete diffusion model". In: arXiv preprint arXiv:2505. 25666 (2025). Sebastian Sanokowski, Sepp Hochreiter. and Sebastian Lehner. "A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization". In: Proceedings the 41st International Conference on Machine Learning. Ed. by Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp. Vol. 235. Proceedings of Machine Learning Research. PMLR, 2024, pp. 43346 43367. mlr. www.ppesssssss..
[SHL24]
URL: https: //rocceecsss..s
http:///////aaoosss1k/a.s sanokowski24s html.
[SME21]
Jiaming Song. Chenlin Meng. and Stefano Ermon. "Denoising Diffusion Implicit Models". In: International Conference Learning Representations. 2021. URL: https:// net/ forumPid-StlgiarCHi.
[Son+21]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. "Score-Based Generative Modeling through Stochastic Differential Equations". In: International Conference on Learning Representations. 2021. URL: https: // /ttp///eppeeeeeeee////ee/ee .eet
[Sun+23a]
Haoran Sun, Katayoon Goshvadi, Azade Nova, Dale Schuurmans, and Hanjun Dai. "Revisiting Sampling for Combinatorial Optimization". In: Proceedings of the 40th International Conference on Machine Learning. Ed. by Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Searlett. Vol. 202. Proceedings of Machine Learning Research. PMLR, 2023, pp. 08859-32874, URL: /ppppe/e/////o/iggg mlr presse /2222/ sun23c. html.
[Sun+23b]
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. "Score-based Contiquous- Discrete Diffussin Models". In: The Eleventh International Conference on Learning Representa- tions. 2023, URL: https: /// http///ppeeeeeeeeeeeeeeee eview. weee//oorumm/ddBBW/essss
[SW86] Robert H. Swendsen and ShangS Wang. "Replica Monte Carlo Simulation of Spin-Glasses"" In: Phys. Rev. Lett. 57 (21 1986). pp.
DoL: 10. w.0./hhyeeeeeet..ec 2607. URL:
https: : :/t::///nnnne... h/tpiikk.pps.... www.appp..o..ooia org/doi/10 11111110/1pp/ www////yseweeee....ooo 572667.
[SW87]
Robert H. Swendsen and Jian-Sheng Wang. "Nonuniversal critical dynamics in Monte Carlo simulations". In: Phys. Rev. Lett. 58 (2 1987), pp. 86-88. DOI: 10. www./haaeeeeeet..tt..m 86.
URL: https: pestpa://limms //link. aps. .aaa..r///oe htt://oi/// / 10. 1101 1103//hysRevLett..c.mm 58.
16

---

flows". In: Statistics and [Tan+25] Sophia Tang, Yuchen Zhu, Molei Tao, and Pranam Chatterjee. "TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion". In: arXiv preprint arXiv:2509.25171 (2025).
[Tou+21] Hugo Touvron, Matthieu Cord, Matthij Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. "Training data-efficient image transformers & distillation through attention". In: Proceedings of the 38th International Conference on Machine Learning. Ed. by Marina Meila and Tong Zhang. Vol. 139. Proceedings of Machine Learning Research. PMLR, 2021, pp. 10347- 10357. URL: https: / /pproceedingg..... mlr. wwwpsrrsssvvveee /v1/s//vv////// wttlv/11//////e /tttoroooonaaa.. a. html.
[TPL24]
Yifeng Tian, Nishant Panda, and Yen Ting Lin. "Liouville Flow Importance Sampler". In: Proceed- ings of Machine Learning Research 235 (2024). Ed. by Ruslan Salakhutdinov, Zico Kolter, Kather- inc Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, pp. 48186- 48210, URL: https: httpas//poccedingeeeeee mlr. (rress/v235/ian///
html.
[TZC25]
Sophia Tang, Yinuo Zhang, and Pranam Chatterjee. "PepTune: Novo Generation of Therapeutic Peptides with Pulli-Objective-Gune Guided Discrete Diffusion".. In: Forty-second International Confer- ence on Machine Learning. 2025. URL: https: http:////peeeeeeeeeeeeee .nw.//e/se.eee net/forum id=FQoy 1Y1Hd8.
[Var+23] Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil D Lawrence, and Nikolas Nüsken. "Bayesian learning via neural Schrödinger- Computing 33.3 (2023). DOI: 10. 1007 s017/2211222-22 000-0022--1-1 10172-5.
[Var+24]
Francisco Vargas, Shreyas Padhy, Denis Blessing. and Nikolas Nüsken. "Transport meets Varia- tional Inference: Controlled Monte Carlo Diffusions". In: The Twelfih International Conference on Learning Representations. 2024. URL: https: / http:///ppp/eviiieeeeeeeeee wwwrriiw.neeffooooo nwtt/eruumr Sid-PPrruue IrudnxiW. Francisco Vargas, Will Sussman Grathwohl, and Amaud Doucet. "Denoising Diffusion Samplers". In: The Eleventh International Conference on Learning Representations. 2023. URl: https
[VGD23]
// //oenreview www//eeenunnneeeeeee htte//rrummnd@@gaflllaaele 1d=8pvnffAbbuf.
[VJ08]
Suriyanarayanan Vaikuntanathan and Christopher Jarzynski. "Escorted Free Energy Simulations: Improving Convergence by Reducing Dissipation". In: Phys. Rev. Lett. 100 (19 2008), p. 190601. URL: https: //1 //win///n////////////////....t aps org/ doi /too///..////////////////// 1103/
DOI: 10. 11t.111001PPh/// www.aeeeeeeeeeee.eeeeeeeeareeeee 100. 190601.
PhysRevLett. 100. 190601.
[Wan+24]
Xinyou Wang. Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. "Diffu- sion Language Models Are Versatile Protein Learners". In: Proceedings the 41st International Conference on Machine Learning. Ed. by Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver. Jonathan Scarlett, and Felix Berkenkamp. Vol. 235, Proceedings of Machine Learning Research. PMLR, 2024, pp. 52309- -20333. URL: https: / /ttp///eeocee.s.... mlr. press/v235 /0m3////aa22cc....mm html.
[Wa++25a]
Chenyu Wang, Masatoshi Uchara, Yichun Hc, Amy Wang, Avantika Lal, Tommi Jaakkola, Sergey Levine, Aviy Regev, Hanchen, and Tommaso Biancalani. TuinnnTTunng Diserete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design", In: The Thirteenth International Conference on Learning Representations. 2025. URL: hatppp://oppeee /ttpe/////peeeeeeeeeeeeeeee net/ forum? forum?id=43280ixxtae
[Wan+25b] Guanghan Wang. Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. "Remasking discrete diffusion models with inference-time scaling". In: arXiv preprint arXiv:2503.00307 (2025).
[Wat+23]
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert Ragotte, Lukas F Milles, al. "De novo design of protein structure and function with RFdiffusion". In: Nature 620.7976 (2023), pp. 1089- 1000. DOI: 10. 1038/ /41///11550000 //////11128000015---0 023-06415-8.
[Wi192]
Ronald Williams. "Simple statistical gradient-following algorithms for connectionist reinforce- ment lcarning". In: Machine learning (1992), pp. 229-256, DOI: 10. 10. BF09999666. Martin J. Wainwright and Michael Jordan. "Graphical Models. Exponential Families, and Variational Inference". In: Foundations and Trends@ in Machine Learning 1.1-2 (2008), pp. 1- 305. ISSN: 1935-8237. DOI: 10. 1561/2200000001. 2200000001.
[WJ08]
URL: http: / /xxx doi www.do.........ooooooooooooooo.. org/10. 1561/
[Xu+24]
Zhe Xu, Ruizhong Qiu, Yuzhong Chen, Huiyuan Chen, Xiran Fan. Menghai Pan. Zhichen Zeng. Mahashweta Das. and Hanghang Tong. "Discrete-state Continuous-tine Diffusson for Graph Generation". In: The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024. URL: https: /ttpaa///ppp http////ppeeeeeeeeee//eeeeeeeeee
wet//formmidd-VsskzE.. dVsSSKZZE.IY..
[ZC22]
Qinsheng Zhang and Yongxin Chen. "Path Integral Sampler: A Stochastic Control Approach For Sampling". In: International Conference on Learning Representations. 2022. URL: https: // /apenreview. .ne./eetttteee fftm/ffrrriddd.ommm wtt.orrid..uccccnaaara...
[Zha+25a]
Leo Zhang. Peter Potaptchik, Jiajun He, Yuanqi Du, Amnud Doucet, Francisco Vargas, Hai-Dang Dau, and Saifuddin Syed. "Accelerated Parallel Tempering via Neural Transports". In: arXiv preprint arXiv:2502.10326 (2025).
17

---

[Zha+25b]
Ruixiang Zhang, Shuangfei Zhai, Yizhe Zhang, James Thornton, Zijing Ou. Joshua M. Susskind, and Navdeep Jaitly. "Target Concrete Score Matching: A Holistic Framework for Discrete Dif- fusion". Forty-second International Conference on Muchine Learning. 2025. URL: https:
//openreview net/ cent/forum?i11 forrmmid-ZZee Ldd-2MrdvSn7x..
[Zhe+24]
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. "A Reparameterized Discrete Diffuusion Model for Text Generation". . In: First Conference on Lunguage Modeling. 2024. URL: https: wewpeeinii-PEEPU-----O
// //openreveew.netee..
[7he+25]
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang. "Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling". In: The Thirteenth International Conference Learning Representations. 2025. URL: https://openreree
net/ net/forum? Citt/iddCCCCcCC CmirNr.
[7hu425a]
Sichen Zhu, Yuchen Zhu, Molei Tao, and Peng Qiu. "Diffusion Generative Modeling for Spatially Resolved Gene Expression Inference from Histology Images". The Thirteenth International Conference on Learning Representations. 2025. URI: https: / http///ppeeeeieeeeeeeeeeeeeee eeeenrevie... net formmfrd PtjLUHyZAO.
[Zhu+25b]
Yuanzhi Zhu, Xi Wang. Stéphane Lathuilière, and Vicky Kalogeiton. "Di[M]O: Distilling Masked Diffusion Models into One-step Generator", In: arXiv preprint arXiv:2503. /5457 (2025). Yuchen Zhu, Tianrong Chen, Lingkai Kong, Evangelos Theodorou, and Molei Tao. "Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups". In: The Thirteenth In- ternational Conference on Learning Representations. 2025. URL: attps://openreee Dorumiid-DDat:jJTDD..
[Zhu+25c]
[Zhu+25d] Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, and Yongxin Chen. "Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization", In: arXiv preprint arXiv:2510.08233 (2025).
18

---

Broader impacts. This work focuses on foundational research in sampling algorithms for discrete distributions where the unnormalized probability mass function is known. Advancements in this area have the potential to accelerate scientific discovery and improve solutions to complex problems in fields like statistical physies and machine learning. Given the specific scope of our proposed method, we do not foresee direct negative societal impacts stemming from its application.
Organization of the appendix. App. A includes a detailed discussion of the related literature. App. B contains the pseudo-code for our proposed algorithm. We review the relevant theory of CTMC and SOC in App. C, and provide all proofs omitted in the main text. In Apps. D and E, we present detailed experimental setups and additional results for learning the Ising model and Potts model, respectively, including ablation studies and further visualizations. Finally, App. F discusses the extension of our SOC framework to the training of uniform diffusion neural samplers (UDNS).
A Related Works
Recently, a growing trend of research directions has been focusing on training neural samplers [Noé+19; He+25], where one amortizes the sampling procedure through learning a neural network. Most of the work has focused on learning distributions over the continuous Euclidean space R.d. One of the main approaches is to learn trainable drift term that drives an SDE to approximate the rime-reversal of a pre- pre-selected noising process that converts any target distribution t to noise distribution. This includes methods like Path Integral Sampler (PIS, [ZC22]). Neural Schrödinger-Föllmer Flows (NSFS, [Var+23]), Denoising Diffusion Samplers (DDS, [VGD23]), time-reversed DIffusion Sampler (DIS, [BRU24]). iterated Denoising Energy Matching (iDEM, [AS+24]), Diffusion-PIND Sampler [Shi+24b], Adjoint Sampling [Hav+25; Liu+25b; Cho+25]. etc. Another popular approach trains the process to transport from a prior distribution To to m along sequence of marginal distributions whose unnormalized densities are all available, such as the geometric interpolation X To = Methods of this type include the paper [MF23], Controlled Monte Carlo Diffusion (CMCD, [Var+24]). Liouville Flow Importance Sampler (LFIS (TPL224)), Non-Equliibrium Transport Sampler (NETS, [AVE25]), Sequential Controlled Langevin Diffusions (SCLD, [Che+25c]), the paper [Che+25a], underdamped diffusion bridges [Ble+25], Accelerated Parallel Tempering (APT, [Zha+25a]), Progessive Tempering with Diffusion (PTSD, [Ris+25]), etc. There are also works motivated by annealed importance sampling [Nea01] and Jarzynski equality [Jar97; VJ08], including Monte Carlo Diffusion (MCD, [Dou+22]) and Langevin Diffusion Variational Inference (LDVI, [GD23]). Despite the progress, these methods are restricted to the setting of Edd and cannot be casily extended to sampling discrete distributions. This distinguishes our proposed MDNS from all the aforementioned approaches.
Neural samplers.
In terms of methodology, MDNS is most relevant to PIS [ZC22]. Despite the difference in the underlying state space of the target distributions, PIS resembles our work in terms of method at high level, since it approaches the sampling problem in Idd also from an SOC perspective. However, PIS training does not suffer from the non-differentiabiiity serious problem in discrete state spaces that we need to address by proposing new learning objectives that do not require trajectory differentiabilityy
underlying sampling dynamics, which is a
In terms of task, MDNS is most connected to LEAPS [HAJ25], which focuses on the same problem of training discrete neural samplers using CTMC. We also note the concurrent work [OZL25]. which appeared after our submission and proposed a framework, DNFS, similar to LEAPS. Unlike our framework, LEAPS takes a measure transport perspective and trains the neural sampler to follow a sequence of pre-defined discrete distributions. However, LEAPS requires special inductive bias called locally equivariant networks to efficiently evaluate the training objectives, which greatly limits the algorithm's 's design flexibility. In contrast, our proposed MDNS does not impose constraints on the choice of score model backbone. Empirically, we also find that secorted-transport-beee approaches like LEAPS tend to be ineffective when learning multimodal, high-dimensional distributions, while MDNS remains competitive.
Discrete diffusion models. Diffusion models [SD+15; HJA20; SME21; Son+21] have been achieving state-of-the-art performances on generative modeling of various data modalities [Wat+23; Zhu+25e; Ess+24; Jin+25; 2hu+25a; Chi+23; HRT24; Roj+25b; Che+25h; Ren+25c]. As a natural generalization of diffusion models to discrete state space, discrete diffusion models [Aus+21; Cam+22;
19

---

CZH23; Sun+23b; LME24; Gat+24; Sha+25] have gradually attracted the attention of the research community as competent method for generative modeling of general discrete sequence data. Discrete diffusion models have seen successful applications in numerous tasks, including image synthesis [Cha+22; Bai+25; Shi+25], text generation [Nie+25; Gon+25; Zhe+24; Arr+25], protein [Gru+23; Ala+23; Hay+25; Wan+24; Che+25d; TZC25], graph [Xu+24; Li+25], combinatorial optimization [SHL24; San+25J etc. Motivated by the huge empirical success of discrete diffusion models in practice, researchers have also worked on the theory of discrete diffusion to understand the key behind its success [Ben+24; CY25; Ren+25a; Fen+25; Kim+25; RRY25; HRT25]. An extensive portion of the literature is also devoted to the development of techniques to enhance further the effectiveness of discrete diffusion models, such as guidance [Sch+25; Guo+24; Nis+25; Roj+25a], distillation [Zhu+25b; DG25], training [Zha+25b; Liu+25a], planning [Par+25; Liu+25c], and inference-time scaling [ [Wan+25b; Tan+25].
Masked discrete diffusion model [Sah+24; Ou+25; 2hi+24a; Zhe+25], as particularly effective variant of discrete diffusion, is most related to our work. Masked diffusion models enjoy many favorable theoretical properties, such as special structure of score functions that enables a simplified, variance-reduced training objective. Masked diffusion also has a natural connection to any-order autoregressive models and can be precisely simulated with fixed number of inference steps. The training of MDNS also benefits from these additional structures, making it more effective than other existing approaches. However, what we propose is indeed general framework and can be applied to other discrete diffusion models as well, such as uniform discrete diffusion. We defer the discussion of such extension to App. F.
Optimal control of stochastic dynamics. The literature of solving SOC problems has rich history, dating back to the work of Richard Bellman [Bel66], which introduces the Hamilton-Jacobi-eell (HJB) equation to characterize the optimal control through PDE. To efficiently solve SOC problems in high dimensions, neural networks are first introduced into the solving process by [EHJ17; HJEI 18] to alleviate the curse of dimensionality suffered by traditional approaches. Later, various learning objectives have been proposed to train the neural network for solving the SOC problems [DE+25; HLVE24]. We refer interested readers to [NR21] and [DE24] for a detailed review of the properties of different objectives.
Although all the works above focus on solving SOC problems in Rdd (i.e., controlling SDEs instead of CTMCs), MDNS is still tightly connected to this literature due to our SOC formulation of the discrete sampling problem. Moreover, in our main training framework, we took a path measure perspective for designing the learning objectives, which generalizes the perspective used in [NR21] to a discrete state space. However, there is notable difference between the SOC for CTMC (considered this work) and the SOC for SDEs (widely studied in the literature) due to the non-differentiable nature of CTMCs, which further limits the feasible choice of learning objectives. These differences separate our paper from the works mentioned above.
B Details of Algorithms
l.d. In Alg. 1, Resample_wt Batplll__wih_h__eeee Maakkha_kkkk means sample random variables (X(rr) hissaassn (ni 1), and for each i and r, obtain ((rr) by randomly masking each entry of X with probability X
~
is defined as in Alg. 2, which is also the inference
algorithm for MDNS.
C Theory of Continuous-time Markoy Chain and Stochastic Optimal Control
C.1
Continuous-time Markov Chain
We refer readers to [Cam+22; LME24; Sun+23b; LME24; CY25; Ren+25a; Ren+25b] for the general theory of CTMC. Below, we present several key lemmas that will be used in our paper.
Lemma 4. Let X be a CTMC with generator Q, and let denote the probubility distribution of X i.e. pr(-) Pr(X, = -). Then p satisfies the following Kolmogorov forward equation:
d.p((z)= Eu(uv = > ((Cr(,,z)piu Qe(z, s)pc(z)), Vx.
(19)
20

---

Algorithm 2 Sample_ Trajetoorsee Sample trajectories and compute weights
Require: score model 8p, reward function r: Ag - IR, batch size B.
Initialize fully masked sequences {_X(+) (M, .... Mss and weights {W 0
2: Sample B i..d. permutations of (1...., D): {II()) = (II) (t) 1 ......, IhsisB
3: for d= D do
4: Call the score model and get all the scores (60(X(()) Issinn
5: For each 1 _ B, sample random integer (1 in [,,.,,.., following the probability
distribution Pr(n (6) n) = so (X " m= and update the M -th entry of X as n()
6: For each 5 < B, update weights W < W + log so so(((()) (n mnt)
7: For each Si S B, update weights with the final reward: W
W + r(x()))
return pairs of sample and weights (x(), W"(X()) hsiss
Moreover the solution p to (19) given boundary condition at either T is unique when [0, T] t + Q E XRXX is continuous function.
Proof. By (1), we have
P++ = L Pr(X I|X, y)p((y) (Lv== AIQi(U,)) + pc())
= =p( + At Eel(,, + +(AE)).
Therefore, by taking the limit At + 0, we have
@xpa(x)= = Σ Q.(..))p(u)
Σ () + ((()
Σ x) uF2
Σ Qc(a,
The uniqueness of the solution follows from the uniqueness of the solution to the linear ODE, by equivalently writing (19) as a.p. = Q Pt. t E [O,T], where pt = (pı(x) x E X) are column vectors in R
Lemma 5. Let X be a CTMC with generator Q. For any bounded test function X R. define (.(-) = Bl(((T)) = Then de satisfies the following Kolmogorov backward equation:
-Orb((x)) = (. - (x, ΦT(x))
((x), Vx X. (20)
y
Moreover, the solution & to (20) is unique when [0, L Qr E XXX is a continuous function.
Proof. By (1), we have
dr(x) =
= Σ (14 + (IQ.(I. ((Q +(AP))
= ++ta(()+ Σ (..
Hence, by taking the limit At O. we have
-z) = Σ (x()),((,,)) = Σ
+ ((t))) = ( (x)) v).
The uniqueness of the solution also follows from the uniqueness of the solution to the linear ODE, by
equivalently writing (20) as - = Q. E [0, 7]; Φr Φ. where &. (+r (x) x E X)) and
= (ó(x) E X) are column vectors in g
21

---

Proof of Lem. 1.
Proof. Let At = T. and to kAt. We have
dlp
duz (Eo)
p2 fe41/te Iia (Stuti + O
log
dIP (5)
log
log
du1
Plat Pletille l (Si (ste)
duz
K-1
log 151211511 + AIQP, (fra-1
151411=a =he + AtQu, (6te-1/St)) + 0(AF)) + +0(at)
log
(Eo)
+
du1
k=0
By the definition of the generator. If Erx.. # Ele+ then
LE Emn
+ +IQU (Erg (Etx)
Pi, (Ster Stal2)
log 11 =b% + AtQ (Eee4s (Exx)
log
Qi, (Ste, Ceees)
If otherwise, by Taylor expansion,
log 11 St + AtQu, (Sta41 15tu)
At(Q (Ete - Q1, (Etx O(
1erce1 =Ei + AIQ), (Eteas (5(e)
At Σ (Qi. (Ete) 2) + O(Dt2)
VF
Hence, taking the limit K -> 0o, we have
dTp
duz
Q²,
log
dp1 (5)
log
(Eo) +
Σ
log Q1, (Etur (1)
dui
+ At Σ
Σ (Q)
Q22 (Etu) +O(At)
=Et MA
d/12 (Eo) (Eo) + Σ
log Qf
T
log d/f11
(()((E))) Q7 s))dt.
1:
USE.
Corollary 1. Given two CTMCs with generators Q' Q² and initial distributions pr, u22 on X, let pl Ip2 be the associated path measures. Then the KL divergence between these two path measures can be written as
KL(P')IP'> = KL(Hel1mm) + Eg S"
Σ 37
Q log
Q? + Q. - QZ -Q?) (Ec, v)d. Q!
Proof. It suffices to take expectation w.r.t. p2 on both sides of (2). For the first term on the rh..s, Egaga log duz duj (Co) = Buo(Es) log dus (Eo) = (L(uellur) for the third term, the expectation dum
is straightforward. To compute the expectation of the second term under p2 start with the
22

---

discrete-time version and take limit when K + 20:
Qi (Eee Sexts) Qi (See Stats)
K-1
log Oi (Gt Sta-rd
Rg-32
log
Renus
Σ
Qi, -tam))
biats AEl
Ribruui #Sen
log Oi (Eeer Since)
Epp, (Sig. )o
Peeeelle (E 1
Ti, (Star Stxis)
LEARI Ahey
K-1
sale (y1Eee) log Qi, (Stx, y)
Epp (Etp.
Σ
p2
Q1 (Eex)1)
K-1
Qi, (Etu. y)
Epp, (5tx)
Σ
AtOP, (Etx. v)
log Qi (Eee
+ +O(AP))
kmO
En
Oi (Lc.))dtt Eswrs 1'
T
Qi log O? (G.,))
ΣΦ
Q7 log Q
Σ
Q
Er
Thus, the proof is complete.
C.2 Stochastic Optimal Control
Here, we give proof of some key lemmas in Sec. 3.1 concerning the SOC for CTMC. See also [Wan+25a] for review of results.
The proof of (8) can be found proving the following lemma.
Lemma 6. The value function satisfies the following Hamitoniia-JJccob-Beee Jelll--elmann (HJB) equation:
0,V((z) = Edl(2 -e = O,eVela)
Q2 (e -e v.(t)). (21)
Proof. By the dynamic programming principle, we have
[1.
-v(xx) = inf EX
Σ N
*X v)ds -((x)) (55)) X1
(+D1)
= At inf
(x + O(At)) - inn 1
= =x],
where *. = =Q Qu Q"+ + Q.. We decompose the last term as
inf (
x]]
inf
Σ Vo+ (I + 1)
inf
- VV+a(()) At Σ (H+A(()) - At Σ Vo
-Ve+ + Atinf
Σ (x. Vetar(u))
+ +
Hence, by taking the limit At 0. have
Q"
V.V((z) = inf voft
Σ
Q
log
Qr + Q.
(I, + (V(()) - (Z(
(22)
23

---

For each pair of * y, by optimizing the r.h.s. of (22) with respect to (z, y), we have
[(
(Q2 log
Qr
(z, = argmin
Qr + + (x, y) + (Vi (z()))
Qi
or (a.)220
Qf(a y)e y)) (ee (((m)
=
and plugging this back into (22) gives the first equality in Lem. 6. The second equality is an immediate consequence of the first one.
Lemma 7. Assume that t. + Q is continuous. Then the following Feynman-Kan Kac formula holds: (V((a) = Epo Jc" (Xx]|x = z].
Proof. This result follows immediately from the second form of the HJB equation (21) and the Kolmogorov backward equation Lem. 5.
Lemma 8. Assume that + Qi is continuous. Then the optimal marginal distribution P? satisfies (x) 1 (z))e where Z Ergg"
Proof. Let (x)) = (P)(e)) (()), and obviously hr P7. Recall from Lem, 4 that (x)) E())(,,))P
Also, from Lem. 6,
Vee) = Σ (i - vi(z)).
Multiplying the first equation by (v(( and the second by (z),, can easily derive
0pha(r)= = I(Qlw.x))aa
(z,
yfr
which is the Kolmogorov forward equation for Ip" driven by generator Q", By the uniqueness result, we conclude that h, (x) = (x))
Note that (z)(() = - (z) = Puii Volo) Thus, P. = Pinit iff. Vo(-) is constant, which can be guaranteed when Xu and Xx are independent under Io due to the Feynman- Kac formula. Otherwise, P. + Pinit leads contradiction, which means there is no solution the SOC problem (6).
Remark 1.
Lemma 9, The following relation between the reference and optimal path measures holds: a (E)
dP"
1 (e((S)) for any & where Z Epg e".
Proof. By (2), (8), and Lem. 8,
dlF*
dP dPp
log (e)
Q. ,()) f. (" E (QPI(E)))) - (Etr
log dIpo (5)
log
(Eo)
1:Se- #6.
UFC.
(Vi(E) Vi(E+-)) + L' Σ (P((E,))( -
Vo(Eo) log Z
Σ I:EE
Since E. is piecewise constant and càdlåg, t + Vz(x)) is continuous for all x, suppose the jump times are 0 < 4 < ... < in < T, and let t0 0, 1n+1 = T, we can write
a
Vi((T) VolEa)
(Vi.s (Ee. (E.))
-Ve.)))
dit1
O.V. +
Σ (Vi(E)) Vi(E.-))
[' a.V. Σ (Vi(E)) V.(E.-))-
24

---

On the other hand, by (8) and (21),
((o((Er. v) (Et, () L.
Of c dt S/" a.Vi(E))dt.
E
Finally, as Vr the proof is completed.
C.3 Omitted Proofs the Main Text
Proof of (9).
Proof. Throughout the proof, assume regularity conditions that guarantee the validity of swapping the position between the integral and the derivative. We also assume Ip" and 1 are mutually absolutely continuous. We have
dlpa
dlea
d
dlp
Vo (P"
PP)
Vo Eps
dIaa
log
dF-
= Eps
o
dpa
log
dF
dPP
dIPu
d
dIPu
Ez30
log
e log
(Chain rule for derivative)
dPa
dle
dF
dP
dF
dFu
dP
dF"
E
V log
log
Ve log
(Gradient of log)
dP
dlPa
dP
dP
dPp
dIp
dPa
dIP
dlp"
dPa
dpa
Eze
V8
log
log
V
log
(When not taking gradient,
1.
dIpa
dP*
dp
a
dlF-
dl
dlp
dPa
+1)
dl"
dF
E
V log
log
Vo log
Vo log dpa )
(23)
dF.
dp.
dp.
The second term is actually zero:
dl
VoldF"/dP"
VoldP"/dl /dF*)
dP"
dP
Epa Vp log
Epe
Eps
(When not taking gradient, dp-
dlp
dlpu /dF
dla /dPp
dp-
dpu
dp
Eg.
Vo
Vo
Ep-
V 0.
dp.
dp-
Therefore, the in (23) can be replaced by any real number C. Swapping the positions of Eps and Va completes the proof.
Proof of Lem.
Proof. Note that under the generator Q: defined as (12), is obvious that each token evolves independently. Therefore, it suffices to consider the transition probability of single token, Qtok (Qiak &R (R(++))(X++)))
Let the mask token M be N + 1. By definition, Qtok can be expressed as Q 7()uv², where
u = O,, 0, 1)T and e = (,...., 1 -1) are both (N + 1))dimmessionl vectors. Thus, by the N: Kolmogorov forward equation (Lem. 4), the transition probability matrix from is given by exp (())dr) ) To compute the matrix exponential, note that (wu u(v u
(-1) uo. and hence
x (-1)
Mun =
I
(wv T) +
m
k!
k!
k=1
k=1
20
(-x)#
=I
uu = I (1 (1-c c>)
k!
k=
Therefore, have Po (n) (M) 1/2 exp (Io w(r)dr)) for E (1,... N). Since the initial sample is completely masked, f-- (r) xo guarantees P. Punir- Moreover, as the initial distribution is point mass, is obvious that Xo and Xy are independent, i.c., p0 is memoryless. D
25

---

Proof of Lem. 3.
Proof. From (8) and (12), we only need to consider the values of (z, otherwise the value is O.
(z, n) for x" M, as
We first compute the conditional distribution of XT given X; E X' under p We know from the proof of Lem. 2 that each token unmasks into .I.... N uniformly and independently, and remains unchanged as long as it is unmasked. Thus, the conditional distribution of XT given X₁ is the uniform distribution over the subset of 2, whose tokens at the unmasked positions of a are the same as æ. i.e.,
P p
154
Nidda MIJ
dex AM
Hence, by Lem. 7, we have
Vi(w)) Ni MJ]
Σ
(
No: Te ad=rdd V
Notably, Vi(x)) is independent of I. Now, suppose x² M, we have
e (v(('="))
Nilddaeen)dd MI) Eietn (e Vdl e
eVelx)
Ni(d:ad-M) MJ] E Licdo: 34=rd, Vd (5)
N Exexo: 3d=rd, Vdda Fl
(Since ((x) ZT((2)
EBENo: 3d=od, Wd (T)) (X'= n|XUM = = MMM,
Punnt (x)
N Pr (X* = X
which, together with (12), completes the proof.
Proof of Prop. 1.
Proof. The claim follows from the data processing inequality KL(pl 2 KL(Tip]TTg) where T is an arbitrary measurable mapping and Tip is the law of T(X) given X ~
Proof of Prop.
Proof. The identity (14) directly implies that Z Ex-pu e ("X), and hence 2 is an unbiased estimation of Z. By the Markov inequality,
dF
dF
Pr
1
2
Ta
2
E
Z
dPP
dPa
dp
dlu
dQ (VQ that dominates both P and P egg. 1(" + P*))
dP
dQ
dPd
dIP*
dQ (Product rule of RN derivatives and 1A S 1, AA)
dQ
dQ
TV(P", P()
(By definition TV distance).
25
By the Pinsker's inequality KL(pllg) > (T.pp, 9),. the probability can thus be bounded by 1. The last claim follows from the median trick, see [GTC25a, Lem. 29].
D Experimental Details and Additional Results: Learning Ising model
D.1 General Training Hyperparameters and Model Architecture
Model backbone. We use vision transformers (VTT, [Dos+21]) to serve as the backbone for the discrete diffusion model. In particular, we use the DeiT (Data-efficient image Transformers)
26

---

framework [Tou+21] ]] with 2-dimensional rotary position embedding [Heo+25], which better captures the 2-dimensional spatial structure of the Ising model.
For I= 4, we use blocks with a 32-diimenioonal hidden space and attention heads, and the whole model contains 26k parameters. For I. = 16, we 6 blocks with 64-dimensional embedding space and 4 attention heads, and the whole model contains 318k parameters.
Training. Among all the training tasks, we choose the batch size as 256, and use the AdamW optimizer [LH19] with a constant learning rate of 0.001. We always use exponential moving average (EMA) stabilize the training, with decay rate of 0.9999. All experiments of learning Ising model are trained on an NVIDIA RTX A6000. The training L = 4 target distributions takes around 10 minutes while the training of I. 16 target distributions takes around 20 hours (or equivalently, 4 hours on an NVIDIA RTX A A100.. For 16 x 16 Ising model, we use FWDCE with a resampling frequency le 10 and replicates R 8 for total of 50k iterations, among which 20k is warm-up training.
Generating baseline and ground truth samples. For the learning- based baseline, we train LEAPS on 16 x 16 Ising model for up to 150k steps for each temperature, which is comparable to, not more this task, to ensure a fair comparison. We also the Metropolis- Hastiigs (MH) algorithm for a sufficiently long time to serve as a baseline, For L = 4 (resp.. batch size 1024, and warm up the algorithm with 210 (resp., 2 burn-in iterations. After that, collect the samples every 200 (resp.. 2 steps to ensure sufficient mixing, and collect for 20 rounds, the Swendsen-War Wang (SW) algorithm
than, requirement
16), we use
so the final number of samples is 2 For L 16, each temperature to generate examples accurately distributed as the ground truth. We use batch size of 128, and warm up the algorithm with 2 burn-in iterations. After that, we collect samples every 128 steps to ensure sufficient mixing of the chain, and collect for 32 rounds to gather total of 212 samples. For L 16, the MH sampling takes around 3 hours while the SW sampling takes around hours on a CPU.
D.2 Evaluation and Additional Results on x 4 Ising model
D.2.1 Definition and Discussion on the Effective Sample Size
The effective sample size is metric commonly used to evaluate the sampling quality and does nor rely on the normalized target probability mass function or ground truth samples.
As in earlier works (e.g., [ZC22; AVE25; HAJ25J), given any control u, suppose we have M trajectories x()).... (M) Lid. p", we can associate each (0) with weight
("((()))
(
M Σラ
e (x66))
so that the weighted empirical distribution D w(X (x serves as a consistent approximation of F as M + co. We define the (normalized) effective sample size (ESS)
M M
ESS
M
>
E
|
(24)
A higher ESS typically means better sampling quality, but this is not always true as ESS is difficult to detect whether the samples miss a mode the target distribution. For instance, suppose F is a uniform distribution two points {a. b) and [ps is delta distribution on a, then sampling from pu would always output a, and thus w (1) = a) = 37 resulting in ESS = 1. This phenomenon is due to the fact IP* is not dominated by (i.c., absolutely continuous with respect to) Pa. In fact, by the strong law of large numbers, can show that
1
ESS No (x
1 M
Epp
de" apw
3.8
MM M (X
Ex
de-
The r.h.s. is ( if P is dominated by IP"., which aligns with the definition of ESS in [ZC22]. But if the learned path measure 1a only covers one of the higbbbrobyyilityy regions in F and misses
27

---

the others, then I may not be dominated by 1"" and ESS does not reveal the correct sampling quality. Therefore, only relying on ESS as the evaluation metric may be problematic.
D.2.2 Evaluation of the 4 Learned Models.
For each step during training, after doing a gradient update to the model and updating the model parameters stored in EMA, we evaluate the current model by sampling using the parameters stored in EMA and computing the weights W along the sampled trajectories. The batch size for sampling is 256. The ESS reported in Tab. 2 are the average ESS of the last 100 steps.
As [o0- 24 65536, the partition function Z can be computed explicitly, and thus we can obtain the whole probability distribution π. We use the random order autoregressive sampler to sample from the learned model, and then compute the empirical distribution of the samples Psamp Recall that for two categorical distributions p and on Xg. the total variation (TV) distance is defined TV(p.g)
2 2ZZe,,
|p(x) {(x),, the KL divergence is defined as (pllg) :=
ExENo: (())>0 P(x) log p(z) q(z)* and the x² divergence is defined as (pllg) == ExEXs: (())>0 ((x)
p(z)
1= -=Er q(=)>0
(p(=)-g(=))² Note that in order to make (L and (2lll) divergences well- ((x)
defined, we require p to be dominated by q, i.e., g(x;) 0 for all E Xo such that p(x) > or
equivalently, q(x) = 0 = p(z)= = 0. Therefore, we do not choose KL(T (PPamp) and x (~ (Psamp))
as the evaluation metrics.
As the partition function Z can be computed explicitly, we can approximate the relative-entropy between the paths, [I(PP"[P))) log Z Ep (X)., by the empirical means of the weights W (X) for X P Finally, as is discussed in Prop. 2, an unbiased estimation of Z can be obtained by 2 = e (X), X
D.2.3 Results of Learning the 4 x 4 Ising model at Other Temperatures
We also provide results for learning distributions at critical (Peritical 0.4407) and low (3 0.6) temperatures in Tabs, and 5. The same model structure, training configurations, and evaluation methods as discussed above are used, except that for learning the distribution at Beritical- we train the model from scratch with 2000 steps, and for learning the distribution at 3low. we train the model for 1000 steps starting from the 1000-step checkpoint for learning the high- highteemperatuee distribution with the LV loss. All four learning objectives are capable of training the model to generate high-quality samples compared with the baseline (MH algorithm).
Table 4: Results for learning 4 x 4 Ising model with J = 1, h = 0.1 and Beritical = 0.4407, best in bold.
Method
ESS f
TV 4
(L(Pmmpllx)) 4
x {f(mmmpl|=)
KL(PPIP")
Abs. of log Ž
RERF
0.8480
0.0841
0.1521
0.1691
0.0565
0.00150
LV
0.9809
0.0801
0.0222
0.0830
0.0083
0.00106
CE
0.9545
U.0454
0.0327
0.1824
0.0259
0.001
WDCE
0.9644
0.0789
0.0375
0.0839
0.0177
0.00010
Baseline (MH)
0.0223
0.0198
0.0615
D.2.4 Ablation Study of the Warm-up Training Strategy
As an ablation study of the effectiveness of the warm-up strategy, in Tab. 5 we also present the results of the model trained from scratch for learning the distribution at 3 and we can see that warm-up generally improves the training quality. Note that at low temperature, high ESS does not necessarily indicate good sampling quality, as the target distribution is highly concentrated on two modes: in fact, the all all-positive configuration has probability 0.7530, the all- all-negative one has probability 0.1104, and all the remaining 216 2 = 65534 configurations occupy the rest portion 0. 0.1366. As result, if the samples are all concentrated on the all- all-positive configuration and do not cover the other mode, then the ESS would be close to 1, yet the overall distribution is far from the target. Thus, we need to rely on a diversified set of evaluation metrics to comprehensively evaluate the learned model, such the
28

---

Table 5: Results for learning 4 x Ising model with J 1, 0.1 and Blow 0.6 (with ablation study of the effectiveness of warm- up in training), best in bold.
Method
warm-up
f
KL(P"(P")
Abs. ca: log 2
RERF
0.9196
0.0320
0.0200
0.4071
0.0263
0.00788
0.9276
8682
2.0487
6.7066
0.0352
2,00995
0.9722
0.0177
0.0098
0.1864
0.0092
0.00267
LV
0.9604
0.1515
0.1619
0.1933
0.0164
0.13500
0.9855
0.0147
0.0138
2,8388
0.0098
0.00259
CE
0.9684
0.0159
0.0287
60.6196
0.0250
0.00887
0.9165
0.0118
0.0282
1,6582
0.0336
0,08373
WDCE
0. 0.9927
0.1365
0.1505
2,4919
0.0057
0.13001
Baseline (MH)
0.0068
0.0044
0.0418
TV distance, KL divergence, and x² divergence reported in Tabs. and 5 as these metries in general do not suffer from the aforementioned problem.
Additionally, it is worth mentioning that this warm-up training strategy is significantly different from the warm-up used in LEAPS. The code implementation of LEAPS also has a warm-up stage in the first 20k steps during the training, where they gradually increase (k)) from 0 to 1 while the step k increases from 0 to the maximum warm-up steps. During the warm-up up phase, the PINN objective of LEAPS is evaluated on time up ta (k). Unlike LEAPS, we do not use partial trajectories for training during the warm-up up phase, which is a major difference.
D.2.5 Ablation Study of the Choice of Number of Replicates R in FWDCE
Bag
022
Auaoal 0.4407
0.5
LTT
0.00
1.30
00100
0.002
0000
K0 and TV
6000
01122
1000
0030
182
0101
115
0.00
Number Replicates
Number Replicates
Number Repiccats
Figure 3: Visualization of learning performance across the number of replicates R for learning 4 x 4 Ising model with J 1 and h 0.1 using the WDCE loss. The metrics reported are TV (Peamp: TT. (L [], and x (Psampll I/=))
We provide an ablation study of the number of replicates R in the WDCE loss in Fig. 3, showing that the performance of training is insensitive to R. Therefore, we only choose a relatively small number R = 8 in learning the 16 x 16 distributions guarantee the efficiency as well as the efficacy of the algorithm.
D.3
Evaluation and Additional Results on the 16 x 16 Ising model
D.3.1
Evaluation of the 16 x 16 Learned Models
As the cardinality of the state space is much larger (21²=256 ~ 1077), we canmot exactly compute the whole probability distribution. Instead, report the values and errors for observables well-studied in statistical physics such as magnetization and 2-point correlation. For any probability distribution v on (+1}A (we recall that A {I,..., ))²), these observables be defined as follows.
Magnetization.
The magnetization of state E A under v is defined as M1) == Eu(w) o². The average magnetization of all states is defined M, Ex) [ E, x']. For our target distribution in (17), M, (i) = 0 for all i A since H (z) H (-xr) when h = 0.
29

---

Moreover, we can define row-wise magnetization (i.e., along the z-axis) Mr column-wise magnetization (i.e., along the y-axis) Mp by summing up the magnetization of states on the subset of specific row or column on the square lattice. They are formally defined as,
{
L 2
M (k) = Σ Mo(x'), M (k)
E Mo(a', iEcol(k)
Vk E
0,
2
(25)
Based on these observables, we define the absolute magnetization error (abbreviated as Mag. in Tab. 1) the learned distribution y compared with the target distribution и as
Σ
[M (k) M.0 (k)l + MM (k) (
(26)
2L
kE
+1
2-Point Correlation. The 2-point correlation of two states i, E A under v is defined as (t, =
Ex(z) [x'x0) Ex(z) [a=]] Bu(z) [a0). For our target distribution in (17), (x(i,) Bn(a) (x' x² due to the zero magnetization.
The average magnetization of all states differing with vector r E { - O,... [E is defined as
Cu(r) T2
Σ E (o' EEpco) (ar) (
where the addition is performed element- wise under modulo L due to the periodic boundary condition.
Similar the magnetization, we can define the row-wise correlation (i.e., along the z-axis) and column-wise correlation (i.e., along the y-axis) by summing up the 2-point correlations between states on the subset of specific pairs of rows or columns. They are formally defined
Cgo (k,
Σ (k)ow(k), JErow((). col
C,(1,),,
(k, =
Σ (Ecol(k), JEcol((), {.j
Co(1,)),
(27)
where k, E (- 0, [ Based on these observables, we can also define the absolute
2-point correlation error (abbreviated as Corr. in Tab. 1) for the learned distribution compared with the target distribution π by
C (k, Crow (ke, (e. + Iceo (ks, Cr (k:,
(28)
1
(K.1)
(a) MDNS (ours)
(b) LEAPS
(c) Ground Truth
Figure 4: Visualization of non-cherry-picked samples from the learned 16 16 Ising model with
J= 1, h = 0. and Brow = 0.6. (a) MDNS. (b) LEAPS. (c) Ground Truth (simulated with SW
algorithm).
30

---

MDNS (ours)
(b) LEAPS
Ground Truth
Figure 5: Visualization of non non-cherry picked samples from the learned 16 x 16 Ising model with J= 1, h = 0, and Peritical = 0.4407. (a) MDNS. (b) LEAPS. (c) Ground Truth (simulated with SW algorithm).
(a) MDNS (ours)
(b) LEAPS
(c) Ground Truth
Figure 6: Visualization of non-cherry-picke samples from the learned 16 x 16 Ising model with = 1, = 0. and Shigh = 0.28. (a) MDNS. (b) LEAPS. (c) Ground Truth (simulated with SW algorithm).
D.3.2 Training Curves for Learning the 16 x 16 Ising Model
The training ESS curves of the score models for learning the 16 x 16 Ising model are provided in Fig. where the models for Peritical and flow are initialized at the 20k-step checkpoint of the model trained for Shigh to implement the warm-up training strategy. Due to the gigantic mismatch between the reward functions used in the warm-up phase and the formal training phase, the ESS experiences sudden drop to near at iteration 20k for 8low and Beritical-
D.3.3 Additional Results for Learning 16 16 Ising Model
To provide more comprehensive evaluation of our learned model for 16 Ising model, we visualize the generated samples in Figs. 4 to 6. From the plots, we can see the high fidelity of our generated samples as they follow statistically similar pattern to the ground truth samples produced by running the SW algorithm. We also plot the 2-point correlation function along the y-axis in Fig. 8. Additionally, we visualize a distribution of absolute error of the estimated Foww and M/col to ground truth along each column superior performance of our proposed MDNS compared to other benchmarks.
of the square lattice in Fig. 9. All of these results demonstrate
D.4 Effects of Preconditioning in MDNS Training
In learning a continuous neural sampler to sample from target distribution & e on Rd, typical approaches such as PIS [ZC22] and DDS [VGD23] leverage the score information logm in the neural network, which is known preconditioning. Efficient preconditioning is shown to facilitate convergence of training and achieve smaller sampling errors [He+25]. In our experiments, already achieve good performance without any information of the target distribution when
31

---

Effective Sample Size (ESS)
0.8
0.6
0.4
Swo
0.2
Beritical
Fnight
warm-up
0.0
10000
20000
30000
40000
50000
Training Steps
Figure 7: Training curves of the effective sample size (ESS) for learning x 16 Ising model at different temperatures. A closer value to 1 generally indicates better performance.
Bow = 0.6
Restiial = 0.4407
Frahh
0.28
1.000
2ppontt Correlation
0.8
0.5
0.5
0.900
0.875
-2
0.2
-2
Distance
Distance
Distance r
MDNS (suis
FAPS
Ground Teith (5W)
Figure 8: Average of 2-point correlation Cool (k, k + r) of samples from 16 x 16 Ising model.
parameterizing the score model; here, explore way do preconditioning for learning to sample from the Ising model, and compare its effectiveness during training.
Recall that in the target distribution of the Ising model (17), the distribution of æt conditional on all
the remaining positions - := (x' j + i) can be computed as follows:
क (tt) xt e 3J (" +'0000
3J( + Sha'
n (x))
BBh at E (+1).
3J Eisni xd +Bh
+e 9J
BJ Eiin
As the model needs to learn the conditional distribution x² given a partially masked x we can naturally treat the mask value as 0 and use the above formula to approximate the conditional
distribution. Specifically, given a partially masked x X (11, -0 with = M, we use the following form of preconditioning:
Pr (x* n n = x. UM ss so(x),,n
X~
= sof coftmax (x + P(x), dim 11)m i A, n {=1}.
where 4y :X + (Lx is free-form neural network and the precondition matrix P X + L)x2 is defined as P(z).-1 log (x= +- 11-1aa 11x= and P(=)., = log k <(-² + 1oc++12=00 1|x-i).
In Fig. 10, we train models from scratch to learn from 16 16 Ising model under 3nigh and Beritical using the WDCE loss, both with and without preconditioning. It obvious that applying precondi- tioning facilitates the convergence of the training in terms of the ESS. However, unlike in the case
32

---

8,0 1.2
Sotiial
1.005
Sushh 0.5
0.020
0.020
0.020
0.015
0.010
0013
0.010
0.010
0010
0,0000
0.005
0.005
0.000
0.000
0.000
6.070
0.030
0.000
0.015
0.015
0.020
0.010
-0
-8
Index
Index
Index
Magenization error un a-anis
Magerization urror on y-azds
Figure 9: Distribution of the the absolute error Mzow (k) and (k)) to the ground values against index & for 16 16 Ising model.
Shigh 0.28
Gettral 0.4407
1.0
Effective Sample Size (ESS)
0.8
0.6
0.2
Preconditioning
Preconditioning
Preconditioning
Presonditioning
10000
20000
1000
50003
10000
20000
30000
10000
50000
Training Steps
Training Steps
Figure 10: Comparison of learning 16 x 16 Ising model with J and using the WDCE loss, with and without preconditioning.
Rd where we can directly leverage the score information V log the above example heavily replies the availability of the closed-form solution of the conditional distributions {π( (-)) Vi E A). which seriously limits the applicability of this preconditioning method. The study of preconditioning methodologies that are capable of dealing with target distributions whose conditional distributions are unavailable is left for future work.
D.5
Failure of DRAKES in Learning Neural Samplers.
Finally, we argue that while DRAKES [Wan+25a] might be useful for fine-tuning a pretrained masked discrete diffusion model to maximize certain reward function, it is not suitable for learning diffusion sampler due to the error of approximating states using the Gumbel softmax trick.
The training dynamics of the 4 x Ising model with J = 1, = 0.1, Shigh = 0.28 using DRAKES are presented in Figs. 11 and 12, where we use the same model structure, training configurations, and evaluation methods as in the experiment of Tab. 2, except that we now use smaller learning rate of 0.0001 with gradient clipping for numerical stability and train for 3000 steps.
test
two different Gumbel softmax temperatures, and 0.1, to show the possible effect of the Gumbel softmax temperature. DRAKES leverages the Fuler sampler [LME22; Ou+25] training, we use 32(= 2L²) steps to generate each sample and let the gradient pack-propagate over the whole
33

---

trajectory with all 32 intermediate states without truncation. As shown in the figures, while the loss decreases and reward (defined by Shigh H (x) (17)) increases during training, indicating that the model is learning, the ESS remains at low level, which means the learned model is not able to sample from the correct target distribution.
Loss Reward
Effective Sample Size (ESS)
1.0
0.8
0.6
16
Reward
Loss
Loss Revord
0.4
0.2
-12
0.0
0
500
1000
1500
2000
2500
3000
0
500
1000
1500
2000
2500
3000
Training Steps
Training Steps
Figure 11: Learning x Ising model with 1, = 0.1, and Buigh 0.28 via DRAKES (with Gumbel softmax temperature 1). Reward corresponds to -niihh (z).
Loss Reward
Effective Sample Size (ESS)
1.0
30
0.8
10
Loss Revord
20
Reward
0.6
-15
16
0,4
20
0.2
-25
0.0
0
500
1000
1500
2000
2500
3000
0
500
1000
1500
2000
25000
3000
Training Steps
Training Steps
Figure 12: Learning x Ising model with = 1, 0.1, and Snigh = 0.28 via DRAKES (with Gumbel softmax temperature 0.1). Reward corresponds to H
E
Experimental Details and Additional Results: Learning Potts model
E.1
General Training Hyperparameters and Model Architecture
Model backbone. Similar the experiments on Ising model, we also use ViT with 2-dimensional rotary position embedding to serve as the backbone for the score model. To ensure enough model representation power, we adopt slightly larger model with a 128-dimensional embedding space, blocks and 4 attention heads, which sum up to 829k trainable parameters in total.
Training. Among all the training tasks for Potts model, we choose the batch size as 256 and use the AdamW optimizer [LH19] with a constant learning rate of 5e-4. Like in the experiment of Ising model, we also use EMA to stabilize the training, with decay rate of 0.9999. All experiments of learning the Potts model are an NVIDIA RTX A A100. For 16 16 Potts model of all three
34

---

Bow 1.2
1.00
Phigh 0.5
2-point Coreelatin
0.2
Distance
Distance
Distance
(surs]
Meinppols-aatting
Teath (SW)
Figure 13: Average of 2-point correlation Cool (k, + r) of samples from 16 16 Potts model.
1.2
Patiial
1.005
0.010
0,000
0.005
0.005
0.000
0.000
0.000
0.005
0.015
0.015
0.000
0.010
0.10
0.06
0.000
0.000
0.000
Index
Index
Index
Magenization error on a-ands
Magenizazion error 3-200s
Figure 14: Distribution of the the absolute error Mrow and (k) to the ground values against index & for 16 x 16 Potts model.
temperatures, we use FWDCE with a resampling frequency k = 10 and replicates It = 8 for a total of 100k iterations, among which 30k is warm-up training. The total training time is about 20 hours.
Generating baseline and ground truth samples. For the learning-based baseline, we train LEAPS on 16 16 Potts model for up to 100k steps for each temperature, which is comparable to, not more than, our requirement on this task, to ensure fair comparison. We also run the MH algorithm for a sufficiently long time to serve as a baseline. For L = 16, we use batch size 32, and warm up the algorithm with 2d burn-in iterations. After that, collect the samples every 26 steps to ensure sufficient mixing. and collect for 128 rounds, so the final number of samples 2 Note that for the Potts model with L 16, the MH algorithm is not capable of sampling from the correct distribution even at the high temperature, and thus we have to resort to the SW algorithm to generate ground truth samples. We
batch size 128, and warm up the algorithm with 2 burn-in iterations. After that, we collect samples every 128 steps to ensure sufficient mixing the chain, and collect for 32 rounds gather final number of total 212 samples. The MH and SW sampling takes about the time as the Ising model on CPU.
E.2 Evaluation and Additional Results on 16 x 16 Potts model
Similar to the case of 16 x 16 Ising model, 16 x 16 Potts model has an intractably large cardinality
of the state space 342=256 s 10122, we cannot compute the exact distribution and thus rely on values of observables in statistical physics to evaluate the performance. As the Potts model is generalization of the Ising model, statistical physics observables such magnetization and 2-point
35

---

MDNS (ours)
(b) LEAPS
Ground Truth
Figure 15: Visualization of concheeryppieeee samples from the learned 16 x 16 Potts model with 1, q 3, and Blow = 1.2. (a) MDNS. (b) LEAPS. (c) Ground Truth (simulated with SW
algorithm).
correlation are also well defined, with different definitions and formulae but share similar idea in spirit.
For any probability distribution on {1,..., g)aA (we recall that A (I,....))))
Magnetization. defined as
The magnetization state A under given n ii.d. samples from is
Mpotts (i) maax (nn//) -1
where nf is the number of samples whose i-th element c. The average magnetization all states is defined as M, :- 7, >, Mpotts (i).
Similar to the case of Ising model, can also define row-wise magnetization (i.e., along the z-axis) or column-wise magnetization (i.e., along the y-axis) by summing the magnetization of states on the subset of specific row column the square lattice. They are formally defined
M (k:)
Σ (E)ow(()
Mpotts (a"), M)
Σ iEcol(k)
Mpotts (x:), (x"),
Vk E
(29)
Based on these observables, define the absolute magnetization error (abbreviated as Mag. in Tab. 3) for learnt distribution by computing the following
Σ
M (k) M1 (k) + |M Mo (k:)l
(30)
2L
ke - 1E....].........111.
1411
where ग is the target optimal distribution and y is the learned distribution.
2-Point Correlation. Potts model has the following definition of 2-point correlation:
1
Cpotts (i.
L Σ Eu(s)
Vi.j e A.
Comparing to the definition for Ising model, the above definition has an additional offset 1 q" which causes the maximal correlation be strictly smaller than 1.
We can define row-wise correlation (i.e., along the æ-axis) or column-wise correlation (i.e., along the y-axis) by summing the correlation between states the subset of specific pairs of columns. It's formally defined as,
Crow (k,
Σ (Erow(k), (I), (. same col
(iotts(i,
C (k,
Σ
Cpntts (i.j),
(31)
(Ecol(k), jEcol!!). i.j same
36

---

MDNS (ours)
(b) LEAPS
Ground Truth
Figure 16: Visualization of non non-cherry- samples from the learned 16 x 16 Potts model with
J= 1, y = 3, and Beritical 1.005. (a) MDNS. (b) LEAPS. (c) Ground Truth (simulated with SW algorithm).
(a) MDNS (ours)
(b) LEAPS
Ground Truth
Figure 17: Visualization of hon-cherry-pickee samples from the learned 16 x 16 Potts model with 1, 9 3, and Shigh 0.5. (a) MDNS. (b) LEAPS. (c) Ground Truth (simulated with SW algorithm).
where k. E (- 4 ..... 0, .... } Based on these observables, we can also define the absolute correlation eтror (abbreviated as Corr. in Tab. 3) for learned distribution by
C (ks,) Czow (k: 01 + |C (kt,) C (k,
(32)
L
(k.)
E.3 Training Curves for Learning x 16 Potts Model
The training ESS curves of the model for learning the 16 16 Potts model are provided in Fig. 18, where the models for Beritical and Slow are initialized at the 30k-step checkpoint of the model trained for Snigh to implement the warm-up training strategy, Again, due the mismatch between reward functions at different temperatures, the ESS experiences sudden drop to near () at iteration 30k for frow and Beritical.
E.4 Additional Results for 16 16 Potts Model
We visualize the generated samples in Figs. 15 to 17. From the plots, we can see the high fidelity of our generated samples as they follow statistically similar pattern the ground truth samples produced by running the SW algorithm. We also plot the average 2- 2-point correlation function along the y-axis in Fig. 13, and visualize the distribution of absolute error of the estimated Mzow and Meal to ground truth along each column or row of the square lattice in Fig. 14. These results suggest that MDNS manages to learn to sample from multimodal, high-dimensional discrete distributions accurately.
37

---

1.0
Effective Sample Size (ESS)
0.8
0.6
0.4
Bra
0.2
Perisical
Brigh
warm-up
0.0
20000
40000
60000
80000
100000
Training Steps
Figure 18: Training curves of effective sample size (ESS) of 16 x 16 Potts model under different temperatures. A closer value to generally indicates better performance.
F UDNS: Extension of MDNS to Uniform Discrete Diffusion Models
In this section, we discuss an extension of our MDNS framework to uniform discrete diffusion models [LME24; Sch+25J. We present the theory in App. F.I and discuss the strategy of preconditioning in App. F.2. Experimental results presented in App. F.3.
F. Theory of Uniform Diffusion Neural Sampler
We choose T and construct the reference path measure To0 by the CTMC that always keeps (r) = ND 1ex at all times t E [0, 1]. This can be achieved by initializing Pg Punif and setting the generator as Q2 d) N Vn * Id. Note that each dimension evolves independently under Io0 ((t Se (ss)ds. One can compute the following transition distribution from time to
D
e-(()) 7(t)
( (-I()))) =
II d=1
e (1))
(33)
N
ijo By similarly assuming 700 oo like in Lem. 2, we can guarantee that T (1)(211(1)) (o0) Punit )
so under po, Xo and X, are independent, making the reference process memoryless. A direct implication is e Vo(+) = Exa (X,, Xo = Epant e" = 2 is a constant, just like the mask case.
Moreover, from Lem. 8 we have (x) = (P( = t e Vi(e), so (8) now reads
(i = z Villamtnn 7(t) P. (( Vn x eVi(u)
P7(x)
We can thus parameterize Q7 (z, x)+")= = - so (x, t) where the neural network so takes x E Xo and E [0, 1] input and outputs nen-negative D x N matrix.
We approximate the continuous-time process by the following (naïve) Euler discretization scheme: for k = 1, .... K - 1, let At Ω and l kAt. To sample from Ipa first sample Xo Punifs and then, for 0, 1, ..., K 1, approximate the transition probability as
P lts (Itaya (xr.)) x 12 + AIQU
Here, we only allow It, and Itess to differ at most one entry in order to correctly compute Qu, (Ity d. ). otherwise the value would be zero. The RN derivative de
(x) is similarly ब
38

---

approximated by
dF log
dP*
dPu
(x) log
apo (x)) log
alpo (a)
puu
K
Pu
Cx+ille lik (Itze (T(x)) (x)x) 10--loo log (34)
Ra r((1)) log
Σ
log
f (Dea+1 (1bz)
Next, if Stuti # Ite.
pu log lstilee (Fturs k( pe Platslte (Itees (1xe) Fea have
Qh Pturr) log
log Qi, (Ity Fea+1)
when Deeti E21 2/ (35)
Otherwise, if Feati
Pu ftttlle (Iee T2x)
1 1+ AIQU, (tee
log ppu
~ log
fatalte (Ity (pt))
1 + AIQU, (Iit Fit.)
1
At Ed=1 Emfert
Qi, (Fte "
log
1 At Ed Empatt Qi (Ite 18 dd
Ar((t)
Ex Empod, s0 Sa Jana
log
(36)
Ab:(0)
Ed E Enjod
Here, unlike in the proof of Lem. we do not use Taylor expansion to approximate log(1 - A(x) by -A* order not to introduce further approximation error. Notably, due to our parameterization of Qr, this only requires architecture such as the locally equivariant network introduced in [HAJ25].
call to the score model without any specifically designed model
Now, with the approximation of RN derivative, we can easily derive the LV, CE and RERF losses just by plugging WXXX into (15). We can similarly derive the WDCE loss as follows. First, as P. = Pumit for all f, the CTMC with path measure po is reversible, and hence we know from (33) that | (z1 (x2) = dj (xa(æ1). Next, due to the property that I (x) 1/2 (=1)ppp from (5), by conditioning on 11 we have
D
=f + -(()))
P P7 (xe(=1) = P (ze|Z1) =
II d=1
e
7(1)
(37)
N
which means we can easily sample from the transition kernel [ (1121) by independently replacing each entry of xi with token from Unif(11.... NJ with probability - e - To derive the WDCE loss, one can first prove the DCE trick [LME24, Thm. 3.4]: for any function f.
(v) X Pi(X.) (Xx,,
Pin ())
(F (XII))
Σ
((X,,t
(X,1x))
Σ
Pip (X,)x.) f(,,t,y)
X
(38)
In fact, simple calculation shows that both sides equal Ex. (X,,))(Xx,,,,))) We can therefore derive the WDCE loss follows:
KL(FI*IP")) Ex-(x)
1''
['
Σ
log log
Qi
Q" - Qi (Xr, (By Cor. 1)
Q
V
-x,P(X) Σ
(Q" - Qi log Q (X,, 3) + const
XxX.
= -E,((X) Σ
Σ (Qr Qi log Q1) (Xe, Xd + const,
d x
39

---

where t Uni 1). and const represents terms that do not depend on 0. Next, we leverage the parameterization of Q" and Q" as well as the DCE trick (38):
(I(PP*IP"))
7(t)
P7 (X
Σ
Σ
So (x(,
log so (X,, t)a..)
+ const
N
PP (X,)
x
7(1)
(, (X//X)
Σ Σ
80 80(X,,))a,
log 80 (Xx, D)n
+ const
N
P.j (X))X1)
d
x
dPP
7(t)
Epa (X)
dPu
X
E
L.
(X,
N
P. (Xplon 1i
]
Σ
80 so(Xe, Dan
log 80 (Xe, Dan
+ const,
X1)
d x
We can simplify the ratio of conditional probabilities according to (37):
Pij (Xp="[Xx)
e X + 1-e 79
1/1 (1)
x1 = + 1-0-7(9)
N
(xx) (X) can be computed via (34), and we can estimate Z by Exapa (x) as in both CE and WDCE losses for MDNS. We thus arrive the WDCE loss for UDNS:
Finally, recall that d
FWDCE (IP ,P), E Xx Z
(x(()
E E.p
7(1)
(x.1x1)
N
c c()
+ 1 1-c 1-e-7(0 T(0)
(39)
ΣΣ
so(Xe, t)dm
log se (Xe, t)d,n
e e-7 1
+ 1-e-7o
X
We summarize the training of the UDNS in Alg. 3. Here, Resample Unif means for sample random variables (f(6rr hiigB,ssrR t.d. Unif(O, 1), and for each and r, first randomly masking each entry of (9 with probability -0-70) and then replacing each masked entry independently with a token from Unif (1,.,,
Algorithm 3 Training of Uniform Diffusion Neural Sampler (UDNS)
Require: score model Sp, batch size B, training iterations K, reward function r Xo + R, learning objective F (FRERF, Fiv, FoE, FWDCE), (number of replicates of each sample R, resample frequency for FWDCE). 1: for step to K do
2: if E [FRERF, Fiv, FCE) then
3:
[X(), ("(X())))sss Compute F with (X(1,, ("("))))))))) else if F FWDCE then
b See Alg. 4. See (15).
4:
5:
6:
if step mod & = then (x (" )esis
A Sample new trajectories every k steps.
7:
fraplee_raajecories_uu Set replay buffer B + (x(), (X (X0), (x(X()) ISEBR Resample_wtn with Unif R). Compute FWDCE with {X (0), (X(
8:
9:
10:
11: Update the parameters / based on the gradient VgF.
refurn trained score model
40

---

Algorithm 4 Sample_ Traill___rjjcttriss___ Unif: Sample trajectories and compute weights for UDNS.
Require:
score model 8p. reward function r : :o + R, batch size B, number of time-intervals K, functions y, early starting parameter 0.
1: Initialize uniformly random sequences X (f) i..d.
to Punif and weights W 0, sisss K, where At xx)0((0)
2: Define + kåt, 0
do not initialize at to 0.
3: for k = to K do
4: Call the score model and get all the scores {so(X ty tx) hsiig.
5: For 1 S - B, sample an update from the approximate transition distribution: X f- lx+1
(x with probability Atr(t)
(X(2) (x)d, for n + (x( and x(() + x(( with 1x
N 8o
probability (tx(()
Ea (xx! so ( by tx),.
N
6: For each 1 S is S B, based on if x(9 f6t1 # x or not, update weights W + W Firtslty Itg (81441/244
log go
according to (35) or (36), respectively.
"etilty (1 [
7: For each S I B, update weights with the final reward: W + W + +(((
return pairs of sample and weights (x(), W"(X() isiss
Table 6: Results for learning 4 Ising model with J 1, h 0.1, and 0.28 using UDNS (with an ablation study of the effect of preconditioning) best in bold.
Method
Use precond.
ESS
TV
KI.G@ampp)(z) 4 x (Rampl)))
KL((P*)PP)) )4
Abs. log 2
0.9671
0.0787
0.0357
0.0738
0.0182
1.45745
FRERF
0.9769
0.0726
0.0341
0.0096
0.0110
1443776
Fix
0.9900
0.0710
0.0332
0.0847
0.0053
0,07752
0,0712
0.1334
0.0678
0.00025
FWDCE
0.9204
0.0878
0.0397
0.0847
0.0450
0.01911 I
0.9301
0.0814
0.1383
0.0811
0.0385
0.00332
Baseline (MH))
0.0667
0.0325
0.0628
F2 Preconditioning
Inspired by path integral sampler [ZC22], we propose to use the following form of preconditioning for learning the score model to sample from
log P (x't-)) (x)
((
log (z, t)... = Do(x, toa (t) log
π(x) Vi (L, LJ2 {=1){(²),
where tb 2 x 1] + (Lx and oo [0, 1] - IR are free-form neural networks. For the Ising model (17),
log (( ((x)
-23J
Σ J:jni
x² 2,3hx',
where -- represents the configuration obtained by flipping the sign of x'. Again, we emphasize that this design of preconditioning replies on the special structure the target probability distribution,
namely the closed-form expression of ((("t")) for all and n. The study of the preconditioning for +(x) more general target distributions is left future work.
F.3 Experiments
We use UDNS to learn the same 4 x 4 Ising model as in Tab. 2 (with J = 1, 0.1, and B = 0.28). and report the quantitative results in Tab. 6. We use the same model backbone and similar training configurations as deseribed in App. D.1, except that (1) the total number of steps is 3000 to ensure all methods fully converge, and (2) for preconditioning, we parameterize oa [) - with multilayer perceptron with hidden sizes 32, 32, 32 and SiLU activation, whose total number of
41

---

parameters is 2k. The number of time points K for discretization is chosen as 50 for both training and evaluation, and the function is chosen as y(t) = J. which implies 7(t) = log and - = t. The number of replicates A. used in FWDCE is set as 32. During training, we find that FeE fails to learn the correct target distribution, and hence its performance is not reported in the table. We can see from Tab. 6 that all three learning objectives are capable of generating samples with quality comparable to the ones generated from the MH algorithm, and the general effectiveness ranking of them is FLV FREEFP FWDCE- Moreover, we find surprisingly that applying preconditioning is generally beneficial when using FLV, but is prone to worsen the performance when using FRERF and FWDCE- However, compared with MDNS, UDNS has higher computational cost in both training and evaluation due to the requirement of fine-discretization, and hence is generally less preferable in practice.
42

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
We believe the main claims made in the paper accurately reflect our contribu- tions.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations in Sec.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
43

---

Justification: All theoretical results have a complete proof in the appendix. We have checked the mathematical derivation and affirmed that they are correct.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All information regarding reproducing the experiments can be found the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
44

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Our experiments do not rely on any data. The code repository can be found in the abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
o icy for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips. ct/pcc/ppl/.o cec/public/guu Seaiiidee/oooeubbi iesisssinnl asssoooliiyyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All experimental details can be found in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[No]
Justification:
We do not report error bars in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
45

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The information on the computing resources for the experiments can be found in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https https://neut /ttpp////////....../// ctt http/uullc/Bttiiee cehiideeenlneeesss
Answer: [Yes]
Justification:
Our research completely complies with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The discussion of broader impacts can be found after the references.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
46

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We believe our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have mentioned the original papers that produced the code in the appendix.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwwode..tcd..... com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
47

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
48

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The the core method development in our paper does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: : /tt///eerippp//e// tweuriip.cc/Coofeeeeeeeee. /LLM) for what should or should be described.
49


Reviews:
[{"review_id": "cxNugnNx0q", "reviewer": "Reviewer_XE64", "summary": "The paper introduces a new diffusion sampler for discrete problems that is based on optimal control of a continuous-time Markov chain (CTMC). To do so, the authors first show that we can obtain a generator by minimizing the KL divergence between the optimal path distribution from a masked diffusion sampler and the generator. The authors discuss and compare different loss functions including REINFORCE loss, log-variance loss and cross entropy loss to solve the SOC problem. finally, they introduce a new weigthed cross entropy loss that is supposely computationally more efficient as it can reuse data. The authors compare their algorithm on challenging Ising models and potts models to a recent SOTA method (LEAPS) and show that it outperforms this SOTA method.", "questions": "- It should be clarified how the optimal control formulation for discrete samplers deviates from the standard optimal control formulation for continuous samplers\n- From the study on the simpler 4x4 task, it seems that the log-variance loss is performing best. However, the paper proceeds with the weighted cross entropy loss (WCE). The advantage of the WCE to logVariance is unclear. Is it computation time due to data reuse? If so, that need to be evaluated as well\n- Its unclear to me why we can not implement a data-reuse mechanism (e.g. with a similar importance weighting) for the other loss functions. Log-Variance loss seems to be better then cross entropy, so  the motivation why this was only done for the cross entropy loss is missing for me", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "strength:\n- interesting new algorithm for discrete sampling\n- sound derivation (even though I could not fully follow everything) based on stochastic optimal control\n- good experimental results with a clear improvement over SOTA on challenging problems:\n\nweaknesses:\n- clarity and motivation. The paper is hard to follow, notation is hard to grasp and it is unclear why certain things are introduced (e.g. Lemma 1). Moreover, basic definitions are hidden in equations and are quite hard to find (e.g. the defintion of W^u, that should be an own equation with W^u = ...) and no intution is given what this terms really mean. \n- Results are not fully conclusive (see questions).", "was_revised": true, "final_justification": "Interesting new algorithm for discrete sampling with strong results. Main weakness was clarity, which was addressed by the authors.", "rebuttal": "We sincerely appreciate the reviewer's time and the constructive suggestions to help us improve the quality of this manuscript. We provide itemized responses below to each of them and hope they can help clarify your questions and concerns.\n\n> **Weakness 1**: Clarity and Motivation\n\nThank you for raising the question and allow us to further clarify. \n\n- The main motivation of our work is to formulate the discrete neural sampler training problem as a path measure matching problem, which naturally connects to SOC for CTMC (Eq. (6)) and gives rises to other training objectives that are more favorable on discrete state spaces. \n- We introduce Lemma 1 as it is an important result for CTMC, which is also crucial for the derivation of our objectives. In Eq. (14), we use Lemma 1 to compute the log RN derivative between the optimal path measure $\\mathbb{P}^ * $ and current path measure $\\mathbb{P}^u$. We will update the manuscript to clearly reflect the use of Lemma 1 in such derivation to improve logical coherence.\n- As of the definition of $W^u$, we agree with the reviewer that it would certainly be clearer to definite it in one separate equation, and we will update it in the revision to clearly emphasize that. As for the meaning of $W^u(X)$, it is the log RN derivative up to an additional log normalizing constant, a weight factor associated with trajectory $X$ and policy $u$ that shows up ubiquitously in the formulation of the four introduced training objectives in this paper. We call it a weight (therefore denoted it with $W$) due to its connection to importance sampling and computation of the effective sample size (see Appendix D.2.1 and Eq. (24)). \n\n> **Question 1**: Connection to SOC for continuous samplers\n\nWe discussed the connection of our method to the related literatures in neural samplers training and SOC in Appendix A. In fact, our formulation of neural sampler training as path measure matching is a general framework motivated by SOC and it indeed can be applied to continuous samplers. For example, when choosing $\\mathcal{F}(\\mathbb{P}^u, \\mathbb{P}^ * ) = \\operatorname{KL}(\\mathbb{P}^u || \\mathbb{P}^ * )$ and applied on continuous state space, we recover the training recipe of Path Integral Sampler [1]. Despite similar in a high-level, theoretical formulation, what distinguishes this work from the literature of SOC-based continuous samplers is the unique challenge of trajectory non-differetiability of CTMC on discrete state space, which renders many existing techniques on continuous samplers training inapplicable. Motivated by this, our work consider only losses that does not requires trajectory differetiability to numerically solve the problem. \n\n> **Question 2**:  Advantage of $\\mathcal{F} _ {\\operatorname{WDCE}}$ to $\\mathcal{F} _ {\\operatorname{LV}}$\n\nThe motivation for introducing WDCE loss is to realize a more *scalable* training in problem with high dimensions, mostly in terms of GPU memory usage. We have explained this motivation in line 211-214 and 224-227. \n\nTo be more specific, for training a discrete samplers of length $D$ with a batch of $B$ trajectories using the log variance loss $\\mathcal{F} _ {\\operatorname{LV}}$, the GPU memory usage is proportional to $BD$, as $\\mathcal{F} _ {\\operatorname{LV}}$ is a trajectory-level objective that requires backpropagating through all the gradient-tracking score values as is defined in Eq. (14). In high dimension scenarios, it's likely that $BD$ is significantly larger than $M$, the maximum possible number of gradient-tracking score values used in each loss evaluation (determined by the maximum GPU RAM). In an extremly high-dim case where $D > M$, $\\mathcal{F} _ {\\operatorname{LV}}$ will be completely inapplicable as it is unafforable to estimate the objective with even one trajectory.\n\nOn the contrary, using $\\mathcal{F} _ {\\operatorname{WDCE}}$ is relatively free from such concern. The importance sampling factor $\\exp(W^{\\overline{u}})$ used in Eq. (16) is detached from the computation graph, and the number of gradient-tracking score values used in loss evaluation can be completely controlled by the remasking phase, and is not influenced of the problem dimension $D$. Besides, $\\mathcal{F} _ {\\operatorname{WDCE}}$ also increases the use rate of score network outputs and further amortize the computation cost through the use of replay buffers. In summary, $\\mathcal{F} _ {\\operatorname{WDCE}}$ is more memory-efficient and applicable even in super high-dim problem settings, and this promotes us to adopt it only for later experiments on $16\\times 16$ distributions from statistical physics with $D = 256$.\n\nIn the revision, we will further highlight this difference in GPU memory consumption and explain clearly why it drives us to only use $\\mathcal{F} _ {\\operatorname{WDCE}}$ in high-dim experiments. Thank you for bringing up this question!\n\n> **Question 3**: Importance Sampling for $\\mathcal{F} _ {\\operatorname{LV}}$ and $\\mathcal{F} _ {\\operatorname{RERF}}$\n\nThank you for the interesting suggestion. We believe it is completely doable to implement data-reuse mechanism for other objectives as well. However, the main reason that we did not discuss them in this work is that importance-weighted versions of $\\mathcal{F} _ {\\operatorname{LV}}$ and $\\mathcal{F} _ {\\operatorname{RERF}}$ do not connect to simulation-free training of discrete diffusion models, which is the DCE loss described in Eq. (4). The main benefit of $\\mathcal{F} _ {\\operatorname{WDCE}}$ is that the simulation-free training with cached importance weights allows a flexible choice of batch size without worrying about being bounded by the problem dimension (see response to previous point). $\\mathcal{F} _ {\\operatorname{LV}}$ and $\\mathcal{F} _ {\\operatorname{RERF}}$ do not seem to enjoy similar connections despite the introduction of data-reuse scheme. \n\nThis indeed brings up a new question of what is the optimal choice for the sampling path measure in the framework. In this work, we mainly explore the on-policy scenario with sampling path measure generated by the current policy $\\overline{u}$. We believe that using importance-weighting to implement the choice of $\\mathbb{P}^ * $ as the sampling path measure would be an intriguing direction to explore for future investigations. \n\n---\nThank you again for your positive review and constructive feedback regarding the paper. We'd appreciate to discuss further if you have any additional questions or suggestions.\n\n\n### References\n\n[1] Zhang and Chen. Path integral sampler: a stochastic control approach for sampling. ICLR 2023."}, {"review_id": "3rISr4D5Qj", "reviewer": "Reviewer_7ypT", "summary": "This paper extends stochastic optimal control for continuous-time Markov chains (CTMCs) to discrete state spaces, introducing the Masked Diffusion Neural Sampler (MDNS)-–a neural sampler tailored to discrete domains. It proposes a weighted denoising cross-entropy loss, a scalable objective for training masked-diffusion models in sampling tasks. The theoretical framework generalizes existing continuous-space neural samplers to discrete settings, while a replay-buffer–based training algorithm improves computational efficiency. Experiments are comprehensive, benchmarking MDNS against both learning-based and classical non-learning Markov-chain methods across diverse probabilistic metrics.", "questions": "1. **Replay buffer**: Does it materially improve training? Have you explored prioritized replay?\n2. **Trajectory length**: Can the algorithm scale to very long trajectories?\n3. **CTMC justification**: Why is a CTMC formulation needed in discrete spaces?\n4. **Variance control**: What techniques could you use to reduce training variance?", "limitations": "No", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n\n1. Strong theory and a solid contribution that extends CTMC-based stochastic optimal control to discrete spaces.\n2. Introduces an efficient loss function and accompanying algorithm.\n\n**Weaknesses**\n\n1. The paper does not report how trajectory-level loss computation affects the variance of learning.\n2. The importance weights may further amplify training variance, yet no variance-reduction estimator is provided.\n3. Motivation is thin: why adopt a CTMC framework in discrete spaces when established approaches like GFlowNets already serve as effective discrete samplers without relying on CTMC dynamics?", "was_revised": true, "final_justification": "This paper bring interesting formulation based on stochastic control into discrete sampling and gives high performances. My questions and concerns are resolved and want to keep high score (5: accept).", "rebuttal": "We sincerely appreciate the reviewer's time and the constructive suggestions to help us improve the quality of this manuscript. We provide itemized responses below to each of them and hope they can help clarify your questions and concerns.\n\n> **Weakness 1, 2 & Question 4**: Variance reduction of training objectives and importance weight estimators\n\nThanks for pointing out the issue for variance reduction. Throughout our experiments, we observe that the variance of the loss is controlled and the training is almost always stable, thus we do not incorporate specific variance reduction techniques.\n\nOne possible trick of variance reduction in the logits could be to leverage the following equivalent formulation of the WDCE loss inspired by [1], i.e.,\n$$\\min _ \\theta \\operatorname{KL}(\\mathbb{P}^ * ||\\mathbb{P}^u) = \\mathbb{E} _ {X\\sim \\mathbb{P}^{\\overline{u}}}\\frac{\\mathrm{d}\\mathbb{P}^ * }{\\mathrm{d}\\mathbb{P}^\\overline{u}}(X)\\mathbb{E} _ {m\\sim\\operatorname{Unif}\\\\{1,2,...,D\\\\}}\\left[\\frac{D}{m}\\mathbb{E} _ {\\mu' _ m(\\widetilde{x}|X _ T)}\\sum _ {d:\\widetilde{x}^d=\\mathbf{M}}-\\log s _ \\theta(\\widetilde{x}) _ {d,x^d}\\right],$$\nwhere $\\mu' _ m(\\widetilde{x}|x)$ means obtain $\\widetilde{x}$ by uniformly selecting a subset of $\\\\{1,...,D\\\\}$ with size $m$ and mask the corresponding positions at $x$. In this way, by directly specifying the number of masks applied onto $x$ instead of independently masking each position of $x$ with probability $\\lambda$, one can possibly reduce the variance of the loss, which has been observed in [2] for evaluating the log likelihood.\n\nFurthermore, one can also use the trick of *antithetic sampling*, i.e., if $m\\ne D$, for each clean sample $x$ one can sample a pair of complementary masked positions whose intersection is empty and union is $\\\\{1,...,D\\\\}$, and compute the logits for these two masked versions of $x$. We believe that these two modifications would further reduce the variance in the logits.\n\nFor the noise in importance weights, we agree that their variance may also amplify the variance in computing the loss. During training, we in fact self-normalize all $W^u(X^{(i)}),1\\le i\\le B$ with softmax in the batch to ensure numerical stability. This is equivalent to first using the weights to estimate the normalizing constant $Z$ then pluging it back into the obejctive.\n\nWe will add more discussion of possible techniques for variance reduction and numerical stability improvement in the revised manuscript. Thank you for the insightful question!\n\n> **Weakness 3 & Question 3**: Motivation of using CTMC\n\nThe use of CTMC originates from the fact that we adopt a **continuous-time** control pespective of the problem. While this is not always necesary to achieve our purpose, taking a continuous-time formulation makes it easier to compute objectives from the path measure perspective and use well-established stochastic analysis tools such as Girsanov's theorem. We believe that our algorithms can also be formulated under the use of a discrete-time dynamics on the finite state space, mimicking the case of GFlowNets. One major difference between our considered SOC approach and GFlowNets objectives is that ours does not require the estimation of normalizing constant during training, while GFlowNets training losses typically involve such process.\n\n> **Question 1**: Use of replay buffer and other choices\n\nThe motivation of using replay buffer is to save the expensive computation costs of simulating the whole trajectories, rather than improve training. In fact, the optimal (yet most expensive) implementation of our algorithm would be not to use replay buffer and re-simulate trajectories for every optimization step, i.e., setting the resample frequency $k$ in Algorithm 1 as one so we are always using the samples from the current policy. Moreover, we only cache the final weights computed from each trajectory and discard the trajectory itself. We periodically drop all the buffered weights and resample new ones using current policy, without considering complicated buffer update algorithms. We did not explore prioritized replay as in fact the algorithms do not require the storage of transitions/trajectories, making prioritized replay largely inapplicable.\n\n> **Question 2**: Scaling to long trajectories\n\nOur algorithms are highly scalable for high dimensional problems, which incur long simulation trajectories. With the WDCE loss, the effectivenss of the training objective is not majorly influenced by the trajectory length. In the numerical experiments of this paper, we present successful results up to trajectory length of $16\\times 16=256$, which is already decently long. We also have additional numerical results showing that the method remains effective for problems with trajectories length up to $24\\times24=576$, which we will include in the revised manuscript. These empirical evidence suggests that our method can scale to long trajectories.\n\n---\nThank you again for your positive review and constructive feedback regarding the paper. We'd appreciate to discuss further if you have any additional questions or suggestions.\n\n\n### References\n\n[1] Ou et al. Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data. ICLR 2025.\n\n[2] Nie et al. Large Language Diffusion Models. ArXiv 2502.09992."}, {"review_id": "N2k1e0eyNw", "reviewer": "Reviewer_Qq8S", "summary": "This paper considers the task of training a neural sampler without data, on discrete distributions by only using the energy function. This paper combines many approaches in continuous space and applies them to CTMC. This paper also proposed an efficient training objection that does not require simulation every gradient update. The authors provided comprehensive evaluation to demonstrate the efficientness of the proposed approach.", "questions": "1. I am not generally familiar with the setup of these systems. Therefore, I might be curious why some of the parameters are set in that way (for example, line 301)? Is it to make the system easy (like reducing phase transitions) or make it harder? Is it a convention?\n\n\n2. I am a bit lost in the equation below line 202 (is there also a typo?). Could you please explain this to me?", "limitations": "Yes,  the authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## Strength:\n\n1. The writing is clear, pleasant to read, and easy to follow. It is a pleasure to review and read this paper.\n\n2. The paper not only proposed one approach. Instead, it provides a systematic recipe with several loss choices one can consider. While most of  these are not novel, it makes the paper complete.\n\n3. The proposed simulation-free loss is quite interesting and effective. \n\n4. Experiments are convincing. Comparing different losses is also helpful. \n\n\n\n## Weakness:\n\n1. I might suggest adding citations and discussion on previous works that use similar ideas to the buffer and weighted loss. As far as I know, FAB[1] and Adjoint Sampler[2]'s buffers have similar ideas. In FAB, one can do IS to reweight samples from the buffer. In AS, one collects buffer using current control and train the network with TSM.\n\n2. By convention, one define potential /  Energy $U$ that $p \\propto \\exp(-U)$. However, in this paper, it was defined as $p \\propto \\exp(U)$ andm hence it might be beneficial to follow that notation. I do not think this is a weakness (but I do not have other places to put this).\n\n\n\n[1] Midgley, Laurence Illing, et al. \"Flow Annealed Importance Sampling Bootstrap.\" ICLR,\n\n[2] Havens, Aaron, et al. \"Adjoint sampling: Highly scalable diffusion samplers via adjoint matching.\" ICML.", "was_revised": true, "final_justification": "The reply has addressed my concern and hence I will support for acceptance.", "rebuttal": "We sincerely appreciate the reviewer's time and the constructive suggestions to help us improve the quality of this manuscript. We provide itemized responses below to each of them and hope they can help clarify your questions and concerns.\n\n> **Weakness 1**: Related works using similar ideas to the buffer and weighted loss\n\nThank you for pointing out related works. Our approach indeed shares similarities with these works, and we will definitely add a discussion in the related work section of a revised version.\n\n> **Weakness 2**: Notation on potential / energy\n\nThank you for your suggestion regarding the notation. We will change it following the convention in the revised manuscript.\n\n> **Question 1**: Choice of hyperparameters\n\nThese hyperparameters are chosen based on the nature of the target distributions. First, for each distribution in statistical physics that is considered in this work, there exists a critical inverse temperature $\\beta _ \\mathrm{critical}$ at which the distribution undergoes a phase transition. The target distributions demonstrate three largely different behaviors at different position across this temperature spectrum, making it an ideal testbed for our algorithms.\n\nGenerally speaking, higher temperatures than the critical threshold lead to more uniform distributions and are thus easier to sample and learn from, while lower temperatures than the critical threshold add challenges in sampling and learning due to mode separation. For both statiscial physics models, we respectively consider the distribution below, at, and above the critical inverse temperature. For Ising model, the critical inverse temperature is $\\beta _ {\\mathrm{critical}} \\approx 0.4407$. For Potts model with $3$-spin state, $\\beta _ {\\mathrm{critical}} \\approx 1.005$. \n\nFor system size, we mostly experiment on the square lattice of $16\\times 16$, which incurs problems of a reasonable high dimension. This setup enables a comprehensive evaluation of algorithm performance. MDNS is successful on both examples across all three temperatures, suggesting its excellent robustness to hyperparameter settings as well as great scalability across dimensions.\n\n> **Question 2**: The equation below line 202\n\nIn the equation below line 202, we are computing the sum of the non-diagonal elements in $Q _ t^u$, which shows up as in the term $\\sum _ {y\\ne X _ t}(Q _ t^u-Q _ t^0)(X _ t,y)$ in the RN derivative $\\log\\frac{\\mathrm{d}\\mathbb{P}^ * }{\\mathrm{d}\\mathbb{P}^u}$ in Eq. (14). \n\nThe first equality is due to the fact that for $x\\ne y$, $Q _ t^u(x,y)\\ne0$ if and only if $x$ and $y$ only differ in one position (say, $d$), and $x^d=\\mathbf{M}\\ne y^d$. The second equality follows from the parameterization due to Eq. (3), where $Q _ t^u(x,x^{d\\gets n})=\\gamma(t)s _ \\theta(x) _ {d,n}$ and $s _ \\theta(x) _ {d,\\cdot}\\in\\mathbb{R}^N$ is a probability vector. We will clarify the purpose of this equation as well as add a short explanation in the revised manuscript, and please do let us know if you have any further questions on this equation.\n\n---\nThank you again for your positive review and constructive feedback regarding the paper. We'd appreciate to discuss further if you have any additional questions or suggestions."}, {"review_id": "7UqdvjZbDB", "reviewer": "Reviewer_ctjr", "summary": "The authors propose to leverage a Stochastic Optimal Control (SOC) of Continuous Time Markov Chains (CTMC) for solving the problem of sampling from an unnormalized density  $\\pi(x) \\propto exp(U(x))$ on a discrete state space. The authors state that sampling from an unnormalized density problem can be formulated as an SOC problem and propose computationally efficient methods to solve the corresponding SOC instance of the problem. In contrast to the previous approach for solving the SOC of CTMC, the authors eliminate differentiation through the sample generation trajectory and provide a total of 4 optimization objectives, i.e., REINFORCE trick-based, Log Variance Divergence-based, two Cross Entropy-based, that require only trajectory samples, but do not differentiate through them. \n\nThe algorithm for solving the SOC of CTMC for sampling with a Masked (and Uniform in Appendix F) prior process is proposed, called MNDS (and UNDS), and tested on a series of statistical physics-based benchmarks. The method is compared to the state-of-the-art neural sampler LEAPS [1] and learning-free Metropolis Hastings baseline, and shows itself on par on high temperature distributions and superior on lower temperature distributions.", "questions": "- It would be nice if the authors added LEAPS to the comparison in theTable 2.\n- Why is WDCE loss used for experiments on 16x16 Ising and Potts models, while the LV loss showed itself to be the best on the training objectives study, see Sec 4.1?\n- In the Proposition 2 main paper, there seems to be a missing expectation. I think it should be $\\hat{Z} = \\mathbb{E}_{X \\sim P^u}e^{W^u(x)}$\n- How do you ensure the memoryless condition $\\int_0^T \\gamma(t)dt=\\infty$ in practice? In the case of masked process MNDS memoryless property seems obvious (since you start from a point), but for the UDLM (Appendix F), in my opinion, it is not obvious how to ensure this condition.", "limitations": "Would like to see the mention of method not being simulation-free in the limitations section.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The proposed approach is quite simple and effective.\n- SOC for CTMC development helps to bridge the gap between CTMC models and diffusion models and, in my opinion, is valuable for the CTMC community.\n- The MNDS method seems to outperform other methods on lower temperature distributions.\n- Theoretical guarantees on sampling quality and normalizing constant are shown.\n- The paper is mainly well-written and easy to follow.\n- Experimental evaluation is through with a lot of quantitative metrics, which is rare for an unnormalized density on a discrete state space sampler paper.\n\nWeaknesses:\n\n- In Section 3.2, when LogVar and CE losses are presented, it is hard to understand where exactly the stop gradient can be placed and why it is so. I think it should be explained more clearly.\n- Would like to see the code in the sake of reproducibility, although sufficient experimental details are described in the Appendix.\n- No information on the computational complexity of the method compared to LEAPS and MH. It would be nice if the authors provided computational complexity analysis and training/inference time comparisons.\n- Performance seems a bit behind the MH on Ising 4x4, see Table 2.\n- I can’t find the preconditioning strategy mentioned in the main text, only in the Appendix. I think it should be mentioned in the Sec 4.\n- The model was evaluated only on the statistical physics benchmarks. It would be nice to see it’s performance on other type of problems, e.g., combinatorial optimization.", "was_revised": true, "final_justification": "The authors have adressed my concerns, but the score of 5 is enough in my opinion. Overall good paper would like to see it accepted.", "rebuttal": "We sincerely appreciate the reviewer's time and the constructive suggestions to help us improve the quality of this manuscript. We provide itemized responses below to each of them and hope they can help clarify your questions and concerns.\n\n> **Weakness 1: Stop gradient for LV and CE losses**\n\nWe apologize for the potential confusion regarding the stop gradient operator. As we use $v=\\overline{u}$ in both losses, we can write the LV loss as $\\mathcal{F} _ \\mathrm{LV}(\\mathbb{P}^u,\\mathbb{P}^ * )=\\mathrm{Var} _ {\\mathbb{P}^{\\overline{u}}}\\log\\frac{\\mathrm{d}\\mathbb{P}^ * }{\\mathrm{d}\\mathbb{P}^u}$, which means we first use the current policy to sample a batch of trajectories $X^{(i)}\\sim\\mathbb{P}^u$, $1\\le i\\le B$, detach them from the computational graph, and then compute all $\\log\\frac{\\mathrm{d}\\mathbb{P}^ * }{\\mathrm{d}\\mathbb{P}^u}(X^{(i)})$ using Eq. (14), and finally do back propagation on the loss $\\mathrm{Var}\\left(\\log\\frac{\\mathrm{d}\\mathbb{P}^ * }{\\mathrm{d}\\mathbb{P}^u}(X^{(i)}):1\\le i\\le B\\right)$. The same applies to the CE loss. The reason for applying the stop gradient operator is to ensure we compute the correct gradient of the loss, and to avoid taking gradient over the sampled trajectories $X\\sim\\mathbb{P}^u$ (which is not differentiable due to the discrete nature of the trajectories). We will provide a clearer explanation of this in the revised manuscript.\n\n> **Weakness 2**: Codes for reproducibility\n\nWe understand the importance of reproducibility and appreciate that you asked. Unfortunately, due to NeurIPS policy, all the texts we post (rebuttal, discussion) cannot contain any links to external pages, and we are unable to upload supplemental materials at this stage. We will make our code publicly available when the anonymous period ends. Thank you for your understanding!\n\n> **Weakness 3**: Algorithm computational complexity compared to LEAPS and MH\n\nWe describe the training details in line 886-891 for Ising model and line 1052-1057 for Potts model. The training time for $16\\times16$ Ising model on an A100 is around 4 hours, and for $16\\times16$ Potts model on an A100 is 20 hours. For generating ground truth samples, we run the MH algorithm for 3 hours to sample from $16 \\times 16$ Ising model, and the SW algorithm for 2 hours to sample from $16 \\times 16$ Potts model on CPU. For the MH baseline in Potts model, we run it for 30 hours on CPU. For all experiments with LEAPS, we run it for a time comparable to that of MDNS to ensure a fair comparison. We will add these information and discussion into the paper during revision.\n\n> **Weakness 4**: Performance on Ising 4x4\n\nWe agree that the TV distance and KL and $\\chi^2$ divergence of the samples generated by MDNS on $4\\times 4$ Ising model are slightly larger than those from MH algorithm. However, we believe this is mostly due to the simplicity of the problem as well as the numerical errors. In general, MDNS performs comparably to MH on Ising models and outperforms it on Potts models (see Fig 2 and Tab 3).\n\n> **Weakness 5**: Preconditioning should be mentioned in the Sec 4\n\nThank you for your suggestion. We will add a brief discussion in the main text regarding the exploration of the preconditioning techniques.\n\n> **Weakness 6**: Performance on other type of problems\n\nThank you for the great suggestion! We will include results on combinatorial optimization problems, such as finding maximum cut and maximal independent set of graphs, in the revision.\n\n> **Question 1**: Add LEAPS to Table 2\n\nThank you for your suggestion. The purpose of Table 2 was mainly to display the ablation results over different choices of training objectives, and therefore we do not include LEAPS in the comparison. We have compared with LEAPS on experiments with a larger scale and results can be found in Figures 1, 2 and Tables 1, 3. We hope these comparisons have sufficiently demonstrated that MDNS outperforms LEAPS in numerical performance.\n\n> **Question 2**: LV loss not used for 16x16 experiments\n\nThe reason that we do not adopt $\\mathcal{F} _ {\\mathrm{LV}}$, $\\mathcal{F} _ {\\mathrm{CE}}$ or $\\mathcal{F} _ {\\mathrm{RERF}}$ for the high dimensional experiments is because they are too memory-intensive. These objectives share the similarity that they all requries using the full trajectory to estimate the objectives, making the loss computation requires at least $BD$ gradient-tracking score values in one step, where $B$ is the number of trajectories and $D$ is the length of the trajectory. For the $16 \\times 16$ experiments, this would be unafforably expensive unless we pick extremely small values for $B$, which potentially would incur large variance detrimental to optimization. On the other hand, using $\\mathcal{F} _ {\\mathrm{WDCE}}$ is free of such concern and has a more controllable, affordable memory usage for high dimensional example. Therefore, we stick with only $\\mathcal{F} _ {\\mathrm{WDCE}}$ in the high-dimensional settings. \n\nIn the revision, we will highlight more this difference in GPU memory consumption and explain clearly why only WDCE is used in large-scale experiments. Thank you for bringing up this question!\n\n> **Question 3**: Missing expectation in Proposition 2\n\nThank you for carefully reading our paper. Following Eq. (14), we know that $Z=\\mathbb{E} _ {X\\sim\\mathbb{P}^u}\\mathrm{e}^{W^u(X)}$. Thus, we can sample $X\\sim\\mathbb{P}^u$, and use the *random variable* $\\widehat{Z}=\\mathrm{e}^{W^u(X)}$ as an unbiased estimator of $Z$. Therefore, we do not add an expectation here.\n\n> **Question 4**: Ensuring the memoryless condition for UDNS\n\nIndeed, we require $\\int _ 0^1\\gamma(t)\\mathrm{d}t=\\infty$ to ensure the memoryless condition (see line 1119-1121). For UDNS, as is discussed on line 1163, we choose $\\gamma(t)=\\frac{1}{t}$ in our experiments. In practice, we use early-stopping (which is a standard practice in diffusion model training) and only train on the interval $[\\epsilon, 1]$. However, by choosing a reasonably small $\\epsilon$, we can ensure the memoryless property to a satisfactory degree that does not affect the correctness of the algorithm.\n\n> **Limitation**: The method is not simulation-free\n\nThank you for pointing this out. Just for us to better answer, may we clarify which \"simulation-free\" were referred to in this comment, as it could mean several things across different communities --\n\n- If being simulation-free refers to the algorithm does not requires differentiate through trajectories to perform optimization, then our algorithm is indeed simulation-free in this sense, as it is a necessary requirement for SOC of CTMCs.\n- If being simulation-free refers to the algorithm does not requries simulating the current policy to generate trajectories, then indeed our algorithm is not simulation-free, and we are happy to discuss this in the limitation section. However, we believe that requiring rolling out policy is a shared limitation of many of the neural samplers. We will commit to explore discrete neural samplers with less reliance on simulations in the future.\n\n---\nThank you again for your positive review and constructive feedback regarding the paper. We'd appreciate to discuss further if you have any additional questions or suggestions."}]

Author rebuttal:

Meta-review: This paper studies diffusion-based samplers of discrete distributions given by an unnormalised probability mass functions. Generalising the long line of existing work on this problem in the continuous setting, objectives for enforcing time reversal of the generative and the forward noising process -- now continuous-time, discrete-space Markov chains -- are derived and some methods for practical training are introduced. This sampler is evaluated on sampling two synthetic distributions from statistical physics (Isings and Potts models).

While the experiments are quite basic, the reviewers appreciate the theoretical/algorithmic contribution and give a unanimous Accept rating, and I agree with their assessment. However, it is strongly recommended that the authors address the following points in the revision:
- Clarify the connection to prior work on sampling via stochastic control in continuous space (e.g., PIS as mentioned by XE64, DIS from reference [11] in the paper).
- Clarify the connection to similar work on amortised sampling with discrete diffusion models (e.g., 7ypT mentions GFlowNet objectives, which have been used for discrete diffusion models -- for example, in [Venkatraman et al., NeurIPS'24](https://arxiv.org/abs/2405.20971) -- and can be trained with the LV objective).
- Improve the presentation of the method, taking reviewers' suggestions and minor corrections into account.
- Include discussion of complexity/cost in comparison to baselines.
Decision: Accept (poster)