Discrete Neural Flow Samplers with Locally Equivariant Transformer
Authors: Zijing Ou
Ruixiang ZHANG
Yingzhen Li
Venue: NeurIPS 2025 poster ()
Keywords: discrete diffusion models
discrete energy-based models
sampling
Primary area: probabilistic_methods
TL;DR: We present an algorithm that learns the rate matrix of a continuous-time Markov chian to satisfy the Kolmogorov equation for sampling from unnormalised discrete distributions.

Abstract: Sampling from unnormalised discrete distributions is a fundamental problem across various domains. 
While Markov chain Monte Carlo offers a principled approach, it often suffers from slow mixing and poor convergence.
In this paper, we propose Discrete Neural Flow Samplers (DNFS), a trainable and efficient framework for discrete sampling. DNFS learns the rate matrix of a continuous-time Markov chain such that the resulting dynamics satisfy the Kolmogorov equation.
As this objective involves the intractable partition function, we then employ control variates to reduce the variance of its Monte Carlo estimation, leading to a coordinate descent learning algorithm.
To further facilitate computational efficiency, we propose locally equivaraint Transformer, a novel parameterisation of the rate matrix that significantly improves training efficiency while preserving powerful network expressiveness.
Empirically, we demonstrate the efficacy of DNFS in a wide range of applications, including sampling from unnormalised distributions, training discrete energy-based models, and solving combinatorial optimisation problems.

Discrete Neural Flow Samplers with Locally Equivariant Transformer
Zijing Ou', Ruixiang Zhang?, Yingzhen Li Imperial College London, *Apple
(z. yingzhen 1im@mmariaa....... ac uk saixiang_hhang2800 com
Abstract
Sampling from unnormalised discrete distributions is fundamental problem across various domains. While Markov chain Monte Carlo offers a principled approach, it often suffers from slow mixing and poor convergence. In this paper, we propose Discrete Neural Flow Samplers (DNFS), a trainable and efficient framework for discrete sampling. DNFS learns the rate matrix of continuous-time Markov chain such that resulting dynamics satisfy the Kolmogorov equation. this objective involves the intractable partition function, then employ control variates to reduce the variance of its Monte Carlo estimation, leading to a coordinate descent learning algorithm. To further facilitate computational efficiency, propose locally equivaraint Transformer, novel parameterisation of the rate matrix that significantly improves training efficiency while preserving powerful network expressiveness. Empirically, we demonstrate the efficacy of DNFS in a wide range of applications, including sampling from unnormalised distributions, training discrete mesey models, and solving combinatorial optimisation problems.
Introduction
We consider the task of sampling from a discrete distribution π(x)
pla)
2 known only up to
normalising constant Z = -E p(z). This problem is foundamental wide range of scientific domains, including Bayesian inference (Murray et al., 2012), statistical physics (Newman & Barkema, 1999), and computational biology (Lartillot & Philippe, 2004). However, efficient sampling from such unnormalised distributions remains challenging, especially when the state space is large and combinatorially complex, making direct enumeration or exact computation of infeasible.
Conventional sampling techniques, such as Markov Chain Monte Carlo (MCMC) (Metropolis et al., 1953) have been widely employed with great success. Nevertheless, MCMC often suffers from poor mixing and slow convergence due to the issues of Markov chains getting trapped in local minima and large autocorrelation (Neal al., 201 These limitations have motivated the development of neural samplers (Wu et al., 2020; Vargas et al., 2024; Máté & Fleuret, 2023), which leverage deep neural networks improve sampling efficiency and convergence rates. In discrete settings, autoregressive models (Box et al.. 2015) have been successfully applied to approximate Boltzmann distributions of spin systems statistical physics (Wu al., 2019). Inspired by recent advances in discrete diffusion models (Austin et al., 2021; Sun et al., 2023c: Campbell et al., 2022), Sanokowski et al. (2024, 2025) propose diffusion-bar based samplers with applications problems. Moreover, Holderrieth et al. (2025) introduces an alternative discrete sampler by learning a parametrised continuous- time Markov chain (CMCT) (Norris, 1998) to minimise the variance of importance weights between the CMCT-induced distribution and the target distribution.
solving combinatorial optimisation
Building these advances, the goal of our paper to develop sampling method for discrete distributions that is both efficient and scalable. To this end, we introduce Discrete Neural Flow
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Samplers (DNFS), novel framework that learns the rate matrix of a CTMC whose dynamics satisfy the Kolmogorov forward equation (Oksendal, 2013). In contrast to discrete flow models (Campbell et al., 2024; Gat et al., 2024), which benefit from access to training data to fit the generative process, DNFS operates in settings where no data samples are available. This data-free setting makes direct op- timisation of the Kolmogorov objective particularly challenging and necessitates new methodological advances to ensure stable and effective training. Specifically, the first difficulty lies in the dependence of the objective on the intractable partition function. We mitigate this by using control variates (Geffner & Domke, 2018) to reduce the variance of its Monte Carlo estimate, which enables efficient optimisation via coordinate descent. More critically, standard neural network parameterisations of the rate matrix render the objective computationally prohibitive. To make learning tractable, a locally equivariant Transformer architecture is introduced to enhance computational efficiency significantly while retaining strong model expressiveness. Empirically, DNFS proves to be an effective sampler for discrete unnormalised distributions. We further demonstrate its versatility in diverse applications, including training discrete energy-based based models and solving combinatorial optimization problems.
2 Preliminaries
We begin by introducing the key preliminaries: the Continuous Time Markov Chain (CTMC) (Norris, 1998) and the Kolmogorov forward equation (Oksendal, 2013). Let x: be a sample in the d-dimensional discrete space f1.....S) S}a by rate matrix R. X x X + R. which captures the instantaneous rate of change of the transition probabilities. Specifically, the entries of R2 are defined by
& X. A continuous- time discrete Markov chain at time is characterised
Pt+ ((+2))(()1z) 1 1==
1, ==
Re(y, lim Ar-
1
(1)
At
0, y
which equivalently yields the local expansion Pt+ (1) = 1y= + Re(y. a) - ot) and the rate matrix satisfies (v, x) 0 if y # and Re x) Eysz R)), Given R., the marginal distribution (z for any e E uniquely determined. Let Toses1 be a sample trajectory. Our goal is to seek rate matrix R₂ that transports an initial distribution po xx 17 to the target distribution P1 x p. The trajectory then can be obtained via the Euler method (Sun et al., 2023c)
Titat ~ Cat (r; 1re+s.=0 + Re(2 x²) At).
20 ~ Po.
(2)
and the induced probability path pt by Rs satisfies the Kolmogorov equation (Oksendal, 2013)
d.pr(x) = Σ (.(x,))p.()) = Σ (. R. ()).
(3)
In this case, we say that the rate matrix R2 generates the probability path pt. Dividing both sides of the Kolmogorov equation by pt leads
po(y) pr(
a log = Re y)
Re.y.x).
(4)
In the next section, we describe how to leverage Equation (4) to learn a model-based rate matrix for sampling from given target distribution π, followed by the discussion of applications to discrete energy-based modelling and combinatorial optimisation.
3 Discrete Neural Flow Samplers
The rate matrix Rt that transports the initial distribution to the target is generally not unique. However, we can select particular path by adopting an annealing interpolation between the prior rj and the target p. defined as Pt x p'n)'-t & Pe (Gelman & Meng, 1998; Neal, 2001). This annealing path coincides with the target distribution π & at time t = 1. To construct an R₂ that generates the probability path Pt, we seek a rate matrix that satisfies the Kolmogorov equation in Equation (4). Specifically, we learn a masel-based rate matrix R2 (f,, parametrised by 0, by minimizing the loss
R² (a, v) Pt Pc(y)
((0) = (w) R,), (x R a logop logg((x)+ Σ
R$ (u.
(5)
Pe (x

---

where di is an arbitrary reference distribution that has the same support as Pt, and w(f) denotes time schedule distribution. At optimality, the condition 6ạ(2; R.)) = O holds for all f. and x E X, implying that the learned rate matrix Ry ensures the dynamics prescribed by the Kolmogorov equation are satisfied along the entire interpolation path. In practice, minimising the loss (5) guides R& to correctly capture the infinitesimal evolution of the distribution pt, enabling accurate sampling from the target distribution via controlled stochastic dynamics.
However, evaluating Equation (5) directly is computationally infcasible due to the intractable summa- tion over y, which spans an exponentially large space of possible states, resulting in a complexity of (S))) To alleviate this issue, we follow Sun et al. (2023c); Campbell et al. (2022); Lou et al. (2024) by assuming independence across dimensions. In particular, we restrict the rate matrix R& such that it assigns non-zero values only states y that differ from x in most one dimension. Formally, Re x) = 0 if & [x),,, where (xx) := {y E X ly: * I; at most one i}. To improve clarity in the subsequent sections, we renotate the rate matrix Ry for y E (x) with yo + H2 as
Rº x)
Ry (3:, and
R (z), = Σ (l..))))
(6)
which yields simplified and more tractable form of the loss in Equation (5)
Pt Pr(y) pa (x)
(x;;; mz) = O (x)) + Σ (V,,
(2x, i ily)
This approximation reduces the computational complexity from O(5)) to O x d). Nonetheless, two main challenges persist. First, the time derivative 0 log (1 remains intractable due to the dependence on the partition function, as it expands to Ô log P. (x) - O log Ze with Z, = Σ fh (x) being intractable. Second, evaluating Equation (7) requires evaluating the neural network [xxx) times, which is computationally expensive for each x. In the following, we propose several techniques address these computational bottlenecks.
3.1 Estimating the Time Derivative of the Log- Partiiatn Function
To estimate the time derivative, note that à log Z, = (x) [0g log pc()r) which can be approximated
via the Monte Carlo estimator θ log Z₁ ~ x Σ Ô log pr )*. However, this approach relies on sampling from Pt. which is typically impractical due to the lack of convergence guarantees for
short-run MCMC in practice and the high variance inherent in Monte Carlo estimation. To address this issue, we leverage a key identity that holds for any given rate matrix R²
Pc(y) ((
ag log Z, = argmin Ep, (E) Re) - 2,)2,E((e R.) A Os log (ox))())) Σ Re 8)
(8)
Because the objective (8) is convex in Ct, the minimizer is given by Og log Z, = Ep, (x; RR). Moreover, in practice, the expectation over pt can be safely replaced by an expec- tation over any distribution qe with the same support, still yielding a valid estimate of the time deriva- tive (see Appendix A. for details). Empirically, we observe that using Equation (8) results in a significantly lower-variance estimator compared to the direct Monte Carlo approach o, log Z, Ep, [0, log pe (=)]. This reduction in variance can lead to improved optimisation performance. To assess this, we conducted experiments on the Ising model (Mézard et al., 1987), min- imising the loss Equation (7) using two different estimators for O log Z,, The standard deviations of both estimators, as well as their corresponding loss values during training, are plotted over training steps in Figure The results demonstrate that the estimator based on E. R.) consistently achieves lower loss values, underscoring the benefits of reduced variance in estimating 0 log Z, for
3.0
EaEibo R.)
2.0
1.5
log
1.0
0.
0
10k
20K
30k
40K
50k
10K
20k
30%
40K
50K
Tramiing Steps
Training Steps
Figure 1: Comparison of std. dev. and training loss for different estimators of θ log Z.. Lower variance estimator exhibits lower training loss, indicating a better learned rate matrix satisfying the Kolmogorov equation in Equation (5).
3

---

improved training dynamics. In Appendix A.2, we provide perspective of control variate (Geffner & Domke, 2018) to further explain this observation. This insight enables coordinate descent approach to learning the rate matrix. Specifically: i) 8 <- argmin, Jo' (E) F - 22 dt; and ii) G. f argming, Eeg (E:(2; R₂) - cq)². Alternatively, the time derivative can be parameterised directly via neural network & allowing joint optimisation of and & through the objective argming.6 = fu (E R2)) - )² dt. This formulation recovers the physics-informed neural network (PINN) loss proposed in Holderrieth al. (2025). A detailed discussion of the connection to the PINN loss is provided in Appendix A.3.
3.2 Efficient Training with Locally Equivariant Networks
As previously noted, computing the ó function in Equation (7) requires evaluating the neural network [V] times, which is computationally prohibitive. Inspired by Holderrieth et al. (2025), we proposed to mitigate this issue by utilising locally equivariant networks, an architectural innovation that significantly reduces the computational complexity with the potential to preserve the capacity of network expressiveness. A central insight enabling this reduction is that any rate matrix can be equivalently expressed as a one-way rate matrix while still inducing the same probabilistic path. This is formalised in the following proposition:
Proposition 1. For rate matrix RG that generates the probabilistic path pt, there exists a one-way rate matrix Qr x)- = R.(y.x) R2 (x, y)
y + x and Qe(z, x) = Eyfz Qr(y,)), x)), that generates the same probabilistic path pt, where [z]+ max(z, 0) denotes the ReLU operation.
po(v]
pr(m)
This result was originally introduced by Zhang et al. (2023b), and we include a proof in Appendix B.1 for completeness. Building Proposition 1, can parameterise R. directly one-way rate matrix. To achieve this, we use a locally equivariant neural network as described by Holderrieth et al. (2025). Specifically, a neural network G is locally equivariant if and only if:
G (T, i|z) G (Ix, i (/Saa (z, i, T)),
... d
where Swap(ar, i, (Z1...., I1-1, T, one-way rate matrix can be defined as (T,
Ia) and T E (1...., S} - S. Based on this, the i|z) _ [G (T, i|z)]++ Substituting this parametrization
S
into Equation (7), we obtain the simplified expression:
1- (1), pr(v) (x(x)
(x; RC) = 0 (o +
Σ (G),,, inet
(10)
This formulation reduces the computational cost from ([NV) to ((), enabling far more efficient training. We term the proposed method as discrete neural flow sampler (DNFS) and summarise the training and sampling details in Appendix C. Nonetheless, the gain in efficiency introduces challenges
in constructing a locally equivariant network (leNet) that both expressive and flexible.
3.3 Instantiation of leNets: Locally Equivariant Transformer
To construct a locally equivariant network, we first introduce hollow network (Chen & Duve- naud, 2019). Formally,
Fit- (II,..., = T..... ...) denote the input with its i-th token set T. A function H : X + dR is termed a hollow network if it satisfies H (xver)); H (z ). VT, E S, where M,. denotes the i-th row of the matrix M. Intuitively, it implies that the output at position i invariant to the value of the i-th input token. Hollow networks provide a foundational building block for constructing locally equivariant networks, as formalised in the following proposition.
Proposition 2 (Instantiation of Locally Equivariant Networks). Let x E X denote the input tokens and H X + RRdxh he a hollow network. Furthermore, for each token T E S, let Wy E Re be a learnable projection vector. Then, the locally equivariant network can be constructed as:
G(r, i|z) (w, - w2,Y
(z))..- (
A rate matrix H is said to be way if M(y, 2) > 0 implies R(x, v) = 0. That is, if a transition from to y is permitted, the reverse transition must be impossible.

---

This can be verified via G(T,i|z) = (w,, -w,)H (Swap(r, r))s. = -G i d/Swap(z, i, r)). Although Proposition 2 offers concrete approach to constructing locally equivariant networks, sig- nificant challenges persist. In contrast to globally equívariant architectures (Cohen & Welling, 2016; Fuchs et al., 2020), where the composition of equívariant layers inherently preserves equivariance, locally equivariant networks are more delicate to design. In particular, stacking locally equivariant layers does not, in general, preserve local equivariance. While Holderrieth et al. (2025) propose leveraging multi-layer perceptions (MLPs), attention mechanisms, and convolutional layers (see Appendix B.2 for details) to construct locally equivariant networks, these architectures may still fall short in terms of representational capacity and flexibility.
Locally Equivariant Transformer.
As shown in Proposition 2, cru- cial component to construct a locally equivariant network is the hollow net- work. Specifically, the key design con- straint is that the network's output at dimension must be independent of the corresponding input value r,, Oth- erwise, any dependence would result in information leakage and violate lo- cal equivariance. However, the i-th output may depend freely all other coordinates of the input token & ex- cept for the i-th entry. This insight motivates the use of hollow transform- ers (Sun et al., 2023c) as a founda- tion for constructing locally equivari- ant networks. Specifically, it employs two autoregressive Transformers (Vaswani et al., 2017; Radford et al., 2018) per layer; one processing inputs from left to right, and the other from right left. In the readout layer, the representations from two directions are fused via attention to produce the output. This design ensures that each output dimension remains independent of its corresponding input coordinate, while still leveraging the expressiveness of multi-layer Transformers. Thereby, the final output (: ilæ:) can be obtained by taking the inner product between the hollow attention output and the token embeddings produced by the projection layer. We term the proposed architecture as locally equivariant transformer (leTF)), and defer the implementation details to Appendix B.3.
(1))))
6f (.,2)))
H HOO)
Attention Readout Layer
Projection Layer
Q
Qu
[V1, VRl
Neural
K-
K₁
Kg
Vp
["2
Kg
Toiee
ridirectional Causal Attention Layer
L2R Causal Attention Layer
R2L Causal Attention Layer
t
X1
X2
X3
X4
t
Figure 2: Illustration of the leTF network.
Comparison of Different leNets. In Figure 3, we compare leTF with other locally equivariant networks for train- ing a discrete neural flow sampler on the Ising model. The results show that leTF achieves lower estimation errors and higher effective sample sizes (see Appendix E. for experi- mental details). leAttn and leMLP, which each consist of only a single locally equivariant layer, perform sig- nificantly worse, highlighting the importance of network expressiveness in achieving effective local equivariance. Although leConv performs comparably to IcTF, its convolutional design is inherently less flcxible. It is restricted to grid-structure readily generalise other data types like text graphs. Additionally, shown in Figure 14, leTF achieves lower training loss compared to leCony, further confirming its advantage in expressiveness.
leTF (ours)
leConv
l0g PAASE
leAttn
1/
W W Mmm
mm
40K
50K
Training Stops
Training Steps
Figure 3: Comparison of log RMSE (L) and ESS (t) for different locally equivariant networks. More expressive net- works achieve better performance.
data, such as images or the Ising model, and does not
Applications and Experiments
To support our theoretical discussion, we first evaluate the proposed methods by sampling from predefined unnormalised distributions. We then demonstrate two important applications to DNFS: i)
5

---

GFlowNet
LEAPS
GelowNet
Energy
Energy
Energy
Oraele
DNFS
Oracie
DNES
Oraele
DNFS
Figure 4: Comparison between different discrete samplers on pre- rre-traned EBMs.
training discrete energy-based models and ii) solving combinatorial optimisation problems. Detailed experimental settings and additional results are provided in Appendix E,
4.1 Sampling from Unnormalised Distributions
Sampling from Pre-trained EBMs. We begin by evaluating the effectiveness of our method by sampling from a pre-trained deep energy-based model. Specifically, we train an EBM on 32- dimensional binary data obtained by applying the Gray code transformation (Waggener & Waggener, 1995) to a 2D continuous plane, following Dai et al. (2020). The EBM consists of a 4-layer MLP with 256 hidden units per layer and is trained using the discrete energy discrepancy introduced in Schröder et al. (2024). Therefore, the trained EBM defines an unnormalised distribution, upon which train discrete neural sampler. We benchmark DNFS against three baselines: (i) long-run Gibbs sampling (Casella & George, 1992) as the oracle; (ii) GelowNet with trajectory balance (Malkin al., 2022); and (iii) LEAPS (Holderrieth et al., 2025) with the proposed leTF network.
The results, shown in Figure 4, demonstrate that the proposed method, DNFS, produces samples that closely resemble those from the oracle Gibbs sampler. In contrast, GFlowNet occasionally suffers from mode collapse, particularly on structured datasets such as checkerboard pattern. Although LEAPS with leTF achieves performance comparable DNFS, it sometimes produces inaccurate samples that fall in smoother regions of the energy landscape, potentially due to imprecise estimation observe that LEAPS with leConv performs poorly in this setting (see
of 02 log Z, Furthermore, Figure 12), reinforcing the limited expressiveness of locally equivariant convolutional networks when applied to non-grid data structures. For on other datasets are provided in Figure 11, further illustrating the effectiveness of our method.
comprehensive evaluation, additional visualisations
Sampling from Ising Models. We further evaluate method on the task of sampling from the lattice Ising model, which has the form of
(11)
where J = AAD with R and Ap being the adjacency matrix of a D x D grid. 2 In Figure 5, evaluate DNFS on a D 10 x 10 lattice grid with a = 0.1, comparing it to baselines methods in terms of effective sample size (ESS) (see Appendix D.1 for details) and the energy histogram of 5, 000 samples. The oracle energy distribution is approximated using long-run Gibbs sampling. The results show that DNFS performs competitively with LEAPS and significantly outperforms GelowNet, which fails to capture the correct mode of the energy distribution. Although LEAPS with leCony achieves comparable effective sample size, it yields a less accurate approximation of the energy distribution compared to DNFS. Furthermore, DNFS attains a lower loss value, as shown in Figure 15, These findings underscore the effec- tiveness of the proposed neural sam- pler learning algorithm and highlight the strong generalisation capability of leTF across both grid-structured and non-grid data. For more comprehen- sive evaluation, we also compare our method with MCMC-based approaches in Figure 13, further demonstrating the efficacy of our approach.
DMFS
0.00
Orade
LEAPS
ONFS
0.15
LEAPS
0.0
GHowhet
8
0.4
67
0.0
0
70k
10x
60K
BOX
100%
0.00
-30
-15
-10
0
16
Training Steps
Energy
Comparison of effective sample size and his- togram of sample energy on the lattice Ising model.
Figure 5:
"Th adjacency matrix is constructed using A D Seraphh.rapphLLLttiiccddii---- DJ, circular=True).

---

Energy
Samples
Data
Energy
Samples
Energy
Samples
Figure 6: Results of probability mass estimation in training discrete EBMs. We visualise the training data, learned energy landscape, and the synthesised samples of DNFS.
4.2 Training Discrete Energy-based Models
A key application of DNFS is training energy- based models (EBMs). Specifically, EBMs define parametric distribution Pa x (x)()),,x))) where the goal is learn an energy function Eg that approximates the data distribution. EBMs are typically trained using contrastive divergence (Hinton, 2002), which estimates the gradient of the log-likelihood as
V lo (x) = (v) |V & - Eg))
(12)
To approximate this intractable gradient, we train a rate matrix Ry to sample from the target Pg. This enables using importance sampling to estimate the expectation (see Appendix D.1 for details).
K
exp(w(())
(k) = =f, Ce(er; R2)) de.
[ [oEs(
jj=1 exp(w6)) V (xEs(
(13)
This
approach is more effective than traditional MCMC methods, in optimal training, it has garuantee to produce exact samples from the target distribution within a fixed number of sampling steps. Moreover, neural samplers arguabelly easier discover regularities and jump between modes compared to MCMC methods, leading to better exploration the whole energy landscape, and thus results in more accurate estimate of the energy function (Zhang al., 2022a). To demonstrate the effectiveness of DNFS in energy-based modelling, we conduct experiments on probability estimation with synthetic data and training Ising models.
Probability Mass Estimation. Following Dai et al. (2020), we first generate 2D floating moating-point data from several wwo-dimensional distributions. Each dimension is then encoded using 16-bit Gray code, resulting in a 32-dimensional training dataset with possible states per dimension.
Figure 6 illustrates the estimated energy landscape alongside samples generated using the trained DNFS sampler. The results demonstrate that the learned EBM effectively captures the multi-modal structure of the underlying distribution, accurately modelling the energy across the data support. The sampler produces samples that closely resemble the training data, highlighting the effectiveness of DNFS in training discrete EBMs. Additional qualitative results are presented in Figure 16. In Table 3, we provide quantitative comparison of our method with several baselines, focusing in particular on two contrastive divergence (CD)- bDDed approaches: PCD (Tieleman, 2008) with MCMC and ED-GFN (Zhang et al., 2022a) with GFlowNet. Our method, built upon the proposed DNFS, consistently outperforms PCD in most settings, underscoring the effectiivnness of DNFS in training energy-based models. While ED-GFN achieves better performance than DNFS, it benefits from incorporating Metropolis-Hastingg
(MH) (Hastings, 1970) correction to sample from the model distribution Pa in Equation (12), which may offer an advantage over the importance sampling strategy used in Equation (13). We leave the integration of MH into DNFS direction for future work.
Training Ising Models. We further assess DNFS for training the lattice model defined in Equation (11). Following Grathwohl et al. (2021), generate training data using Gibbs sampling and use these samples to learn a sym- metric matrix Jo to estimate the true matrix in the Ising model. Importantly, the training algo- rithms do not have access the true data- aatt--enerating matrix J, but only the synthesised samples.
Truth
Initialised
Learned
Figure 7: Results on learning Ising models.
In Figure 7, we consider a D 10 x 10 grid with a = 0.1 and visualise the learned matrix Jg using a heatmap. The results show that the proposed method successfully captures the underlying pattern of the ground truth, demonstrating the effectiveness of DNFS. Further quantitative analysis across various configurations of D and is presented in Table 5.

---

Table Max independent set experimental results. We report the absolute performance, approxima- tion ratio (relative to GUROBI), and inference time.
1-16
ER32-40
ER64-75
METHOD
SIZE t
DROP 1
TIME
SIZE f
DROP 4
TIME
SIZE t
DROP :
TIME 1
GUROBI
8.92
0.00%
4:00
11.62
0.00%
4:03
20.55
0.00%
4:10
RANDOM
5.21
41.6%
0:03
6.31
56.8%
0:06
8.63
58.0%
0:09
DMALA
8.81
1.23%
0:21
14.02
4.10%
0:22
19.54
4.91%
0:24
GFLOWNET
8.75
1.91%
0:02
13.93
4.72%
0:04
19.13
6.91 %
0:07
DNFS
8.28
7.17%
0:03
13.18
9.85%
0:06
18.12
11.8%
0:09
DNFS+DMALA
8.91
0.11%
0:10
14.31
2.12%
0:15
20.06
2.38%
0:22
4.3 Solving Combinatorial Optimisation Problems
Another application of DNFS is solving combinatorial optimisation problems. As an example, we describe how formulate the maximum independent set sampling problem.
Maximum Independent Set as Sampling. Given a graph G (V, E), the maximum independent set (MIS) problem aims to find the largest subset of non-adjacent vertices. It can be encoded 1 if vertex is included in the set, and
binary vector E {0, 1)llI, where
0 otherwise.
The objective maximise Σ F; subject to z,z, 0 for all (i.j) E E. This can be formulated as sampling from the following unnormalised distribution:
V Σ
p(x) x exp
T
Tq
Σ (UJE
(14)
where 0 is the temperature and À > 1 is a penalty parameter enforcing the independence constraint. As T - 0, p(x) uniformly concentrates on the maximum independent sets.
Therefore, we can train DNFS to sample from p(z), which will produce high-quality solutions the MIS problem. To enable generalisation across different graphs G. we condition the locally equivariant transformer on the graph structure. Specifically, we incorporate the Graphformer architecture (Ying et al., 2021), which adjusts attention weights based on the input graph. This allows the model to adapt to varying graph topologies. refer to this architecture as the locally equivariant Graphformer (leGF), with implementation details provided in Appendix D.4.
Settings. In this experiment, we apply our method to solve the Maximum Independent Set (MIS) problem, with other settings deferred to Appendix E.3. Specifically, we benchmark MIS on Erdōs-Rényi (ER) random graphs (Erdos, 1961), comprising 1,000 training and 100 testing instances, each with 16 to 75 vertices. Evaluation on the test set includes both performance and inference time. We report the average solution size and the approximation ratio with respect the pest-performing mixed-integer programming solver (GuROBI) (Gurobi, 2023), which serves as the oracle.
Experimental
Results & analysis. We compare our method against baselines: an annealed MCMC sampler (Sun et al., 2023b) using DMALA (Zhang et al., 2022b), and a neural sampler based on GFlowNet (Zhang al., 2023a). Additionally, without training, which serves as an estimate of the task's intrinsic difficulty. As shown in Table 1, DNFS after training substantially outperforms its untrained counterpart, highlighting the effectiveness of our approach. While the MCMC-based method achieves the strongest overall performance, it requires longer inference time. Compared to Get, another neural sampler, DNFS performs slightly worse. This may be attributed to the fact that GFlowNet restricts sampling to feasible solutions only along the trajectory, effectively reducing the exploration space and making the learning problem easier. Incorporating this inductive bias into DNFS is promising direction for future work. Nevertheless, a key advantage of our method is that the unnormalised marginal distribution ps. known, allowing us to integrate additional MCMC steps to refine the sampling trajectory. As shown the last row of Table 1, this enhancement leads to a substantial performance gain. Further analysis of this approach is provided in Table 7 in the appendix.
include results from randomly initialised version of DNFS

---

5 Related Work
CTMCs and Discrete Diffussin. Our work builds on the framework of continuous-time Markov chains (CTMCs), which were first introduced in generative modelling by Austin et al. (2021); Sun et al. (2023e); Campbell et al. (2022) under the context of continuous-time discrete diffusion models, where the rate matrix is learned from training data. This approach was later simplified and generalised to discrete-tine time masked diffusion (Shi et al., 2024; Sahoo et al., 2024; Ou et al., 2024), demonstrating strong performance across a wide range of applications, including language modelling (Lou et al., 2024; Zhang et al., 2025), molecular simulation (Campbell et al., 2024), and code generation (Gat et al., 2024; Gong et al., 2025). However, these methods require training data and are inapplicable when only an unnormalised target is given.
MCMC and Neural Samplers. Markov chain Monte Carlo (MCMC) (Metropolis et al., 1953) is the de facto approach to sampling from target distribution. In discrete spaces, Gibbs sampling (Casella George, 1992) is a widely adopted method. Building on this foundation, Zanella (2020) improve the standard Gibbs method by incorporating locally informed proposals to improve sampling efficiency. This method was extended to include gradient information to drastically reduce the computational complexity of flipping bits in several places. This idea was further extended by leveraging gradient information (Grathwohl et al., 2021; Sun et al., 2022a), significantly reducing the computational cost. Inspired by these developments, discrete analogues of Langevin dynamies have also been introduced to enable more effective sampling in high- nigh-dimensional discrete spaces (Zhang et al., 2022b; Sun et al., 2023a). Despite their theoretical appeal, MCMC methods often suffer from slow mixing and poor convergence in practice. To address these limitations, recent work has proposed neural samplers, including diffusion-based (Vargas et al., 2024; Chen et al., 2024; Richter & Berner, 2024) and flow-based (Máté & Fleuret, 2023; Tian et al., 2024; Chen et al., 2025) approaches. However, the majority of these methods are designed for continuous spaces, and there remains a notable gap the literature when it comes to sampling methods for discrete distributions. A few exceptions include Sanokowski et al. (2024, 2025), which are inspired by discrete diffusion models and primarily target combinatorial optimisation problems. A concurrent work, MDNS (Zhu et al., 2025), introduces a masked diffusion neural sampler grounded in stochastic optimal control theory (Berner et al., 2022). LEAPS (Holderrieth et al., 2025) and our method DNFS are more closely related to discrete flow models (Campbell et al., 2024: Gat et al., 2024), as both can be view as learning a CTMC to satisfy the Kolmogorov forward equation. While LEAPS parametrise O log Z. using a neural network, DNFS estimates it via coordinate descent.
Discrete EBMs and Neural Combinatorial Optimisation. Contrastive divergence is the de facto approach to train energy-based models, but it relies on sufficiently fast mixing of Markov chains, which typically cannot be achieved (Nijkamp et al., 2020). To address this, several sampling-free alternatives have been proposed, including energy discrepancy (Schröder al., 2023; Schröder et al., 2024), ratio matching (Lyu, 2012), and variational approaches (Lázaro- GreeillGree et al., 2021). More recently, Zhang et al. (2022a) replace MCMC with GFlowNet, a neural sampler that arguably offers improvement by reducing the risk of getting trapped in local modes. Our work follows this line of research by using DNFS as a neural alternative to MCMC for training energy-based models. Sampling methods are also widely used to solve combinatorial optimisation problems (COPs). Early work (Sun et al., 2022b) demonstrated the effectiveness of MCMC techniques for this purpose. More recent approaches (Zhang et al.. 2023a; Sanokowski et al., 2024, 2025) leverage neural samplers to learn amortised solvers for COPs. In this paper, we further show that DNFS is well-suited for combinatorial optimisation tasks, demonstrating its flexibility and broad applicability.
6 Conclusion and Limitation
We proposed discrete neural flow samplers (DNFS), a discrete sampler that learns a continuous- time Markov chain to satisfy the Kolmogorov forward equation. While our empirical studies demonstrate the effectiveness of DNFS across various applications, it also presents several limitations. A natural direction for future work is to extend DNFS beyond binary settings. However, this poses significant challenges due to the high computational cost of evaluating the ratios in Equation (7). As demonstrated in Appendix C.2, a naive Taylor approximation introduces bias into the objective, resulting in suboptimal solutions. Overcoming this limitation will require more advanced and principled approximation techniques. Additionally, we find that the current framework struggles to

---

scale to very high-dimensional distributions. This difficulty arises mainly from the summation over the ratios in Equation (7), which can lead exploding loss values. Designing methods stabilise this computation represents promising avenue for future research. Finally, extending DNFS to the masked diffusion setting offers another compelling direction, with the potential to support more flexible and efficient sampling over structured discrete spaces.
Broader impact. This paper aims advance machine learning research. While there may be potential societal impacts, none require specific mention at this time.
Acknowledgements
Zo is supported by the Lee Family Scholarship. We would like to thank Tobias Schröder and Sangwoong Yoon for their valuable discussions the early draft of this work. ZO also thanks Peter Holderrieth for generously sharing his insights on LEAPS during their discussion at ICLR 2025.
References
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and van den Berg, R. Structured denoising diffu- sion models in discrete state-spaces. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 17981-17993, 2021. URL attps://proceedeee www.eeeddinnn..uuis...... 0B565066446666869726666666666883 20-11677706677s--bbtaaaatt
Abstract html.
Barabási, A.-... and Albert, R. Emergence of scaling in random networks. science, 286(5439):509-512, 1999.
Bemer, J., Richter, L., and Ullrich, K. An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:2211.0133 2022.
Box, G. E., Jenkins, G. M., Reinsel, G. C., and Ljung, G. M. Time series analysis: forecasting and control. John Wiley & Sons, 2015.
Campbell, A., Benton, J., Bortoli, V. D.. Rainforth, T., Deligiannidis, G., and Doucet, A. A continuous time framework for discrete denoising models. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Confer- ence on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, Novem- ber 28- December 9, 2022, 2022. URL http:/ http://papers wwwp.c//paaaa..aeeeeee happppppr___a (illes/paper/2222 aaa522777aa3t5baaaffeeeaaaaaa Abstractt-bbbttttt---eee-------e Conffrenceeeeeeeeee oot html.
Campbell, A., Yim, J., Barzilay, R., Rainforth, T.. and Jaakkola. T. S. Generative flows on discrete state- spaces: Enabling multimodal flows with applications protein co-design. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL httphtttt//opereeeeeee wtwporiieeeeeeeeeteeeeeeee wwt.toetteoreeeeeet hett/foruumdddkuuuSrrrrrrr dQxQSS00E...
Casella, G. and George, E. I. Explaining the gibbs sampler. The American Statistician, 466 1992.
Chen, J., Richter, L., Berner, J., Blessing, D., Neumann, G., and Anandkumar, Sequential controlled langevin diffusions. ArXiv preprint, abs/2412.070811, 2024. URL https: //axxiv////xvvv.////a/// wttpi//.rrraab////2/2/
Chen, т. and Duvenaud, D. Neural networks with cheap differential operators. Advances in Neural Information Processing Systems, 32, 2019.
Chen, W. Ou, Z., and Li, Y. Neural flow samplers with shortcut models. ArXiv preprint, abs/2502.00337,, 2025. URL pestpa://arxee wwwaiiivoorr...............eaoo. org.otgg///.o.eoooooooo. aaaaaabb/.22000//
Cohen, T. and Welling, M. Group equivariant convolutional networks. In Balcan, M. and Weinberger, K. Q. (eds.). Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 2990-2999. J.... 2016. URL http: http:::////occcingg.....m mlr. press/ htt///////hheeee coh.ncc....6...../// html.
Dai, H., Singh, R., Dai, B., Sutton, C., and Schuurmans, D. Learning discrete energy-based models via auxiliary-variin variaale local exploration. Larochelle, H., Ranzato, M.. Hadsell, R., Balcan, M., and Lin, (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 12, 2020, viriual, 2020. URL htttps://ptt //prc/eedinc/ss.. neurips. wwtpcccppppcpceeeeccccccee.eeeet
Abrtt--bsttttttt html.
10

---

Del Moral, P. and Penev, S. Stochastic processes: From applications to theory. Chapman and Hall/CRC, 2017.
Erdos, P. On the evolution of random graphs. Bulletin of the Institute of International Statistics, 38:343-347, 1961.
Ferdinand, A. E. and Fisher, M. E. Bounded and inhomogeneous ising models. specific-heat anomaly of finite lattice. Physical Review, 185(2)832, 1969.
Fuchs, F. Worrall, D. E., Fischer, V., and Welling, M. transforanfoomm:::::: 3d roto-translation equivariant attention networks. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (cds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: httppp://ppoceeeiii pptppeeieceeeaaa. seneeeedinnn.euuripps..e www.nmuurcc.oo cc/ paper haper/2020/hash/ /ww.///ahhha//aae htt00/00haah//553a7/4ebaeae 35231a77c4ba7799411772222282
Abstracttttaat....o.om .html.
Gat, I., Remez, T., Shaul, N., Kreuk, F., Chen, R. T. Q., Synnaeve, G., Adi, Y., and Lipman. Discrete flow matching. In Globersons, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tom- czak, J. M., and Zhang. C. (eds.), Advances in Neurul Information Processing Systems 38: Annual Conference Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 15, 2024, 2024. URL http://papers nips cwtppper.pppppeppeeeeeeee htc//aper_fileee//eeeeeeeeeeeeee fel.ffillllppppe fiwes/paper///e/eeeeeeeeeaapaeet 2024/hashs 0119294734556442771bbbea
Batrbbbbddd-bbbracttttee -bsstract-Coffeeeeeee Conference html.
Geffner, T. and Domke, J. Using large ensembles of control variates for variational inference. In Bengio, S., Wallach, H. M., Larochelle, H.. Grauman, K., Cesa-Bianche N., and Garnett, R. (eds.). Advances in Neural Information Processing Systems 31: Annual Conference
Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 9982-9992, 2018. URLhttps: //tpr//edingg.......e wwstt//raaaaaattt..aaaaaaaaaa... Abstract htmml.
neuripe- cc/paper///eeeeeepepceceeee.ee / 2018/hashh hwsh/ dead35f 1551222866733110093277777222222
Gelman, A. and Meng, X.L. Simulating normalizing constants: From importance sampling to bridge sampling to path sampling. Statistical science, pp. 183-185,7 1998.
Gong. S., Zhang, R., Zheng, H., Gu, J.. Jaitly, N., Kong, L., and Zhang. Y. Diffucoder: Understanding and improving masked diffusion models for code generation. arXiv preprint 2rXhv:2506.2633, 2025.
Grathwohl, W., Swersky, K., Hashemi, M., Duvenaud, D., and Maddison, C. J. Oops I took A gradient: Scalable sampling for discrete distributions. Meila, M. and Zhang. T. (eds.). Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 3831 13811 PMLR, 2021. URL http::/ttocceeeeeee v139/ grathwoh12aa. html.
mlr press/
Gurobi, O. Llc gurobi optimizer reference manual. http:://ww..gumbiiccm// 2023.
Hastings, W. K. Monte carlo sampling methods using markov chains and their applications. 1970.
Hinton, G. F. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8): 1771-1800, 2002.
Holderrieth, P., Albergo, M. S., and Jaakkola, T. Leaps: A discrete neural sampler via locally equivariant networks. 2025. URL https::tpppp: www.etiisscooii....ooooo..oooo.e org/ (ttoorpuaID:2744888555
Kahn, H. Random sampling (monte carlo) techniques in neutron attenuation problems. i. Nucleonics (US) Ceased publication, 6(See also NSA 3-990), 1950.
Lartillot, N. and Philippe. H. bayesian mixture model for across-site heterogeneities in the amino-acid replacement process. Molecular biology and evolution, 21(6):1095-1109 2004.
Lázaro- -Geeila,,,,,l,, M., Dedieu, A., and George, based models. In Ranzato, M., Beygelzimer. A., Dauphin, Y. N., Liang, P. and Vaughan. J. W. (eds.). Advances in Neural Information Processing Systems 34: An- nual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 928-940, 2021. URL Abrrrrbbbtrtt..... html.
Perturb and and-max product: Sampling and leam-
ing in discrete energy-based
Lee, C. K., Joha, P., Frellsen, J., Lio, P., Albergo, M. S., and Vargas, E. Debiasing guidance for discrete diffusion with sequential monte carlo. ArXiv preprint, 0000799020607, 2025. URL https: // //txxiv.ivv///////////////aaaala org/ a/aabb//// 2502.. 06079.
Liu, J. S. and Liu, J. S. Monte Carlo strategies in scientific computing, volume 10. Springer, 2001.

---

Lou, A., Meng, C., and Ermon, S. Discrete diffusion modeling by estimating the ratios of the data distribution. In Forty-first International Conference on Machine Learning, ICML 2024. Vienna, Austria, July 21-27, 2024. OpenReview.nen 2024. URL httpttttt////ereeettee http://errrriieeeeeetooorraaaaaa www.://ee///nnneeeeeeeee (it.ffrrrrn..--0M.cRIIIII...
Lyu, S. Interpretation and generalization of score matching. arXiv preprint arXiv: /205.22220 2012.
Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Y. Trajectory balance: Improved credit as- signment in gflownets. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NewrIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http://papers
ccwpaper.ppp..ppce/ee/ee Staper_filesseeeee wwwes//pppr////222 hash/ Conferrnce .html.
Máté, B. and Fleuret, F. Learning interpolations between boltzmann densities. ArXiv preprint, abs/2301.07388. 2023. URL https: /arxiv/rxvvvvvv wwwaaiiii.....ioeoooooooe /bs/ 07388.
Metropolis, N., Rosenbluth, A. W., Rosenbluth, M. N., Teller, and Teller, Equation of state calculations by fast computing machines. The journal of chemical physics, 21 1887-1997 1953.
Mézard, M., Parisi, G., and Virasoro, A. Spin glass theory and beyond: An Introduction to the Replica Method and Its Applications, volume 9. World Scientific Publishing Company, 1987.
Murray, L., Ghahramani, Z., and MacKay, D. Meme for doubly distributions. arXiv preprint arXiv: 206.6848, 2012.
Neal, R. M. Annealed importance sampling. Statistics and computing, 1 11:125-139, 2001.
Neal, R. M. et al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Newman, M. and Barkema, G. T. Monte Carlo methods in statistical physics. Clarendon Press, 1999.
Nijkamp, E., Hill, M., Han, T., Zhu, S., and Wu, N. On anatomy of meme-hased maximum likelihood learning of energy- based models. In The Thirty-Fourr Fourth AAAI Conference on Artificial Intelligence. AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY. USA, February 7-12, 2020, pp. 5272-5280. AAAI Press, 2020. URL httph:/tp/////t view/5972
org/ jooo/////nddeee https/idee...pd./daaeeepp php/ AAAT/ AAAT/ArTiacaaeeeoco
Norris, J. R. Markov chains. Number 2. Cambridge university press, 1998.
Oksendal, B. Stochastic differential equations: an introduction with applications. Springer Science Business Media, 2013.
Ou, J., Nie, S., Xue, K., Zhu, F., Sun, J., Li, Z., and Li, C. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. ArXiv preprint, abs/2406.003766, 2024. URL https: //arxiv. org/abs//////oooooooooooooee abs//b//// /400. .3736.
Pani, C., Ou, Z., and Li, Y. Test-time alignment of discrete diffusion models with sequential monte carlo. arXiv preprint ar05v: 22524, 2025.
Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language understanding by generative pre-training. 2018.
Ranganath, R., Gerrish, S., and Blei, D. M. Black box variational inference. In Proceedings of the Seventeenth International Conference Artificial Intelligence and Statistics, AISTATS 2014, Reykjavik, Iceland, April 22-25, 2014, volume 33 of JMLR Workshop and Conference Proceedings, pp. 814 822, JMLR.org, 2014. URL htatt:://proceedinngeeoo mml. www.ppreeaavve.e h//v//33/33aaaaaaa/a/ hhml.
Richter, L. and Berner, J. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024, URL https: wwwpppeeeeeeeeeeeeeeoeeeeeeeeeee htt.:ffoormiiiddddiiii 1d=hh4
Sahoo, S. S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J. T., Rush, A., and Kuleshov, V. Simple and effective masked diffusion language models. In Globersons, A., Mackey, L., Belgrave, D., Fan,
A., Paquet, U., Tomezak. M., and Zhang, (eds.), Advances in Neural Information Processing Systems
38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024, URL http:://tapppp //ppp//pppeeeee/eea///eeeeeeeee/ nips. cc/paper__p. pupperffllee waleeffleesppppee papp//224/// hash/
Abstract- wwnetttrrrc---fffeeeeeeeeee coofference html.
12

---

Sanokowski, S., Hochreiter, S.. and Lehner, S. A diffusion model framework for unsupervised neural combinato- rial optimization. In Forty-first International Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27. 2024. OpenReview.net, 2024, URL https: /ttp///eppeeeeeeee///e/e/ee het//orruTideAAPEIIEE..
Sanokowski, S., Berghammer, W., Ennemoser, M., Wang, H. P., Hochreiter, S., and Lehner, S. Scalable discrete diffusion samplers: Combinatorial optimization and statistical physics. ArXiv preprint, tbs/2502.08690 2025. URL httt:///a org /et.b////222cce0 08696.
Schröder, T., Ou, Z., Lim, J., Li, Y., Vollmer, S., and Duncan, A. Energy discrepancies: a score-independent loss for energy-bass based models. Advances in Neural Information Processing Systems, 00:45300-453333 2023.
Schröder. T., Ou, Z., Li, Y., and Dunean, A. B. Energy-based modelling for discrete and mixed data via heat equations on structured spaces. In Globersons, A., Mackey, L.. Belgrave, D., Fan, A., Paquet, U.. Tomczak, J. M., and Zhang, C. (cds.), Advances in Neural Information Processing Systems 38: Annual Conference Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024. URL http://papers sesaaprssnipp.... wwwp.cc/pape...eeeeeeeeppppppppp cc/paper__.... 0112//48886662226661888002222220
www.-abtt-tttttereaaaeeaeeee-eee wwtpccceee.leeeeeeeeeeeee www.ss/pppe/e////0/// 22024/hash/
.ttml.
Shi, J., Zhou, Y. Hwang. J. Titsias, M. K. and Mackey, L. Gradient estimation with dis- crete stein operators. In Koyejo, S.. Mohamed, S.. Agarwal, A., Belgrave, D., Cho, K., and Oh, (eds.), Advances Neural Information Processing Systems 35: Anmual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL. attp::/pappee nips. wttpp.c/pppprr..../.// wwwes//pppee//eeee/ee 2222/haash/ Abstracctbttttttteeoaaeaoaaaatte Conference html.
Shi, J.. Han, K.. Wang. Z. Doucet, A.. and Titsias, K. Simplified and generalized masked dif- fusion for discrete data. In Globersons. A., Mackey, L., Belgrave, D., Fan, A., Paquet, U. Tom- czak, J. M., and Zhang. C. (eds.), Advances in Neural Information Processing Systems 38: Annuul Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC. Canada, December 10 15, 2024, 2024. URL attp://papez: wwp.. cwwppper...eee.pceeeeeeoo wttpppeeeeeepppeeeeeeeeeepeeeeee pttpe///pppee//2////aaaa Mad230998090019eed6666666 -at.ccoo0eeeoto---baaatattt Abstract -bstractt-Coffereneee html..
Sun, H., Dai, H., Xia, W., and Ramamurthy, A. Path auxiliary proposal for MCMC in discrete space. In The Tenth International Conference on Learning Representations, ICLR 2022, Viriual Event, April 25-29, 2022.
OpenReview.nen 2022a. URL https: / //ppn//vvie/eeeeeeee/meeeer http://paaeeeeeeeeeeeeeeorrriaa -i-.I.--D------.
Sun, H., Guha, E. K., and Dai, H. Annealed training for combinatorial optimization graphs. ArXiv preprint, 0is/2207.111444 2022b. URL aattps://arzie org/abs/ 2207 11542.
Sun, H., Dai, H., Dai, B., Zhou, H., and Schuurmans, D. Discrete langevin samplers via wasserstein gradient flow. In Ruiz, F. J. R., Dy, J. G., and van de Meent, J. (eds.), International Conference on Artificial Intelligence and Statistics, 25-27 April 2023, Palau de Congressos, Valencia, Spain, volume 206 of Proceedings of Machine Learning Research, pp. 6290-6313, PMLR, 2023a. URL httpss://ppoceeediinngg///// .ml..ngsm........mmm.mm.....m..m press/ v206/
v206/sun23f sun23f. html.
Sun, H., Goshvadi, K., Nova, A., Schuurmans, D., and Dai, H. Revisiting sampling for combinatorial op- timization. In Krause, A., Brunskill, E.. Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.). International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA. volume 202 Proceedings of Machine Learning Research, pp. 32859-32874. PMLR, 2023b. URL mlr.press ppees//202// v202 v202/sun23c st///cnn..cc...m html.
httpa: httpae//ppocceeeeee
Sun, H., Yu, L.. Dai, B.. Schuurmans, D., and Dai, H. Score-based continuous- time discrete diffusion models. In The Eleventh International Conference on Learning Representationss ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023c. URL /estps://openreree/ net/ wttppppwwww.......eoooooaa/aaa/a BT
Theis, L., van den Oord, A., and Bethge, M. A note the evaluation of generative models. In Bengio, Y. and LeCun, Y. (cds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http::tt///xiii orgg///ioooroeeeoo wrg./org///...............oo.... abs//511.. 01844.
Thiébaux, H. J. and Zwiers, F. W. The interpretation and estimation of effective sample size. Journal of Applied Meteorology and Climatoloyy, 2305)8800-111, 1984.
Tian, Y., Panda, N., and Lin, Y. T. Liouville flow importance sampler. In Forty-first International Conference on Machine Learning ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: // /ppenreviee /tt/forummt////// 10OOMKM221HH...
13

---

Tieleman, T. Training restricted boltzmann machines using approximations to the likelihood gradient. In Cohen, W. w. McCallum, A., and Roweis. S. T. (eds.), Machine Learning, Proceedings of the Twenty- Fifth International Conference (ICML 2008), Helsinki, Finland, June 5-9, 2008, volume 307 of ACM International Conference Proceeding Series, pp. 1064-1071. ACM, 2008. doi: 10. 10.1145/1390166.1 1390290.. URL https:// /top::://do...o/ org/10.////oo.o. 10. 1115/ 11155/399156... 1390290.
Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Transport meets variational inference: Controlled monte carlo diffusions. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.nee 2024. URL https: http////ppeeeeeeeeeeeeeeeeeeeeee
1rudnxiv..
Vaswani, A.. Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,. Gomez, A. N., Kaiser, L.,, and Polosukhin, I. Attention is all you need. In Guyon, I., von Luxburg. U.. Bengio. S. Wallach, H. M.. Fergus, R., Vishwanathan, S. V. N., and Gamett, R. (eds.). Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-6000 6008, 2017. URL httpe: / http///pooeeeeeee///e////e/ ppooeedding..neurip.... httpp//ccppppppe///ee//pp /21h/hashh Abstract html.
Sf5ee243547 (ee4443444dee9111 aareaaaseeeeeee1eaaaaaaaeeeeeeee www.:aaaa0aa3aaaeae5e Aaa.aaaaaaaaA---aaaatttaaa
Waggener, B. and Waggener, W. N. Pulse code modulation techniques. Springer Science & Business Media, 1995.
Wu, D., Wang, L.,, and Zhang, P. Solving statistical mechanics using variational autoregressive networks. Physical review letters, 12800),08062,, 2019.
Wu, H., Köhler, J., and Noć, F. Stochastic normalizing flows In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Sys- tems 33: Annual Conference on Neural Information Processing Systems 2020, NewrlPS 2020, De- cember 6-12, 2020, virtual, 2020. URL httpss//proceedeie www.a//aaaaaaaaaaa.tttea.a Abstract..ttml
neurips cc/paper 22020/hash//
Ying, C., Cai. T., Luo. S.. Zheng, S., Ke, G.. He, D., Shen, Y.. and Liu, T. Do transformers re- ally perform badly for graph representation? In Ranzato, M.. Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (cds.), Advances in Neural Information Processing Systems 34: An- nual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6- 6-14, 2021, virtual. pp. 28877-28888, 2021. URL https: http:://cccceece.....m seneeeddinnsn.uuris... wwwrips.cc/ppaee.ee cc/paper/ 2222/hashhh www.taa-.batttt................. html.
Zanella, G. Informed proposals for local meme in discrete spaces. Journal of the American Statistical Association, 115(530):852 865, 2020.
Zhang. D., Malkin, N., Liu, Z., Volokhova, A., Courville, A. C., and Bengio, Y. Generative flow networks for discrete probabilistic modeling. In Chaudhuri, K., Jegelka, S., Song. L.. Szepesvári, C., Niu, G., and Sabato, Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland,
S. (eds.). International Conference USA, volume 162 of Proceedingy of Machine Learning Research, pp. 16412-26428. PMIR, 2022a. URL httpa:/ hatpaa//ppooceeiiig
mlr. wwwpprreaaavaae wwr12//haagg2.....m catpnn22..tmmm.... html.
Zhang, D., Dai, H., Malkin, N., Courville, A. C., Bengio, Y., and Pan, L. Let the flows tell: Solv- ing graph combinatorial problems with gflownets. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10- 16. 2023, 2023a. URL http::/papers nips. cat.secccpppe httppppppeeepppppppeeae ptt.ee//pppr///23///// hhash/ 27571b7466c6500ueb6ccc www.:eefcc38819588 1837913388 Abstractt--atttttaeeeeaa--a--ete Conferennoeeeeeeeeeeeeeeeeeeeeee .ttll.
Zhang, P.. Yin, Li, C., and Xie, X. Formulating discrete probability flow through optimal transport. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA. USA, December 10- 16, 2023, 2023b. URL http: / /apprrs wwws.cccpppp...ctte wwtppper.ppppppeeeeeee paler_files/ paper/ poger//023//ae 22222//ahh//aae/ e6e (07776554412211866266... www.1/ffooob/odoooo Abssract--Coffreennce .htll.
Zhang, R., Liu, X., and Liu, Q. A langevin- like sampler for discrete distributions. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato, S. (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 26375-26396. PMLR, 2022b. URL https: //tppe/ed///ooooogg/.//. . mll., press/ /etee//vvvv//vvv v162/zhang22t html.
Zhang, R., Zhai, S., Zhang, Y.. Thornton, J., Ou, Z., Susskind, J., and Jaitly, N. Target concrete score matching: A holistic framework for discrete diffusion. arXiv preprint arXiv:2504. 16431, 2025.
Zhu, Y., Guo, W., Choi, J., Liu, G.-H., Chen, Y., and Tao, M. Mdns: Masked diffusson neural sampler via stochastic optimal control. arXiv preprint 2it::2508..08844 2025.
14

---

7 NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims are supported by the experimental results in Section 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in 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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long 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 discuss the limitation of our work in Section 6.
15

---

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 on 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 a 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 low or images are taken in low lighting. Or speech-to-tex 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
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]
Justification:
All theoretical results are correct with the proof given in the appendix.
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 provide short proof sketch to 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
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Our code will be released in the open-source git repository. In addition, we provide implementation details in the appendix to reproduce results.
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

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 to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to 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 good way accomplish this, but reproducibility 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, other means that are appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibilityy which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our code and data will be fully open-sourced. In the interim, we provide the code for sampling from Ising models in the supplementary materials whilst we prepare our code for the full release.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. ca/
Baailc/gguides/CoeSSubmi. wwtssSummissioooooooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is 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. wcc/public/ggicee
for more details.
The authors should provide instructions on 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. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
17

---

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:
Details of experimental settings are provided in Appendix E.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to 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
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: To reduce the impact of randomness, we repeated each experiment multiple times and report aggregated results.
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 run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to 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 or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma 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 not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables 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 exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
All experiments are conducted on a single Nvidia RTX A6000 GPU.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
18

---

The paper should provide the amount of compute required for each of the individual experimental runs as well 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 failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: :ttps:///neue /ttp///eeer..p.eeeee wwtppcccccceeeeeeeeeeee ctwiiccCueeliness
Answer: [Yes]
Justification: We have been strictly observing the relevant code of ethics, including but not limited to the NeurIPS one.
Guidelines:
The answer NA means that the authors have not 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 regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Our neural sampler does not pose any foresceable societal impact outlined in the guidelines, as the specified concerns are either irrelevant or not applicable to our work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 a 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 as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility 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: This does not apply to our case.
19

---

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 challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
Licenses for existing assets
Question: 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?
Answer: [Yes]
Justification: Relevent open-sourced code and models are properly cited in our paper. Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a 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 (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi wwwoehe com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- 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 creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This does not apply to our case.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/model 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 an 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 details about compensation (if any)?
20

---

Answer:
[NA]
Justification:
This does not apply to our case.
Guidelines:
The answer NA means that the paper does 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 Ethies, 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 equivalent for research with human subjects
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: This does not apply to our case.
Guidelines:
The answer NA means that the paper does 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 as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if 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.
Answer: [No]
Justification: The core method development in this research 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, non non non-standard components.
Please refer to our LLM policy (https //eurippp.....// for what should or should not be described.
/LLM)LLM)
21

---

Appendix for "Discrete Neural Flow Samplers with Locally Equivariant Transformer"
Contents
Variance Reduction and Control Variates
22
A.1 Proof of Equation (8)
22
A.2 Discrete Stein Control Variates
23
A.3 Connection LEAPS
24
B Derivation of Locally Equivariant Transformer
25
B.1 Proof of Proposition
25
B.2 Locally Equivariant Networks
26
B.3 Locally Equivariant TransFormer (leTF)
27
Details Training and Sampling of DNFS
28
Training and Sampling Algorithms
28
Efficient Ratio Computation
28
Applications to Discrete Neural Flow Samplers
29
D.1 Importance Sampling
30
D.2 Training Discrete EBMs with Importance Sampling
33
D.3 Combinatorial Optimisation as Sampling
33
D.4 Locally Equivariant GraphFormer (leGF)
34
Details Experimental Settings and Additional Results
35
E.1 Sampling from Unnormalised Distributions
35
E.2 Training Discrete EBMs
37
E.3 Solving Combinatorial Optimisation Problems
39
A Variance Reduction and Control Variates
In this section, we analyse the variance reduction phenomenon shown in Figure I through the lens of control variates. We begin with the proof of Equation (8) and the discussion of control variates, then establish its connection to LEAPS (Holderrieth et al., 2025).
A.1 Proof of Equation (8)
Recall from Equation (8) that the following equation holds:
Pr(y) Σ Re 8)
0 log Z, argmin Ep, (E:(:; Rτ) c.), (,;;,, & 0 log
pr(x)
We now provide detailed proof of this result, beginning with two supporting lemmas. Lemma (Discrete Stein Identity). Given a rate matrix R(y, that satisfies R(x, x))
Σ Cazt R(y, I),, we have the identity By(x) Ey R(x, u) ply) = 0.
p(m)
22

---

Proof. To prove the result, notice that
p(u) p(z) =
Ep(x) Σ
R(x, v)
ΣP(u) ((,,)) EΣ I y
p(u) Σ V
R(a, y) = 0,
1
which completes the proof. For a more comprchensive overview of the discrete Stein operator, see Shi et al. (2022).
Lemma 2. Let ci argming, Ep, (Ee(x;; Re) -a)?, then ot Ep, Re).
Proof. To see this, we can expand the objective
[(a)) argmin Ep. (Et (z; R₂) - aa)² oî 20, 51 Re) + cc= (ce Ep, R.))² d'
where the final expression is minimized when of = (x; Re). We are now ready prove Equation (8). Specifically:
& = argmin p. (E:(0; Re)) - ct)²
4
E.p Ce(x; Rs))
= EEp ú, (c)) Ep. Σ Re y) Pc(y)
(x(
M
Ex Og (x()) = 0g log Ze.
where the second and third equations follow Lemmas and respectively.
A.2 Discrete Stein Control Variates
To better understand the role of control variates (Geffner & Domke, 2018) in variance reduction as shown in Figure I, let us consider standard Monte Carlo estimation problem. Suppose our goal is to estimate the expectation μ = If where f(x is a function of interest. A basic estimator for
is the the Monte Carlo average j Σ Ek= [xx(k)), x;() s. Now, suppose we have access to another function ( known as a control variate, which has a known expected value 7 = E, [g()]] We can the use g(x) to construct new estimator: jı K Σ (f(z(1)) (z(xx(k)) + By. This new estimator já is unbiased for any choice of 3, since [l = = [[](())] (B([((x) + 87 = μ. The
benefit of this construction lies not in bias correction but in variance reduction. To see this, we can compute the variance of ãr:
1 (V[) [VU]) 2BCov(/) (o))) VIM K
(15)
This is quadratic function of 8, and since it is convex, minimum can be found by differentiating w.r.t. B and setting the derivative to zero. This yields the optimal coefficient 3 Cov (ʃ: ([]] Substituting it back into the variance expression (15) gives:
V vI Cor
(16)
This result shows key insight: the effectiveness of a control variate depends entirely on its correlation with the target function f. As long f and are correlated (positively negatively), the variance of j strictly less than that of f.. The stronger the correlation, the greater the reduction. In practice, the optimal coefficient 3 can be estimated from the same sample used to compute the Monte Carlo estimate, typically with minimal additional cost (Ranganath et al., 2014). However, the main challenge lies in selecting or designing suitable control variate g that both correlates well with and has tractable expectation under æ. For an in-depth treatment of this topic and practical considerations, see Geffner & Domke (2018).
Fortunately, Lemma 1 provides a principled way construct a control variate tailored our setting Ep, [f(x)) A Ep, [0, log (x)) R K Σ O log pe (x)()), where (kk) Pt- To reduce the variance of this estimator, we seek control variate g(æ) whose expectation under Pt is known. Inspired by
23

---

the discrete Stein identity, we define g(z) - Ey (z, v) Pr(2) which satisfies Ep, [x() = by Lemma 1. This makes g valid control variate with known mean. Using this construction, we can define a variance-reduced estimator as:
K
d, Po 8" (
12)
Pc(y) Pr (x))
(R) Pp
(17)
This estimator remains unbiased for any B, but with the optimal choice 3", can substantially reduce variance. Moreover, in the special case where the parameter 0 is optimal (in the sense that the objective in Equation (5) equals zero), an even stronger result emerges: the control variate becomes perfectly (negatively) correlated with the target function. That is, g(z) = f()) + c. where is constant independent of the sample z, leading Corr(f.g) a zero- vero-variance estimator, rare but highly desirable scenario. For more comprehensive discussion of discrete Stein-based control variates and their applications in variance reduction, we refer the reader to Shi et al. (2022).
In this idealised case, becomes
A.3 Connection to LEAPS
Equation (8) provides natural foundation for learning the rate matrix R. using coordinate ascent. This involves alternating between two optimisation steps:
i) Updating the rate matrix parameters 0 by minimising the squared deviation of (x Rt)) from a baseline ct. averaged over time and a chosen reference distribution 9e
i argmin L. (E Re())) c²)
ii) Updating the baseline o to match the expected value of &. under the true distribution pr (x)
Ce t- E.. (E, cx)²
Alternatively, instead of treating e as scalar baseline, we can directly parametrize as neural network & This allows us to jointly learn both 9 and @ by solving the following objective:
Pe(y) pc(x))
0*, Φ0 argmin (x(())
Or log og Ene
Re V)
(18)
0.0
which matches the Physics- Informed Neural Network (PINN) objective as derived in (Holderrieth al., 2025, Proposition 6. At optimality, this objective recovers two important conditions: The learned network of recovers the true derivative Ô log Ze: and ii) The rate matrix RÇ satisfies the Kolmogorov forward equation.
key insight is that even though Equation (8) formally holds when the expectation is taken under Pt- the training objective in Equation (18) remains valid for any reference distribution as long as it shares support with pt. This because, at optimality, the residual
Po(y) px(x))
oy (o) - of - -E R& (x.
0, Vr.
Integrating both sides with respect to Pt, and invoking the discrete Stein identity (Lemma 1), we find:
Ep T (o
4
Ep. ( log] = θ log Z,.
This also naturally admits coordinate ascent training procedure, where @ and & are updated in turn:
- - Σ (x, v) (c())
0 t argmin (m(().(()))
0 logg(()))
(c(x)
Rd y) (v) Pr (x)
c0 - (w(t))((())
0 logg
Σ
24

---

In this light, the term the term Ey Ry (2, y) pe(y) pr(
serves as a control variate for estimating 0 log Z
effectively reducing variance in the learning signal. Under optimal training, of accurately captures the log-partition derivative, confirming the correctness of the learned dynamics.
Since NDFS closely resembles LEAPS, summarise the main distinctions below to highlight the unique contributions of our work:
While LEAPS and DNFS yield similar objective functions, they are derived from different perspectives. DNFS derives the objective by learning the rate matrix to satisfy the Kol- mogrove equation, whereas LEAPS learns the rate matrix by minimising the importance weights. Perhaps surprisingly. these two perspectives lead similar objectives. However, the new perspective from the Kolmogorov equation offers a new insight for future research: leveraging more accurate estimators of Ô log Ze to further improve performance, which is not evident from the LEAPS framework.
The success of DNFS highly depends the proposed Locally Equivariant Transformer (leTF). Compared to the locally equivariant networks in LEAPS, leTT offers greater model capacity and improved adaptability, making more suitable for diverse modalities and complex input structures. We hope this architectural advancement will inspire future developments, which are essential for advancing both LEAPS and DNFS.
Unlike LEAPS, which is only evaluated on synthetic Ising and Potts models, DNFS is tested on broader applications, including sampling from Ising models, training EBMs, and solving combinatorial optimisation problems. We hope this wider empirical scope will inspire further research into additional applications of discrete neural samplers.
B Derivation of Locally Equivariant Transformer
B.1 Proof of Proposition
It is worth noting that Proposition subsequently utilized by Campbell et al. (2024) to construct the conditional rate matrix, as well as by Holderrieth et al. (2025) in the development of the locally equivariant network. For completeness, provide a detailed proof of Proposition I in this section.
first introduced in (Zhang al., 2023b, Proposition 5), and
We begin by formally defining the one-way rate matrix
Definition 1 (One-way Rate Matrix). rate matrix R is one-way if and only if R(y, 0 R(x, y) = 0. In other words, if a one-way rate matrix permits a transition from x to y. then the transition probability from y to x. must be zero.
We then restate Proposition and provide detailed proof follows.
Proposition 1. For a rate matrix R₂ that generates the probabilistic path ps, there exists a one-way rate matrix 2) Rz(y.z) R& (z, y) pe(v]
pr( i + and Qc(z,z) Evtz Q1,,, that generates the same probabilistic path pt, where [z]+ max(z, 0) denotes the ReLU operation.
Proof. We first prove that the one-way rate matrix Qe generates the same probabilistic path as Re:
Σ az( v)p(()))))
Σ Qi(z, v)pr(v) - Qe(y, (x)
Σ
- (R( - (x(x,))):(v)
Σ Re(s, u)p(()) ().x)
Σ
(a, v)p((1) = (x).
y
25

---

which completes the proof. We then show that Qi is one-way:
Q+(v.z) > = x) - Re(E, v) p+(U)
Pc(x) pc(y)
> 0 = Ro( - Ru 2)
<0
pz(x))
= +Q(zz v) [R( v) - Re(y,T) (((x)
= 0.
(19)
(z
Thus, Qt(y, = Qx y) = 0. which completes the proof.
B.2
Locally Equivariant Networks
Based on Proposition I, we can parametrise R2 one-way rate matrix, which is theoretically capable of achieving the optimum that minimizes Equation (5). Although the one-way rate matrix restricted subset of general rate matrices and thus offers limited flexibility, it enables efficient computation of the objective in Equation (5). To see this, we first formally define the local equivariant network, originally proposed in Holderrieth et al. (2025).
Definition 2 (Locally Equivariant Network). A neural network G is locally equivariant if and only if
Ge(T, f|.) = GG( Swapp( (z, i)) i= 1,...,
(20)
where Swap(x, i, = (Z I1-1, #1+1 Id) and T E S}.
We can then parametrise the one-way rate matrix RS using a locally equivariant network G₁:
(T, i|zr = (G((T,(z=))+ if r * 24 and Re x) = =Endz (y, This construction ensures that R, is a one-way rate matrix. To see this, consider state y (I1 (I1...., T, ..., E2). If
Rg (y. x) > 0, then
G.(T,ilz) > 0 -Ge //Sapp( i,)) > + (z, y) 0,
(21)
demonstrating the one-way property. With this parameterisation, the objective function in Equation (5) can be computed as
Pr(y) pt(x:)
8 R₂) 0, logge(z))
+
Σ
(,,, R2 (mi, ily)
(u. [x)) (T, aly)] Pe(y)
=0ggg +
Σ
(e())
pc(y)
= dg +
Σ (v,, - 1- (v. i1 (z)
where the final expression only requires single forward pass of the network G to compute the entire sum, significantly reducing computational cost. To construct a locally equivariant network, we first define its fundamental building block, the hollow network, as follows
T,..., x.a E X denote the input tokens with its i-th component set to r. A function H X' + Rdxh is called hollow network if satisfies the following condition
Definition 3 (Hollow Network). Let Tie-r = (I1,.. ........
H (zier)); = - (
for all 7, E,
where Md,: denotes the i-th of the matrix M. In other words, the output at position i is invariant to the imput at position i; that is, the i-th input does not influence the i-th outut.
Inspired by Holderrieth al. (2025), we then introduce the following proposition, which provides concrete method for instantiating a locally equivariant network.
Proposition 2 (Instantiation of Locally Equivariant Networks). Let x. E X denote the input tokens and H X dx be hollow network. Furthermore, for each token S, let w, g be learnable projection vector. Then, the locally equivariant network can be constructed as:
G (w, H (2)
26

---

Proof. We verify local equivariance by showing:
G(T, f/xx) = (wr wx.) (z));
= -ww, - w (Swap(zz, i, T))s:
-G /Sapswa i))),
where the second equality follows from the definition of the hollow network in Definition 3, which ensures that the i-th output invariant to the changes in the i-th input. This confirms that G satisfies the required local equivariance condition as Definition 2
Based on Proposition 2, we present two locally equivariant architectures introduced in Holderrieth al. (2025), followed by our proposed locally equivariant transformer.
Locally Equivariant MLP (leMLP) (Holderrieth et al., 2025). Let i E Rdxh be the embedded imput data. To construct a locally equivariant multilinear perceptron (MLP), we first define hollow MLP as (M) Σ oWW* + b²) where each W E Raxd is a weight matrix with zero diagonal entries (i.e., W for all i). by E R is the bias term, and o denotes an element-wise activation function. A locally equivariant MLP can then be defined as G(T, i/)) = Wa,)' "",,)"HMLEE (M)L(
(w, -w,,
Locally Equivariant Attention (leAttn) (Holderrieth et al., 2025). Let (z1...., Rd) E Rdx be the embedded input data. Similarly, we first define a hollow attention network as
exp(k(x.)) g(xx))
H A i:)
Σ E (x,,)((,)') g(())) v(x.),
where k, denote the query, key, and value functions, respectively. Thus a locally equivariant attention network can be defined as G(T,ilz)= = (wi - w., )a Atta
While these two architectures offer concrete approaches for constructing locally equivariant networks, their flexibility is limited, as they each consider only a single layer. More importantly, naively stacking multiple MLP attention layers violates local equivariance, undermining the desired property.
B.3 Locally Equivariant TransFormer (leTF)
In this section, we present the implementation de- tails of the proposed Locally Equivariant Transformer (leTF), an expressive network architecture designed to preserve local equivariance. As introduced in Sec- tion 3.3, 1cTF is formulated as
oulput
Attention ad Layer
Feed Fo forward
Add Norm
(r, a/0) = (w, - w2,)" (uTT(())..
(22)
Attention
where w denotes the learnable token embeddings pro- duced by the projection layer (illustrated Figure 2), and HHTF represents the Hollow Transformer mod- ule. In the following, we focus on the implementation details of the hollow transformer HUTF.
on Causal
Layer
Add & Norm
Add & Norm
Feed Forward
Feed Forward
L
X
As illustrated in Figure 8, the hollow transformer comprises L bidirectional causal attention layers fol- lowed by single attention readout layer. For clar- ity, we omit details of each causal attention layer, as they follow the standard Transformer architecture (Vaswani et al., 2017). We denote the outputs of the final left-to-right and right-to-left causal attention layers as
Add & No
Add & Norm
L2R Causal Attention
R2L Causal Attention
positional encoding
input
KL, VL CausalAttnnaa QR, KR. VR = CausalAttnRaI (input).
Figure 8: Illustration of the hollow trans- former.
'Note that Holderrieth et al. (2025) also introduces locally equivariant convolutional network: we refer interested readers to their work for further details.
27

---

Algorithm Training Procedure of DNFS
Input: initial rate matrix R, probability path Pt, time spans { outer-loop batch size M, inner-loop batch size N, replay buffer B
B 0
> Initialise replay buffer
2: while Outer-Loop do
{x(t (")))) M.K
I 1===,k=0 ~ T.0.1..
3:
> Generate training samples
4:
G Em O log Pi (m/)) ) Ey R² (x,
(m))
.3)
pe(u) pe(a, mm))
V E {tx)
0 log Ze
5:
B + B-B BU((UA,
T1x (mm)
} m==1kk=0
M.K
D update replay buffer
6:
while Inner-Loop do
7:
[] ~ U(B)
Uniformly sample from buffer
< 0- log Pe (zt (m)) Ey (x(") p(()) pe(a)"))
8:
9:
(( + N En (c(mm -2,
D. Compute training loss Perform gradient update
10:
+ apptiiize_step(( vo((O)
11: end while
end while
Output: trained rate matrix R
Because causal attention restricts each token to attend only to its preceding tokens, the resulting outputs Q, inherently satisfy the hollow constraint, In the readout layer, we first fuse the two query representations by computing Q = Q2 + QR. and then apply a masked multi-head attention to produce the final output
QI KnOMA]
softmax(
VVZ. VRl.
/2dx
where denotes the element-wise product, d, is the dimensionality of the key vectors, and Mr. and MR masks that enforce the hollow constraint by masking out future-token dependencies the left-to-right and right-to-left streams, respectively.
C Details of Training and Sampling of DNFS
C.1 Training and Sampling Algorithms
The training and sampling procedures are presented in Algorithms 1 and 2. For clarity, the rate matrix RC is parametrised using the proposed locally equivariant network, defined as R2 (y, & [G] (ys, i[z)]+, where y and æ only differ at the i-th coordinate. To initiate training, we discretise the time interval [0, into set of evenly spaced time spans {tx 1 and 2t, fx+1 t--1 for all valid indices k.
o satisfying 0 = <.. tk 1
In each outer loop of training, we generate trajectory samples by simulating the forward process under the current model parameters. Specifically, samples are drawn from the probability path
T defined by the initial distribution Pu and the current rate matrix R!" where fag denotes stop_gradiee tro_ggrddient((t This forward trajectory is simulated using the Buler-Maruyant detailed in Algorithm 2, and stored in replay buffer for reuse. During the inner training loop, draw mini-batches uniformly from the buffer and compute the training loss based on Equation (10). The model parameters then updated via gradient descent, as described in steps 6-11 of Algorithm 1.
method, as
C.2 Efficient Ratio Computation
Although the ratio
P((z) (E) can be computed in parallel, it remains computationally expensive
po(y)
in general. However, in certain settings, such as sampling from Ising models and solving combinatorial optimisation problems, the ratio can be evaluated efficiently due to the specific form of the underlying
distribution. In these cases, the unnormalized distribution takes quadratic form:
p(x) p(x) exp(x" Wa + hTx), E (0. 1]d, W E ddxd h E Rdd
(23)
28

---

Algorithm 2 Sampling Procedure of DNFS
Input: trained rate matrix Ry, initial density pu, # steps
1: Po, At + R t+0
Initialisation
2: for k O...., K - - do
3:
I++i: + Cat( (1a441= + ()
Euler-Maruyama update
4: ++ At
5: end for
Output: generated samples 21
where the neighbourhood (x) is defined as the 1-Hamming Ball around æ, i.e., all vectors differing from in exactly one bit. Let be such a neighbor obtained by flipping the i-th bit: Yi -2; and Uj 2y for j # i. The log of the unnormalised probability can be decomposed as:
log (())
Σ
Fa Fa Tb
Σ
War
Ti Σ Fa Wai 20 WwwW
Σ Matia-
(24)
byf's
agis
Thus, the change in bog-probability when flipping bit i is:
p(y) logp(x) (1 2x,))
Σ Wa
Σ DaW + Wie + h4
(25)
light
This expression can be vectorised to efficiently compute the log-ratios for all neighbours:
p(v)
log
(z) @ ((W + W jz diag (W) + h),
(26)
p(x)
(EA((z)
where @ denotes element-wise multiplication. Furthermore, in special cases, such as combinatorial optimisation where W is symmetric with zero diagonal, the log-ratio simplifies to
p(y) p(z) vEN(2)
log
= (1 2x) o (2Wx + h).
(27)
In more general settings, where is categorical variable and the energy is non-quadratic, there is no flosed-form solution for efficiently calculating the likelihood ratio. However, it can be approximated through first-order Taylor expansion (Grathwohl et al., 2021). Specifically,
log (z log (z)) = (y xx) V log)
(28)
which gives the following approximation:
Pt (1)
log
logp( - x? [V, log
(29)
p=(x))
where [log p=(w) (m)
and we take the fact that y and differ only in one position. This approximation requires computing V log (x) just once to estimate the ratio for the entire neighbourhood, thus improving computational efficiency. However, introduces bias into the training objective in Equation (5), potentially leading to suboptimal solutions.
log Pe
1.j
Pe(a)
To illustrate this, train DNFS to sample from pre- pre-trained deep EBMs minimising the objective
(x))) exact computation in parallel; and ii) an approximation via the Taylor expansion in Equation (29). As shown Figure 9, using the approximate ratio leads to inaccurate samples that tend lie in overly smoothed regions of the energy landscape. This degradation in sample quality is likely caused by the bias introduced by the approximation. Consequently, in deep EBM settings where the energy function non-quadratic, opt to compute the exact ratio in parallel, leaving the development of more efficient and unbiased approximations for future work.
in Equation (5) using two methods for computing the ratio p=(v).
D Applications to Discrete Neural Flow Samplers
In this section, we introduce details of two applications discrete neural flow samplers: training discrete energy-bar based models and solving combinatorial optimisation problems.
29

---

Figure 9: Illustration of the effect ratio computation in Equation (5). Top row: results using the Taylor expansion approximation; bottom row: results using exact computation. The approximated ratio yields less accurate samples, likely due to the bias introduced in the training objective.
D.1 Importance Sampling
We begin by reviewing the fundamentals of importance sampling, followed by a proof based on the Radon-Nikodym derivative. This approach was also introduced in Lee al. (2025); Holderrieth et al. (2025); Pani et al. (2025). For completeness, we include a brief recap to make the paper more self-contained.
- where p(x) 2 is the un- normalised distribution and Z Ex p(a:) denotes the normalising constant, which is typically intractable. For test function ((r) of interest, estimating its expectation under n through direct sampling can be challenging. Importance Sampling (IS) (Kahn, 1950) addresses this by introducing a proposal distribution (xx) that is easier to sample from. The expectation under π can then be re-expressed as
Importance Sampling. Consider a target distribution <(x) = p(x)
[o(x)] 1
p(x)
E By(x))
p(a)
o (((r)
Z (x))
g(x)) ((x))(0)
6(x))
(30)
Ey(x)
p(z)
g(()
This leads to the Monte Carlo estimator:
(A)
[x(=)(((x)]
Eド= w() (zx(())
x(1) g(x),
(31)
where (+)
q(2(ary denotes the importance weight. Although this estimator is consistent as K + oc, it often suffers from high variance and low effective sample size (Thiébaux & Zwiers, 1984), especially when the proposal q is poorly matched the target æ. In theory, the variance of the estimator is minimized when (Ix) x p(x))(x)), yielding a zero- variance estimator. While this condition is rarely attainable in practice, provides useful guideline: a well-designed proposal should closely approximate the target distribution, i.e., q(x) T(x).
(4))
CTMIT-Inppred Importance Sampling. As noted previously, an ideal proposal should should closely approximate the target, i.e., ((x) R (2)) This motivates the use of continuous-time Markov chains (CTMCs) to construct the proposal. Specifically. let Re(y) x)) denote rate matrix that defines forward CTMC with initial distribution Po oc 17, generating a probability path denoted by R. To complement this, we define a backward CTMC with initial distribution pl x p and interpolated marginals Pt X P: == -pn1²' The backward process is governed by the rate matrix (1, R2 (3x, y) Pc(y) leading the backward path distribution D...R This construction yields (x)
the following importance sampling identity:
Ean (ani(zoSeS)) 6(21)
[=(=)(((z) E dogigin 5- nt [d(x1)] = E 2ogigi Tn.
(32)
(10S(1)))
which can be approximated via Monte Carlo sampling:
K
(k(k)
(1 (k)
[x(x)
~
k= Σ w(i) (x((t)),
(8) P0gig1 ~ T.R. (x())
(33)
n. (2.1) (1(2)))))
30

---

This estimator is consistent for any choice of Rg, and it becomes zero variance when RG satisfies the Kolmogorov equation in Equation (4), condition that can be approximately enforced by minimizing T
the loss in Equation (5). Before delving into the computation of the importance weights introduce two key lemmas that underpin the derivation.
n.
Lemma 3 (Radon-Nikodym Derivative (Del Moral & Penev, 2017)). Let po and pt two initial distributions; Rg and R2 be two rate matrices, which induce the forward and backward CTMCs Q R.o.R and Qu.. over the time interval [0, l, respectively, Then,
= .. (x.: - (2 ds I3)
R2 Rg R2 I
log
Σ log
(34)
R. R.
8,82 #
which induces that
P...R. pe(x;) exp (Jo R²,, (s))
R2,15, sts) Re(Tes 25
II
(35)
Po.Re (20) exp (So R2 R.( Rx) ds)
ds
a,n tx,
Proof. For a comprehensive proof, we refer readers to (Camphell al., 2024, Appendix C.I), which follows the exposition in Del Moral & Penev (2017).
Lemma 4 (Fundamental Theorem of Calculus). Let f 10,T] + R be piecewise differentiable function on the interval T). Suppose that f is differentiable except at a finite set of discontinuity points { C [o, TJ, where the left-hand limit f(s,) and the right-hand limit f( at each 8; exist but are not necessarily equal. Then, the total change of f can be expressed as
f(t)-((0) = f000=(00 d())
E (($)))
(36)
for all t E [O, where f'(s) denotes the derivative of f at points where f is differentiable.
Now, it is ready compute the importance weight. Specifically,
log
logg
Pt
+ +log
Zo +
R(( (x),
Σ log
R,(
2 R..
Pu
Z,
R. (I 23
Zo
O log ds+ ( (2( (ast))--a(a,,
R' (IT Nt )D
= llog
+
x. ds + Σ
log R R(Fs, Ne (x(e
Z
8.24 $2.
Zo
/ o log ds + I (xs (1
log
Zz
Zo
o log do + [0 Σ
Dolu) (o(xr)
log
+
Ro(w y)
Rs ds
Z
vsx.
7o
+ {
o, L. Σ Ra y) Py(y)
log
ds,
Zi
ps (zi.)
where the first equation follows from Lemma 3, and the second from Lemma 4, leveraging the fact that ++ log şr piecewise differentiable. The final equality holds by the detailed balance condition satisfied by the backward rate matrix, namely (2 (v. (x) = R2 (x, v)px(y). Importantly, although the partition function Z, is generally intractable, cancels out practice through the use of the self-normalised importance sampling estimator in Equation (32).
31

---

Free Energy and Internal Energy Estimation. To estimate the log- peg-parttiion function log Zt we consider the following lower bound
Ze Zo.
minni Ze
log Z₁
log Pop
log
sogage~ vo...
a Ro Zo
Ze
log
+ +log
Zoc.se Tp
Po.Ri
Z
Io O2 10g1,,
Po(y) (x(
-
Σ Ro y)
ds
(37)
I0 Ro
where we assume that Zo = 1. Considering the test function & = log Pt. one can also estimate the negative entropy, which is related the internal energy:
D. Bro.R. (or)
K ~
exp(w!!))
Ep, [log p.] =E
log log
(38)
ΣØ= exp(wi))
where (K) "DSuSe"
n. and wi (K) (R) (ROSSg)) So' 0g log (x.
Ey Ro (M)
pe(y) ds. (m(m)))
Effective Sample Size (ESS). The Effective Sample Size (ESS) (Liu & Liu, 2001, Chapter 2) quantifies how many independent and equally weighted samples
importance- weighted Monte Carlo samples is effectively worth. It reflects both the quality and diversity of the weights: when most weights are small and few dominate, the ESS is low, indicating that only small subset of samples contributes meaningfully to the estimate. Formally, consider the importance sampling estimator:
exp(w(()) exp(w(l)
(xx(k)), w(()) log ((x((k) q(x(k)))
((k) ((x),
(39)
K Lタ=1
where q the proposal distribution. The normalised ESS is given by:
1 E (x))
ESS
(40)
K E (X))
K E
where (M)
ER ())( denotes the normalised importance weight. In CTMC-based importance sampling, the estimator takes the form:
(x)
exp(wx()) (*))
Ep, [p(=)]
j= exp(w.!)) (1)
(t)
ToSase
(k) R..
(41)
with log-weight wi (k) ESS is computed as
So 0. log (og Ey Ra(x (k)
w)
poly) ds. Thus, the corresponding (1
exp(ui (w (K)
ESS
w
(42)
K KEK k= (w) EX (k))))
ΣØ exp(wi!) ()))
(r, (|ا) be learned rate matrix, Ti-1, 1, T, I4+1.-.., T.)- This matrix parameterizes forward that progressively transforms noise into data via Ruler-Maruyama discretization:
Log-LLikelihood Estimation. Let Re (v, A (T 1/x) = (T,
where y = (#1, .........--..
TH Cat (1m. + Re (A))). xo Po
(43)
The corresponding reverse-time rate matrix is given by
pr(y) ( ily)]+
p((y)
pe(y) px(x)
R ()) 12,( 2)
(v,, i|x]]+
(44)
(x
pe(x)
32

---

which enables simulating the reverse CTMC that maps data back into noise:
Cat 111-A + (T,, (Ix, Te- Pe Att)
Tt
Pr(
Dt)
31 ~
Pi-
(45)
This reverse process enables estimation a variational lower bound (ELBO) on the data log- likelihood:
log Epanea [x)] log E
3 1o po(21))
Pogici
J n
E
rppi Rogtein Taples
log
log logo)(
log
+ log (o (,)
(K) TUsisi ~ (x1) (fl))
K
n can be evaluated using Lemma 3
where the log-ratio log
T T l=1
0
Rd (Tar 25
log
Re (I,, - (r,, v.) ds
Σ
log
[
R (1o Fis)
8,2a $o,
Pa(y))
Ps (2(sss)
Σ
2g, (x v)
ds
log (o(x5))
(2(x
v
8,8x f..
(1(1)
log (x(xx)
Σ
(( ( G-G i|1,s)]
ds + Σ
(x(x))
Po (x
8,0, #x.
D.2 Training Discrete EBMs with Importance Sampling
To train a discrete EBM Ps (x) x exp(-F.gh (x))))), we employ contrastive divergence, which estimates the gradient of the log-likelihood as
Tgll (w) [log = (x) V [x)]]] - E (2) [ (z)],
(46)
where the second term can easily approximated using the training data with Monte Carlo estimation. To estimate the intractable expectation over Pa. MCMC method is typically used. However, for computational efficiency, only a limited number MCMC steps are performed, resulting in a biased maximum likelihood estimator and suboptimal energy function estimates (Nijkamp et al., 2020), To address this issue, we replace MCMC with the proposed discrete neural flow samplers. Specifically, train rate matrix RV sample from the target EBM pa- The expectation over p&can then be estimated using CTMT-inspired importance sampling, as described in Appendix D.I:
K
exp(w(k))
( (oLE = E an BP..
VoEs(T)
2
V (z
j=1 exp(w(i))
where (*) = Ir (xy; R$) dt. To summarise, we jointly train the EBM Ps and the DNFS by alternating the following steps until convergence:
I) Updating the rate matrix parameters 0 using the training procedure described in Algorithm 1;
2) Updating EBM P& via contrastive divergence, as defined in Equation (46).
D.3 Combinatorial Optimisation as Sampling
Consider general combinatorial optimisation problem of the form min f(xx) subject to <(x) = This problem can be reformulated as sampling from an unnormalised distribution p(z) exp(-E(x)/T)), where the energy function is defined as E(x) f(x) + Ac(x) (Sun et al.,
33

---

2023b). As the temperature T - oc, p(z) approaches the uniform distribution over X, while as T 0, p(z) concentrates on the optimal solutions, becoming uniform over the set of minimisers. In this paper, we focus on two combinatorial optimisation problems: Maximum Independent Set (MIS) and Maximum Cut (MaxCut). Below, define the energy functions used for each task following Sun et al. (2022b). Given graph G (V. we denote its adjacency matrix by A, which is symmetric and zero-diagonal binary matrix.
MIS. The MIS problem can be formulated as the following constrained optimisation task:
min (0,1)
Zis s.t. 2,,, 0, V(i.)) E.
(47)
We define the corresponding energy function in quadratic form as:
7 Ax
lo) B(x) = -1
(48)
Thus, the eeg--rbbability ratio between neighboring configurations E (x)), differing from æ by single bit flip, has a closed-form expression
P(U) p(x) vEN(2)
log
(1- 2x) (x AAz).
(49)
Following Sun et al. (2022b), we set À = 1.0001. After inference, we apply post-processing step to ensure feasibility: we iterate over each node F;, and if any neighbour Ty = 1 for (æ,, E, set æjj O. This guarantees that the resulting configuration is valid independent set.
Maxcut. The Maxcut problem can be formulated
min x
Σ Auj
(50)
(-1,))
ijER
We define the corresponding energy function
logg(x) x F(x) = 7"' Ar.
(51)
This leads to the following closed-form expression for the log-ratio
p(U)
1
Ax Since any binary assignment yields a valid cut, no post-processing is required for MaxCut.
log p(x) vEN(x)
(52)
2
D.4 Locally Equivariant GraphFormer (leGF)
To train an amortized version of DNFS for solving combina- torial optimization problems, it is essential to condition the model on the underlying graph structure. To this end, we inte- grate Graphormer (Ying et al., 2021) 1) into our proposed Locally Equivariant Transformer (leTF), resulting in the Locally Equiv- ariant Graphformer (leGF).
MatMul
SoftMax
The leGF architecture largely follows the structure of leTF, with the key difference being the computation of attention weights, which are modified incorporate saph--seciffic structural bi- ases. Following Ying et al. (2021), given graph G (V, E). we define (v,, U;) as the shortest-path distance between nodes and vj if a path exists; otherwise, we assign it special value (e.g.. 1). Each possible output of & is associated with a learnable scalar byy)). which serves structural bias term in the self- self-attention mechanism. Let Ar, denote the (i, jith
Scale
MatMul
V
Figure 10: Illustration of the graph- aware attention mechanism. The figure is adapted from (Ying et al., 2021, Figure 1)
34

---

GelowNet
GelowNet
SS
S
Energy
Energy
Energy
SC
Oracle
DNFS
Oraeie
DNFS
Oraele
DNFS
S
S
Gelow
LEAPS
GelowNet
LEAPS
GFlowNet
LEAPS
Energy
Energy
Energy
DNES
Ornele
DNFS
Oraele
DNFS
Figure 11: Additional visualisation of sampling from pre-trained EBMs.
element of the Query-Key interaction matrix. The attention weights are then computed as:
Q + bb(os)n)))
A.j softmax(
vdl
where (,() is shared across all attention layers. An illustration of this graph-aware attention mechanism is shown in Figure 10. For further details, refer the reader Ying et al. (2021).
E Details of Experimental Settings and Additional Results
In this section, we present the detailed experimental settings and additional results. All experiments conducted on a single Nvidia RTX A6000 GPU.
E.1
Sampling from Unnormalised Distributions
E.1.1
Experimental Details
Sampling from Pre-trained EBMs. In this experiment, we adopt energy discrepancy to train an EBM, implemented as 4 layer MLP with 256 hidden units and Swish activation. Once trained, the pretrained EBM serves as the target unnormalized distribution, with the initial distribution Po set uniform. probability path is then constructed using linear schedule with 128 time steps. To parameterise the rate matrix, we employ the proposed locally equivariant transformer, which the causal attention block consists of multi-head attention layers, each with 4 heads and 128 hidden units. The model is trained using the AdamW optimizer with learning rate of 0.0001 and a batch size of 128 for 1,000 epochs (100 steps per epoch). To prevent numerical instability from exploding loss, the log-ratio term log (x]
Pc()) is clipped to a maximum value of
Sampling from Ising Models. We follow the experimental setup described in (Grathwohl et al., 2021, Section F.1). The energy function of the Ising model is given by log p(x) о ax Jæ + 6Fx. In Figure 5, we set a = 0.1 and b = 0. The probability path is constructed using a linear schedule with 64 time steps, starting from uniform initial distribution. The le model comprises bidirectional causal attention layers, each with 4 heads and 128 hidden units. Training is performed using the AdamW optimizer with learning rate of 0.001, batch size of 128, and for 1,000 epochs (100 steps per epoch). To prevent numerical instability, the log-ratio term is clipped to a maximum value of
Experimental Setup for Figures and 3. In this experiment, we set probability path is defined via a linear schedule over 64 time steps, starting from a uniform initial distribution. The leTF model is composed of 3 bidirectional causal attention layers, each with heads
0.1 and b 0.2. The
'https: hattps://github comet/hhb.oomm-_....o..o han/J-zin/// http:/disseeeeeeeeeeeeee energy www.crcpaan//trre//eesttteeeeeee
35

---

and 64 hidden units. Training is conducted using the AdamW optimiser with a learning rate of 0.001, a batch size of 128, and for 500 epochs (100 steps per epoch). To mitigate numerical instability, the log-ratio term is clipped to maximum value of 5.
E.1.2 Additional Results
Additional Results of Sampling from Pre-trained EBMs. We present additional results comparing DNFS to baseline methods on sampling from pre-trained EBMs in Figure 11. The results demonstrate that DNFS produces samples that closely resemble those from the oracle distribution.
LEAPS with leCony and leTE. While leConv performs comparably to DNFS on Ising models, its convolutional architecture may struggle to adapt to non-grid data struc- tures. In this experiment, we use the LEAPS algorithm (Holderrieth et al., 2025) to train a neural sampler with two different locally equivariant architectures: leConv (Hold- errieth et al., 2025) and leTF (ours). Given the data has 32 dimensions, grid to make it compatible with leConv. As shown in Figure 12, LEAPS with leConv fails to achieve meaning- ful performance, whereas the proposed transformer-biue based architecture leT' performs comparably to DNFS. This result highlights the limitations of leCony, whose expres- siveness is constrained non-grid data, while leTF offers greater flexibility and gencralisation across diverse input structures.
LEAPSS 12TF LEAPS + IeCom
pad 36 and reshape it into 6 6
Comparison of leConv and lcTF on sampling from pre-trained EBMs with LEAPS.
Figure 12:
Comparing to MCMC Methods. Compared to MCMC
methods, key advantage of neural samplers is their ability to guarantee convergence when trained optimally, whereas MCMC methods often suffer from slow mixing and poor convergence. To highlight this benefit, we adopt the same setting as in Figure using DNFS with 64 sampling steps. For a fair comparison, we evaluate MCMC baselines with the same number of steps, including Gibbs sampling (Casella & George. 1992), Gradient with Gibbs (GwG) (Grathwohl et al., 2021), and the Discrete Metropolis- adjusted Langevin Algorithm (DMALA) (Zhang et al., 2022b). We present energy histograms based on 5,000 samples in Figure 13, with long-run Gibbs sampler serv- ing as the oracle. The results show that short-run MCMC methods struggle to produce accurate samples, although gradient-based variants (GwG and DMALA) outperform conventional Gibbs sampling. In contrast, the energy distribution produced by DNFS closely matches the oracle, demonstrating the effectiveness of the proposed neural sampler.
0.20
Oracle
DNFS
0.15
Gibbs
GwG
DMALA
0.05
0.00
-20
-15
-10
-5
Energy
Figure 13: Histogram of sample energy
for different sampling methods.
Training Loss Comparison. We train our model by minimising the loss function defined in Equa-
leTF (ours)
DNFS
7.5
leConv
LEAPs
tion (5), where smaller loss
8
5.0
leAttn leMLP
value reflects better fit to the pa-
2.5
8
0.0
log
rameterised rate matrix. As such,
1
the loss value serves as a proxy for training quality and model convergence. In Figure 14, we compare the training loss across different locally equivariant net- work architectures. The results show that the proposed locally equivariant transformer achieves the lowest loss value, indicating its superior capacity to fit the target
-7.5
10K
20K
J0K
40K
50K
Z0K
AUK
60K
100K
Training Steps
Training Steps
Figure 14: Training loss of dif- ferent locally equivariant net- works (ref: Figure 3).
Figure 15: Training loss parison between DNFS and LEAPS (ref: Figure 5).
httppe://gttabb..c////bbbgg////
36

---

Table 2: Comparison of the estimated free energy F/D, internel energy E/D, entropy S/D, and effective sample size for Ising models on = 10 x 10 grids at different temperatures.
Method
ESS
Free Energy J/D
Internal Energy E/D
Entropy S/D
Optimal Value
3.6727
-0. 00.282
0.6489
LEAPS
0.9956 0.0001
-3.6709 + 4.9182e
-0.4262 0.0034
0.6489 0.0006
DNFS
0.9985 + 4.6912e-5
-3.6709 + 6.2463e
-0.4271 + 0.0010
0.6488 0.0002
Optimal Value LEAPS
1
-2.1242
-1.4763
0,2853
0.22305
0.3631
184
-2.1011 + 0.0002
- -1.4493 + 0.0130
0.2873 +: 0.0057
DNFS
0.9685 0.0010
-2.1120 1 6.6258e
-1.4743 1 0.0068
0.2811 + 0.0030
Table Experiment results of probability mass estimation on seven synthetic datasets. We display the negative log-likelihood (NLL) and MMD (in units of 1 x 10 10_
Metric
Method
2spirals
8gaussians
circles
moons
pinwheel
swissroll
checkerboard
PCD
20.094
19.991
20.565
19.763
19.593
20.172
21.214
ALOE+
20.062
19.984
20.570
19.743
19.576
20.170
21.142
ED-Bern
20.039
19.992
20.601
19.710
19.568
20.084
20.679
NLL
EB-GFN
20.050
19.982
20.546
19.732
19.554
20.146
20.696
EB-DNFS
20.118
19.990
20.517
19.789
19.566
20.145
20.682
PCD
2.160
0.954
0.188
0.962
0.505
1.382
2.831
+LOE+
0.149
0.078
0.636
0.516
1.746
0.718
12.138
ED-Bern
0.120
0.014
0.137
0.088
0.046
0.045
1.541
MMD
EB-GFN
0.583
0,531
0.305
0.121
0.492
0.274
1.206
EB-DNFS
0.603
0.070
0.527
0.223
0.524
0.388
0.716
rate matrix. Furthermore, Figure 15 compares the loss values between DNFS and LEAPS, showing that our approach again outperforms the bascline. This demonstrates the combined effectiveness of our coordinate descent learning algorithm and mensffrmer-bbase architecture.
Quantatitive Results. Recall Equation (11), where the Ising model is defined as
p(œ) x exp(-E(x)) exp(or" ADæ), x {-1,11
(53)
where R is the temperature and Ap denote the adjacency matrix of the lattice graph. We evaluate our method by comparing the estimated free energy F= 1 and entropy S 20r [ F) with their theoretically optimal values derived in Ferdinand & Fisher (1969). The free energy and internal energy are estimated using Equations (37) and (38), respectively. We compare DNFS with LEAPS under the latter corresponds to the critical temperature and thus presents more challenging sampling problem. The result is reported in Table 2, where use 2, 048 Monte Carlo samples to estimate the values and report the mean and standard deviation averaged over 10 independent runs. The results show that DNFS provides accurate estimates of the free and internal energies, closely matching the theoretical value with low variance. In contrast, LEAPS exhibits significantly lower ESS and larger deviations in other metrics under the critical temperature setting, which may indicate insufficient mode coverage. These comparisons highlight the robustness of DNFS and suggest it is less prone to mode collapse across different temperatures.
2g log Z, internal energy & E, [x)),
temperature settings, 0.1 and 0.22305, where
E.2
Training Discrete EBMs
E.2.1 Experimental Details
This experiment follows the setup of Dai et al. (2020). We first sample 2D data points ftl E IR? from a continuous distribution p. and then quantise each point into a 32-dimensional binary vector {0. 32 using Gray code. Formally, the resulting discrete distribution follows p(z) x p[GradyTTollon ToFloat( Grady GradyToFloat(22ppp5
Probability Mass Estimation.
Following Schröder et al. (2024), we parameterise the energy function using 4-layer MLP with 256 hidden units and Swish activation. The leTF model consists of 3 bidirectional causal attention layers, each with 4 heads and 128 hidden units. The probability path is defined by linear schedule over 64
37

---

Table 4: Experimental Results of EBM and DNFS on probability mass estimation: Rows labelled
'EBM' represent metrics evaluated using the trained EBM model, while rows labelled "DNFS' represent metrics evaluated using the trained DNFS.
Metric
Method
2spirals
8gaussians
circles
moons
pinwheel
swissroll
checkerboard
EBM DNFS
20.118
19.990
20.517
19.789
19.566
20.145
20.682
NLL
20.947
20.948
21.043
20.908
21.011
20.899
21.106
EBM DNFS
2.553
1.429
0.897
2.808
1.733
0.731
6.168
MMD
0.603
0.070
0.527
0.223
0.524
0.388
0.716
Data
Energy
Samples
Data
Energy
Samples
Data
Energy
Samples
Data
Energy
Samples
Data
Energy
Samples
Data
Energy
Samples
Figure 16:
Additional qualitative results in training discrete EBMs. We visualise the training data, learned energy landscape, and the synthesised samples of DNFS.
time steps, starting from a uniform initial distribution. Both the EBM and DNFS are trained using the AdamW optimiser with learning rate of 0.0001 and batch size of 128. Notably, each update step the EBM is performed after every 10 update steps DNFS. To ensure numerical stability, the log-ratio term is clipped maximum value of 5.
After training, we quantitatively evaluate all methods using negative log- log-likelihood (NLL) and maximum mean discrepancy (MMD), as reported in Table 3. Specifically, the NI.I is computed using the trained EBM on 4,000 samples drawn from the data distribution, with the normalization constant estimated via importance sampling using 1, 1,000,000 samples from variational Bernoulli distribution with p = 0.5. For the MMD metric, we follow the protocol in Zhang al. (2022a), employing an exponential Hamming kernel with bandwidth of 0.1. All reported results are averaged over 10 independent runs, where each run uses 4,000 samples generated by the trained DNFS.
Training Ising Models. Following Grathwohl et al. (2021); Zhang et al. (2022b), we train learnable adjacency matrix Jg to approximate the true matrix in the Ising model. To construct the dataset, we generate 2,000 samples using Gibbs sampling with 1,000,000 steps per instance. The leTF model consists of three bidirectional causal attention layers, each with four heads and 128 hidden units. The probability path is defined by linear schedule over 64 time steps, starting from uniform distribution. Jø is optimised using the AdamW optimiser with learning rate of 0.0001 and batch size of 128. To promote sparsity, we follow Zhang et al. (2022a) and apply 4 regularisation with a coefficient of 0.05. DNFS is trained separately using AdamW with a learning rate 0.001 and the same batch size. To ensure numerical stability, the log-ratio term is clipped at maximum value of 5. We train both models iteratively, performing one update step for Jø for every ten update steps of DNFS.
E.2.2 Additional Results
Additional Results of Probability Mass Estimation. We compare DNFS to various baselines. including PCD (Ticleman, 2008), ALOE+ (Dai et al., 2020), ED-Bemn (Schröder et al., 2024), and EB-GFN (Zhang et al., 2022a). As shown in Table 3, DNFS outperforms both Persistent Contrastive Divergence (PCD), which relies on conventional MCMC methods, and the variational approach ALOE+, demonstrating the effectiveness of our method. Additional qualitative results are provided in Figure 16, where DNFS consistently produces accurate energy landscapes and high-quality samples that closely resemble the training data. Furthermore, Figure 17 illustrates the sampling trajectory of DNFS alongside marginal samples from long-run Gibbs sampling. The comparison shows that DNFS produces samples that closely resemble those from Gibbs, demonstrating ability to approximate the target distribution with high fidelity.
38

---

Notably, since the EBM and the sampler are trained jointly, we can evaluate the negative log- likelikoehood (NLL) of the data using the trained CTMC, as described in Appendix D.I, and assess sample quality via the maximum mean discrepancy (MMD) using samples generated by the trained sampler. In Table 4, we report results under four evaluation settings: i) NILL (EBM): using the trained EBM with importance sampling; ii) NLL (DNFS): using the trained CTMC following the method in Appendix D.I; iii) MMD (EBM): samples drawn via Gibbs sampling from the EBM; and iv) MMD (DNFS): samples generated by DNFS. It can be seen that samples generated by DNFS achieved lower MMD compared to those generated by Gibbs sampling, demonstrating the superiority of the learned sampler in capturing the target distribution and producing higher-fidelitg samples. However, we observe that DNFS yields less accurate likelihood estimates compared to importance sampling performed with the trained energy function. This is likely because DNFS is trained to satisfy the Kolmogorov equation rather than explicitly optimising the evidence lower bound, as is common in other discrete diffusion models for generative modelling (Shi al., 2024; Sahoo al., 2024). Moreover, it is noteworthy that the sample quality and likelihood are not necessarily consistent (Theis al., 2016, 6, Section 3.2), and DNFS does not directly optimise the likelihood. Thus, the performance of DNFS on NLL is not guaranteed.
Additional Results of Training Ising Models. We further provide a quan- titative comparison against baselines for training Ising models. Following Zhang et al. (2022a); Schröder et al. (2024), we evaluate on D = 10 x 10 grids with
Table 5: Mean negative log-RMSE (higher is better) between the learned connectivity matrix Js and the true matrix J for different values of D and o.
D
10a
D
g?
Method \ a
0.2, .0. 0.5 and D = 9 x 9 grids with a = -0.1, -.2.. Performance measured by the negative log-RMSE between the estimated Jo and the true adja- cency matrix J. As shown in Table 5, while our method underperforms ED- GFN, which is also neural sampler, it achieves results comparable to Gibbs, GwG, and ED-Bern in most settings, demonstrating its ability to uncover the underlying structure in the data.
0.1
0.2
0.3
0.4
0.5
-0.1
0.2
Gibbs
4.8
4.7
3.4
2.6
2.3
4.8
4.7
GwG
4.8
4.7
3.4
2.6
2.3
4.8
4.7
ED-Bem
5.1
4.0
2.9
2.6
2.3
5.1
4.3
EB-GFN
6.1
5.1
3.3
2.6
2.3
5.7
5.1
DNFS
4.6
3.9
3.1
2.6
2.3
4.6
3.9
E.3
Solving Combinatorial Optimisation Problems
E.3.1
Experimental Details
This experiment follows the setup in (Zhang et al., 2023a), where we train an amortised combinatorial solver using DNFS on 1,000 training graphs and evaluate it by reporting the average solution size over 100 test graphs. To be specific, we use Erdös-Rényi (ER) (Erdos, 1961) and Barabási-Albert (BA) (Barabási & Albert, 1999) random graphs to benchmark the MIS and MCut problems, respectively. Due to scalability limitations of our current method, we restrict our evaluation to small graphs with 16 to 75 vertices, leaving the exploration of more complex graphs for future work.
We parameterise the rate matrix using the proposed locally equivariant GraphFormer (leGF), which consists of 5 bidirectional causal attention layers, each with 4 heads and 256 hidden units. Training is performed using the AdamW W optimiser with a learning rate of 0.0001 and a batch size of 256. The log-ratio term is clipped to a maximum value of 5 to ensure stability. More importantly, we find that the temperature T plays a crucial role in performance. Fixed temperature values generally lead to suboptimal results. Therefore, we adopt a temperature annealing strategy: starting from an inverse temperature of 0. and gradually increasing it to a final value of A more comprehensive annealing strategy, such adaptive schedules based loss plateaus, may further improve performance by better aligning the sampling dynamics with the learning process. Exploring such adaptive annealing schemes is a promising direction for future work.
E.3.2 Additional Results
Additional Results on Maximum Cut. We further evaluate DNFS on the maximum cut problem. As shown in Table 6, the trained DNFS significantly outperforms its untrained version, underscoring the effectiveness of our approach. While DNFS slightly lags behind the baselines DMALA and
39

---

Table 6: Maximum eut experimental results. We report the absolute performance, approximation ratio (relative to GUROBI), and inference time.
BA16-20
BA32-40
BA64-75
METHOD
SIZF
DROP 4 TIME 4
SIZR
DROP
TIME 4
SIZE f
DROP 4 TIME. +
GUROBI
40.85
0.00%
0:02
93.67
0.00%
0:05
194.08
0.00%
0:14
RANDOM
25.70
37.1%
0:03
46.19
50.7%
0:05
81.19
58 58.2%
0:08
DMALA
40.32
1.30%
0:04
93.47
0.21%
0:06
192.33
0.90%
0:07
GELOWNET
39.93
2.25%
0:02
90.65
3.22%
0:04
186.60
3.85%
0:07
DNFS
39.60
3.06%
0:03
88.64
5.37%
0:05
181.75
6.35%
0:08
DMALA
40.76
0.22%
0:08
93.63
0.01%
0:12
192.30
0.92%
0:17
Table 8: Results for the maximum independent set problem on RB32-40 graphs with varying parameter Reported values include the solution size (larger is better) and the percentage drop in performance relative to GUROBI (lower is better), indicated in brackets.
0.1
0.3
0.5
0.7
0.9
1.0
GUROBI 8.52(0.00%)
8.24(0.00%)
7.06(0.00%)
7.89(0.00%)
(.05(0
(.63(0
Random
3.38(60.3%)
4.59(4..2%)
4.25(39.8%)
6.29(21.8%) 5.54(29.7%)
6.93(19.6%)
DMALA
8.50(0.23%)
8.16(0.97%
(.01().70%)
8.01(0.12%) (.8.6)
8.63())00%)
GflowNet
8.20(3.75%)
7.93(3.76%)
6.83(3.25%)
7.78(..39%)
8.03(0.62%)
8. 8.600%))0%)
DNFS
8.00(6. 10%)
7.65(7.16%)
7.97(0.99%) 7.66(2.91%) (.67()
8. 8.60()
DNFS+DMALA
(..1().
7.02().56%% 8.19().60%)
(.85().50%)
8.04(0.12%)
8. 8.63()
GFlowNet, its MCMC-refined variant achieves the best overall performance, closely approaching the oracle solution provided by Gurobi.
MCMC Refined DNSF. As previ- ously discussed, a key advantage of DNFS is its ability to incorporate addi- tional MCMC steps to refine the sam- pling trajectory, thanks to the known marginal distribution Pt- To vali- date the effectiveness of this MCMC- refined sampling, we conduct an ex- periment the MIS problem us- ing the ER16-20 dataset. this ex- periment, compare two methods:
Table Comparison between DNFS and its DLAMA-
refined version by solving MIS on the ER16-20 dataset.
DLAMA
DNFS+DLAMA
STEPS
SIZE f
TIME (S)
SIZE f
TIME (S) $
8.33
1.75
8.37
4.53
2
8.53
3.19
8.63
5.89
3
8.71
4.70
8.75
7.62
4
8.77
6.12
8.84
8,85
5
8.80
7.74
8.91
10.47
DMALA (Zhang et al., 2022b) and
DNFS combined with DLAMA, both sampling from the same interpolated distribution pt OC Pb p1i- For reference, the average solution size obtained by DNFS without DLAMA refinement is 8.28. As shown Table 7, applying DLAMA refinement significantly boosts performance, with improvements increasing more refinement steps are added. More importantly, integrating the proposed neural sam- pler (DNFS) with the MCMC method (i.e., DNFS + DLAMA) outperforms the standalone MCMC baseline (i.e., DMALA), demonstrating that the learned sampler provides strong initialisation that guides the refinement process toward better solutions. This result confirms the synergy between neural samplers and MCMC refinement in solving challenging combinatorial problems.
Benchmarking on the RB Graphs. Following Sanokowski al. (2024), we further evaluate our method on the Maximum Independent Set problem using the RB32-40 graphs, varying the parameter, which controls the problem's 's difficulty. SSpecfiically, higher p values of yield easier instances, while lower values result in more challenging graphs. For each setting, we report the solution size, along with the percentage performance drop relative to GUROBI. As shown in Table 8, the results are consistent with the observation in Table 1. On easier instances (i.e., higher values of p), our method performs competitively, approaching the performance of the oracle solver GUROBI. On more challenging instances, however, DNFS exhibits a performance gap relative to GelowNet. This gap arises because GFlowNet restricts its sampling trajectories to the feasible solution, effectively narrowing the exploration space. Introducing such inductive bias into DNFS represents promising direction for future work. Nevertheless, despite DNFS underperforming GFlowNet in its pure form, it offers distinct advantage: the intermediate target distribution is explicitly known. This property enables integration with MCMC-based refinement methods (e.g., DNFS+DMALA), which significantly improves performance.
40

---

Noise to Data (t: : 0 - 1)
S
5
S
S
$
S
S
5
5
S
S
5
S
S
S
5
S
S
S
S
S
S
Data to Noise (L: (1:1+0)
S
S
S
S
S
9
S
S
5
S
9
S
S
S
Marginal Samples with Gibbs Sampling (t (t:00 0 + 1)
S
S
5
S6
9
S
S
S
S
S
S
S
S
S
SS
S
S
S
S
S
S
S
Figure 17: The sampling trajectory of DNFS in discrete EBM training. Top: noise to data trajectory;
Middle: data to noise trajectory; Bottom: marginal samples with Gibbs sampling.
41


Reviews:
[{"review_id": "E42QSBy1Hp", "reviewer": "Reviewer_ameA", "summary": "The paper introduces a new algorithm called Discrete Neural Flow Samplers (DNFS) to perform sampling from unnormalized discrete distributions. DNFS works by learning the rate matrix of a continuous-time Markov chain (CTMC).", "questions": "1. why does DNFS suffer less from mode collapse?\n\n2. Is there any convergence guarantee for the coordinate descent used in the paper? how is it compared with joint optimization as in LEAPS\n\n3. The method uses a specific annealing path to connect the initial and target distributions. How sensitive is the performance of DNFS to different interpolation schemes? is it related to question 1\n\n4. can you comment on why it encounters issues in scaling to higher dimensions\n\n5. in prop 1, does the simplification to a one-way rate matrix impose any practical limitations on the learning dynamics", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1. It proposes a new, trainable framework for sampling from discrete distributions without access to training data. It learns a rate matrix for a CTMC by optimizing an objective based on the Kolmogorov forward equation. The algorithm is efficient using coordinate descent.\n2. the introduction of Locally Equivalent Transformer is novel and interesting and worth further studying.\n3. The paper demonstrates the effectiveness of DNFS across a broad range of applications, including training discrete energy-based models and solving combinatorial optimization problems like the Maximum Independent Set\n\nWeaknesses:\n\n1. It only works in binary settings.\n2. Scaling issue for high-dimensional distributions\n3. performs less favorably in combinatorial optimization tasks; cannot incorporating constraints", "was_revised": true, "final_justification": "I think the authors addressed most of my concerns and provided discussions for the rest. I raised my evaluation by 1 point.", "rebuttal": "Thanks for your constructive comments.\n\n> weakness: only works in binary settings;scaling issue for high-dimensional distributions; performs less favorably in CO tasks\n> \n\nWe thank the reviewer for the thoughtful summary of the limitations of our method, which we have also acknowledged in the conclusion section. Regarding the weaknesses pointed out, we would like to provide the following clarifications to motivate future research:\n\n- Our method supports both binary and categorical distributions, although it may still face challenges due to the computational cost and instability of ratio evaluations in high dimensions or with large categorical vocabularies. Addressing these issues could make DNFS more scalable. Additionally, while our current work focuses on uniform discrete diffusion, extending it to masked discrete diffusion is a promising direction. This extension could enable the incorporation of inductive constraints, an important feature for solving combinatorial optimisation problems.\n\nDespite these limitations, we believe our work still offers valuable insights. We would like to reiterate our main contributions:\n\n1. A new discrete neural sampler motivated by the Kolmogrove equation:\n    \n    We propose a method for learning a discrete neural sampler by training a rate matrix that satisfies the Kolmogorov equation, where we estimate the intractable term $\\partial_{t} \\log Z_{t}$ using Monte Carlo with control variates. This approach offers a distinct perspective from LEAPS, which focuses on minimising the variance of importance weights. Our formulation not only provides a complementary viewpoint but also highlights a promising future direction: leveraging more accurate estimators of $\\partial_t \\log Z_t$ to further enhance performance, an insight that is not evident from the LEAPS framework.\n    \n2. A more expressive locally equivariant network\n    \n    We propose locally equivariant transformer (leTF), a novel architecture that enables highly efficient objective computation while preserving model expressiveness. Compared to the locally equivariant networks used in LEAPS, leTF offers greater model capacity and improved adaptability, making it more suitable for diverse modalities and complex input structures. We hope this design also inspires further research into more effective locally equivariant architectures.\n    \n3. Broader empirical scope and generality\n    \n    Unlike LEAPS, which is only evaluated on sampling from Ising models, our method is tested across a broader range of applications, including sampling from Ising models and pretrained EBMs, training deep EBMs, and solving combinatorial optimisation problems. These experiments highlight the versatility and generality of our approach across diverse domains. We believe this broader empirical scope establishes a clear foundation for future research toward more effective, scalable, and adaptable discrete neural samplers for real-world applications.\n    \n\n> Why does DNFS suffer less from mode collapse?\n> \n\nWe believe you are referring to Fig.4, where GFlowNet exhibits mode collapse on the checkerboard dataset, whereas DNFS does not. We would like to clarify this difference as follows:\n\n1. Why can DNFS cover all modes\n    \n    DNFS learns a geometric interpolation between the base and target distributions. In principle, the base distribution can be arbitrary, which provides the flexibility to incorporate useful inductive biases into the neural samplers. A more structured base makes it easier for the model to learn the interpolation. In our experiments, we use a simple uniform prior, which, despite its simplicity, covers all modes at the initialisation. This property helps DNFS avoid mode collapse.\n    \n2. Why does GFlowNet suffer from mode collapse\n    \n    In contrast to DNFS, GFlowNet begins from a point mass (i.e., mask tokens) and learns a stochastic policy to generate the target samples. This setup can lead the model to converge prematurely on narrow, high-reward regions, resulting in reduced diversity and mode collapse [1]. To mitigate this, additional techniques, such as loss-guided GFlowNets [1], are often required to encourage more thorough exploration of the sample space.\n    \n\n[1] Loss-guided auxiliary agents for overcoming mode collapse in gflownets\n\n> Is there any convergence guarantee for the coordinate descent? how is it compared with joint optimization as in LEAPS\n> \n\nYes, it has convergence guarantees. As detailed in Appendix A.3, our alternating optimisation scheme consists of two steps. The 2nd step involves minimising a convex and smooth function, which ensures convergence to a global minimum. The 1st step also has convergence guarantees, supported by the universal approximation theorem for neural networks. Therefore, alternating between these two steps reliably leads to convergence to a stationary point of the joint objective function.\n\nTo compare with LEAPS, we first highlight the key difference: LEAPS estimates $\\partial_t \\log Z_t$ by learning a neural network, whereas DNFS computes it by solving Eq. 8, which is convex and admits a closed-form solution. We argue that our method is preferable, as it directly leverages the closed-form optimal solution, avoiding the approximation errors that may arise during neural network training. Moreover, as outlined in lines 691–694, when the model is optimally trained, DNFS is guaranteed to recover the true value of $\\partial_t \\log Z_t$. In contrast, LEAPS may incur additional errors due to its decoupled modeling strategy. To validate this empirically, we estimate the free energy $\\frac{1}{2*\\sigma} \\log Z_t$ for the Ising model defined in Eq. 11 under two settings: $\\sigma = 0.1$ and $\\sigma = 0.22305$, the latter corresponding to the critical temperature. The results, reported on a scale of $10^4$, are presented in the table below.\n\n|  | $\\sigma=0.1$ | $\\sigma=0.22035$ |\n| --- | --- | --- |\n| optimal value | -3.67265 | -2.12417 |\n| DNFS | -3.67091 | -2.11201 |\n| LEAPS | -3.67084 | -2.10111 |\n\nThe results show that DNFS provides a more accurate approximation of $Z_t$. This also highlights a promising direction for future research: improving the estimation of $\\partial_t \\log Z_t$ during training can potentially lead to better performance. This insight is not immediately apparent from the LEAPS perspective, which derives its objective by minimising the variance of importance weights. In contrast, it emerges naturally in DNFS, which learns a rate matrix to satisfy the Kolmogorov equation, requiring the approximation of the intractable term $\\partial_t \\log Z_t$. We believe this new insight motivates future work to explore more accurate estimators, such as integrating SMC or applying more advanced variance reduction techniques beyond control variates.\n\n> How sensitive is the performance of DNFS to different interpolation schemes?\n> \n\nThe choice of interpolation scheme is indeed important. We adopt the commonly used geometric interpolation of the form $p_t \\propto \\rho^{\\alpha_t} \\eta^{1-\\alpha_t}$, where $\\alpha_t$ denotes the noise schedule. Specifically, we use a linear schedule with $\\alpha_t = 1 - t$ throughout the paper. However, alternative schedules, such as those used in diffusion generative models, can also be applied. Below, we report the performance of DNFS when sampling from Ising models under different interpolation schemes:\n\n|  | $\\alpha_t=$  | ESS | Free energy |\n| --- | --- | --- | --- |\n| optimal value | - | 1 | -3.67265 |\n| linear | $1-t$  | 0.9985 | -3.67091 |\n| Polynomial ($w=0.5$) | $1-t^w$  | 0.9969 | -3.46192 |\n| Polynomial ($w=2$) | $1-t^w$  | 0.9992 | -3.61088 |\n| cosine | $1-\\cos(\\frac{\\pi}{2}(1-t))$  | 0.9993 | -3.67156 |\n| geometric ($\\beta_{min}=10^{-5},\\beta_{max}=20$) | $\\exp(-\\beta_{min}^{1-t} \\beta_{max}^t)$  | 0.9978 | -3.67255 |\n\nThe results show that all interpretation schemes perform well, confirming the robustness of our method. However, we note that this observation is task-dependent. We hypothesise that for more complex tasks, the choice of interpolation scheme will play a more critical role in ensuring stability and achieving optimal performance.\n\n> can you comment on why it encounters issues in scaling to higher dimensions\n> \n\nThe difficulty is twofold:\n\n1. High computational cost\n    \n    When the ratio $\\frac{p(y)}{p(x)}$ lacks a closed form (e.g., Eq27), it must be computed in parallel over all $y$ (see Eq5). This is prohibitive in high-dimensional spaces or with large categorical vocabularies, especially when $p$ involves neural models like deep EBMs or reward networks.\n    \n2. Instability of ratio values\n    \n    Early in training, samples $x \\sim q_t$ (from $R_\\theta$) may lie in low-density regions of $p_t$, making $\\frac{p(y)}{p(x)}$ unstable or divergent. This is exacerbated in high-dimensional settings with rugged, sparse energy landscapes.\n    \n\nDespite these challenges, extending DNFS to large categorical vocabularies and high-dimensional problems remains a promising research direction.\n\n> in prop 1, does the one-way rate matrix impose any practical limitations on the learning dynamics\n> \n\nYes, although prop1 shows that every rate matrix has an equivalent one-way rate matrix that induces the same probability path, the restricted expressiveness of the one-way rate matrix makes it more difficult to learn compared to an unconstrained rate matrix.\n\nHowever, in the context of sampling, there is a trade-off between expressiveness and computational efficiency. The vanilla (fully unconstrained) parameterisation is impractical due to its prohibitive computational cost. As a result, we must rely on the more efficient one-way rate matrix. To enhance its expressiveness, we introduce the locally equivariant transformer. As shown in Figures 3 and 14, adopting a more expressive architecture leads to improved performance. We hope this finding offers valuable insights for future research on designing more powerful locally equivariant networks."}, {"review_id": "NqlLTZB6lu", "reviewer": "Reviewer_i34y", "summary": "The paper proposes a method for learning the rate matrix of a CTMC to sample from discrete distributions with large or complex state spaces, useful for Monte Carlo inference. To guide learning, the authors define a probability path using an annealing schedule and formulate a loss based on the Kolmogorov forward equation of the CTMC. To reduce computational complexity, the loss is restricted to transitions between states that differ in at most one dimension. The method employs a variance-reducing approximation equivalent to control variates to estimate this loss efficiently. Additionally, they parameterise the rate matrix using a local equivariant network built with a hollow transformer architecture, which they prove it replicates the same probability paths and empirically show it leads to better approximations at lower computational cost. The approach is evaluated through empirical comparisons on discrete energy-based models and combinatorial problems, demonstrating competitive performance relative to existing methods.", "questions": "What is the initialisation strategy for the rate matrix across experiments?\n\nHow sensitive are the results to the initialisation of the rate matrix?\n\nPlease elaborate on the treatment of the reference distribution q_t and its empirical impact.\n\nHow are MCMC refinement steps (e.g., DMALA) integrated into the sampling procedure?\nIt might be useful to release code and add implementation details for the combinatorial optimisation experiment.", "limitations": "Yes, except for the initialisation strategy mentioned before.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n    Technically sound, well-grounded in prior work, and theoretically justified.\n\n    Empirically demonstrates improvements over strong baselines using equal architecture.\n\n    Novel composition of ideas for a new, effective method.\n\nWeaknesses:\n\n    Insufficient discussion of the role, initialisation, and impact of the reference probability path q_t, which may significantly affect performance.\n\n    Lack of implementation details for certain experiments (e.g., combinatorial examples), reduces reproducibility.\n\n    Potential sensitivity to initialisation of the rate matrix or probability path is not addressed.\n\nQuality:\nThe authors build on a set of well-established ideas to produce a novel algorithm for sampling from discrete distributions. Each of the components is rooted in prior literature and the proposed Algorithm 1 is empirically and theoretically justified. The experimental section compares fairly to relevant baselines using a shared architecture. However, while the overall methodology is rigorous, certain choices critical to the algorithm's success (particularly the handling and role of the reference distribution q_t) are under-explained and raise questions regarding robustness and reproducibility.\n\nClarity:\nThe paper is well written and structured, with clear explanations of the algorithmic design. While the authors state that any q_t with support matching p_t would suffice in theory, the experimental results suggest it may have a nontrivial empirical effect, especially in the combinatorial optimisation example where using DMALA improves performance. My understanding is that the learned probability path is updated in the coordinate descent approach and used to approximate the expectation in (8), but its initialisation strategy lacks discussion or supporting code for some experiments. Clarifying these aspects, especially the empirical sensitivity of results to different initialisations would significantly improve both clarity and replicability.\n\nMinor Corrections:\nL33-34: CMCT -> CMTC\nL152: x -> x_{i \\leftarrow \\tau}\nL264: Arguably\n\nSignificance:\nThe paper contributes a new method that improves empirical performance compared to baselines, particularly under equal architectural constraints. This has practical implications and it may inspire follow-up work that modifies components of the proposed algorithm. The empirical results are promising and demonstrate potential for broad applicability. However, the significance of the contribution could be strengthened by a more thorough investigation of the algorithm’s sensitivity to the initialisation of the probability path.\n\nOriginality:\nWhile each individual methodological component is drawn from prior work, the paper presents an original and thoughtful combination of these ideas into a novel algorithm. The objective can be viewed as a novel variant of the PINN loss, which is shown empirically to outperform comparable baselines when using the same architecture.", "was_revised": true, "final_justification": "Despite the concerns on clarity and significance mentioned above, the core methodological contribution is solid, and the empirical results are promising. My current evaluation stands as is.", "rebuttal": "Thanks a lot for your valuable comments. Your suggestions are very helpful in further improving the work.\n\n> While the overall methodology is rigorous, certain choices critical to the algorithm's success (particularly the handling and role of the reference distribution q_t) are under-explained and raise questions regarding robustness and reproducibility. Please elaborate on the treatment of the reference distribution q_t and its empirical impact.\n> \n\nThank you for raising this concern. While in theory the model is guaranteed to converge for any reference distribution $q_t$ that shares the same support as $p_t$, the choice of $q_t$ can significantly affect training stability in practice.\n\nTo illustrate this, consider the case where $q_t$ is poorly aligned with $p_t$. In such situations, it is likely to sample points $x$ that lie far from the high-density regions (modes) of $p_t$. At these points, $p_t(x)$ can be extremely small, causing the ratio $\\frac{p(y)}{p(x)}$ to become very large or even numerically unstable, which leads to exploding loss values and unstable training dynamics.\n\nTherefore, a well-chosen reference distribution should balance **exploration** (ensuring coverage of the full support) and **exploitation** (favouring regions with higher probability mass). This trade-off helps stabilise training by avoiding extreme ratios while still providing meaningful learning signals. In our implementation, the reference distribution $q_t$ is generated by a rate matrix $R_\\theta$, which is updated during training to adaptively improve alignment with $p_t$ over time. This strategy is commonly used in training continuous neural diffusion [1] and flow-based [2] samplers, and we find it also empirically effective in our discrete setting. We will add more discussion on this point in the revision.\n\n[1] Akhound-Sadegh, Tara, et al. \"Iterated denoising energy matching for sampling from boltzmann densities.\" *arXiv preprint arXiv:2402.06121* (2024).\n\n[2] Tian, Yifeng, Nishant Panda, and Yen Ting Lin. \"Liouville flow importance sampler.\" *arXiv preprint arXiv:2405.06672* (2024).\n\n> While the authors state that any q_t with support matching p_t would suffice in theory, the experimental results suggest it may have a nontrivial empirical effect. My understanding is that the learned probability path is updated in the coordinate descent approach and used to approximate the expectation in (8), but its initialisation strategy lacks discussion or supporting code for some experiments. Clarifying these aspects, especially the empirical sensitivity of results to different initialisations would significantly improve both clarity and replicability.\n> \n\nThanks for pointing this out. You are correct, we adopt an “on-policy” training strategy, where the reference distribution $q_t$ is defined as the probability path induced by the rate matrix $R_\\theta$, which is updated throughout training.\n\nTherefore, the initialisation of $q_t$ is inherently determined by the initialisation of the rate matrix $R_\\theta$, which is parameterised by a neural network. We adopt PyTorch’s default parameter initialisation and observe that our method is robust to this choice. Specifically, we repeat the experiment five times with different random seeds, training DNFS to sample from the Ising models defined in Eq. 11 ($\\sigma=0.1$), and report the mean and standard deviation of the effective sample size (ESS) and free energy (under the scale of $10^4$).\n\n|  | ESS | Free Eenrgy |\n| --- | --- | --- |\n| Optimal value | 1 | -3.67265 |\n| DNFS | 0.99851 ± 4.6912e-05 | -3.67091 ± 6.2463e-05 |\n\nThe results exhibit minimal standard deviation, indicating the robustness of our method. Furthermore, we provide code for sampling from Ising models in the supplementary material, which may help clarify the training process. We will include this discussion in the revised version of the paper as well.\n\n> Potential sensitivity to initialisation of the rate matrix or probability path is not addressed. What is the initialisation strategy for the rate matrix across experiments? How sensitive are the results to the initialisation of the rate matrix?\n> \n\nThe rate matrix is parameterised by a neural network, initialised using PyTorch’s default weight initialisation scheme. Empirically, we find that our method is robust to this initialisation. To illustrate this, we replicate the previous experiment under a more challenging setting: sampling from Ising models at the critical temperature $\\sigma = 0.22305$ (see Eq. 11), suggested by reviewer prrE. We repeat the training five times with different random seeds, training DNFS to sample from the Ising models defined in Eq. 11. We report the mean and standard deviation of the effective sample size (ESS) and free energy (scaled by $10^4$).\n\n|  | ESS | Free Eenrgy |\n| --- | --- | --- |\n| Optimal value | 1 | -2.12417 |\n| DNFS | 0.9685 ± 0.0010 | -2.11201 ± 6.6258e-05 |\n\nThese results show that despite variations in initialisation, the model consistently converges to high-quality solutions with stable performance. This suggests that the training process is not overly sensitive to the choice of initial weights.\n\n> How are MCMC refinement steps (e.g., DMALA) integrated into the sampling procedure?\n> \n\nTo perform MCMC refinement, we adopt the predictor-corrector approach commonly used in diffusion generative models [3]. Specifically, before applying the transition $p(x_{t-\\Delta t} \\mid x_t)$ as defined in Eq.2, we refine the distribution of $x_t$ by performing one or more steps of MCMC at the fixed time step $t$. This correction step brings the distribution of $x_t$ closer to the target distribution $p_t$. Since $p_t \\propto \\rho^{1-t} \\eta^t$, we have access to its unnormalised form, allowing us to apply any suitable MCMC algorithm for refinement.\n\n[3] Song, Yang, et al. \"Score-based generative modeling through stochastic differential equations.\" *arXiv:2011.13456* (2020).\n\n> Lack of implementation details for certain experiments (e.g., combinatorial examples), reduces reproducibility. It might be useful to release code and add implementation details for the combinatorial optimisation experiment.\n> \n\nThank you for raising this concern. The training algorithm is summarised in Appendix C, and experimental details are provided in Appendix E. However, we agree that it may still be difficult to fully grasp the implementation details from the text alone. To address this, we have included code in the supplementary materials that reproduces the results of training DNFS for sampling from Ising models. We also plan to open-source the code for all experiments upon acceptance of the paper."}, {"review_id": "wC0i3h0969", "reviewer": "Reviewer_VDCK", "summary": "The authors propose a method to learn a Continuous-Time Markov Chain (CTMC) for sampling from an unnormalized density on a discrete state space. They craft a continuous in time process given an unnormalized density  $\\hat{p}(x_t) \\propto p(x_t)$ and learn a rate matrix $R_t$ to satisfy the Kolmogorov equation. However, evaluation of the Kolmogorov equation proves to be challenging due to the need to evaluate the time derivative of the probability density ($\\partial_t p_t$ or $\\partial_t \\log p_t$) and rate matrices $R_t(y|x)$ at many points $y$ and $x$. Authors address the first challenge with a coordinate descent algorithm and the second challenge by employing a special parameterization of the rate matrix with a special kind of neural networks with local equivariance.\n\nIn the practical part, the authors show the performance of their method on a series of benchmarks for sampling from unnormalized densities on discrete state spaces: sampling from a neural network trained EBMs, sampling from Ising models, solving combinatorial problems, i.e., maximum independent set problem. In addition, the authors utilize the byproduct of the coordinate descent algorithm for the training of energy-based models and test this approach on 2D real space grids quantized to the discrete state space and Ising models.", "questions": "- I ask the authors to clearly outline what is the difference of their method w.r.t. LEAPS [1], see weaknesses.\n- Since your method is so close to LEAPS, can DNFS employ the proactive importance sampling (Sec 4, 5 [1])?\n- In Table 1 Size, Drop and Time measurements are reported, but what do they mean? What is the size, drop, and in what units is time measured? I think such an explanation should be added to the Table 1 caption to improve clarity.\n- In Sec 4.3 the authors use the combination of DNFS and DMALA. How, in particular, does it work? Can the authors provide an algorithm or a clear step-by-step explanation?\n- In the conclusion and limitations section, the authors say that it is a problem to extend DNFS beyond binary settings. Why is that? To me it seems like the computational complexity will grow linearly with the number of categories.\n- What are the possible applications of learning an unnormalized density (energy function) on the discrete state space?", "limitations": "Mostly yes, for others see the Strengths and Weaknesses section.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n\n- The quality of the paper text and narrative is good, and most of the concepts are well presented. I enjoyed reading this paper.\n- The design of the locally equivariant transformer is interesting and could be a good contribution to the locally equivariant neural networks family.\n- The experimental part is strong, with a lot of experimental setups and comparisons. It was especially interesting to see the method's performance on solving combinatorial problems.\n- The way to learn an energy function on a discrete state space through contrastive divergence and log partition function, i.e., $\\partial_t \\log Z_t$, estimation, though solving the additional optimization problems seems quite novel to me.\n\nWeaknesses\n\n- Originality. My biggest concern is the originality of the proposed method. Seems like the core of the DNFS method was already presented in LEAPS [1], and the authors add very little on top. From the paper, it is hard to say whether there is some difference between the methods; it feels like the methods are almost the same.\n    - The leTF neural network architecture seems to me to be the only one notable contribution, w.r.t sampling from unnormalized densities.\n    - Coordinate descent presented in Sec 3.1 seems equivalent to PINN objective (Proposition 6.1 [1]), moreover authors dedicate an Appendix A.3 to the comparison of these objectives, and I cannot see a clear conclusion out of it.\n    - “One-way” rate matrix and locally equivariant neural networks parameterization of rate matrices has also been proposed in Sec 8 [1].\n    - Overall, most of the methods seem to have already been proposed in LEAPS [1] and hence DNFS is mostly not original.\n    - I encourage authors to clearly explain ***all*** the differences between DNFS and LEAPS, and incorporate it into the paper.\n- Method computational complexity. Does one have to do the sampling from the learned backward process during the training procedure? Seems like this is the case, and Algorithm 1 supports this guess, but then it adds a lot of computational burden on top, and it should be mentioned in the main part of the paper and in the limitations.\n- Practical performance. The performance of pure DNFS falls a bit behind the GFlowNets on Maximum Independent Set problem evaluations.\n\n[1] Holderrieth, P., Albergo, M. S., & Jaakkola, T. LEAPS: A discrete neural sampler via locally equivariant networks. In *Frontiers in Probabilistic Inference: Learning meets Sampling*.", "was_revised": true, "final_justification": "Many of my concerns were adressed, but the main one about the originality, unfortunately, still stands. Although the paper seems like overall of a decent quality, I don't feel like pushing my rating over 3, because of the originality concerns.", "rebuttal": "We appreciate the reviewer's constructive feedback. Regarding the weaknesses pointed out, we would like to provide the following clarifications:\n\n> My biggest concern is the originality. I encourage authors to clearly explain all the differences between DNFS and LEAPS.\n> \n\nTo address the reviewer’s concern, we summarise the main differences below\n\n1. A new perspective from the Kolmogrove equation\n    \n    While DNFS and LEAPS yield similar objective functions, they are derived from fundamentally different perspectives. LEAPS trains the rate matrix by minimising the log-variance of the importance weights. In contrast, DNFS takes a more straightforward approach: it learns a rate matrix by enforcing it to satisfy the Kolmogorov equation, where we estimate the intractable term $\\partial_t \\log Z_t$ using Monte Carlo with control variates. This formulation provides a complementary viewpoint and also a more intuitive explanation of the required estimation of $\\partial_t \\log Z_t$.\n    \n    Perhaps surprisingly, even though DNFS and LEAPS are derived from two different perspectives,  DNFS has the same objective as LEAPS if $\\partial_t \\log Z_t$ is estimated by training a neural network regressor. However, we argue that DNFS is preferable, as it can directly leverage the closed-form optimal solution, avoiding the approximation errors that may arise during neural network training. Moreover, as outlined in lines 691–694, when the model is optimally trained, DNFS is guaranteed to recover the true value of $\\partial_t \\log Z_t$. In contrast, LEAPS may incur additional errors due to its decoupled modelling strategy. From this perspective, DNFS highlights a promising direction for future research: leveraging more accurate estimators of $\\partial_t \\log Z_t$ to further improve performance, such as incorporating SMC methods or using more advanced variance reduction techniques beyond control variates. These insights are not evident from the LEAPS framework.\n    \n2. More advanced Model Architecture and Expressiveness\n    \n    We further propose the locally equivariant transformer (leTF), a novel architecture that achieves highly efficient objective computation without sacrificing model expressiveness (see Figures 3 and 14 for empirical validation). Unlike the locally equivariant networks (leNets) instantiated in LEAPS, which are limited to simple architectures such as MLPs, single-layer attention, and convolutional networks, leTF provides greater model capacity and enhanced adaptability, making it better suited for diverse modalities and complex input structures. We believe that this transformer-based design significantly broadens the applicability of DNFS across a wider range of tasks. We hope this architectural advancement will serve as a foundation for further research into expressive yet efficient parameterisations of leNets.\n    \n3. Broader Empirical Scope and Generality\n    \n    Unlike LEAPS, which is only evaluated on Ising models, DNFS is tested across a broader range of applications, including sampling from Ising models and pretrained EBMs, training deep EBMs, and solving combinatorial optimisation problems. These experiments highlight the versatility and generality of our approach across diverse domains. We believe this broader empirical scope lays a clear foundation for future research aimed at developing more effective discrete neural samplers that are both scalable and adaptable to real-world applications.\n    \n\nWe hope this clarification helps address the reviewer’s concerns regarding originality. We believe that the new perspective of framing the problem in terms of learning a rate matrix that satisfies the Kolmogorov equation, along with the introduction of a more expressive locally equivariant transformer architecture and broader applications and evaluations of discrete neural samplers, may offer useful insights and contribute meaningfully to the ongoing research in this area.\n\n> Does one have to do the sampling from the learned backward process during the training procedure?\n> \n\nYes. In DNFS, the reference distribution $q_t$ is defined as the probability path induced by the rate matrix, and we need to sample from the learned backward process during training.\n\nWe acknowledge that this sampling step introduces computational cost. However, this cost is inherent in training neural samplers where direct access to data samples from the target distribution is unavailable. In such cases, it is necessary to generate “synthetic” samples from the reference distribution $q_t$ to compute the training objective (see Eq. 5).\n\nWhile the model is theoretically guaranteed to converge for any $q_t$ that shares the same support as the target distribution $p_t$, the practical choice of $q_t$ has a significant impact on training stability. A well-designed reference should strike a balance between exploration (ensuring support coverage) and exploitation (focusing on high-probability regions), which helps avoid extreme loss values and stabilises learning. In our approach, $q_t$ evolves with the model via the rate matrix $R_\\theta$, which is updated during training to better align with $p_t$ over time. This adaptive, on-policy strategy is standard practice in training continuous diffusion samplers [1], flow-based samplers [2], and LEAPS, and we find it empirically effective in our method as well.\n\nWe will include a discussion in the revised version.\n\n[1] Iterated denoising energy matching for sampling from Boltzmann densities\n\n[2] Liouville flow importance sampler\n\n> In Sec 4.3 the authors use the combination of DNFS and DMALA. How, in particular, does it work?\n> \n\nTo perform MCMC refinement, we adopt the predictor-corrector strategy, a technique commonly used in diffusion-based generative models [3]. Specifically, before applying the forward transition $p(x_{t-\\Delta t} | x_t)$ (see Eq2), we insert MCMC-based correction steps to refine the current sample $x_t$ so that its distribution is closer to $p_t$. This correction is possible because we know its unnormalised form $p_t \\propto \\rho^{1-t} \\eta^t$, which allows us to use MCMC such as DMALA to refine toward $p_t$.\n\n[3] Score-based generative modeling through stochastic differential equations\n\n> The performance of pure DNFS falls a bit behind GFlowNets.\n> \n\nAs noted in the conclusion, a promising direction for future work is to extend DNFS from uniform diffusion to *masked* diffusion processes. This would allow us to incorporate constraints directly into the sampling trajectory, similar to GFlowNet, where sampling is restricted to the feasible set only. By doing this, it can narrow the search space,  significantly improving sampling efficiency and performance.\n\nNevertheless, although pure DNFS underperforms GFlowNet, it has a key advantage: the intermediate target distribution $p_t$ is explicitly known. This allows us to combine DNFS with MCMC refinement (e.g., DNFS+DMALA, see Table 1), which significantly boosts performance and enables DNFS to outperform GFlowNet.\n\n> Can DNFS employ the proactive importance sampling as in LEAPS?\n> \n\nYes, we can. The sampling procedure is independent of the training method. We can train the rate matrix using either LEAPS or DNFS. At inference time, we can simulate CTMC using the Euler method and apply proactive importance sampling as done in LEAPS, which is essentially a type of sequential Monte Carlo.\n\n> What are Size, Drop and Time in Table 1?\n> \n\nSize refers to the number of nodes in the solution. A larger size indicates better performance. Drop denotes the percentage drop in performance relative to GUROBI, which we treat as the oracle. A smaller drop is therefore better. Time represents the inference time. Lower values indicate greater efficiency.\n\n> Why is extending DNFS beyond binary settings problematic?\n> \n\nThe difficulty comes from two-fold:\n\n1. High computational cost of computing the ratio $\\frac{p(y)}{p(x)}$\n    \n    When this ratio lacks a closed-form expression, as in Eq27, it must be computed in parallel for all possible values of $y$ (see Eq5). This becomes computationally prohibitive in high-dimensional spaces or when dealing with large categorical vocabularies, particularly if the target distribution involves neural networks, such as deep EBMs or deep reward models.\n    \n2. Potential instability of ratio values\n    \n    During training, samples $x$ are drawn from the reference distribution $q_t$, which is defined by the probability path induced by the learned rate matrix $R_\\theta$. In the early stage of training, when $R_\\theta$ is still inaccurate, the sampled $x$ may lie far from the high-density regions. In such cases, $p_t(x)$ can be extremely small, making the ratio $\\frac{p(y)}{p(x)}$ explode. This instability is further amplified in high-dimensional settings, where the energy landscape tends to be rugged and the modes sparse.\n    \n\nDespite these challenges, extending DNFS to large categorical vocabularies and high-dimensional problems remains a promising direction for future research.\n\n> What are the applications of learning EBMs on the discrete state space?\n> \n\nTraining discrete EBMs has been extensively studied in the literature [4,5], and they have demonstrated broad applicability across various domains. For example, EBMs have been applied to language modeling [6] and enhancing diffusion language models [7]. Additionally, they support compositional generation [8] and offer a flexible framework for decision-making and reasoning [9].\n\n[4] oops i took a gradient scalable sampling for discrete distributions\n\n[5] Generative Flow Networks for Discrete Probabilistic Modeling\n\n[6] Residual Energy-Based Models for Text Generation\n\n[7] Energy-Based Diffusion Language Models for Text Generation\n\n[8] Compositional generation with energy-based diffusion models and mcmc\n\n[9] Learning Iterative Reasoning through Energy Diffusion."}, {"review_id": "Mp7OOFVRQB", "reviewer": "Reviewer_prrE", "summary": "This paper builds upon continuous-time discrete flow samplers and introduces a novel locally equivariant Transformer architecture to efficiently parameterize the rate matrix. Additionally, the authors propose an alternative method for estimating the time derivative of the partition sum, which exhibits lower variance compared to previous estimators. The paper includes ablation studies and evaluates the method on sampling from pre-trained Energy-Based Models (EBMs), the Ising Model, and Combinatorial Optimization problems. However, the system sizes considered are quite small.", "questions": "### Questions:\n- Why does the method not scale to large dimensions, and what is the bottleneck?\n- How many integration steps are used in the method for training and evaluation?\n- The Ising Model is considered at \\(\\sigma = 0.1\\). Does this correspond to the critical temperature? If not, how does the model perform at the critical temperature? I am particularly interested in the log Z, entropy, and internal energy estimations compared to ground truth values from theory, as these are good indicators of mode collapse.\n- In line 1134, how many Combinatorial Optimization problem instances were used to condition the neural network per update step?\n- The authors might also consider benchmarking on the RB-model. I would be interested in how the model performs for different hardnesses (see [1] Fig. 2 for reference).\n- I find section C.2 in the appendix rather weird. Here, log ratios are computed by using a taylor approximation. However this assumes that $x$ is continuous and not discrete which is not the case in discrete optimization. Could the authors comment on that? Also this computation is only necessary for EBMs, right?\n\n\n[1] Sanokowski, Sebastian, Sepp Hochreiter, and Sebastian Lehner. \"A diffusion model framework for unsupervised neural combinatorial optimization.\" Proceedings of the 41st International Conference on Machine Learning. 2024.", "limitations": "Limitations are properly adressed.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths:\n- The paper is well-presented and clearly written.\n- It introduces some novel contributions.\n- Method is evaluated on a decent amount of experiments\n- The limitations of the study are appropriately addressed.\n\n### Weaknesses:\n- The authors mention in the appendix the use of variational neural annealing, and it would be appropriate to cite [1] in this context.\n- The paper does not present log Z values for the Ising Model or compare them to theoretically known baselines, making it unclear whether mode collapse occurs. It would be beneficial to include a comparison with theoretically known values. A script to compute these theoretical values for finite-size Ising models is available [here](https://github.com/ml-jku/DiffUCO/blob/main/IsingTheoryBaselines/IsingTheory.py).\n- The Ising Model does not appear to be evaluated at the critical temperature, where solving it is particularly challenging.\n- The experimental setup for the Ising Model, as described around line 285, is insufficiently explained and thus difficult to understand.\n\n### References:\n[1] Hibat-Allah, Mohamed, et al. \"Variational neural annealing.\" *Nature Machine Intelligence* 3.11 (2021): 952-961.", "was_revised": true, "final_justification": "I think this is a technically solid paper, with a large amount of experiments and openness about the limitations of their method.\nAlso a lot of work was put into the rebuttal.\nI therefore recommend acceptance.", "rebuttal": "Thanks for your constructive comments. Here are our responses to your concerns.\n\n> Does $\\sigma = 0.1$ correspond to the critical temperature? If not, how does the model perform at the critical temperature?\n> \n\nThank you for raising this point. The value 0.1 is not the critical temperature. In our experiments, we use 2D lattice Ising models, for which the critical temperature is $\\sigma = \\frac{\\log (1+\\sqrt{2})}{4} \\approx 0.22035$ based on Onsager’s solution. \n\nTo assess our model's performance at the critical temperature, we conduct additional experiments sampling from the Ising model at $\\sigma = 0.22035$. We estimate the free energy, internal energy, and entropy using the importance sampling method described in Appendix D.1. Each experiment is repeated five times with different random seeds, and we report the mean and standard deviation of the effective sample size (ESS), free energy, internal energy, and entropy (scaled by $10^4$)\n\n|  | ESS | Free Energy | Internal Energy | Entropy |\n| --- | --- | --- | --- | --- |\n| Optimal Value | 1 | -2.12417 | -1.47632 | 0.28550 |\n| DNFS | 0.9685 ± 0.0010 | -2.11201 ± 6.6258e-05 | -1.47425 ± 0.0068 | 0.28106 ± 0.0030 |\n\nThe results show that DNFS performs well even at the critical temperature, with high effective sample size and accurate estimates of free energy, internal energy, and entropy. The small deviations from the optimal values, along with the low standard deviations across runs, further demonstrate the robustness and reliability of our method under challenging sampling conditions.\n\n> The paper does not present log Z values for the Ising Model or compare them to theoretically known baselines.  I am particularly interested in the log Z, entropy, and internal energy estimations compared to ground truth values from theory.\n> \n\nThank you for your insightful suggestions and for sharing the code scripts. Following your recommendations, we evaluated DNFS and LEAPS under the setting $\\sigma = 0.1$ and $\\sigma = 0.22305$, the latter corresponding to the critical temperature. We report estimates of the log partition function (free energy), internal energy, and entropy—key indicators of mode coverage—alongside effective sample size (ESS). The results, scaled by $10^4$, are summarised below.\n\n$\\sigma = 0.1$\n\n|  | ESS | Free Energy | Internal Energy | Entropy |\n| --- | --- | --- | --- | --- |\n| Optimal Value | 1 | -2.12417 | -1.47632 | 0.28550 |\n| DNFS | 0.9685 ± 0.0010 | -2.11201 ± 6.6258e-05 | -1.47425 ± 0.0068 | 0.28106 ± 0.0030 |\n| LEAPS | 0.3631 ± 0.1184 | -2.10111 ± 0.0002 | -1.44926 ± 0.0130 | 0.28727 ± 0.0057 |\n\n$\\sigma = 0.22305$\n\n|  | ESS | Free Energy | Internal Energy | Entropy |\n| --- | --- | --- | --- | --- |\n| Optimal Value | 1 | -2.12417 | -1.47632 | 0.28550 |\n| DNFS | 0.9685 ± 0.0010 | -2.11201 ± 6.6258e-05 | -1.47425 ± 0.0068 | 0.28106 ± 0.0030 |\n| LEAPS | 0.3631 ± 0.1184 | -2.10111 ± 0.0002 | -1.44926 ± 0.0130 | 0.28727 ± 0.0057 |\n\nThe results show that DNFS provides accurate estimates of the partition function and internal energy, closely matching the theoretical value with low variance. In contrast, LEAPS exhibits significantly lower ESS and larger deviations in other metrics under the critical temperature setting, which may indicate insufficient mode coverage. These comparisons highlight the robustness of DNFS and suggest it is less prone to mode collapse across different temperatures.\n\n> The authors might also consider benchmarking on the RB-model to see how the model performs for different hardnesses.\n> \n\nThank you for your valuable suggestions. We conducted additional experiments on the Maximum Independent Set problem using the RB32-40 graphs, varying the parameter $p$, which controls the problem’s difficulty. Specifically, higher values of $p$ yield easier instances, while lower values result in more challenging graphs.  For each setting, we report the solution size (larger is better), along with the percentage performance drop relative to GUROBI (lower is better), shown in brackets below:\n\n| p= | 0.1 | 0.3 | 0.5 | 0.7 | 0.9 | 1.0 |\n| --- | --- | --- | --- | --- | --- | --- |\n| GUROBI | 8.52 (0.00%) | 8.24 (0.00%) | 7.06 (0.00%) | 7.89 (0.00%) | 8.05 (0.00%) | 8.63 (0.00%) |\n| Random | 3.38 (60.3%) | 4.59 (44.2%) | 4.25 (39.8%) | 5.54 (29.7%) | 6.29 (21.8%) | 6.93 (19.6%) |\n| DMALA | 8.50 (0.23%) | 8.16 (0.97%) | 7.01 (0.70%) | 7.84 (0.63%) | 8.04 (0.12%) | 8.63 (0.00%) |\n| GFlowNet | 8.20 (3.75%) | 7.93 (3.76%) | 6.83 (3.25%) | 7.78 (1.39%) | 8.03 (0.62%) | 8.63 (0.00%) |\n| DNFS | 8.00 (6.10%) | 7.65 (7.16%) | 6.67 (5.52%) | 7.66 (2.91%) | 7.97 (0.99%) | 8.63 (0.00%) |\n| DNFS+DMALA | 8.51 (0.11%) | 8.19 (0.60%) | 7.02 (0.56%) | 7.85 (0.50%) | 8.04 (0.12%) | 8.63 (0.00%) |\n\nOverall, these results are consistent with the observation in Table 1 of the paper. On easier instances (i.e., higher values of $p$), our method performs competitively, approaching the performance of the oracle solver GUROBI. On more challenging instances, however, DNFS exhibits a performance gap relative to GFlowNet. This gap arises because GFlowNet restricts its sampling trajectories to the feasible solution, effectively narrowing the exploration space. Introducing such an inductive bias into DNFS represents a promising direction for future work.\n\nNevertheless, despite DNFS underperforming GFlowNet in its pure form, it offers a distinct advantage: the intermediate target distribution $p_t$ is explicitly known. This property enables integration with MCMC-based refinement methods (e.g., DNFS+DMALA), which significantly improves performance.\n\n> Why does the method not scale to large dimensions?\n> \n\nThe scalability challenge arises from two factors:\n\n1. High computational cost of evaluating $\\frac{p(y)}{p(x)}$\n    \n    When this ratio lacks a closed-form expression, as in Eq. 27, we must compute it across all possible values of $y$, which becomes computationally prohibitive in high-dimensional spaces. This issue is especially pronounced when the target distribution is defined by a neural network, such as in deep EBMs or deep reward models. For example, computing the loss in Eq. 10 requires evaluating the reward model $\\mathcal{O}(S \\times d)$ times, where S is the vocabulary size and d is the input dimension, making the computation scale poorly with dimensionality.\n    \n2. Potential instability of the ratio values\n    \n    During training, we sample $x$ from the reference distribution $q_t$, which is induced by the probability path of the rate matrix $R_\\theta$. In the early stages of training, when $R_\\theta$ is still inaccurate, the sampled $x$ may lie far from high-probability regions. As a result, $p_t(x)$ can be extremely small, making the ratio $\\frac{p(y)}{p(x)}$ explode. This issue is further amplified in high dimensions, where the energy landscape is often not smooth and modes are sparse.\n    \n\n> In Appendix C.2, log ratios are computed via taylor approximation. However this assumes that is continuous and not discrete which is not the case in discrete optimization. Is this computation only necessary for EBMs?\n> \n\nYes, this computation is only required for deep EBMs, where the ratio $\\frac{p(y)}{p(x)}$ does not admit a closed-form expression. The idea of using Taylor approximation to estimate the difference between the log-probability of two different inputs $y$ and $x$ was first introduced in [1], where it was shown to improve the convergence of Gibbs sampling. This approach was later adopted in discrete Langevin samplers [2].\n\nThe key insight is that most energy functions are continuous and differentiable with respect to real-valued inputs, even though they are only evaluated over a discrete subset of their domain (see Section 3 in [1]). This allows the use of gradients and Taylor expansions to approximate likelihood ratios effectively.\n\nHowever, while this approximation improves computational efficiency, it introduces bias into our training objective, which can degrade the performance of the learned rate matrix (see Figure 9). As such, developing more accurate and scalable approximation methods remains an important direction for future work, particularly for applications in large-scale and high-dimensional settings.\n\n[1] Oops I Took A Gradient: Scalable Sampling for Discrete Distributions\n\n[2] A Langevin-like Sampler for Discrete Distributions\n\n> The experimental setup for the Ising Model, as described around line 285, is insufficiently explained.\n> \n\nSorry for the confusion. In this experiment, the goal is to learn the adjacency matrix of lattice Ising models from samples synthesised via long-run Gibbs sampling. Thus, this is framed as an EBM training task, rather than a sampling problem. Moreover, the training algorithm has no access to the true adjacency matrix, but only the synthesised samples. We will clarify this in the revised version.\n\n> In line 1134, how many Combinatorial Optimization problem instances were used to condition the neural network per update step?\n> \n\nAs outlined in line 1134, we use a batch size of 256. The training procedure follows Algorithm 1, with the only difference being that the inputs are of the form $(t, x_t^{(n)}, G^{(n)})$. This means that, in the best case, up to 256 distinct problem instances may be included in each update step. However, it is also possible for all $x_t^{(n)}$ to belong to the same graph $G$, in which case the update step corresponds to a single instance of the combinatorial optimisation problem.\n\n> How many integration steps are used in the method for training and evaluation?\n> \n\nSorry for the confusion. As described in Appendix E, we use 64 integration steps for both training and evaluation. We will make it clearer in the revised manuscript.\n\n> The authors mention in the appendix the use of variational neural annealing, and it would be appropriate to cite [1] in this context.\n> \n\nThanks for the reference. We will consider citing it in the revised version."}]

Author rebuttal:

Meta-review: The main concern about this paper is the novelty of the DNFS method compared with LEAPS. But the reviewers acknowledge that the paper offers practical advantages over LEAPS, both in terms of performance and the thoroughness of the empirical evaluation.
Decision: Accept (poster)