LEAPS: A discrete neural sampler via locally equivariant networks
Authors: Peter Holderrieth
Michael Samuel Albergo
Tommi Jaakkola
Venue: ICML 2025 poster (2025)
Keywords: CTMC
Bayesian inference
MCMC
sampling
Primary area: probabilistic_methods
TL;DR: A discrete neural sampler for unnormalized densities via locally equivariant networks.

Abstract: We propose *LEAPS*, an algorithm to sample from discrete distributions known up to normalization by learning a rate matrix of a continuous-time Markov chain (CTMC). LEAPS can be seen as a continuous-time formulation of annealed importance sampling and sequential Monte Carlo methods, extended so that the variance of the importance weights is offset by the inclusion of the CTMC. To derive these importance weights, we introduce a set of Radon-Nikodym derivatives of CTMCs over their path measures. Because the computation of these weights is intractable with standard neural network parameterizations of rate matrices, we devise a new compact representation for rate matrices via what we call \textit{locally equivariant} functions. To parameterize them, we introduce a family of locally equivariant multilayer perceptrons, attention layers, and convolutional networks, and provide an approach to make deep networks that preserve the local equivariance. This property allows us to propose a scalable training algorithm for the rate matrix such that the variance of the importance weights associated to the CTMC are minimal. We demonstrate the efficacy of LEAPS on problems in statistical physics. We provide code in https://github.com/malbergo/leaps/.

LEAPS: A discrete neural sampler via locally equivariant networks
Peter Holderrieth
Michael S. Albergo 23 Tommi Jaakkola
Abstract
Truth
LEAPS
LEAPS uncorrected
We propose LEAPS, an algorithm to sample from discrete distributions known up to normalization by learning a rate matrix of continuous-time Markov chain (CTMC). LEAPS can be seen as continuous-time formulation of annealed im- portance sampling and sequential Monte Carlo methods, extended so that the variance of the im- portance weights is offset by the inclusion the CTMC. To derive these importance weights, we introduce a set of Radon-Nikodym derivatives of CTMCs over their path measures. Because the computation of these weights is intractable with standard neural network parameterizations of rate matrices, we devise new compact rep- resentation for rate matrices via what we call locally equivariant functions. To parameterize them, we introduce a family locally equivari- ant multilayer perceptrons, attention layers, and convolutional networks, and provide approach to make deep networks that preserve the local equivariance. This property allows us to pro- pose a scalable training algorithm for the rate matrix such that the variance of the importance weights associated to the CTMC are minimal. We demonstrate the efficacy of LEAPS on problems in statistical physics. We provide code in https: //gittubb wttp://www..a..o//////aaaa leaps//
Po
P1
Categories
Figure 1. Illustration of the LEAPS algorithm. LEAPS allows to learn dynamical transport discrete distributions from = to t (bluc). Sample are generated via the simulation of Continuous-time Markov chain (yellow). Further, importance sampling weights allow correct training errors trading off sample efficiency (red).
the molecular dynamics of chemical compounds (Berendsen et al., 1984; Allen & Tildesley, 2017), and computational ap- proaches to statistical and quantum physics (Wilson, 1974; Duane et al., 1987; Faulkner Livingstone, 2023).
The most salient approach to such sampling problems Markov chain Monte Carlo (MCMC) (Metropolis al., 1953; Robert et al., 1999), in which randomized process is simulated whose equilibrium is the distribution of interest. While powerful and widely applied, MCMC methods can be inefficient as they suffer from slow convergence times into equilibrium, especially for distributions exhibiting multi- modality. Therefore, MCMC often combined with other techniques that rely on non-equilibrium dynamics, e.g, annealing from a simpler distribution with annealed impor- tance sampling (AIS) (Kahn & Harris, 1951; Neal, 2001) or sequential Monte Carlo methods (SMC) (Doucet et al., 2001). Even then, the variance of these importance weights may be untenably large, and making sampling algorithms more efficient remains an active area of research. Inspired by the rapid progress in deep learning, there has been extensive
1. Introduction
A prevailing task across statistics and the sciences is to draw samples from probability distribution whose probability density is known up to normalization. Solutions to this prob- lem have applications in topics ranging across Bayesian un- certainty quantification (Gelfand & Smith, 1990), capturing
'Equal contribution 'Massachusetts Institute of Technology "Society of Fellow, Harvard University 'Institute for Artificial Intelligence and Fundamental Interactions. Correspondence to: Peter Holderrieth <phold@ mit.edu>, Michael Albergo <mal- bergo@ fas. harvard.edu>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

LEAPS: A Discrete Neural Sampler
interest in augmenting contemporary sampling algorithms with learning (Noé al., 2019; Albergo et al., 2019; Gabrié et al., 2022; Nicoli al., 2020; Matthews et al., 2022), while still maintaining their statistical guarantees.
2. Setup and Assumptions
In this work, we are interested in the problem of sampling from target distribution Pi on a finite state space S. We refer to pt by its probability mass function (pmf) given by
Recently, there has been rapid progress in development of generative models using techniques from dynamical mea- sure transport, i.c. where data from base distribution is transformed into samples from the target distribution via flow diffusion processes (Ho et al., 2020; Song et al., 2020; Albergo & Vanden-Eijnden, 2022; Albergo et al., 2023; Lipman et al., 2023; Liu et al., 2022). More recently, these models could also be developed for discrete state spaces (Campbell et al., 2022; Gat et al., 2025; Shaul et al., 2024; Campbell et al., 2024) and general state spaces and Markov processes (Holderrieth al., 2024).
1 exp(-Ui(x)) (ar E S)2 P: (1( = Z1
(1)
where we assume that we do not know the normalization constant Z1 but only the function potential U1. Our goal is to produce samples X ~ Pl. To achieve this goal, is common to construct a time-dependent probability mass function (pmf) (Pt)osesı over which fulfils that Po has distribution from which we can sample easily, e.g. Pu Unifg, and Pr is our target of interest. We write Pr as:
While there have been various developments on adapting these non-equilibrium dynamics for sampling in continu- ous state spaces (Zhang & Chen, 2022; Vargas et al., 2023;
Px(x) Ze exp(-U;(z)),
(2)
Z, Σ expí (V)(),,()), F. = -- -log Ze yES
Máté & Fleuret, 2023; Tian et al., 2024; Albergo & Vanden- Eijnden, 2024; Richter & Berner, 2024; Akhound-Sadeee et al., 2024; Sun et al., 2024), there a lack of existing liter- ature on such sampling approaches for discrete distributions. However, discrete data are prevalent in various applications, such as in the study of spin models in statistical physics, protein and genomic data, and language. To this end, we provide a new solution to the discrete sampling problem via CTMCs. Our method is similar in spirit to the results in (Albergo& & Vanden Rijddnn,,, 2024; Vargas et al., 2024) but takes the necessary theoretical and computational leaps to make these approaches possible for discrete distributions. Our main contributions are:
where Z. (or equivalently F)) is unknown. The value F, is also called the free energy. Throughout, we assume that U. is continuously differentiable in t. For example, we can set Uz(x) = (U'(() (x) so that Pa = Unifs and we get that P: x exp(-tU (z)) that can be considered a form of temperature annealing.
3. Sampling with CTMCs
In this work, we seek to sample from #1 using continuous- time Markoy chains (CTMCs). A CTMC (Xt)osıg1 is given by a set of random variables X, E $ (0 S S 1) whose evolution is determined by a time-dependent rate matrix Qr x) E R (0 S S 1, E S) which fulfills. the conditions:
We introduce LEAPS, a non-equilibrium transport sam- pler for discrete distributions via CTMCs that com- bines annealed importance sampling and sequential Monte Carlo with learned measure transport.
Qc(y: 20
(for y # x)
(4a)
To define the importance weights, we derive a Radon- Nikodym derivative for reverse-time CTMCs, control of which minimizes the variance of these weights.
Qx(x;))
Σ 2)(,22) wylar
(for S)
(4b)
The rate matrix Qr determines the generator equation
We show that the measure transport can be learnt and the variance of the importance weights minimized by optimizing physics-infonnee informed neural network (PINN)
P[X++A y| = x] 1-=9 + hQr(y, x) + o(h) (5)
loss function.
for all x, y E and h > where ohh) describes an error function such that lim o(h)/h O. Because this equa- h-0 tion describes the infinitesimal transition probabilities of the CTMC, following iterative Fuler scheme:
We make the computation of the PINN objective seal- able by introducing the notion of a locally equivariant network. We show how to build locally equivariant versions of common neural network architectures, in- cluding attention and convolutions.
sample from X, approximately via the
(6) where [f] describes a valid probability distribution for small h by the conditions we imposed on Qe (see (4)).
Xeth ~ P((X == (1x,=y + hu, Xi))
We experimentally verify the correctness and efficacy of the resulting LEAPS algorithm in high dimensions via simulation of the Ising model.
2

---

LEAPS: A Discrete Neural Sampler
Our goal is find Qe that is solution to the Kolmogorov forward equation (KFE)
proposed scheme in the following paragraph. Intuitively, the log-weights A, should accumulate the deviation from the true distribution of X₁ to the desired distribution Pt. We can rephrase this as "accumulating the error of the KFE" that one may want to write as the difference between both sides (7):
Depe(a) = (i(n)))))) Pt=0 - po-
BES
Fulfilling the KFE is a necessary and sufficient condition to ensure that the distribution of walkers initialized as Xo ~ po and evolving according to (6) follow the prescribed path Pt. in particular such that Xt=1 Pt.
(x) - (elr,
yES
Remark 3.1. While the problem we focus on this work is sampling, the contributions of this work hold more generally for the problem of finding a CTMC that follows prescribed time-varying density Pt. i.c. that the condition in (7) holds.
As we do not know the normalization constant Ze, it is intuitive divide by pe(x:) get
depr(a) (t(x)
Pr(y) (x(
Σ BES
Ei( 1)
4. Proactive Importance Sampling
Using that we obtain after dropping the unknown time-dependent con- stant F.F;
= O log (r())) = 0, 0,F,
In general, the CTMC (Xe)ogegı with arbitrary Qs will have different marginals than Pt- To still obtain an unbi- ased estimator, is common to use importance sampling (IS) to reweight samples obtained while simulating X, or sequential Monte Carlo (SMC) (Doucet et al., 2001) to re- sample the walkers along the trajectory. Here, introduce time-evolving set of log- weights A, E R for 0 S 1 to re- weight the distribution of X2 to a distribution μr defined IR
- -,U((() - 2e((I,2) 2) Poly) p+(x)
KGp((0) =
(10)
yES
where we defined new operator KePe. Intuitively, the oper- ator Kr, measures the violation from the KFE in log-space (up to Ft) and it is intuitive to define A, as the accumulated error of that violation, i.e. as the integral
such that for all h:
Elexp( Ar (x)))
Ea-p (M(x))
EE. [[epp(AA (A))
f (apa(X.)ds
E E(exp( AA) Xx, z]
A
(11)
(e(x)
Σ Elexp( (At)),
DES
To simulate A, alongside Xx, we the approximate update role:
where [E[-] denotes expectation over the process (X,, A,). Intuitively, the distribution μr is obtained by re-weighting samples from the current distribution of X, This ef- fectively means that from a finite number of samples (X), A!),..... (x,, A,), we can obtain a Monte Carlo esti- mator via
A+ta A, + (X,)
We call this proactive importance sampling (proactivate IS) as the update operator Kr, anticipates where X, is jump- ing to. We next provide a rigorous characterization of A, defined in this manner.
exp(Ai) 1(x))
Exnna [bxx)]
exp( A()
5. Proactivate IS via Radon-Nikodym Derivatives
Our goal is to find a scheme of computing A, such that its reweighted distribution coincides with the target densities Pe:
A priori, it not clear that the log-weights A₁ that we obtain via the proactive rule fulfil the desired condition in (9) (i.e. provide a valid scheme). Beyond showing this property, we show that there are many possible IS schemes but the proactive update rule is optimal among natural family of IS schemes. To do so, we present set of Radon-Nikodym derivatives in path space.
= 10
(9)
In particular, this would mean that (8) is good approxima- tion for large n.
Proactive importance sampling. We next propose an IS scheme of computing weights A,. Before we provide a formal derivation, we provide heuristic derivation of our
Let X' be a state space and P. Q be two probability measures over X. Then the Radon-Nikodym derivative (RND) aF allows to express expected values of Q via expected values
3

---

LEAPS: A Discrete Neural Sampler
of P. Specifically, the RND is a function d2 : that for any (bounded, measurable) function G: X holds that:
R such R it
(X) gives valid set of importance weights. Note that this holds for arbitrary Q1-
i.c. the RND a'e..a d
Exxo|G(X)] = -xxee ((X)
dQ (x)
However, to sample efficiently, it is crucial that the IS weights have low variance. Therefore, we will now derive the optimal IS scheme of this form. Ideally the weights will
dP
have zero variance in other words the RND d'='1.e' (X) will be constant = 1. This is the case if and only if the path measures are the same, i.e. if the CTMC reverse time is a time-reversal of the CTMC running in forward time. is well-known that this is equivalent to:
The state space that we are interested in is the space X' of CTMC trajectories. Specifically, for a trajectory we denote with X,- lim X the left limit and with Xx+ Jime Xy the right limit. The space X' of CTMC trajectories is then defined
X' = (X: [0, 1] - s|X,- exists and Xx+ Xt):
(:l)) (x())
Pil. x)) Qz v)
for all y + (13)
L.c all trajectories that are continuous from the right with left limits. Such trajectories are commonly called cådlăg trajectories. In other words, jumps (switches between states) happen if and only if X,- * X,. We consider path distribu- tions (or path measures), i.e. probability distributions over trajectories. For CTMC X (Xt)osts) with rate matrix Qi and initial distribution μl. we denote the corresponding - path distribution as T.. where the arrow P denotes that we go forward in time. Similarly, we denote with F a CTMC running in reverse time initialized with v. We present the following proposition whose proof can be found in Appendix A:
where q denotes the true marginal of X,, i.c. X, ~ q. As we strive to make g: = Pe. it is natural to set qt = pr (13) and define Q2 = Qi as
pr(y) pc(x)
Q1(1,x) (Z
(2, v)
for all y fx
(14a)
Σ Qc(x, v) pr(v)
Q((x,)== --
(14b)
(x(
VES,yşe
Let now return the proactive update that we defined in (11). We can now rigorously characterize it. Plugging in the definition of Q, we can use Proposition 5.1 to obtain the main result of this section:
Proposition 5.1. Let μ, v be two initial distributions over S. Let Qu, Q2 be two rate matrices. Then the derivative of the corresponding path distributions running in opposite time over the time interval [0, 8] is given by:
Theorem 5.2. For the proactivate updates A, as defined in (11) and Qi defined in (14), it holds:
aF (X) log((((.)) log(u(Xo))
log
d Q
d dFuq
A₁ + F. Fo = loog
(X)
(15)
aP l
+ f -x( - Qo(x,,
This implies that we obtain valid IS scheme fulfilling:
Σ log (2((x,
((x Bznp, [M(x)] = E((ex((A))
(051S1) (16)
(x,,
*X.
i.e. (9) holds. Further, As will have zero variance for every
where we over all points where Xs jumps in the last term.
0 < L 1 if and only if X, ~ Pt for all 0 S $111 1.
A proof can be bound in Appendix B. Note that Theorem 5.2 is useful because we can, in principle, compute Ar, i.e. there are no unknown variables, and that this holds for arbitrary Q.. This theorem can be seen as discrete state space equivalent of the generalized version of the Jarzynski equal- ity (Jarzynski, 1997; Vaikuntanathan & Jarzynski, 2008) that has also recently been used for sampling in continu- ous spaces (Vargas et al., 2024; Albergo & Vanden-Eijnden, 2024). Finally, it is important to note that the IS scheme will not have zero variance if it does not hold that X, ~ Pt-
Let us now revisit our goal of finding an IS scheme to sample from the target distribution P1. The key idea is to construct a CTMC running in reverse-time with initial distribution and then use the RND from Proposition 5.1. For a function h: S + R, we can then express its expectation under Pt as:
Expp, [k(x)] = EE xa
Fr [Xx))
d P
(12)
= EE =Ix~ZpO
A(X,)
(X)
d p

---

LEAPS: A Discrete Neural Sampler
6. PINN Objective
As next step, we introduce a learning procedure for learn- ing an optimal rate matrix of a CTMC. For this, we denote with Qf a parameterized rate matrix with parameters 0 (e.g. represented in neural network). Our goal is to learn Q such that X; Pt is fulfilled for all 0 S L < 1. By Theo- rem 5.2 this equivalent to minimizing the variance of the IS weights. To measure the variance the weights, it is common to use the log variance divergence (Nüsken & Richter, 2023; Richter& & Berner, 2023) given by
(+1,444)))
y
G (+1,3
G (11.3l
1 2 3
4
5
Clog-var (0; =Vx xx>gxlll flog I Qo.Q0
(X)
=Vxmol4 [4 + F. Fo)
Figure 2. Visualization of local equivariance. Two tokens T {-1,+1) and d = 6. Local equivariance means that the "flux" to transition neighbor is the negative of the flux of transitioning from that neighbor back.
[Ad)
where Q is a reference measuring whose support covers the support of p and F r= and where we used that Fo, Fr are constants. The above loss is tractable but we can bound it by loss that is computationally more effi- cient. To do so, we use an auxiliary free energy network F : R-- + R with parameters 6. Note that F is a one- dimensional function and therefore induces minimal addi- tional computational cost. As before, let the operator KÇ" be defined as:
MCMC scheme with learned transport. Most of the time, MCMC schemes are formulated as discrete time Markov chains. Hence, we first describe how they can be formu- lated as CTMCs. For every fixed 0 S / S 1, an MCMC scheme for the distribution Ps is given by a stochastic matrix Me x), i.e. Me x) 2 with Eves Me x) = 1. These are constructed to satisfy the corresponding detailed balance condition:
(.p = 0,U,(1) El y) Pc(y)
(17)
Pc(x))
Me(y,2)pa(l) Mc(2,y)pe00
(18)
yES
Proposition 6.1. For any reference measure Q the PINN- objective defined by
We can convert this into an annealed CTMC scheme by only accepting updates with a probability scaled by a parameter
E, 2 0:
L(0, oo t) = Esn-Unif = Q.
0
QCMMC (v, = (,,,(,,
if
has unique minimizer (0*, )²) such that Q! satisfies the KFE and F." = Fe is the free energy. Further, this objective is an upper bound to the cor-variaece divergence:
x)) yf==
The rate matrix Q.CM satisfies equation 18 so that the second term of the K--operator (see equation 10) vanishes:
Clog var (0:t) S 1 L(0,0;t)
OMMOM (x. Pc(V)
-Q (g.z)=0 =0
MCMC
In particular, if L(0, Φ; t)= = O, then also Llog-sor (l)))) = 0 and the variance of the IS weights is zero.
p((z)
WES
where we used (4). This implies that the rate matrix
A proof can be found in Appendix C. Note that we can easily minimize the PINN objective via stochastic gradient descent (see Algorithm 2). It is "off-policy" as the reference distribution Q is arbitrary. This objective can be seen as the CTMC equivalent of that in (Máté& & Fleuret, 2023; Albergo & Vanden-Eijnden, 2024; Tian et al., 2024; Sun et al., 2024).
of x) + Q QMCM
will have the same PINN loss and the same IS weights for the same trajectories because the X² remains unchanged (Note that while the RND in (15) is the same, the path measures do change). Specifically, this means that we can sample and compute the weights via Algorithm 1. The parameter 6, controls "how much local MCMC mixing" we want to induce. Note that the IS weights can be used both for reweighting at the final time or in addition to resample
7. Adding Annealed IS and SMC
is possible to add an arbitrary MCMC scheme to the above dynamics. This effectively combines an "unlearned"
5

---

LEAPS: A Discrete Neural Sampler
Algorithm LEAPS Sampling
for jumps to neighbors, i.e. (y x) 0 if N(x). One then neural network Of represented by the function
1: Require: N time steps, M walkers, model F2, replay buffer B3 MCMC kernel M,, density Pt. coeff. 2 0,
resample thres. 0 S ô - 1
0 : (R (BN-1)d x x
2: Init: X ~ Pu. Ab = 0
(m = 1,............... .,....M)
3: Set h = =
4: for = 0 N- do
i.e. the neural network is given by the function x -> (-12r)
5: for m = 1 M do
6:
X² Xi"..X(") with prob. he, else X!"
that returns for very dimension a value for every token different from Ty We then parameterize a rate matrix via
7:
ith ~ (1x,=y hQt Xt))
8:
4hth = Ap + +CC (X."")
if y & N()) else if yj + æj (19) if =
9:
end for
Viw Σ (T,
10:
+-t+A
11:
if (A,)(A.) < then
12:
X₁ resample(X., At) (m 1,...,,100 M)
13:
A, 0
(m = 1, 1....,,,.. M)
This parameterization is commonly used in the context of discrete markov models ("discrete diffusion models") (Campbell et al., 2022; 2024). With that, the operator K;; in (10) becomes:
14:
end if end for
15:
16: Optional: Store {(X,", ,)),t)m,, m in
(x)((x)) + (.U((r)
the walkers along the trajectories. connecting it to the SMC literature (Doucet et al., 2001). We specify that one may want to do this whenever the effective sample size (ESS) (see Appendix G) drops below a threshold.
Qi ily) Pc(y) Pr pr(x)
Σ
The key problem with the above update is that it requires us to evaluate the neural network (z)) times (for every neighbor y). This makes computing KC" computationally prohibitively expensive. Hence, with the standard rate ma- trix parameterization, the proactive IS sampling scheme and training via the PINN-objective is very inefficient.
Generalization of AIS and SMC. For Q the above dynamics describe continuous formulation of AIS that can be simulated approximately via Algorithm 1. In particular, this means that the algorithm presented here is strict generalization of AIS and SMC (Neal, 2001; Doucet et al., 2001). Note that in the / + 0, - do limit, LEAPS would recover the exact distributions pr with AIS (i.e. even with Q = 0). Of course, this asymptotic limit not realizable in practice with finite number of steps, and the inclusion of Q allows sample much more efficiently while still maintaining statistical guarantees.
To make the computation of KCŞ efficient, choose to in- duce an inductive bias into our neural network architecture to compute KGς with no additional cost. Specifically, we introduce here the notion of local equivariance. A neural network G represented by the function
G S NRgd-
8. Efficient IS and Training via Local Equivariance
x + (G. (T, (()
is called locally equivariant the following condition holds:
We now turn to the question of how to make the above train- ing procedure efficient. Note that for small state spaces S we could rely on analytical solutions to the KFE (Camp- bell et al., 2022; Shaul et al., 2024). In many applications, though, the state space S is so large that we cannot store [S] elements efficiently in a computer. Often state spaces S are of the form S Tdd where T (I,..., NN is a set of N tokens. We use the notation r for token, i.e. element E T. One then defines notion of a neighbor of x, i.c. an element
(T, = G/ (([aa/,(SSapp (i = 1, 1.... d) where Swap(æ, (Z1 (21...., I1- T,41, Id))
In other words, the function Go gives the "flux of probabil- ity" going from x to each neighbor. Local equivariance says that the flux from to its neighbor is negative the flux from the neighbor to æ (see Figure 2).
(yI,..., Ma) that differs from x in at most one dimension (i.e. y4 # I; for at most one i). We denote as (ω)) the set of all neighbors of z. We then restrict functional form of the rate matrices to only allow
Therefore, every coordinate map F3 is equivariant with re- speet to transformations of the j-th input ("locally" equiv- ariant). Note that we do not specify how F, transforms for
6

---

LEAPS: A Discrete Neural Sampler
i A j under transformations of zy. This distinguishes it from "full" equivariance as, for example, used in geometric deep learning (Bronstein et al., 2021; Weiler & Cesa, 2019; Thomas et al.. 2018). We can use locally equivariant neural network to parameterize a rate matrix via:
sequence of vectors: (21: Wii...,,20 F.)) E (R as the
input the neural network. The specific designs discuss
here are based on two components: (1) locally-invariant prediction head given by
H : (RR (R
@ (G (T,],)))+
(20)
where [2]+ m2xx(2.0) 0) describes the ReLU operation. This representation is not restriction (see Appendix D for a proof):
i.e. 11 fulfills that Hi (i/Swaa(zz,,, (z)),,),))))))) - (ff((()) (i/x) for all
x, i. (2) This is combined with small network P. T R that we call token projector. The full network is then given by
Proposition 8.1 (Universal representation theorem). For any CTMC as in (19) with marginals Pt, there is a cor- responding CTMC with the same marginals Pt and a rate matrix that can be written as in for a locally equivariant function G.-
((T -( (I;)) (1()
In Appendix E, we verify that G? defined in this way is locally equivariant. We discuss specific instances of this design.
Note that this representation implies that rates only go from I to y or from y to x, and in that sense are kinetically favorable (Shaul al., 2024). Crucially, this representation allows to efficiently compute XCg in one forward pass of the neural network:
this paragraph for readability. Let W W W E R be a set of weight matrices with zero diagonal, i.e. Wa 0 for i = 1,... I,.... d. Further, let R + R be an activation function and wi, E IR* be a learnable projection vector for every token E T. Then define the map:
Multilayer perceptron (MLP). Let us set Cn
K. (,U(())
[G G Pr(y) Σ i=l.....
(z(
* (((T,)l)== wo (W 1=1
(EN(
With this, we can efficiently compute the proactive IS up- date A, and evaluate the PINN-objective. Therefore, this construction allows for scalable training and efficient proac- tivate importance sampling. We call the resulting algorithm LEAPS (Locally Equivariant discrete Annealed Proactivate Sampler). The acronym also highlights that we use a Markov jump process to sample (i.e. that takes "leaps" through space).
where aWVxx) denotes the j-th element of the vector ob- tained by multiplying the vector x with the matrix W and applying the activation function componentwise. This is a locally equivariant MLP with one hidden layer.
-aaally-euuvarriant attention (LEA) layer. Let us con- sider self-attention layer operating on keys ky ky (myy), queries Øj = q (x,)) and values vj (z(z)) each of which is function of element Zy. We define the locally equivariant attention layer then as:
Remark 8.2. It is important to note that with the above con- struction, we would need to naively evaluate (z( as often as d times for a single computation of KÇ. However, note that the sum goes over all neighbors over æ. Therefore, this can be a considered as computing a discrete gradient. Such ratios can often be computed efficiently, c.g. for many scien- tific and physical models it is often only 2x the computation compared to a single evaluation of pi (2:).
exp(k? 43) Σ exp(ki (o. gj)
GT( = (wr -2,) )T
It can be shown that this layer is locally equivariant if the queries 1j are independent of the sign of xy (i.e. (, = (-xx))) which can be easily achieved. By stacking across multiple attention heads, one can create a locally equivariant MultiHeadAttention (LEA) with this.
9. Design of Locally Equivariant Networks
remains be stated how construct locally equivariant neural networks question we turn to in this section. We will focus on three fundamental designs used throughout deep learning: Multilayer perceptrons (MLPs), attention layers, and convolutional neural networks. Usually, tokens are embedded as token vectors e, E Rº where Gja is the embedding dimension. We therefore consider the embedded
Hierchical local equivariance. Local equivariance is dif- ferent from "proper" cquivariance in that the composition of locally equivariant functions is not locally equivari- ant in general. Therefore, we cannot simply compose lo- cally equivariant neural network layers as we would do

---

LEAPS: A Discrete Neural Sampler
Input
Embedding
(:)
22
ho
h (z)
Convolution h2(21)
Const
Convolution
Convolution
ReLU
Convolution
(W)) * x
hut1 (z)
Token
Tokenn
Neural
h, (z)
A
Wx(
Tokan
Token Projector
P
Figure 3. Overview of locally equivariant convolutional neural network architecture.
with "proper" equivariant neural networks. In particular, the above MLP and the attention layers cannot simply be composed as their composition would violate the local equiv- ariance. This fundamentally changes considerations about how compose layers and how construct deep neural networks. We will now illustrate this for the case of convo- lutional neural networks.
Transport and sampling, Over the past decade there has been continued interest in combining MCMC and IS with learning transport maps. non-parametric version of this described (Marzouk et al., 2016), and parametric ver- sion through coupling-based normalizing flows was used study systems in molecular dynamics and statistical physics (Noé et al., 2019; Albergo et al., 2019; Gabrié et al., 2022; Wang et al., 2022). These methods were extended to weave normalizing flows with SMC moves (Arbel et al., 2021; Matthews et al., 2022). More recent research focuses on replacing the generative model with a continuous flow dif- fusion (Zhang & Chen, 2022; Vargas al., 2023; Akhound- Sadegh et al., 2024; Sun et al., 2024). Our method is inspired by approaches combining measure transport with MCMC schemes (Albergo & Vanden Vanden-Eijnden, 2024; Vargas et al., 2024) and other samplers relying on PINN-based objectives in continuous spaces (Mâté & Fleuret, 2023; Tian et al., 2024; Sun et al., 2024).
Locally-qquivarian convolutional (LEC) network. To construct a locally equivariant convolutional neural net- work (LEC), we assume that our data lies on grid. A convolutional layer is characterized by a matrix W E (222-1))(2X-1) and its operation is denoted via k(W) a where & denotes the convolutional kernel with weights W. Here, we set the center of W to zero: Wa = 0 (i.e. such that corresponding location is effectively ignored). To stack such layers, we can make the output of the previous layer feed into the weights of the next layer:
Wi o (Ashi + b:) + oj
(i 1,..., L)
Discrete Neural samplers. The primary alternative what we propose is to correct using importance weights arising from the estimate of the probability density computed using an autoregressive model (Nicoli et al., 2020; Wu et al., 2019; McNauggton et al., 2020). However, the computational cost of producing samples in this case scales naively as O(d), whereas we have no such constraint a priori in our case so long the error in the Euler sampling scheme is kept small. Other work focuses on discrete formulations of normalizing flows, but the performant version reduces to an autoregres- sive model (Tran et al., 2019). Recent work has considered using CTMCs for sampling by parameterizing their evolu- tion operators directly via tensor networks (Causer et al., 2025) as opposed neural network representations of rate matrices here. Finally, discrete diffusion models have al- ready been explored in the context of unsupervised neural optimization (Sanokowski et al., 2024; 2025).
1441 (W,) x
(i (i= 1, L)
H (x)==ı,
where A, E Rdxxd. by Rdd, E R.d. are learnable ten- sors which operate on each coordinate independently (i.e. a 1x1 convolution) and a I + is an activation function to make it non-linear. With this construction, one can stack deep highly complex convolutional neural networks. Note that this convolutional neural network has two (separate) symmetries: it is geometrically translation equivariant and locally equivariant in the sense defined in this work.
10. Related Work
CTMCs. CTMCs (Campbell et al., 2022) have been used for various applications generative modeling ("discrete diffusion" models), including text and image generation (Shi al., 2024; Gat et al., 2025; Shaul et al., 2024; Sahoo et al., 2024) and molecular design (Gruver et al., 2023; Campbell et al., 2024; Lisanza et al., 2024). While here we use a RND for CTMCs running in reverse time, one recovers the loss functions these generative models considering a RND of two worward time CTMCs (see Appendix A).
11. Experiments
As a demonstration of the validity of LEAPS in high di- mensions, we benchmark it on models of statistical physics, namely the Ising model and the Potts model. We apply
8

---

LEAPS: A Discrete Neural Sampler
LEAPS on each model 15 15 lattice corresponding to ad 15x15 225 dimensional space. To construct Pt. use temperature annealing making the inverse temperature 3t a linear function of time.
10-1
soore)
poOW 910
A Ablation- Ising model. To establish the validity of our method also experimentally, we compare our results against ground truth of long-run Glauber dynamics, an efficient al- gorithm for simulation in this parameter regime. In Figure 8, one can see that LEAPS recovers the expected physical observables Further, in Figure 8 we compare three dif- ferent realizations of our method, one using LEA, and the other two using deep LEC that vary in depth. The deep LEC network performs best and we use it for all subsequent experiments.
Figure 5. LEAPS MCMC samplers on Ising model (DISCS benchmark (Goshvadi et al., 2023)). Note that the comparison has limitations: Function call for LEAPS neural network evaluations and energy calls for MCMC samplers. Further, ESS is measured differently for MCMC samplers.
Effective Sample
AIS Steps
0.7
LEAPS steps
importance weights to recover exact observables.
Effective Sample Size vs AIS Steps
0.5
LEAPS 100 steps
0.4
0.175
0.150
0.2
0.100
0.1
A LEAPS
0.075
0.050
10a
10a
10
10*
AIS steps
0.025
Figure 4. LEAPS AIS Ising model. Comparison across dif- ferent number of simulation steps for AIS. LEAPS is run for 100 steps. Note, however, that LEAPS requires neural net evaluations.
0.000
LEAPS
AIS steps
Benchmarks- Ising model. We then apply LEAPS on the parameters corresponding the corresponding critical tem- peratures of the Ising model, i.e. the "hardest" temperature to sample from. Here, LEAPS achieves a high effective sam- ple size (ESS) of ~ 68%. In Figure 4, we compare LEAPS against anncaled importance sampling (AIS) and see that AIS needs around 10t integration steps to achieve a similar ESS (vs. 100 integration steps from LEAPS). Next, we use the DISCS benchmark (Goshvadi al., 2023) to benchmark LEAPS against MCMC samplers (see Figure 5). The ESS of LEAPS is significantly higher, even after accounting for the number of function evaluations (NFEs).
Figure 6. LEAPS vs AIS on Potts model for critical temperature. Comparison across different number of simulation steps for AIS. LEAPS is run for 100 steps.
12. Discussion
In this work, we present the LEAPS algorithm to learn to sample from discrete distributions via continuous-tim Markov chains ("discrete diffusion"".. Crucially, we trans- lated the sampling problem into a neural network design problem. While we presented first locally equivariant neu- ral network designs, anticipate that future work will focus on building more scalable and more expressive locally equivariant network architectures. Another direction of fu- ture work will be to connect the ideas presented here with guidance and reward fine-tuning of generative model a problem strongly tied to sampling. Further, neural samplers like LEAPS can be used sample across whole family of distributions as opposed to only for single target.
Benchmarks Potts model. We also apply the LEAPS method on the Potts model following Goshvadi et al. (2023). LEAPS achieves a high ESS (~~ 20%) in only 100 integra- tion steps while AIS has lower ESS even for 10e integra- tion steps (see Figure 6). Also on the DISCS benchmark, LEAPS shows higher ESS than MCMC samplers.
Overall, the above results demonstrate the validity of LEAPS and that tostiiraining, it is able to generate real- istic samples with few simulation steps, while leveraging
9

---

LEAPS: A Discrete Neural Sampler
Impact Statement
Berendsen, H. J. C.,, Postma, J. P. M., van Gunsteren, W. F,
DiNola, A., and Haak, J. R. Molecular dynamics with coupling to external bath. The Journal of Chemical Physics, 81 (1)3366844)369 10 1984. ISSN 0021-9606. doi: 10.1063/1 18 URL https: //ooi. /tggg
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
10. 10../1631..... 448118.
Acknowledgements
Bronstein, M. M., Bruna, J., Cohen, T., and Veličković, P. Geometric deep learning: Grids, groups, graphs, geodesies, and gauges. arXiv preprint arXiv: 104. 13478, 2021.
We thank Eric Vanden Eijnden for useful discussions. PH acknowledges support from the Machine Learning for Phar- maceutical Discovery and Synthesis (MLPDS) consortium, the DTRA Discovery of Medical Countermeasurres Against New and Emerging (DOMANE) threats program, and the NSF Expeditions grant (award 1918839) Understanding the World Through Code. MSA is supported by Junior Fel- lowship at the Harvard Society of Fellows as well as the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelli- gence and Fundamental Interactions, http://iaffi.org/).
Camphell, A., Benton, J., De Bortoli, V., Rainforth, T., Deli- giannidis, G., and Doucet, A. A continuous time frame- work for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-28277 2022.
Campbell, A., Yim, J., Barzilay, R., Rainforth, T., and Jaakkola, Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:240..04997, 2024.
References
Causer, L.., Rotskoff, G. M., and Garrahan, J. P. Discrete gen- erative diffusion models without stochastic differential equations: a tensor network approach. Physical Review E. 111(2):025302, 2025.
Akdound-Sadegh, T., Rector-Brooks, J., Bose, A. J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y., et al. Iterated denoising energy matching for sampling from boltzmann densities. arXiv preprint XrXiv:2402.06/2/1 2024.
Doucet, A., de Freitas, N., and Gordon, N. J. (eds.). Se- quential Monte Carlo Methods in Practice. Statis- tics for Engineering and Information Science. Springer, 2001. ISBN 978- 4418--1--1-228--0.
Albergo, M. S. and Vanden-Eijndee E. Building normaliz- ing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2022.
doi: 10.1007/
978-1-4757-3437-99 URL https: //tp:i://d.... org/10. 1007 1007/978-1-4755 1-778-1-4777-3433
Duane, s., Kennedy, A. Pendleton, B. J., and Roweth, D. Hybrid monte carlo. Physics Let- ters B, 195(2):216-22) 1987. ISSN 0370-2693, doi: URL http:://ww..ccennedd.ceee com/ calence//attiille////////2222//
Albergo, M. S. and Vanden-Eijnden, E. Nets: A non-equilibrium transport sampler. arXiv preprint arXiv:2410.02711,22 2024.
Albergo, M. S., Kanwar, G., and Shanahan, P. E. Flow- based generative models for markov chain monte carlo in lattice field theory. Phys. Rev. D. 100: 034515, Aug 2019. doi: 10.1103/PhysRewD.1000 034515. URL https: /ttppp:://inkk...// 1103/PhysRevD. 100. 034515.
Faulkner, M. F. and Livingstone, S. Sampling algorithms in statistical physics: a guide for statistics and machine learning, 2023. URL 2208. 04751.
www.kiinn...........mmmmmm.....a aps www.ikkkas..o////////// owt..o////oool/
abs/
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.088. 2023.
Gabrié, M., Rotskoff, G. M., and Vanden- E. Adaptive monte carlo augmented with normalizing flows. Proceedings of the National Academy of Sci- ences, 119(10):e209420111 2022, doi: 10.1073/pnas. 2109420119. URL https: /www. www.ppaaas....... www..nag.aa..p.pp............... org/doi/ abs /1t//1.1100 1073/ //a///////s/../ 2109420113
Allen, M. P. and Tildesley, D. J. Computer simulation of liquids. Oxford university press, 2017.
Arbel, M., Matthews, A. G. D. G., and Doucet, A. An- nealed flow transport monte carlo. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, 18-24 Jul 2021.
Gat, I., Remez, T., Shaul, N., Kreuk, F., Chen, R. T., Syn- naeve, G., Adi, Y., and Lipman, Y. Discrete flow match- ing. Advances in Neural Information Processing Systems, 37:133345 133385, 2025.
10

---

LEAPS: A Discrete Neural Sampler
Gelfand, A. E. and Smith, A. F. M. Sampling-based approaches to calculating marginal densities. Journal of the American Statistical Association, 85(410):398- 409, 1990. ISSN 01621459, 1537274X. URL http: //www. wwt.ssstooo..........
10. 10.10188.41537--22-223999- URL https: : /doi.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209. S00::229903003,, 2022.
Goshvadi, K., Sun, H., Liu, X., Nova, A., Zhang, R., Grath- wohl, W., Schuurmans, D., and Dai, H. Discs: a bench- mark for discrete sampling. Advances in Neural Informa- lion Processing Systems, 36:79055-79066, 2023.
Marzouk, Y., Moselhy, T., Parno, M., and Spantini, A. Sam- pling via measure transport: An introduction. Handbook of uncertainty quantification, 1:2, 2016.
Matthews, A., Arbel, M., Rezende, D. J., and Doucet, A. Continual repeated annealed flow transport Monte Carlo. In Chaudhuri, K., Jegelka, S., Song, L., Szepes- vari, C.,, Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Conference on Machine Learn- ing, volume 162 of Proceedings of Machine Learn- ing Research, pp. 15196-15219, PMLR. 17-23 Jul 2022. URL http:s:://occedddeee www.hngg.............eem presss/ v1 #162/maatthews22..tttt
Gruver, N., Stanton, S. D., Frey, N. C., Rudner, T. G. J., Hotzel. I., Lafrance- Vanasse, J., Rajpal, A., Cho, K., and Wilson, A. G. Protein design with guided discrete diffusion. In seirty- Conference on Neural In- formation Processing Systems, 2023. URL ht https: // //openreview net forum? AfE
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. In Larochelle, H., Ranzato, M., Hadsell, R., Balean, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc., 2020. URL https: neurip.. ca httpppppr_fillleeppeer////////// acbbofee350af0066711111117766669 pdf.
ml
McNaughton, B., Miloševié, M., Perali, A., and Pilati, S. Boosting monte carlo simulations of spin glasses using autoregressive neural networks. Physical Review E, 101 (5):053312, 2020.
Metropolis, N., Rosenbluth, A. W., Rosenbluth, M. N., Teller, A. H., and Teller, E. Equation of state calcu- lations by fast computing machines. The Journal of Chemical Physics, 21(6):0087-1092 06 1953. ISSN 0021-9606. doi: 10.1003//.16991 14. URL https: / //dbi.org/10
Holderrieth, P., Havasi, M., Yim, J., Shaul, N., Gat, I., Jaakkola, T. Karrer, B., Chen, R. T., and Lipman, Y. Generator matching: Generative modeling with arbitrary markov processes. arXiv preprint arXiv:2410.20587 2024.
Máté, B. and Fleuret, F. Learning interpolations between boltzmann densities, 2023. URL https:
Jarzynski, C. Nonequilibrium equality for free en- ergy differences. Phys. Rev. Lett., 78:2690-2693, Apr 1997. doi: 10. 10.1109Physkevleett8...666 URL /ttps://link: wwt.kkapp..rg//ddeeeee www./o/g/ooooooooo http//.o111.001 1103/
org / abs 2301. 07388.
Neal, R. M. Annealed importance sampling. Statistics and computing, 11:125-139, 2001.
PhysRevIett 7690.
Nicoli, K. A., Nakajima, S., Strodthoff, N., Samek, W., Müller, K.-R., and Kessel, P. Asymptotically unbiased es- timation of physical observables with neural samplers. Phys. Rev. E. 101:023304 Feb 2020. doi: 10.1103/ PhysRevE. 101 1010033041... org/ /ww//0......... 111111111110110 1113/Physsev3...... 101. 023304.
Kahn, H. and Harris, T. E. Estimation of particle transmis- sion by random sampling. National Bureau of Standards applied mathematics series, 12:27-30, 1951.
URL https: //tp:k::///ikkkpppp
Lipman, Y., Chen, R. T. Q., Ham-H H., Nickel, M., and
Le, M. Flow matching for generative modeling. In The Eleventh International Conference Learning Represen- tations, 2023. URL https: / opwnreview.e..e//////p/pe....wwe www.eiieeeeeee. n forum? di PPP.MPMDTT....
Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltz- mann generators: Sampling cquilibrium states of many-body systems with deep learning. Science, 365(6457)xeaaw 147, 2019. doi: 10.1126/science. aaw 147. URL https: /www. sct eece..nn...... .rg/ http:///....///c/eeeeeoooo
Lisanza, S. L., Gershon, M., Tipps, S. W. K., Sims, J. N., Arnoldt, L., Hendel, S. J., Simma, M. K., Liu, G., Yase, M., Wu, H., Tharp, C. D., Li, X., Kang, A., Brackenbrough, E., Bera, A. K., Gerben, S., Wittmann, B. J., McShan, A. C., and Baker, D. Multistate and functional protein design using rosettafold sequence space diffusion. Nature Biotechnology, 2024. doi:
Nüsken, N. and Richter, L. Solving high- high-dimensional bamilton-jacobi-eelmee pdes using neural networks: per- spectives from the theory of controlled diffusions and measures on path space, 2023, URL https: //arxiv. aa/////s/0055. 05409.

---

LEAPS: A Discrete Neural Sampler
Richter, L. and Berner, J. Improved sampling via learned diffusions. arXiv preprint arXiv:2307.01 198, 2023.
on Machine Learning, volume 235 of Proceedings of Ma- chine Learning Research, pp. 48186-48210, PMLR, 21- 27 Jul 2024. URL https: ttttps::/paaee /ttp:://ppoooee /ttp//eeoooeeeeseeeeeeee .mlr. press Deess/v235/tian24...oo html.
Richter, L., and Berner, J. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview
Tran, D., Vafa, K., Agrawal, K. K., Dinh, L.,, and Poole, B. Discrete flows: Invertible generative models of discrete data, 2019. URL //atps:://rxiv.. org/ abs/ 1905. 10347.
Robert, C. P., Casella, G., and Casella, G. Monte Carlo statistical methods, volume 2. Springer, 1999.
Vaikuntanathan, S. and Jarzynski, C. Escorted free energy simulations: Improving convergence by re- ducing dissipation. Phys. Rev. Lett., 100:190601, May 2008. doi: 10.1103/PhyseevLet. 100.eee99...1001190000 URL https://lind PhysRevIett 100. 190601.
Sahoo, S. S., Arriola, M.. Gokaslan, A., Marroquin, E. M., Rush, A. M., Schiff, Y., Chiu, J. T., and Kuleshov, V. Simple and effective masked diffusion language mod- els. In The Tiihty-eighth Anmual Conference on Neural Information Processing Systems, 2024. URL https: //openreview wwt//frruuee wawaammmid---uaaaa LAaARAAMM
aps. .pw.arg/..aa................ www.ooog///////ooooooo/o//oo/e/t http:/oo/11111.11
Vargas, F., Grathwohl, W. S., and Doucet, A. Denoising diffusion samplers. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: / //openreview. net/ B
Sanokowski, S., Hochreiter, S., and Lehner, S. A diffusion model framework for unsupervised neural combinatorial optimization. arXiv preprint arXiv:2406.016: 2024.
Sanokowski, S., Berghammer, W., Ennemoser, M., Wang, H. P., Hochreiter, S., and Lehner, S. Scalable discrete diffusion samplers: Combinatorial optimization and sta- tistical physics. arXiv preprint arXiv:2502. 2025.
Vargas, F., Padhy, S., Blessing, D., and Nusken, N. Trans- port meets variational inference: Controlled monte carlo diffusions. In The Twelfth International Conference on Learning Representations: ICLR 2024, 2024.
Shaul, N., Gat, I., Havasi, M., Severo, D., Sriram, A., Hold- errieth, P., Karrer, B., Lipman, Y., and Chen, R. T. Flow matching with general discrete paths: A kinetic-optimal perspective. arXiv preprint arXiv:2412.03487 2024.
Wang, Y., Herron, L., and Tiwary, P. From data to noise to data for mixing physics across temper- atures with generative artificial intelligence. Pro- ceedings of the National Academy of Sciences, 119(32):2200366611,, 2022. doi: 10.1073/pnas. 2203656119. URL https: http://www.p..ttt//////.//.s pwwp.org/ds.......... abs/10 11111.10000 1073.1077/pss.s 2203656119.
Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. Simplified and generalized masked diffusion for discrete data. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL https: / ///penrevvew.nng net /a///forooe forum? www.mmidddxxssexoo 4cqSOfHt
Weiler, M. and Cesa, G. General e (2)-equivariant steerable cnns. Advances in neural information processing systems, 32, 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011. 1.200011114455 13456, 2020.
Wilson, K. G. Confinement of quarks. Phys. Rev. D, 10:2445-2459, Oct 1974. doi: 10.1 10.1103/PhysRevD. 10.2445. htts: /ttp///nnnn/p//////e wtw....p..oog///11100/
1103/ P133/hhysReDD.... 10. 2445..
Sun, J., Berner, J., Richter, L., Zeinhofer, M., Müller, J., Az- izzadenesheli, K., and Anandkumar, A. Dynamical mea- sure transport and neural pde solvers for sampling, 2024.
Wu, D., Wang, L., and Zhang. P. Solving statistical mechan- ics using variational autoregressive networks. Physical review letters, (2806000 2019.
URL https:: //t:://rxxa /ttpaaawwrgg/a//a///////../ooott abs 07873.
Zhang, Q. and Chen, Y. Path integral sampler: A stochastic
Thomas, N., Smidt, T., Kearnes, S., Yang, I.., Li, L.., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and and translation-equii equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.082) 2018.
control approach for sampling. In International Confer- ence on Learning Representations, 2022. URL https: / //openreview, net/ forum? cetdddd____uuu______
Tian, Y., Panda, N., and Lin, Y. T. Liouville flow impor- tance sampler. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 4 1st International Conference
12

---

LEAPS: A Discrete Neural Sampler
L= =0,
Histogram of M(x) M
Connected 2-pt Correlation G(r)
1.0
1.0
Ellectiv Sample Size (ESS)
0.8
0.0
0.4
-0.5
0.4
0.8
0.2
0.0
0.0
3000
5030
0.6
0.4
0.2
0.0
0.2
0.4
0.6
10
Training
Distance (r)
LEC (depth 5)
LEC (depth
LEA
no transport
truch
Figure Ablation experiments simple Ising model with L 15, 0.4, 0.7 with LEA and LEC architectures. Left: Effective sample size of LEAPS samplers over training. Green area denotes an annealing phase during training where is increased from to 1. Increasing
depth of LBC significantly improves performance. For locally equivariant attention (LEA) mechanism, we use 40 attention hcads, each with query, key, and value matrices of dimension 40x40. As such, there about 350,000 parameters in the model. In addition, the locally equivariant convolutional net (LFCC of depth three uses kernel sizes [5, 7, 15), while the depth version uses [3, 5. 7, 9, 15]. amounting around 100,000 parameters. Center: Difference in the histograms of the magnetization M (x) of configurations as compared to the ground truth set attained from a Glauber dynamics run of 25,000 steps, labeled as M', We denote by "no transprr" the case of using annealed dynamics with just the marginal preserving MCMC updates to show that the transport from Q: essential in our construction. Right: Comparison of the 2-point correlation function for the LEA and LEC samplers against the Glauber dynamics ground truth.
Connected 2-pt Correlation G(r)
Isapa
Histogram of Magnetizations
8 1
Distance (r)
Magnetizations
Figure Repeat experiments from Figure 8 for Ising model with critical temperature. Left: Comparison of the 2-point correlation function against the Glauber dynamics ground truth. Right: Histograms of the magnetization ((z) of configurations as compared to the ground truth attained from Glauber dynamics of 25,000 steps. LEAPS recovers the expected observables.
13

---

LEAPS: A Discrete Neural Sampler
ESS
ESS function
10
10
ESSINFE (leg scale)
ESS (lop scale)
10-4
10-4
randammnnit
tanmingthal
dincc
grog
blechgibis
dinais
Heaos
hanmsngbbll
beckgibis
dincc
randerrar
grog
giibls
dmsite
auailary Mappa
Figure 9. LEAPS vs MCMC samplers on Potts model (DISCS benchmark (Goshvadi et al., 2023)). Note that the comparison has limitations: Function call for LEAPS is neural network evaluations and energy calls for MCMC samplers. Further, ESS is measured differently for MCMC samplers.
Algorithm LEAPS training with optional replay buffer
Require: batch size, N time steps, model Gf, free energy net F. learning rate 1), replay buffer B3
1: while not converged do
2: if use buffer then
3:
(Xm Aln. (m))m Im......, SampleBatch(B)
4: else
5:
(X.m
A (m) mm)m.l..... +- Algorithm
6:
end if
7:
[(0 Φ) = B Σn 1K Pt.m (x - Ocm tm
8:
0 11 (oL(0,0)
9: ot 6+6 Φ n V O)
end while
14

---

LEAPS: A Discrete Neural Sampler
A. Proof of Proposition 5.1
Without loss of generality, we set the final time point to be = 1. We compute for bounded continuous function
④ X R:
Bx- Fu. [(X)]
n-+00 Xx Fu.o
lim E
( Xi/ns x,.,, xI-,,
I..0 X1/m: X1)
lim E xx Fu
& ((Xu, ////////////1
X X1)
puol (Xu, X Xi/ns X2/n: ..., X X X1)
Fuo (Xoth
X X1) ((Xo)
= lim By Fn n-t00
$ (Xo, 1/nn X2/m
v(11)
puo (X, (XglXath)
==0, 0=0.1/2,//n....
((X) (((o))
Qu(Xsthi
lim 8x~ Fule! n-+00
v(Xi)
exp
Σ (o(X..X.) (ota(XoxXs)
*.Xo+a#X, (otn Xatn)
II
(X(o)
(()x(tx..))) -((x:: X,)ds
Q,(X,, 9xx,-XX (X X,)
=B =Fx=
((X)
exp
II
xXx)
where we used the definition of the rate matrix Qt, Qú, the continuity of Q in and the faet that the left and right Riemann integral coincide. As was arbitrary, the RND is given by:
dErQ Q log
v. (X) = log(u((X))) - log(v(X)))+ / - (XX,,,X))ds -X,
Qx(Xx.X,)
Σ log
(X,,,,,,)
s,X, HX.
B. Proof of Theorem 5.2 Specifically, we use Proposition 5.1 compute
log dF dF P. (X) log(p((X.)) - -
d d....
Qx(X...)) (o(x,
Σ log
Qx(X, Xx))
Qx X²)
s,X xX,
F. - Fo - Ux,) (o0o + (x((x,,))) -x,xx, +
Σ
log
Qx (Xo)
Qx X,
8,X, XX.
Note that the function + UX,),) is piccewise differentiable function. Therefore, we can apply the fundamental theorem on every differentiable "piece" and get:
U.(X,) Uo(Xo) / D,Ui(Xi)ds+ +
Σ - U2(x5)
log Px(X5) Px(Xx)
Q,U,(X.)ds+ Σ
N,X, *X.
15

---

LEAPS: A Discrete Neural Sampler
Next, we can insert the above equation and get:
d P..0.
log
(X)
d Po.Q.
/
(x(X5,X.))
=F - Fo- Ux(x,) + Uo(o) +
Σ log
(X,,
8,X, *X,
Pa(X;)
Qa(Xx
=F - Fo-
(X,)ds
E
log
Pa(Xx) + (o( (x,,
Σ log
(x(
*XXxX
8,77 +Xx
Q, x,) Px(
=Fi - Fo-
(x,
-x) Qo(x,, 8, *X,
Σ log
Qo Xo Po (x 5
Σ (o(Xxs)) po(y) Pc(),))
=F: Fo-
(XU,(X))ds +
Qs(Xs, X,)ds +0
Pe(y) ds
=F - Fo +
(X.)(s
Σ
Qx(
(X,)
VES
=F Fo + A
where used the definition of A, in (11) and the definition of Q in (14). Note that for 1, get that
exp(Fi - Fo) = Expp, [m(=)] = [A + Fi Fo)] = [x.))((A:] [xxp
because FG, Fo constants. Therefore, in particular [l(exp Ax)] exp(Fo F() Z Note that assume that Zo 1 as we know Pu- Therefore, E[cxp(A,)] = Z₁. This proves (16).
C. Proof of Proposition 6.1
We can use the variational formulation of the variance as the minimizer of the mean squared error to derive a computationally more efficient upper bound, i.e. we can re-express for every s S 1:
log var (0: t)
=V [X1l
min E Ex- F F,ER
=t²
min 8, ER.
Ex-Q
- (x.)
<i²
min
Ex-e
(x.)
a, Fy ER.DDSE
1
=t2
min
E Be-
(x,) -
0, F.R.ER
where we used Jensen's inequality and denote with Q. the marginal of Q at time 8. We now arrive the result by replacing the above with the free energy network F. Further, note that the above bound tight for Q-almost every X:
(X.) O,F. C(Xo:)
a constant time s. However, this constant may depend on X.
16

---

LEAPS: A Discrete Neural Sampler
D. Proof of Proposition 8.1
Before we prove the statement, we prove an auxiliary statement about one-way rate matrices. We call a rate matrix Qr a one-way rate matrix if
# Q((1,y) = 0 for all æ + 2) =0 or Q.(x, y) for all + y
Intuitively, a rate matrix Qi is a one-way rate matrix we can always only go from x - or from y - 2. The next proposition shows that there no problem restricting ourselves to one-way rate matrices.
Lemma D.1. For every CTMC with rate matrix Qt and marginals qt, there is a one-way rate matrix Qs such that its corresponding CTMC X, has marginals q: if Xo ~ qo is initialized with the same initial distribution. Furhter, if Quly, x) = for y + x, then also Qe x)
Proof. Let Qt be rate matrix defining CTMC with marginals qr. Then
0.q((x)
Σ
Qe(a, (y)
VES
Σ
Qe(2, v))(
(x(x)
Σ
Qi(x, v) Qr(,, x)
(c(v) (v()
Qr(x, v) (((,,)) (x)
Qe(y. x) 2((x)
Σ
(x(y)
Qx(x, (c())
(:(
(o(1)
Qc(x, v) - ())..)) (c(x))
Qe(y. x)) - Q2( v) (:(v)
Σ
(e(v)
(x(x) (x())
gc(v))
Σ
Qi(x, (v) Q,(,,)))((()
u
Σ Q1 ())
where we defined
[Qnlu.z) - Qc(2, 3) (c(3))
+
9((2))
Q+(y,))=
Σ Q0(2,z) 24/2
y=
Note that
Q+(u,)) >0
(1())
Qi(y,z) -Qi(I v)
(1(x)
2((,, (((x)
(z.
((())
(Q(xs,)) - 2)
-2((,) x) (x(x) =0
(y)
Q. y) =0
Therefore, we learn that Q fulfils the desired condition and fulfils the KFE. Therefore, we have proved that we can swap out Q for Q and we will have an CTMC with the same marginals.
17

---

LEAPS: A Discrete Neural Sampler
Now, let return the proof of Proposition 8.1. Given rate matrix Q₁, the same marginals and define function:
one-way rate matrix Qt with
((]r) = x))
if ()) >
where y = Swap(a,,,,7
Qc
otherwise
By construction, holds that P,(T,i/2) is locally equivariant and that [F((T, Qe(..)). This finishes the proof.
E. Local equivariance of ConvNet
To verify the local equivariance, one can compute
G (T, 1|z) =Pp (e+) P. H((a))
(x:) - (e=)) (1lz) (i/z)
(P& (2ri) (e+))) )) (i(Swa(z,,,,)) :/Swap((,p i, i,),
where we have used the invariance of the projection head H (ilzr) to changes in the i-th dimension. This shows the local equivariance.
E.1. Universal representation of decomposition into locally-invariant prediction head
In this section, we answer the question: To what extend do we reduce the expressive power of the neural network by imposing the above conditions? To show universal representation theorem, define the polynomial basis features
P(x) (1,x*, _² xx x ....
i.c. we assume infinite width of the neural network output here.
Proposition F.1. Any continuous function F C TRd defined on compact set C Rdd be approximated as follows: For any there is a function F (F1, Fa) C RD that can be decomposed in
F((x) P(x)² H:(x)
for locally-invariany function H:C- yd (i.e. H1(x)) is independent of x;)) for some lo such that
sup [(() - F(+) <e
nec
In other words. for infinite width this representation universal.
Proof. Further, let F: C ded be continuous function defined compact set C. Further, write F == (F.,.... Fa) for the coordinate functions. Then fix j. By the Seie--Wiessssasss theorem, there polynomial p² (I1,..., (a) of the form
p (z
Σ
o x
ka.....k20
where Kk. F only for finite number of coefficients such that
(1)(2)) =ec
sup
pd (x)
Then we can re-express that polynomial via
1 (II...,, xz))
2U
k220,
1#1
Pk. (a.
4,20
=P(x,)" (H)
=(((x)
18

---

LEAPS: A Discrete Neural Sampler
Note that H (H Ha) is locally invariant. Therefore, conclude that F (F15' .. Fa) fulfills the desired
conditions:
sup
[[ - F(x))l S
Σ sup P((x) F.(x)) < Σ6 i=1 ZEC
sEC
As 0 was arbitrary, this finishes the proof.
F. Recovering loss functions for CTMC models via RNDs
We discuss here in more detail how the Radon-Nikodym derivatives (RNDs) presented in Proposition 5.1 relate to the construction loss function for CTMC generative models, also called "discrete diffusion" models. The connection lies in the fact that the loss function of these models relies on RNDs of two CTMCs running both in forward time. We can prove the following statement:
Proposition F.1. Let μs, v be two initial distributions over S. Let Qs, Q2 be two rate matrices. Then the Radon Gadon-Nikodym derivative of the corresponding path distributions in forward time over the interval [0, I] is given by:
dFFQ
log
(X)
dFuq
log dp dy (Xu) +
(xo, Xx,),s
log
Q',
8,X, #X.
where we sum over all points where N, jumps in the last term.
The proof the above formula is very similar to the proof of Proposition 5.1 and an analogous formula also appeared in (Campbell et al., 2024), for example. The above proposition allows us by compute the KL-divergence:
(PR |F <DKi(Fus puoe)
dFme =Ex-pmo log dFu. (X)
(x(,,
=Dxi(ullv) + B Bnw
(r(,, - Ci(x, X,)dt+
Σ log
Xx,,
4,X, *X.
-Dxi(Rllv) + 1| 1Qe(X, - 2((x,, Xi)l
Σ log (1((
+ F
Qi(X(,
1.X, #X.
=Dxi(ullv) + L . 1. Fu. (Qr(X,, xx(X,,,x,)
+ /
E Fue Σ
Qr(y; log
Qe(y: Xi)
) dt
Qi(y, X,)
D#X.
Qr(v; Xt) Q2 X,)
=Dxx(ullp) + +R N1 Ner
Σ
Cllv, Xt) Q.(v, X,) + Qi(y: X.)
V
19

---

LEAPS: A Discrete Neural Sampler
where we have used the data processing inequality in the first term. Having a parameterized model Qi Qi this leads the following loss:
L(0) =DKL
=DKz(ullv) + E Em X,x pua
Σ (1 X
(v. +C
where Q: is some reference process. The above recovers loss functions in the context of CTMC and jump generative models (Campbell et al., 2022; Gat et al., 2025; Shaul et al., 2024; Campbell et al., 2024) and Euclidean jump models (Holderrieth al., 2024, Section D.I.). Note that the above loss cannot be used for the purposes of sampling straight-forwan manner because we do not have access to samples from the marginals of the ground reference Fu.Q.
G. Numerical experiments
code for experiments can be found under the following link: https: httpsss//githubb.cmm/m/ cowm/maaleergg/eeapppeie
We briefly explain the Ising model here. The Potts model is similar and refer to details for the Potts model to the DISCS benchmark (Goshvadi et al., 2023). Configurations of the L L Ising lattice follow the target distribution
pi (1x) = ( where 3 is the inverse temperature of the system, Fi the free energy, and H(x) {-1, + R
is the Hamiltonian for the model defined as
H(x) JJ Σ TaE + Σ Dz. (3.3)
(21)
Here, is the interaction strength, (i denotes summation over nearest neighbors of spins æi, Ty and μu is the magnetic moment. Neighboring spins are uncorrelated at high temperature but reach critical correlation when the temperature drops below certain threshold. As observables, we use the histograms the magnetization M(x) = Σ; 2g compared ground truth, well as the two point connected correlation function
Geonn (r) [[]
E [r:) [[p++r]-
(22)
The latter measure of the dependency between spin values a distance lattice separation.
For evaluation, the self-normalized definition of the effective sample size such that, given the log weights A, associated to N CTMC instances, the ESS at time t in the generation is given by:
N Σm exp (4:)
ESS,
(23)
N 1
E exp (2A:)
20


Reviews:
[{"review_id": "DDxJVrIkXV", "reviewer": "Reviewer_YU3Q", "summary": "The authors propose an algorithm for sampling from discrete distributions by combining importance sampling with a learned continuous-time Markov chain (CTMC). They derive the importance weights via a Radon–Nikodym derivative for CTMCs and introduce locally equivariant neural architectures to ensure tractable learning of the rate matrices. Empirical results on a 2D Ising model demonstrate the sampling efficiency of the algorithm.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback on our work. We are glad to read that the reviewer considers our work a \"sound\" and \"novel\" contribution. Below, we provide answers to your questions and concerns.\n\nWe understand that the reviewer sees our experiments as the main area of improvement. Addressing this, we obtained new experimental results - as possible within the limited time frame. The results can be found under this anonymous link, in which we show that our method performs well at the critical phase of the Ising model and compares favorably against existing samplers (see more below): https://tinyurl.com/leapsicml\n\nWe would also like to emphasize that the we consider the main contribution of this work the introduction of a new paradigm for learning to sampling from discrete distributions. We therefore stress our methodological and theoretical contributions:\n- A new derivation of Radon-Nikodym derivatives of path measures\n- Proactive importance sampling as an IS scheme for CTMCs\n- A PINN-objective for discrete samplers by bounding the variance of the IS weights\n- Local equivariance as a symmetric constraint to enable scalable proactive IS\n- Locally equivariant neural network architectures such as convolutional neural networks\n\n**Additional benchmarks (1) - LEAPS vs. MCMC.** Thank you for sharing the DISCS benchmark [1] for discrete samplers with us. To use this benchmark, we adapt the parameters of the Ising model to match theirs, i.e. we use the parameters for the critical temperature for $\\beta=0.4407$. Note that this makes the problem harder (both for traditional methods and for LEAPS). We then ran benchmarks of LEAPS vs other discrete samplers (including the ones in the DISCS benchmark). The results can be found in figure 1 in the google drive folder. As one can see, LEAPS achieves an effective sample size (ESS) of $\\sim 70$% vs ESS$<0.1$% of MCMC-based samplers. This shows how LEAPS effectively converts a simple distribution into a highly complex distribution that can only be sampled with many MCMC runs. In figure 3, we also show that the samples recover known physical observables of the Ising model. Finally, we note that there are significant difficulties in creating a fair assessment between non-equilibrium (LEAPS) and equilibrium methods (MCMC):\n- The ESS is measured differently  in DISCS because it focuses on equilibrium MCMC methods, rather than non-equilibrium sampling  methods such as LEAPS.\n- The ESS quantities for non-equilibrium methods such as LEAPS do not take into account the number of steps used for annealing (i.e. for simulation of the continuous-time Markov chain). To account for this, we also plot the same quantities normalized by the number of function evaluations in Figure 1 (ESS/NFEs). However, LEAPS does *not* require evaluations of the energy at inference time but only of the neural network, while MCMC methods require energy evaluations (see metrics used in DISCS [1]). The best way we account for that now would be to normalize by the NFEs as measured in neural net evaluations. However, note that this does not measure efficiency or wall clock time. The main goal of LEAPS is to solve the sampling problem, not efficiency.\n\nWe hope that our new experiments showcase the efficacy of the LEAPS algorithm, while also explaining the difficulty in benchmarking.\n\n**Additional benchmarks (2) - LEAPS vs. AIS.** We also benchmark LEAPS against annealed importance sampling (AIS). AIS is a standard non-equilibrium sampling technique. Note that here, ESS is measured in the same way. In Figure 2, we plot the number of AIS steps vs the ESS of AIS. As one can see, even with $100k$ simulation steps, AIS only achieves an ESS of $<30$% vs an ESS of $69$% of LEAPS with 100 simulation steps. We note that LEAPS is a true extension of AIS (as well as SMC) and we can run AIS with the learned transport without changing the weights. However, for the purposes of benchmarking we turn off the AIS (here and above) to show that the learnt transport via the neural network is what enables the performance boost.\n\n**Other benchmarks.** Thank you for bringing up other possible benchmarks. We note that combinatorial optimization does not aim to faithfully recover a distribution or any observable thereof, where importance weights could be used. This puts our method fundamentally at a disadvantage. Please let us know if there are other benchmarks that come to mind that we can run to convince you of the potential of the method. \n\n**References.** Thank you for highlighting the references that we included in the updated draft of our work. We have benchmarked against the suggested methods (see above).\n\nWe thank the reviewer again for the insightful response. If these results have addressed your concerns, we would greatly appreciate any increase in score.\n\n[1]. Goshvadi K, Sun H, Liu X, et al. DISCS: a benchmark for \ndiscrete sampling[J]. Advances in Neural Information Processing Systems, 2023, 36: 79035-79066.", "claims_and_evidence": "Most of the theoretical claims in the paper are supported by clear proofs or sound intuitions. However, while the empirical results on the 2D Ising model are promising, further evidence on additional benchmarks or experiments could strengthen the overall support for the method’s general applicability. (Please see \"Experimental Designs Or Analyses\" part)", "methods_and_evaluation_criteria": "The proposed methods are well-suited to the problem at hand. The derivation of importance sampling weights via the Radon–Nikodym derivative is both theoretically sound and intuitively reasonable, and the parametrization of the rate matrix aligns with common practices in discrete Markov models. Moreover, the use of effective sample size (ESS) as the evaluation criterion in the 2D Ising model experiment is standard for assessing sampling efficiency, providing a clear metric for evaluating the algorithm’s performance.", "theoretical_claims": "I checked the proof of Proposition 5.1 and Theorem 5.2 roughly. Both proofs appear to be logically structured and correct.", "experimental_designs_or_analyses": "There are two main concerns regarding the experiment part of this paper:\n\n\n1. Insufficient benchmarks: Considering the many effective discrete sampling algorithms proposed in recent years, the paper would be more convincing if the authors compared their method against these state-of-the-art techniques. For example, they could refer to [1], which presents a set of advanced discrete samplers and provides useful comparisons between these samplers.\n\n\n2. Lack of more experiments: The evaluation is limited to the Ising model. To more comprehensively demonstrate the algorithm’s effectiveness, the authors should consider additional experiments on other tasks, such as other classical graphical models, combinatorial optimization problems, or generative tasks. The settings and benchmarks in [1] could serve as a useful guide for expanding the experimental section.\n\n\nReferences:\n\n\n[1]. Goshvadi K, Sun H, Liu X, et al. DISCS: a benchmark for discrete sampling[J]. Advances in Neural Information Processing Systems, 2023, 36: 79035-79066.", "relation_to_broader_scientific_literature": "The paper provides its contributions to discrete sampling problems. In particular, it builds on prior advances in annealed importance sampling, sequential Monte Carlo methods, and recent discrete diffusion models that employ CTMCs. The introduction of locally equivariant neural architectures for parameterizing the rate matrix extends existing ideas from neural parameterizations in discrete Markov models.", "essential_references_not_discussed": "It would be a good idea if the authors could discuss the family of newly proposed discrete sampling methods mentioned in [1], like Locally Balanced [2], Gibbs with Gradients [3],  Path Auxiliary\nSampler [4], Discrete Metropolis Adjusted Langevin Algorithm [5],\nDiscrete Langevin Monte Carlo [6].\n\n\n\n\nReferences:\n\n\n[1]. Goshvadi K, Sun H, Liu X, et al. DISCS: a benchmark for discrete sampling[J]. Advances in Neural Information Processing Systems, 2023, 36: 79035-79066.\n\n[2]. Zanella G. Informed proposals for local MCMC in discrete spaces[J]. Journal of the American Statistical Association, 2020, 115(530): 852-865.\n\n[3]. Grathwohl W, Swersky K, Hashemi M, et al. Oops i took a gradient: Scalable sampling for discrete distributions[C]//International Conference on Machine Learning. PMLR, 2021: 3831-3841.\n\n[4]. Sun H, Dai H, Xia W, et al. Path auxiliary proposal for MCMC in discrete space[C]//International Conference on Learning Representations. 2021.\n\n[5]. Zhang R, Liu X, Liu Q. A Langevin-like sampler for discrete distributions[C]//International Conference on Machine Learning. PMLR, 2022: 26375-26396.\n\n[6]. Sun H, Dai H, Dai B, et al. Discrete langevin samplers via wasserstein gradient flow[C]//International Conference on Artificial Intelligence and Statistics. PMLR, 2023: 6290-6313.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper is highly motivated, contributing to the challenging and important topic of discrete sampling, which has significant implications in statistics and machine learning.\n\n2. The derivation of the proposed algorithm is intuitive and theoretically sound, and the approach is novel to me.\n\n\nWeaknesses:\n\nMy main concern lies with the experimental section. Expanding the benchmarks and providing more extensive empirical comparisons, as discussed in the “Experimental Designs or Analyses” section, would strengthen the paper. With improvements to the experimental evaluation, the paper would make a stronger contribution.", "other_comments_or_suggestions": "There are several typos and unclear notations that need to be corrected:\n\n1. **Line 110 (right):**  \n   \"As we do not the normalization constant\" should be \"As we do not know the normalization constant.\"\n\n2. **Equation 12 (second line):**  \n   It appears that $\\mathbf{Y}$ might be a typo; \n\n3. **Proposition 6.1:**  \n   The sampling notation should likely be $s \\sim \\mathrm{Unif}[0,t]$ instead of $s \\sim \\mathrm{Unif}[0,1]$.\n\n4. **Equation 16 (second line):**  \n   The notation $y^j$ should be corrected to $y_j$.", "questions_for_authors": "Please see \"Experimental Designs Or Analyses\" and \"Other Comments Or Suggestions\"", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uF9ZpsmQv4", "reviewer": "Reviewer_YMxM", "summary": "The paper proposes locally equivariant functions, a compact neural parameterization of rate matrices in continuous-time Markov processes over discrete state spaces.  This effectively allows them to use recent \"discrete diffusion\" models as proposals in annealed importance sampling and sequential Monte Carlo over discrete spaces.  They train these via a variational upper bound on the log-variance divergence. The experiment with a 2D Ising model that has a computable ground truth shows an improvement in performance using the resulting proposals over proposals designed without any measure-transport or diffusion machinery.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback on our work. Below we address your questions and provide new information on additional experiments. \n\n**Updates to experiments: harder sampling problem and comparison benchmarks**\nWe obtained new experimental results and benchmarks as possible within the limited time frame. The results can be found at this anonymous link, in which we show that our method performs well at the critical point of the Ising model and compare it extensively against other methods in this hard regime (see more below): https://tinyurl.com/leapsicml\n\n\n**Additional benchmarks (1) - LEAPS vs. MCMC.** We use the DISCS benchmark [1] for discrete samplers to compare LEAPS to various MCMC samplers. We also adapt the parameters of the Ising model to the critical temperature for $\\beta=0.4407$. Note that this makes the problem harder (both for traditional methods and for LEAPS). The results can be found in figure 1 in the google drive folder. As one can see, LEAPS achieves an effective sample size (ESS) of $\\sim 70$% vs an ESS of $<0.1$% of MCMC-based samplers. This shows how LEAPS effectively converts a simple distribution into a highly complex distribution that can only be sampled with many MCMC runs. In figure 3, we also show that the samples recover known physical observables of the Ising model. Finally, we note that there are significant difficulties in creating a fair assessment between non-equilibrium (LEAPS) and equilibrium methods (MCMC):\n- The ESS is measured differently  in DISCS because it focuses on equilibrium MCMC methods, rather than non-equilibrium sampling  methods such as LEAPS.\n- The ESS quantities for non-equilibrium methods such as LEAPS do not take into account the number of steps used for annealing (i.e. for simulation of the continuous-time Markov chain). To account for this, we have also plotted the same quantities normalized by the number of function evaluations in Figure 1 (ESS/NFEs). However, LEAPS does *not* require evaluations of the energy at inference time but only of the neural network, while MCMC methods require energy evaluations (see metrics used in DISCS [1]). The best way we account for that now would be to normalize by the NFEs as measured in neural net evaluations. However, note that this does not measure efficiency or wall clock time, which would highly dependent on implementation and is usually not measured in the literature on neural sampling methods. The main goal of LEAPS is to solve the sampling problem, not efficiency.\n\nWe hope that our new experiments showcase the efficacy of the LEAPS algorithm, while also explaining the difficulty in benchmarking LEAPS vs traditional equilibrium methods.\n\n**Additional benchmarks (2) - LEAPS vs. AIS.** We also benchmark LEAPS against annealed importance sampling (AIS) [2]. AIS is a standard non-equilibrium sampling technique. Note that here, ESS is measured in the same way and the two methods are exactly comparable. In Figure 2, we plot the number of AIS steps vs the ESS of AIS. As one can see, even with $100k$ simulation steps, AIS only achieves an ESS of $<30$% vs an ESS of $69$% of LEAPS with 100 simulation steps. We note that LEAPS is a true extension of AIS and we can run AIS with the learned transport without changing the weights. In addition it can be turned into a sequential monte carlo sampler by performing resampling along the trajectory. However, here, for the purposes of benchmarking we turn off the AIS (here and above) to showcase that the learnt transport via the neural network enables the efficiency boost.\n\n**Comparison to other works cited:** Many of the works that we cite in our work are about the mathematical realization of discrete diffusion models and sampling algorithms for continuous variables. To the best of our knowledge, we do not know of another paper which learns a neural sampler for discrete distributions via Continuous Time Markov Chain. Therefore, we cannot benchmark against them. If you know of any, please let us know.\n\nFinally, we would like to emphasize that the we consider the the main contribution of this work the introduction of a new paradigm for learning to sampling from discrete distributions. We therefore emphasize again our methdological and theoretical contributions:\n- A new derivation of Radon-Nikodym derivations of path measures\n- Proactive importance sampling as an IS scheme for CTMCs\n- A PINN-objective for discrete samplers by bounding the variance of the IS weights\n- Local equivariance as a symmetric constraint to enable scalable proactive IS without losing representational power\n- Locally equivariant neural network architectures such as LE-convolutional neural networks\n\nThanks again for your valuable feedback. If these results have fulfilled your questions and concerns, we would greatly appreciate any increase in score.\n\n[1]. Goshvadi K, Sun H, Liu X, et al. DISCS: a benchmark for \ndiscrete sampling[J]. Advances in Neural Information Processing Systems, 2023, 36: 79035-79066.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Proofs are in the appendix and were not checked, but the theorems and propositions themselves do look sensible to me.", "experimental_designs_or_analyses": "There was only one experiment, and its design looks fine.  In the course of reviewer discussion, the authors have added additional benchmarks and experiments, which appear quite thorough.", "relation_to_broader_scientific_literature": "The paper contextualizes itself adequately within the literature.", "essential_references_not_discussed": "I could not tell the authors if they were missing a discrete diffusion paper.", "other_strengths_and_weaknesses": "The paper's main weakness is that it gives exactly one experiment in which the only baseline against something other than a locally-equivariant sampler is \"no transport\".  Could methods from other works cited not be compared to the present method?\n\nAfter review the paper has become significantly stronger through the addition of more experiments.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NNQGKqMqF8", "reviewer": "Reviewer_esu2", "summary": "The goal of this paper is to draw samples from a distribution $\\rho_1$, known up to a normalization constant, over a discrete space.\n\nOne way to do so is to simulate a prescribed path of marginal distributions ${(\\rho_t)}_{t \\in [0, 1]}$ that ends with the desired target. \n\nTo do so, the authors introduce a generic process ${(p_t)}_{t \\in [0, 1]}$ driven by a Markov operator $Q_t$,\n\nand then consider the reweighted process $(w_t p_t)_{t \\in [0, 1]}$, \n\nwhere the reweighting is done by the probability distribution $w_t(x) \\propto \\mathbb{E}_{A_t | X_t = x}[\\exp(A_t)]$. \n\nThis reweighted sampling process simulates exactly the prescribed path of marginals, if either:\n\n1. the Markov operator $Q_t$ verifies (Eq 7): $\\partial_t \\rho_t(x) = \\sum_{y \\in S} Q_t(x, y) \\rho_t(y)$\n\n2. the log-unnormalized weights $A_t$ verify (Eq 11): $\\partial_t A_t = (1 / \\rho_t(X_s)) \\big( \\partial_t \\rho_t(X_s) - \\sum_{y \\in S} Q_t(X_s, y) \\rho_t(y) \\big)$\n\nThe authors propose to enforce the equation on $Q_t$ by minimizing its violation, which can be written as a PINN loss. In that loss, the authors propose an computationally efficient parameterization of $Q_t$ based on equivariant neural networks.\n\n## update after rebuttal\n\nI appreciate that additional experiments were added, where the authors compare against baseline methods that were lacking in the original submission. Sampling in discrete spaces is a rapidly evolving field with a broad literature, ranging from statistical physics to mainstream machine learning, so finding all the relevant literature is not so obvious. It seems like the authors agreed to cite the references that were rightly brought up by many reviewers. I agree with reviewer 8trc that further clarity would be welcome in the experiments (using other metrics than ESS) and the text (claims about auto-regressive models and the log-variance loss). Yet overall, the paper is already an interesting contribution to the field. So I will maintain my positive score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback on our work. Below, we provide answers to your questions and concerns.\n\n**Additional benchmarks and experiments.** We obtained new experimental results. The results can be found under this anonymous link, in which we show that our method performs well at the critical phase of the Ising model and compares favorably against existing samplers (see more below): https://tinyurl.com/leapsicml\nWe use two benchmarks:\n1. **LEAPS vs. MCMC.** We benchmark LEAPS against a diversity of MCMC samplers on the critical temperature of the Ising model. The results can be found in figure 1 in the google drive folder. As one can see, LEAPS achieves an effective sample size (ESS) of $\\sim 70$% vs ESS$<0.1$% of MCMC-based samplers.\n2. **LEAPS vs. AIS:** We also benchmark LEAPS against annealed importance sampling (AIS), see figure 3. As one can see, even with $100k$ simulation steps, AIS only achieves an ESS of $<30$% vs an ESS of $69$% of LEAPS with 100 simulation steps.\n3. \nPlease see our response to YU3Q for a thorough discussion of these results. \n\n**Questions.** We address your remaining questions and comments in the following:\n- ```What does \"no transport\" refer to?```: This refers to not using the neural network at all. We show via this ablation that it is truly the neural network (and not the AIS) that enables the sampling from the distribution with only a few samples. We ran additional benchmarks on this that we explained above (AIS vs LEAPS) not using AIS at all for LEAPS.\n- ``LEAPS uncorrected vs. LEAPS``: For the simulation of the continuous-time Markov chain, we always use the rate matrix $Q_t$ represented by our network. Due to the local equivariance, we get importance weights for free. \"LEAPS uncorrected\" means that we do not use the importance weights (at perfect training, this would not be necessary). \"LEAPS corrected\" means that we use the importance weights to reweight samples towards the correct distribution. Both parts are essentials as our ablations show.\n- ```More datasets and benchmarks:``` We provided other benchmarks, in particular against a series of MCMC samplers. We further actively run experiments on the Potts model, another model from statistical physics. The only other established benchmark for discrete sampling is combinatorial optimization. However, we note that metrics used for combinatorial optimization do not aim to faithfully sample from a distribution or any observable thereof, where importance weights could be used. Rather it tries to find low energy states. This puts our method fundamentally at a disadvantage.\n- ```Clarifying equivariance```: We note that for an input $x$ the neural network outputs $F_t^\\theta(y^i,i|x)$ for every $y^i,i$ for token $y^i$ and index $i$ (this is also true for discrete diffusion models, it returns a rate per neighbor). Therefore, this requires only one forward pass. Thank you for bringing it up that this needs further clarification. We will highlight the distinction between input and output of the neural network more in future versions.\n- ```Connection with NETS/ALD```: Here is the analogy. Because ALD is like running Langevin dynamics locally on some $\\rho_t \\sim e^{-U_t}$ for fixed $t$ and then iterating $t \\rightarrow t + dt$, so too can we add any local MCMC kernel to the dynamics of the CTMC that preserves detailed balance. This would \"simulate approximately the prescribed path\" just like Langevin does in the continuous setting. Does that make sense? We will add a remark about this in the text.\n- ```Question about the PINN loss```: The goal of LEAPS (and of all importance sampling schemes) is that the weights have as low variance as possible. In our case, this would mean that $A_t$ is just a constant in time and space. In other words, this means $\\partial_sA_s=0$. As we show in the paper, this is equivalent to $\\mathcal{K}_s\\rho_s=\\partial_sF_s$. Thank you for bringing this up. We will highlight this in future versions of our work.\n- ```Interpolation schemes```: The interpolation scheme is given by a time-dependent Ising model for coupling constant $J_t=tJ_1$ where here $J_1=1$. As you point out, this is equivalent to temperature annealing, i.e. $\\rho_t(x)\\propto \\exp(-tU_1(x))$. As also pointed out, discrete diffusion models use arithmetic means (or probabilistic mixtures) between a uniform (or mask) and the target distribution. In principle, one could use such an interpolation. We focuses on temperature annealing because it is common in the sampling literature and it allows to sample from the Ising model for higher temperatures on the fly (by simply stopping at earlier time points). Therefore, this is a natural and physical way. We have explored various annealing speeds and schedules (i.e. time parameterizations) for the annealing that is happening during training and found the schedule of $t\\to \\sqrt{t}$ to perform best.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The evaluation is encouraging.", "theoretical_claims": "Not in detail but they seem coherent with previous literature.", "experimental_designs_or_analyses": "I looked at the reported results which are encouraging. I have a few questions, however:\n\n- **What does \"no transport\" refer to?**\n\nIn Figure 4, what does \"no transport\" refer to? I read in the caption that it denotes the case \"of using annealed dynamics with just the marginal preserving MCMC udpates to show that the transport from Q_t is essential\". I don't understand what this corresponds to exactly in the text. Could the authors provide more detail?\n\n- **More datasets**\n\nThe experiments on the Ising model, while principled (because the ground truth is known) and encouraging, are limited. Did the authors benchmark against other target distributions, such as quantized MNIST as in [1].\n\n- **LEAPS uncorrected vs. LEAPS**\n\nIn Figure 1, the authors contrast LEAPS (perfect transport and reweighting) with uncorrected LEAPS (reweighting only). Which $Q_t$ is used for the uncorrected LEAPS? Because the distinction between LEAPS and uncorrected LEAPS is central to the authors' paper, it would be useful to see both these results in Figure 4 as well. \n\n\n[1] Discrete Langevin Sampler via Wasserstein Gradient Flow. Haoran Sun, Hanjun Dai, Bo Dai, Haomin Zhou, Dale Schuurmans. 2022.", "relation_to_broader_scientific_literature": "This paper relates to the sampling literature.", "essential_references_not_discussed": "Related work is discussed.", "other_strengths_and_weaknesses": "-", "other_comments_or_suggestions": "-", "questions_for_authors": "1. **Confused about motivation of the equivariant parameterization of the rate matrix**\n\nIn section 7, the authors first write an equation involving $Q_t^{\\theta}(y^i, i | x)$ and $Q_t^{\\theta}(x^i, i | y)$. If I correctly understand, each of these terms should be evaluated for each neighbor $y$, so roughly 2 |N(x)| times. \n\nAfter enforcing equivariance, the equation now only involves $F_t^{\\theta}(y^i, i | x)$. As I understand it, this term should also be evaluated for each neighbor $y$ so |N(x)| times. What am I missing?\n\n2. **Connection with NETS**\n\nThe authors mention that their paper can be understood as an extension of [1]. The comparison is\n\n| NETS    | LEAPS |\n| -------- | ------- |\n| ALD  | CTMC |\n| ALD + perfect additional transport (Eq 16) |  CTMC + perfect rate matrix  |\n| ALD + reweighting (Eqs 10, 11) |  CTMC + reweighting |\n\nOne difference is that the vanilla ALD defines a specific, tractable Markov kernel, that does not require learning and simulates *approximately* the prescribed path. \nIn contrast, CTMC is defined by a general Markov kernel that does not necessarily simulate approximately the prescribed path. \n\nDo the authors have an idea of how to choose a CTMC that would be a closer comparison to ALD? Specifically, do they have an idea of what would be a default, tractable choice of $Q_t$, that does not require learning and simulates *approximately* the prescribed path? \n\n3. **Question about the PINN loss**\n\nThe PINN loss in Proposition 6.1. encourages $K_s \\rho_s = \\partial_s F_s$. Shouldn't it encourage $K_s \\rho_s = \\partial_s A_s$ from Eq 11.\n\n4. **Interpolation schemes**\n\nThe authors define a general interpolation in Eq 2. In Figure 4, the authors use an interpolation that is specific to the Ising model. In Figure 1, which interpolation do the authors use? In the text after equation 3, the authors say that when the initial distribution is uniform, then we get $\\rho_t \\propto exp(-t U_1(x))$ but this is the case when the interpolation is a geometric mean. In discrete diffusion models, the interpolation is an arithmetic mean between the start and end distributions. Can the authors clarify which interpolations they use? Also, have the authors tried different interpolations or have any thoughts on how that choice affects the learning of $Q_t$?\n\n[1] NETS: A Non-Equilibrium Transport Sampler. Michael S. Albergo, Eric Vanden-Eijnden. 2024.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OcF09WvPRn", "reviewer": "Reviewer_8trc", "summary": "This paper introduces a continuous-time diffusion sampler that operates along an annealed energy path in the discrete domain. The approach is trained using a PINN-based objective as an upper bound for the Log Variance Loss. The authors propose locally invariant neural network architectures for parametrization. The method is experimentally validated on a 15x15 Ising Model.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and valuable feedback on our manuscript. Below we itemize and address your comments and concerns. \n\n**Further experiments and benchmarks:**\n- We have included additional experiments, as possible within the limited time frame. Specifically, we run additional tests of the Ising model at the critical parameters you suggested, and benchmarked this setup against the samplers that reviewer YU3Q suggested (DISCS) and annealed importance sampling (AIS). The results are provided in the anonymous google drive https://tinyurl.com/leapsicml where we see that the LEAPS performs favorably in this setting. To save space to address your other concerns, **please see our response to YU3Q for a thorough discussion of these results.**\n\n**No experiments with locally equiv MLP**\n\n- We have tested the MLP architecture and it has performed significantly worse. The convolutional architecture is geometrically translation equivariant, which is known in deep learning to outperform MLPs if the data (e.g. the statistical physics models we consider) is symmetric. We therefore do not consider this a noticable benchmark.\n\n**Experimental details.** Thank you for letting us know about missing details for replication. Here we supply an exhaustive list:\n- Learning rate=5e-3\n- Hyperparameter selection: We peform manual optimization of the neural network architectures and hyperparameters.\n- Batch size 256 walkers, stored in a replay buffer. New walkers are added to the buffer every 50 training steps\n- Training iterations: 30k\n- Ising model: periodic lattice.\n- External magnetic field: $\\mu=0$. \n- Comparison to theoretically available solutions: The theoretically available solutions are infinite grid sizes. We run Glauber dynamics that simulate the actual physics of the system for a sufficiently long time to recover the statistics that we compare against.\n\nWe will add the above details to the paper. Please let us know if you are missing other details.\n\n**Essential References Not Discussed:**\n- Thank you for pointing us to these citations. We will happily add them to the text when we are allowed to edit it. We note that the paper by Tian et al is cited already in multiple locations of the manuscript.\n- One distinguishing feature of our approach to autoregressive models is that we can draw samples with fewer than d steps where d is the dimension. For example, we achieved an ESS of 69% in $<100$ steps for the Ising model. An autoregressive could only achieve the same results in $d=L^2$ steps. Therefore, our effective ESS will be higher.\n\n**Non-locally invariant architecture clarification**\n\n- Thank you for bringing this up. As stated in section 7, a non-locally equivariant architecture would require $\\mathcal{O}(N*d)$ evaluations of the network for computation for $N$ the number of neighbors and $d$ the dimension. This number is prohibitively large for any $d$ that is interesting for applications, e.g. $d=225$ in the Ising experiment. We do not consider a need to give experiments to show that this is unfeasible. Of course, one could *train* a neural network architecture to be locally equivariant via an additional loss but then lose the computation of *exact* importance weights, which is one of the main goals of our paper.\n\n**Log-variance.** \n- We thank the reviewer for bringing up the log-variance loss. Evaluating the log-variance divergence is computationally much more expensive here, which is why we do not benchmark against it. Specifically, the log-variance divergence is $K$ times more expensive in memory, where $K$ is the number of simulation steps. The reason for that is as follows: The log-variance loss is not valid when sampled pointwise (it requires the whole trajectory). In contrast, the PINN loss can be evaluated pointwise. However, the PINN objective generally requires a “discrete divergence” computation (the $\\mathcal{K}_t$ operator). We get rid of the cost of this discrete divergence with the locally equivariant networks. Note that this is not necessarily true in the continuous case, where the log-var can avoid the computation of a divergence via Ito-integrals.\n\n**Proposal distribution.**\n- We use the model itself as a proposal distribution.\n\nFinally, we like to emphasize that the we consider the the main contribution of this work the introduction of a new paradigm for learning to sampling from discrete distributions. We therefore stress our methdological and theoretical contributions:\n- A new derivation of Radon-Nikodym derivatives of path measures\n- Proactive importance sampling as an IS scheme for CTMCs\n- A PINN-objective for discrete samplers by bounding the variance of the IS weights\n- Local equivariance as a symmetric constraint to enable scalable proactive IS\n- Locally equivariant neural network architectures such as convolutional neural networks\n\nThank you again for your valuable review. We hope that this addresses your questions. We would accordingly appreciate any increase in rating you see fit.", "claims_and_evidence": "While the approach is theoretically sound, the empirical evidence is limited. The study includes only one experiment on the 15x15 Ising Model and lacks ablation studies for the proposed methods. Additionally, crucial details for reproducing the experiments are missing. Although a locally invariant MLP architecture is proposed, no experiments using this architecture are presented.", "methods_and_evaluation_criteria": "The model is evaluated based on effective sample size, two-point correlation function, and magnetization.\nThis makes sense, however, an estimation based on log Z, ELBO, internal Energy and entropy could be added, especially when $\\mu = 0$, where these values are theoretically available.", "theoretical_claims": "Theoretical claims were partially checked.", "experimental_designs_or_analyses": "The experiments are not reproducible due to several reasons:\n- No information is provided about the hyperparameter selection process or the final hyperparameters used, such as learning rate and the number of samples during training or evaluation.\n- It is unclear whether the Ising Model is defined on a periodic or non-periodic grid.\n- The value of $\\mu$ in Eq. 18 is not specified, and there is no comparison to theoretically available solutions for the Free Energy, Entropy, and Internal Energy when $\\mu = 0$.", "relation_to_broader_scientific_literature": "The paper provides a broader relation to scientific literature but omits some relevant papers, particularly those related to discrete domain diffusion samplers.", "essential_references_not_discussed": "- [1] and [2] should also be cited in this context. [1] is seminal work in discrete samplers and [2] should be cited together with Nicoli et al., 2020.\n- [3] and [4] are highly relevant discrete diffusion sampler papers. The Annealed Noise Distribution in [1] resembles a discrete variant of the transport proposed in this paper. [4] applies diffusion samplers in the discrete domain in statistical physics.\n- [5] also proposes a PINN-based loss to learn a transport between a prior and target distribution in the continuous domain\n\n**References:**\n\n1. Wu, Dian, Lei Wang, and Pan Zhang. \"Solving statistical mechanics using variational autoregressive networks.\" Physical review letters 122.8 (2019): 080602.\n2. McNaughton, B., et al. \"Boosting Monte Carlo simulations of spin glasses using autoregressive neural networks.\" Physical Review E 101.5 (2020): 053312.\n3. Sanokowski, S., Hochreiter, S., & Lehner, S. A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization. In *Forty-first International Conference on Machine Learning*.\n4. Sanokowski, Sebastian, et al. \"Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics.\" arXiv preprint arXiv:2502.08696 (2025).\n5. Tian, Yifeng, Nishant Panda, and Yen Ting Lin. \"Liouville Flow Importance Sampler.\" *arXiv preprint arXiv:2405.06672* (2024).", "other_strengths_and_weaknesses": "**Strengths:**\n- The proposed framework and locally invariant architectures are interesting.\n\n**Weaknesses:**\n- Many ablation studies are missing, such as comparisons between LV loss vs. PINN loss and locally invariant networks vs. vanilla architectures.\n- The experiments are scarce and not reproducible with the available information.\n- There is no comparison to other discrete domain diffusion samplers or autoregressive samplers.", "other_comments_or_suggestions": "This paper has the potential to be very impactful, but it requires more experimental evaluation and detailed descriptions of the experiments.\nExperiments on Spinn Glasses as in [4] would also be nice to have.", "questions_for_authors": "- How does a non-locally invariant architecture compare to the proposed architectures?\n- How was the proposal distribution $\\mathbb{Q}$ chosen?\n- Why were $\\beta = 0.7$ and $J = 0.4$ specifically chosen? The correlation length seems quite small for this choice (see Fig. 4 right). This might have the consequence that the problem is not hard for this choice of parameters.\n- Why is there no comparison on the ising model, with $\\mu = 0$, $J = 1$ and $\\beta = 0.4407$ at the critical temperature where the correlation length does not decay and theoretical solutions can be used as a baseline?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes LEAPS, a novel discrete sampler that uses locally equivariant neural network architectures to parameterize the rate matrix of a continuous-time Markov chain (CTMC). The method is designed for efficient sampling from unnormalized discrete distributions. It integrates ideas from annealed importance sampling, non-equilibrium transport, and Radon–Nikodym derivatives over path measures. The experiments are conducted on a 2D Ising model.

All reviewers found the core idea compelling, particularly the use of importance sampling and locally equivariant networks within the CTMC framework. However, several reviewers noted that the empirical evaluation, limited initially to the 2D Ising model without any baselines, was too narrow. In response, the authors provided ESS comparisons on both the Ising and Potts models with several existing discrete samplers. While this significantly improves empirical support, relying solely on ESS is insufficient to demonstrate sampling performance. Moreover, the paper lacks ablation studies to isolate the contributions of each proposed component.

Overall, the paper presents a theoretically sound and promising approach to discrete sampling using locally equivariant neural networks. The methodology represents a valuable contribution to the field. However, due to the limited empirical evaluation, it remains difficult to assess the true impact of the proposed method. The authors are strongly encouraged to expand the empirical evaluation by including more comprehensive comparisons on different tasks, particularly on real-world datasets, and against a broader set of discrete sampling baselines.
Decision: Accept (poster)