Proximal Diffusion Neural Sampler
Authors: Wei Guo
Jaemoo Choi
Yuchen Zhu
Molei Tao
Yongxin Chen
Venue: ICLR 2026 Poster (2026)
Keywords: Neural sampler
proximal gradient descent
diffusion models
discrete diffusion models
cross entropy
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We propose a framework that enables efficient and robust training of neural samplers via proximal point method on the space of path measures.

Abstract: The task of learning a diffusion-based neural sampler for drawing samples from an unnormalized target distribution can be viewed as a stochastic optimal control problem on path measures. However, the training of neural samplers can be challenging when the target distribution is multimodal with significant barriers separating the modes, potentially leading to mode collapse. We propose a framework named **Proximal Diffusion Neural Sampler (PDNS)** that addresses these challenges by tackling the stochastic optimal control problem via proximal point method on the space of path measures. PDNS decomposes the learning process into a series of simpler subproblems that create a path gradually approaching the desired distribution. This staged procedure traces a progressively refined path to the desired distribution and promotes thorough exploration across modes. For a practical and efficient realization, we instantiate each proximal step with a proximal weighted denoising cross-entropy (WDCE) objective. We demonstrate the effectiveness and robustness of PDNS through extensive experiments on both continuous and discrete sampling tasks, including challenging scenarios in molecular dynamics and statistical physics. Our code is available at https://github.com/AlexandreGUO2001/PDNS.

Published as a conference paper at ICLR 2026
PROXIMAL DIFFUSION NEURAL SAMPLER
Wei Guo", Jaemoo Choi", Yuchen Zhu, Molei Tao", Yongxin Chen Georgia Institute of Technology {wei. jchoi843, yzhu738, mtao, yongchenj@@attech.edu
ABSTRACT
The task of learning diffusion-based neural sampler for drawing samples from an unnormalized target distribution can be viewed as a stochastic optimal con- trol problem path measures. However, the training of neural samplers be challenging when the target distribution multimodal with significant barri- separating the modes, potentially leading mode collapse. We propose framework named Proximal Diffusion Neural Sampler (PDNS) that addresses these challenges by tackling the stochastic optimal control problem via proxi- mal point method on the space of path measures. PDNS decomposes the learn- ing process into a series of simpler subproblems that create path gradually approaching the desired distribution. This staged procedure traces progressively refined path the desired distribution and promotes thorough exploration across modes. For a practical and efficient realization, we instantiate each proximal step with a proximal weighted denoising cross-entropy (WDCE) objective. We demonstrate the effectiveness and robustness of PDNS through extensive exper- iments on both continuous and discrete sampling tasks, including challenging scenarios in molecular dynamics and statistical physic.. Our code is available at httpps://itthue www../////ww.eeeeooocccoooo http///lelaaeeet000///P0..
INTRODUCTION
Sampling from an unnormalized Boltzmann distribution is a fundamental task across various domains, including computational statistics (Liu 2008 Brooks et al., 2011), Bayesian inference (Gelman et al.) 2013), statistical mechanics (Landau & Binder 20144) etc. Formally, we aim to draw samples from target distribution π on the state space X,, whose probability density or mass function specified by an energy function E X + IR and an inverse temperature 3 >> 0:
π (z)
1 -E(())
where Z Z:
(E(x) dr or Σe -BE(a)
(1)
Z
REX
Classical Markov chain Monte Carlo (MCMC) methods may suffer from slow mixing when in high
dimensions or when the target distribution multimodal (Latuszyński et al., 2025 He & Zhang 2025), motivating neural samplers that transport known reference distribution the target via score-based diffusion or normalizing flows (Gabrié et al. 2022 Zhang & Chen 2022).
Recently, diffusion-baaae neural samplers have been extensively studied in learning high-dimensional complex distributions (Zhang & Chen 2022 Vargas et al. 2023 PPillipp et al.) 2024 Havens et al. 2025 Liu al., 2025). These approaches typically parameterize the control of diffusion process and aim to transport a simple initial distribution toward a prescribed target distribution æ by solving a stochastic optimal control (SOC) problem. However, when the model initially captures only few modes in the early stages of training. the global loss tends to reinforce those modes, thereby concentrating probability mass to only several modes, resulting in mode collapse. This observation motivates replacing one-shot global minimization with sequence of gradual and constrained steps that improve control while preserving mode coverage.
In this paper, propose the Proximal Diffusion Neural Sampler (PDNS), unified framework of diffusion-basse sampling for both continuous and discrete domains, which leverages proximal
"Equal contribution. Joint mentorship. Corresponding author.

---

Published as a conference paper at ICLR 2026
iterations over path measures. Starting from given initial diffusion process, PDNS iteratively addresses a local sub-problem and updates the control of the dynamics so that the terminal marginal distribution progressively approaches the target distribution я. This proximal structure yields stable optimization and mitigates mode collapse multi-modal settings. As a one instantiation of PDNS, develop proximal algorithmic variants of cross-entropy (CE) and weighted denoising cross-entropy (WDCE) losses (Phillips et al., [024 Zhu et al. (025), which are widely adopted and computationally efficient. Finally, evaluate our PDNS in the multi-modal and high-dimensional benchmarks for both continuous and discrete target distributions. We compare against recent works, including the original WDCE variant, achieving sfate-of tetthehart art for many benchmarks. Our experimental results show the effectiveness and robustness of our methods.
Contributions Our key contributions are summarized as follows. In Sec. 2 by using a path measure formulation, we integrate continuous and discrete SOC-based samplers into the unified framework. In Sec. [ starting from this unified viewpoint, we demonstrate the problem of mode collapse in the training of diffusion neural samplers. Then, we introduce Proximal Diffussion Neural Sampler (PDNS), which performs proximal steps in path space to stabilize learning and mitigate collapse. 3. In Sec. we develop proximal weighted denoising cross-entropy (proximal WDCE) for both continuous and discrete settings as a concrete example of PDNS. We then outline principled strategies for selecting the proximal step size. 4. Finally, in Sec. 5 provide extensive experiments across continuous and discrete sampling tasks, demonstrating that PDNS improves stability, mode coverage. and overall sampling quality relative strong baselines.
Notations Let X be the state space and T > 0 be the terminal time. A path measure is a probability measure P.: :: [0, T] x X' + on the space of all paths, which is subset of all functions from [0, 2] to X P All path measures considered in the paper start from the same initial law μt, i.e., Po u.. We denote pref as reference path measure and denote its terminal marginal by v == PT The target distribution is π on the state space X' whose probability density function (p.d.f.) or probability mass function (p.m.f.) is specified up to a normalizing constant Z by an energy function E X - T. and inverse temperature 3 >> O, i.e., π (x) x e -BE((x), Expectation under a path measure and its marginals are written as E3 [-] and Ep, [.]. respectively. 14 {0, 1) is the indicator of statement or event A, which is 1 if.f. A holds or is true. For x (x' ...., x²)) E IR., the notation 2 means the vector obtained by replacing the i-th entry of x with
2 UNIFIED FRAMEWORK FOR SOC-BASED NEURAL SAMPLER
In this section, we discuss the key concepts and methods that form the foundation of our approach, including stochastic optimal control (SOC) interpretation of diffusion samplers (Zhang & Chen 2022
Vargas et al. 2023 Domingo-Enrich et al. 2025 Liu et al. 2025 Choi et al., 2025) and the weighted cross-entropy-basse approaches (Phillips et al. 2024 Zhu al. 2025). We defer detailed discussions of the related works to Sec. A
2.1 LEARNING NEURAL SAMPLER AS A STOCHASTIC OPTIMAL CONTROL PROBLEM
Unifying Framework for SOC-based Neural Samplers One approach to building diffusion samplers is to formulate them through a stochastic optimal control (SOC) perspective. We aim to fit a neural-network- paraalletttwrr-pprrreeeeerede parameterized path measure p6 to optimal one T whose terminal distribution is the target л. Let pref j0. 7] x X - IR be a reference path measure" where Poof and Ip = v. We assume that the reference path measure is memoryless (Domingo-Ennn Enrich al. 2025): the initial and terminal distributions are independent, i.e.,
0.7 Xx, XT) pae (Xo) (Xr) u(Xo)u(XT).
With this memoryless condition, the following holds:
For instance, paths induced by Brownian motions are a.s. -Hi-Höl continuous, and paths induced by contimnous-time Markov chains are a.s. piecewise constant, right-continuous with left limits. Throughout this paper, path measure refers to the distribution of a stochastic process indexed by E [0, T] and takes values X.
2

---

Published as a conference paper at ICLR 2026
IP* == argmin
[- E (x) + KL(PP
pe: P8
PT (X)
pref (x))((r)) r -3 log Po
The proof is provided in Sec. C other words, the optimal path measure bridges between the initial distribution of the reference path measure pref to our target distribution We remark that the memoryless condition on ppref necessary for this framework, and defer further discussion after the proof of (3) in Sec. C Many SOC-based neural samplers (Zhang & Chen 2022, Vargas et al.) 2023) solve 3 or equivalent problem (Richter & Berner 2024; Zhu et al. 2025). In the following
paragraphs, we provide continuous and discrete instantiations of the optimization problem (3)
SOC Approaches for Continuous Diffusion Samplers tic differential equation (SDE) on the state space X' F
Consider the following controlled stochas-
dX;- (X,) + o,uu (X.)) + o,dWe, Xo μt,
T) x X X' and o 10 T) + IR are the base drift and noise schedule, respectively, and W is the Brownian motion (BM). Let po be the path measure induced by (4). The reference path measure pref is the path measure induced by (4) with v = O. By applying Girsanov's theorem (Särkkä & Solin 2019), the optimization problem can be reduced to the following SOC problem:
where b: [0,T)
min E Ex=
2 (x))l [ r(()) -l(())))
Under the memoryless condition 2 the solution of (5) then yields a controlled process whose terminal marginal equals æ. Such memoryless references can be constructed by suitable choices of (b. jz). For example, Path Integral Sampler (PIS, Zhang & Chen (2022) and Adjoint Sampler
(AS. Havens et al. (2025) use μ = 60 and b, O. Denoising Diffusion Sampler (DDS) (Vargas
et al. 2023) takes (x)
2 z. vou, and (O, 7) where exp (- --o" ordl & 0. This
scheduling choice is the same as Variance Preserving SDE (VPSDE) in Song et al. (2021b). In these cases, the reference terminal law is tractable Gaussian, so r(-) in (5) is available in closed form.
SOC Approaches for Discrete Diffusion Samplers A counterpart of SDE in a discrete state space X' is the continuous-time Markov chain (CTMC), which is stochastic process X (Xt) taking values in X. The distribution of CTMC is determined by its generator Q (Qa)teju,7) defined as Q (3r, y) == lima--0 7. (Pr(X+++ X æ²) 1==y) for y. E X.
Existing SOC approaches for training discrete neural samplers mainly rely on masked discrete
diffusion (Zhu et al. 2025). We suppose the target distribution is supported on Xo - 21,2,.... ... Nj and let the mask-augmented state space be X' {1, 2, ..., N, The typical choice of the reference path measure pref has generator Quef defined by Qi (x,,""=)== 1 M
where ? [0,7]-1 R.+ can be any function satisfying J 7())dt co. By construction, ipge = fa is the delta distribution on the fully masked sequence, denoted Prnas (which guarantees memorylessness) and paoe is the uniform distribution over Xo. denoted Punif. SOC-based mask discrete diffusion neural samplers (Zhu et al. 2025) solve the following problem akin to (5):
7(!)
L'Σ
E
Qf @f log Qia
Oi
-Qt + Qioe) (Xe, v)d r(XT)
min Ex.rs 0
VXX,
s.t. X (Xt)ue10.11 is a CTMC on X' with generator Q", Xo Pmask:
where r has the same definition as in (5). The generator Q" is parameterized by Qft (z, f) where sy sa(x) is a d x matrix whose each row a probability vector.
2.2 WEIGHTED CROSS ENTROPY SAMPLERS
Cross-entropy-based Sampler The optimization problem 3 can be viewed minimizing the relative-entropy (RE) KL [P*). Another popular choice of the loss functional F is the dF' cross-entropy (CE) loss (L(P [P*) Ep- log - The CE loss requires importance sampling to evaluate, we have access the ground-truth path measure P A common choice for the
3

---

Published as a conference paper at ICLR 2026
Magnetization
Run
Run
-0.5
-1.0
Step
(a) Magnetization
(b) Run 1
(c) Run 2
(d) Ground Truth
Figure 1: An example of mode collapse in WDCE-based method where the target distribution is an Ising model under low temperature with two runs. (a) The magnetization (i.c., average spin) of the generated samples during training. The ground-truth value is O. We plot two independent runs, namely Run I and Run 2. (b, c) Visualization the generated samples at the final training step for Run and Run 2. (d) Visualization of the ground-truth samples.
proposal path measure is the detached path measure pd for faster convergence. In this case, the CE loss be expressed follows:
KL(P*)|P)) Eps
dp+
log
dF
E
Z dpó (X)e
1 dlPref
log
dpref dlpp (X) + + (6)
X
where C' a term independent of
Weighted Denoising Cross-entropy (WDCE) Samplers A drawback of the RE and CE losses is that they need to store the whole path of X for computing loss. A more memory-efficient alternative the weighted denoising cross-entropy (WDCE) loss which replaces negative log likelihood term in (6) with a denoising score (or bridge) matching objective (Phillips et al. 2024 Zhu et al. (025). This objective theoretically equivalent to the CE formulation and therefore implements the same reverse- KI. objective up to additive constant. Practically, WDCE allows to retain only the terminal samples Xy from rollouts X ~ IP. The pathwise weight wxX) x dud (X) can be accumulated online during simulation, training proceeds from storing pairs (Xr, w(X)), without storing full trajectories. See Sec. and Sec. for a detailed CE and WDCE formulation and its instantiation.
dF'
PROXIMAL DIFFUSION NEURAL SAMPLER: A GENERAL FRAMEWORK
3.1 TACKLING MODE COLLAPSE OF SOC-BASED SAMPLERS
Although SOC-based neural samplers theoretically grounded, these methods suffer from mode collapse in the early stages of training, when μo is far from the target IP.. These samplers attempt to optimize a given SOC problem using the trajectories generated by the current model pd, When the distributional mismatch is large, only small subset of trajectories carry a meaningful signal (e.g., high likelihood or high importance weight under IP*), causing the objective to focus on this narrow subset of the trajectories. As result, the losses are dominated by few high-weight paths, producing unstable updates and poor coverage
target. Moreover, due to the lack of exploration of the entire state space, these SOC-based methods tend to fit the modes that the existing trajectories have already reached, and reinforce this memory without attending to the rest of the possibilities, potentially leading mode collapse.
To illustrate this effect, we train a WDCE sampler [Zhu et al. 2025 on a discrete benchmark: 24 24 lattice Ising model under low temperature i i~j means i, are adjacent. At this temperature, the target distribution is bimodal, concentrating around the two ferromagnetic states (all-positive spins or all-negative spins) separated by a large energy barrier. As shown in the training dynamics in Fig. II WDCE rapidly collapses into single mode and subsequently reinforces this mode through self-generated samples.
0.6, where E (r) Σ Gunj 2. {+1] 24x222 and
"aample the path from Td and detach from the computational graph.
4

---

Published as a conference paper at ICLR 2026
3.2 PROXIMAL DIFFUSION NEURAL SAMPLER
We develop the Proximal Diffusion Neural Sampler (PDNS) to mitigate the mode collapse observed in Sec. 3.1 Rather than directly solving the global SOC objective in 33. PDNS proceeds by incrementally solving local sub-problems that progressively move the current model toward the target distribution. Coneretely, we apply proximal point method (Parikh & Boyd 2014) in the space of path measures, progressively driving Pe toward IPt step-by-step while preserving mode coverage.
Proximal Point Method for Diffusion Samplers To transform (3) into a local sub-problem, we add proximity term by KL divergence from the previous iterate po,-1. Formally, we define the proximal point sub problem with proximal step size Ti E (0, 20] as follows:
P argmin (x) r(Xr) KL((P") + - _ p
(7)
Rather than solving the difficult SOC problem (corresponding to infinite step size m)). the problem
(7 searches for the minimizer of the KL.ppealized optimization problem. This KL--penalization term ensures that IP remains close to the previous iterate p thereby resulting more robust updates and simplifying the optimization due to the dominant regularization effect. In practice, for given
path, the proximal objective effectively tempers the importance weights d= compared to that AP" of non-proximal counterpart dr reducing the dominance
single mode, thereby mitigating
mode collapse. Theoretically, one can show that F + I as k oo. Moreover, when pdi prref the optimal Fot can be written as geometric interpolation between ppee and IP*. We summarize in the following result, with proof in Sec. C
Proposition 3.1. I. The optimal solution of is
dl
T x (100-1)
x
(8)
dPP
dp
2. Assume for all k 2 1, the subproblems are solved to optimality and ler pdo & pref Denote p as the corresponding path measure poi which satisfies Pp x (Fk-1) TITI (IP*)
We thus have
ph (Iprey prreetya (Pr)I-de, where An
TJ +1
This implies Pa converges to PM if Ax + O. Moreover,
(x() prof (x))(tAAar((X (x)r(XT) prof (X) T Ae
(XT),
(10)
ie., the terminal distribution of pk is Pt x π 1-1-x vux, the geometric interpolation between v and m.
In other words, iteratively solving the local optimization problems yields sequence of path measures (IP%) [ that converges TP* For simplicity, we use
Pab" denote the optimal solution the
k-th proximal step, which can be chosen as either 1o1 ( or pale 9 We summarize the meanings of the notations of path measures in Tab. and present the overall PDNS framework in Alg
Here, the proximal step size Th is a crucial hyperparam- eter: smaller 77 increases the relative weight of the KL regularizer, resulting in more conservative updates that better preserve mode coverage. Instead, it has slower convergence to optimum. Larger 77s weakens the regularization, enabling faster conver- gence, but risking mode collapse. Note that it is the step size πp affects performance due to this trade-off between convergence and mode coverage. We further discuss the possible proximal step size scheduler in the last paragraph of Sec. 4
Table 1: Explanation of the path measures (P.M.).
General notation for P.M. of a neural sampler parameterized by
paot P- po.-1 pot
Reference and optimal P.M.s defined in (3]
P.M. of the neural sampler the end of the (k 1)-th proximal step
Optimal P.M. 1-th proximal step, starting from pb_-1
ps pk
Exact solution P.M. for k-th proximal step, i.c. and 10
cither pst or Pt training objective at the k-th proximal step
Proximal CE for Diffusion Samplers By reversing the arguments in the KL divergence of obtain cross-entropy
3 we have obtained cross-entropy counterpart (6). Likewise, counterpart to the proximal subproblem (7). Given the current iterate Fd, we define the proxi-
mal cross-entropy update as minpo KL (F") [[P"") or mingo KL (P*IP""). By leveraging the right- hand side of (3) and (8), the former object reduces weighted negative log-likelihoods under
5

---

Published as a conference paper at ICLR 2026
Algorithm Proximal Diffusion Neural Sampler (PDNS)
Require: Model parameters 0, target distribution b x e schedule of step sizes (14)221 reference path measure Ipref, initial path measure pdo, loss function F.
1: for k in 1. 2... do
2: Set [p* E 1I , the optimal solution to the current subproblem, following 8 or (1
3: for step in 1,2, do
4: Compute the loss F(Pe; Tpk") using samples from poz-1 and update O.
5: end for
6: Set _od 4- 1po
7:
end for
the previous iterate, up to positive multiplicative factors and additive constants independent of po:
KL (P%
|
(X)
log
dP
dIP (X)
E
(d (x)
dTp
"k log
ya d
(X)
dipb_-,
X~p"k
X~P
x - X ((xr) dlpref
mu+T
d (x)
log (X) + +C,
(11)
Furthermore, leveraging
we obtain
KL [Pd)) - Ex-pa-1 ((1-)e)r(XF) de
dpk-r (X) log po (X) + C,
(12)
up to some multiplicative normalizing constant and some irrelevant constant C. Note that both objectives [II and (1) are reduced to weighted oggliikllhood objective.
Proximal WDCE for Diffusion Samplers As described in Sec. [.2] the CE loss (6) can be approximated by the WDCE loss by simply replacing the negative log-likelihood of the CE term into the denoising score matching loss. Similarly, we can also further analyze the proximal CE losses (T)-(12) and derive proximal weighted denoising cross-entropy (proximal WDCE) losses, denoted as ming F(IP"; F ming ; depending on which local optimum we are targeting. In the next section, we provide explicit forms of the proximal WDCE objective F in both continuous and discrete settings.
PROXIMAL DIFFUSION NEURAL SAMPLER IN PRACTICE: CONTINUOUS AND DISCRETE CASES
Among various PDNS instantiations, we focus on a proximal WDCE variant PDNS, selected for its practical effectiveness: it avoids storing entire trajectories and enjoys efficiency of (discrete) score matching objectives (Zhu et al. 2025). Now, we derive explicit forms of the proximal WDCE training objective FF in Alg. for both continuous and discrete diffusion settings. Furthermore, we also provide principled guidelines for selecting the proximal step sizes { ] k=1- Complete algorithms for proximal WDCE and scheduler policy are provided in Sec. B and additional implementation details are provided in Secs. and
Proximal WDCE for Continuous Diffusion Samplers By applying Diffusion Schödinger Bridge Matching (DSBM) (Sh et al. 2024), the proximal WDCE objective is written as follows:
(L [P=))
E Unif(( X-DaA
|lu Ilu (X.) V log Proff (X+|X.)11
dlik"
(13)
dp (1) (l(((X)) - V.V log pere (X xx))2
U. X
Here, the first line is the original DSBM loss with target path measure pk" and the second line we use importance sampling to rewrite the expectation under the detached path measure po. obtained from the previous iteration. We sample N trajectories { li=l pā and compute
the corresponding weights (RN derivative) { (w(()) = ar* /ap" (x())] which will be discussed

---

Published as a conference paper at ICLR 2026
in the next paragraph. Then, store pairs (Xp) w()) for (I....,N) NJ in buffer B. In the training iteration, we sample
(Xr, B, Xo~p p. to Unif T):
and sample X, using reciprocal property (Léonard et al. 201 4), i.e., 1p peof to compute
=jo.T
60.77 0
the matching objective in the equation. The detailed discussion provided in Sec. C.2
We now introduce closed form to compute the weight (RN derivative) control that induces the path measure [ Using the Girsanov theorem Ksendal 2003).
13 Let n (x)
be the
dirref
14 B (x)ll
uh 8 Ou-1 (X,) dW,
dP (X) = exp
(14)
2
so the RN derivative
app
in (3
where pp
E (Fot pu,, can be computed as follows:
dp
dPp (x(Xx) dlpref d=ox=1
dlk dP
deref
dp (X) x
(x)
or
e(1-Aa)r(X))
dp (X).
(15)
The detailed discussions are provided in Sec.
Proximal WDCE for Discrete Diffuusion Samplers
The WDCE objective, following Ou al.
(2025a) Zhu et al. (2025), can be written as
disk"
KI.(P |TPC) = X de (X) Ex-Uni((O.1)
x Bux(F)
Σ
log sp (e()),
M
where (x(-(a) (r) means independently masking each entry of x with probability X, and 89 is the parameterization of Ip introduced in Sec. 2.1
We adopt (15) to compute the weight
dpt dP"c- for discrete cases, where pk E (P J. It suffices
obtain the target weight. Through the Girsanov theorem of CTMCs (Ren et al. 2025a Zhu et al. 2025). for any trajectory X, we could obtain the RN derivative follows:
to compute the following RN derivative
disref
dipref
1/N
dPe (X) = exp
Σ
log
(16)
80
(Xt x000)
1:X,_
#X.
where assume the jump time happens at the -(h)-h entry, and the total number of jumps is the sequence length d.
Design of Schedulers for the PDNS framework, which we refer to as the scheduler because it governs the trade-off between convergence speed and stability. larger ты enables faster convergence but risks difficulty from larger distributional gaps, whereas smaller 77 yields conservative yet more stable progress.
As discussed in Sec. 3.2, the design of the step sizes {mw Ja21 is important
The simplest design is choose a predefined schedule for пі or Åg. leveraging the relationship in (9) and making sure A - O. Unlike directly training the model to fit the target distribution æ, the knowledge of the optimal solution of each subproblem (Prop. B.T h PDNS provides information on how the samples fit the local target distributions. A demonstration can be found in Fig. [2
We also propose an adaptive strategy that selects Vik automatically based on the current state of the model. The key idea is ensure that the next target Ipk
E (I* Fd} is not too far from the
current one p0a e.g., through controlling (L((pt-)) !ppk" < e. Detailed description of the criteria
deferred to Sec. and refer readers Secs. D.3 and E.3 for ablation study the schedulers.
EXPERIMENTS
We now demonstrate the effectiveness of our proposed PDNS framework learning continuous and discrete target distributions. All implementation details and additional experimental results can be found in Secs. and F
7

---

Published as a conference paper at ICLR 2026
Table 2: Evaluation on the synthetic energy functions. Following Chen et al. (2025); Choi al. (2025). we report Sinkorn distance (1) and MMD (4). The best results and the second best results are highlighted.
MWS4 (d (d=b, 5)
Funnel (d
GMM40 (d 50)
MoS 50)
Method
Sinkhorn
MMD
Sinkhorn
Sinkhorn
MMD
Sinkhom
SMC
Del et 2006
20.71 +.5.5
149.35 4.755
46370.34
137,79
3297.2 28 2184.444
SMC-
ESS
Bacthole
201
1.11 1.015
117.48 9:99
24240.68
66.66
1477.0k 044
a 135.90
CRAFT
2011
11:47 5:000
0.115
134.34
28960.70
0.257
1918.14
DDS
Vampa
0.63
0.172
142.89
5435.18 177.0000
0.131
0.000
2154.88
13366
Zhang
0.42 0.00
10405.75 00.41
0.218
000000
2113.17
51.17
0.32 0.000
18984.21
0.210
0.000
2178.60
1.54.0
Viggaset 2024
0.57 0.000
0.095
513.33
+ 122.4
22132.25
18418.80 19.88.89.
+ 5322200
CMCD-LV Vagas et 2024
0.51 0.00
139.07
+:935
4258.57
4-737.15
19:15.71
48.79
SCLD
Chenerall
0.44 2.000
134.23
8:800
3787-73
656.10
:s 99.997
NAAS
Clas stal. 3025
0.10 0.0075
0.076
0.00%
132.30
5.5.87
496.48
27:00
0.113
0.070
394..5
20.30
PDNS
0.08
0.086
129.54
327.83
0.059
353.05
Table 3: Evaluation on the particle-base energy functions. Following Havens et al. (0225), we report the Wasserstein-2 distances w.r.t samples, W2(4), and energies, W( The best results and the second best results are highlighted.
2025
: Liu et al.
DW44 4 (d 8)
LJ-13 (d 39)
LJ-55 (d = 165)
Method
W2 4
E(-) W/ 4
W2 4
E(-) W2
W2 4
EX() We +
PDDS
Phillips
2024
0.92 0:08
0.58
4.66 0.97
56.01
10.00
SCLD Chen
2025
1.30 0:64
0.40 2:19
2.93 0:19
27.98
PIS
Zrong Chan.
0.68 0.28
0.65 0.25
1.93
0.07
18.02
4.79
0.45
228. 70
131.27
DDS
Vargas ot al. 2025
0.92 4: 0.11
0.90 0.37
1.99 = 0:18
24.61
8.99
4.60
0.00
173.09
18:01
LV-PIS
Richer Burner
1.04 1: 0:29
1.89 0.89
iDM
Aklouad Sadegh or
2024
0.70 + 0.06
0.55
0.14
1.61
0.01
30.78
24.46
4.69
1.62
93.53
16.31
AS Hayens al. 2025
0.62 + 0.08
0.55
0.12
1.67
+ 0.01
2.40
1:25
4.50
0.00
58.04 + 20.98
ASBS
Liu et all 20255
0.38 0.08
0.19 0.08
1.59 =: 0.00
1.28
8: 0.22
4.00
0.08
27.69
3.86
PDNS
0.51 0.04
0.21
1.57 0.00
1.01
3.95
21.97
3.14
5.1 LEARNING CONTINUOUS TARGET DISTRIBUTIONS
Overview We evaluate our model on the following continuous target benchmarks:
Continuous synthetic energy functions Following (Choi et al. 2025). ) we consider four different
multi-modal synthetic benchmarks; the 32-well many-well potential (MW-54, SD), Funnel (5D), a 40-
component Gaussian mixture (GMM40, 50D), and a mixture of Student's 's distribution (MoS, 50D)
are considered. We evaluate with two metries; the maximum mean discrepancy (MMD) (Akhound 'Sadegh et al. 2024) and the Sinkhorn distance (Sinkhorn) with a small entropic regularization (10->) (Chen et al. 2025).
Particle potentials We consider classical continuous potentials where the energy E(x) is analytic and depends on pairwise distances in an n-particle system. Following Akhound-Sadegh al. (2024) we consider DW-4 (2D, 4 particles), LJ-13 (3D, 13 particles), and LJ-55 (3D, 55 particles). We use the equivariant Wasserstein distance and 2- Wasserstein distance between the ground-truth and the generated samples, following (Akhound-Sadege) et al.) 2024 Liu et al. 2025). For the ground-truth samples, we use MCMC samples from Klein et al. (2023).
Alanine Dipeptide (AD) The AD system contains 22 atoms in diree-dimensionll space. Following prior work (Midgley et al., 2023), we simulate the molecule using the OpenMM package (Eastman et al. 2017) and parameterize its geometry in terms of internal coordinates, resulting in 60- dimensional representation. For evaluation, we compare the distribution of generated samples against the reference distribution using the KL divergence computed from 10a samples. Our evaluation metric is KI. divergence (DKL) between generated and reference samples for the backbone angles d. v and the methyl rotation angles 71, 72, 73- In addition to these quantitative metrics, we also visualize the energy histogram, torsion plot, and Ramachandran plot for both generated and reference samples to assess structural fidelity, We use 10 samples for torsion plot and Ramachandran plot.
8

---

Published as a conference paper at ICLR 2026
Table 4: Comparison between diffusion samplers sampling the molecular Boltzmann distribution of the alanine dipeptide. We report the KL divergence KL for the ID marginal across five torsion angles.
KL each torsion's is marginal
Method
71
332
Zhang
0.05
5.61
DDS
Nargas alg
0.03
0.16
2.44
0.03
0.03
AS Haxens et
0.17
ASBS
2025
0.02
0.01
0.03
0.02
0.02
Figure 2:
51
Energy his-
X
PDNS (Ours)
0.02
0.04
0.03
0.02
0.02
tograms for AD.
P
100
density
0.30
Figure 3: Comparison of torsions between PDNS and reference.
Detailed benchmark descriptions, experimental settings, and evaluation metrics are provided in Sec. D In Sec. D.3] also compare the WDCE baseline, the main baseline of our models to proximal WDCE. Moreover, in Secs. D.3 and D.4] we further present qualitative and quantitative results, including ablations over hyperparameters and schedulers, and sample visualizations.
Main Results Quantitative results are demonstrated in Tabs. 2 and 3 Among seven benchmarks, PDNS (with proximal WDCE) attains the best scores on five tasks, with the only exceptions being Funnel and DW4, where performance is comparable but slightly behind recent baselines. Notably, PDNS excels on the more challenging targets MoS (50D) with heavy tails and the L.I-13/LJ-55 particle systems with steep, rugged energy surfaces. We attribute these gains to the proximal step, which tempers trajectory weights and constrains each iteration to a local move in path space. This prevents the model from committing to single model early, promotes gradual exploration across separated modes, and stabilizes optimization high dimensions. As illustrated Sec. D.3 removal of the proximal term induces mode collapse, whereas the PDNS maintains coverage. Finally, as shown in Tab. 4 Figs. 2 and 3. our method achieves strong performance on the alanine dipeptide task. Quantitatively, the KI. divergences of the torsion angles comparable to those reported by the current state-of-the-an method, ASBS (Liu et al., 2025), and the qualitative assessments including energy histograms and torsion plots, demonstrate that our model captures the main conformational structures of the system.
5.2 LEARNING DISCRETE TARGET DISTRIBUTIONS
Overview We evaluate our method on the following discrete sampling tasks:
Learning distributions from statistical physics Following Holderrieth et al. (2025): Zhu et al. (2025), we consider the Ising and Potts models on a square lattice with L sites per dimension, which are two classical models from statistical physics that exhibit phase transitions under different inverse temperatures [Onsager] 1944; Beffara & Duminil-Copin) 2012). To emphasize the advantage of our method in learning distributions with mode separation, we mainly consider temperatures around or below the critical one: for Ising model, we choose L = 24, and consider Beritical = log(1 + 22)/2 /2- = 0.4407 and Blow = 0.6; for Potts model, we choose L = 16, number of states q 4, and consider Beritical = log (1 + v 1.0986 and Blow 1.3. We compare our method with LEAPS (Holderrieth et al. 2025) and the Settroooli--Hasstinn -anns
(MH) algorithm, and use the Swendsen- Wang (SW) algorithm (wendssen & Wang 1986 1987) to generate ground-truth samples for evaluation. As the original WDCE loss cannot lear the correct target distribution (Fig. we do not include its performance here.
Amortized sampling for combinatorial optimization Recently, there is a trend of leveraging sampling methods for solving combinatorial optimization problems (Sun et al. 2022 Sanokowski et al. 2024 Ou et al. 20256), by writing the loss function as the energy E and sampling from the distribution TT X under relatively large inverse temperature (e.g., s.... 5), that ny concentrates on the

---

Published as a conference paper at ICLR 2026
Ising, 24, Benna 0,4407
Ising, 24,
Potts, q- Peena 1.0986
Ports,
2.iint Conrelation
66-
-18.58.5.44 Distance
-32-32-4
10
Distance
Distance
Distance
Haw PONG Isuns
104'5
-8h 101
-
Groed IFM
Figure 4: Average 2- 2-point correlations in both vertical and horizontal directions of samples from learned Ising and Potts models at different inverse temperatures.
global minimum of E, thus improving the optimization accuracy. This low-temperature nature of the target distributions serves perfect example demonstrate the effectiveness of PDNS.
Main Results For sampling from Ising and Potts models, we report quantitative results, including
error in magnetization and 2-point correlation with respect to the ground samples in Tab. 5 and visualize the learned average 2-point correlations at states differing at different distances r in Fig. We observe that our method significantly outperforms LEAPS. While the MH algorithm can achieve good accuracy in Ising model, its performance degrades significantly in Potts model due to the larger number of spin states, yet our method still achieves good performance. For combinatorial optimization, in Tab. 6 we use PNDS to solve the maximum cut problem and test the performance on two classical random graphs. We report the solution compared with the ground-truth value solved by Gurobi (Gurobi Optimization, LLC 2025), showing PDNS's s capability of reaching comparable performance. We defer detailed explanation of the distributions and settings to Sec. E.2
Table 5: Results for learning Ising and Potts models at different inverse temperatures, the best result highlighted. Mag. and Corr. represent absolute errors of magnetization and average 2-point correlations to ground-truth values given by the SW algorithm.
Distribution
bing model, L = 24, =1
Potls model. L = 16, J 1. =
Inv. Temp.
Pemmiai 0/4407
Puoee 0.6
Pertiaal
1.0886
Puw 1.3
Metrics
Mag.
Com J
ESS f
Mag.
Cor
ESS f
Mag.
Corr.
ESS
Mag.
Corr.
ESS
PONS (ours)
1.2g 2
5.6e 3
0.903
9.0e 3
1.7e 3
0.950
5.2e 3
4.8e 4.8e-3 3
0.948
8:4e 4
6.1e 4
0.978
LEAPS
5.5e 2
2.80 1
0.020
3.0 2
6.50
0.001
3.0
2.60 1 0.112
3.60
3.50
0.021
Baseling (MH)
2.3e
1.9e-3
1.6e
5.3e
10e-1
7.6e
Table 6: Results for finding maximum cut of Barabási-Albert (BA) and Frdös-Rényi (ER) graphs by PNDS sampling. The evaluation metrics are maximum and average cut size, i.c., the maximum / average of the size of all generated samples for each graph, taken average over all graphs in the test set. Higher is better.
Num. of Vtx.
16
20
32
40
64
75
100
128
Graph Type
BA
ER
BA
ER
BA
ER
BA
ER
Random
0.680;
0.865
0.725; 0.873
0.685; 0.806
0.780; 0.861
0.678; 0.700
0.836 0.879
0.676; 0.740
0.865; 0.894
PDNS tours)
0.947: 0.991
0.962; 0.905
0.939; 0.973
0.966; 0.989
0.918; 0.947
0.970: 0.983
0.882; 0.915
0.973; 0.983
CONCLUSION
We propose Proximal Diffusion Neural Sampler (PDNS), a unified framework for SOC-based training of diffusion neural samplers that perform proximal point method on the path measure space. We instantiate PDNS with the proximal WDCE objective, and evaluate our method on various continuous and discrete target benchmarks where PDNS shows competitive performance. Our paper has two main limitations: we investigate a single instantiation of PDNS, and our experiments focus on controlled benchmarks rather than real-world, large scale applications. Potential directions include exploring SOC problem with memoryless reference dynamic pref, other instantiations of PDNS (e.g., other proximal objectives or priors), and extending the framework for distributions on more general state spaces and neural samplers beyond the SOC formulation.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
WG thanks Zijing Ou and Sebastian Sanokowski for discussions on combinatorial optimization problems. We thank Petr Molodyk and Bo Yuan for discussions during the early stage of this work. WG, JC, and YC acknowledge supports from NSF Grants ECCS-1942523 and and MT are thankful for partial supports by NSF Grants DMS-1847802, DMS-2513699, DOE Grants NA0004261, SC0026274, and Richard Duke Fellowship.
YZ
REFERENCES
Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh. Gauthier Gidel, Yoshua Bengio, et al. Iterated denoising energy matching for sampling from boltzmann densities. Proceedings of the 41st International Conference Machine Learning, pp. 760-786, 2024.
Michael Samuel Albergo and Eric Vanden-Eijnden. NETS: non-equilibrium transport
pler. In Forty-second International Conference on Machine Learning, 2025. httpst
/pppnrevieww.oo net forum? Sitimmid-Qooee QQGW9StPbQ
Brian DO Anderson. Reverse-time diffusion equation models. Stochustic Processes and their Applications, 12(3):313-326, 1982.
Michael Arbel, Alex Matthews, and Arnad Doucet. Annealed flow transport monte carlo. In International Conference on Machine Learning, pp. 318-330. PMLR, 2021.
Vincent Beffara and Hugo Duminil-Copin. The self-dual point of the two-dimensional random-cluster model is critical for g 2 1. Probability Theory and Related Fields, 153(3):5 2012. doi:
11100775004440001--03333
Denis Blessing. Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Beyond clbos: a large-scale evaluation of variational methods for sampling. In Proceedings of the 41st International Conference Machine Learning, pp. 4205-4229, 2024.
Denis Blessing, Julius Berner, Lorenz Richter, Carles Domingo-Enrich, Yuanqi Du, Arash Vahdat, and Gerhard Neumann. Trust region constrained measure transport in path space for stochastic optimal control and inference. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. BtLhttps:// //tp:rev//pperevieee wtapeerreiieeeeee/ooutrttttt
6R1b0EcOS4
Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of Markov chain Monte Carlo. CRC press, 2011.
Alexander Buchholz, Nicolas Chopin, and Pierre E Jacob. Adaptive tuning of hamiltonian monte carlo within sequential monte carlo. Bayesian Analysis, 16(3):745-771, 2021.
Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anand- kumar. Sequential controlled langevin diffusions. In International Conference on Learning Representations (ICLR), 2025.
Ricky Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrödinger bridge. SIAM Review, 63(2):249-33 2021.
Jaemoo Choi, Yongxin Chen, Molei Tao, and Guan-Horng Liu. Non-equilibrium annealed adjoint sampler. In The Tiihh Annual Conference on Neural Information Processing Systems, 2025. URL https: //tprreiie.e..eeeeeeeeeeeee net. www.:neett.oooe forum? i id=ay 7WDSqCKb)
Nicolas Chopin. A sequential particle filter method for static models. Biometrika, 89(3)5539-552, 2002.
Valentin De Bortoli, Michael Hutchinson, Peter Wirnsberger, and Amaud Doucet. Target score matching. arXiv preprint wxXv:2242.008677 2024.

---

Published as a conference paper at ICLR 2026
Pierre Del Moral, Amaud Doucet, and Ajay Jasra. Sequential monte carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, (033:411-4336 2006.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv: 140..85 16, 2014.
Carles Domingo-Enrien Jiequn Han, Brandon Amos, Joan Bruna, and Ricky TQ Chen. Stochastic optimal control matching, Advances in Neural Information Processing Systems, 37:112459-111111 12459-1411111 112504, 2024.
Carles Commingo-Enriche Michal Drozdzal, Brian Karrer, and Ricky I. Q. Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In The Thirteenth International Conference Learning Representations, 2025. URL https
//openreview.nee net www.:nnnee//ooutieoomoooooo forummrriooooooooo ?id
xQBRrt XQB
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL httts::/ fest/fooum?2 Ieforummid=Yiie id=YicbFdVTTT
Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, x3(7);e 1005659, 2017.
Marylou Gabrié, Grant M Rotskoff, and Eric Vanden-Fijnden. Adaptive monte carlo augmented with normalizing flows. Proceedings of the National Academy of Sciences, 119(10)e210942011 19, 2022.
Tomas Geffner and Justin Domke. Meme variational inference via uncorrected hamiltonian annealing. Advances in Neural Information Processing Systems, 34:639-651. 6619, 2021.
Andrew Gelman, John Carlin, Hal S. Stern, and Donald B. Rubin. Bayesian data analysis. Chapman and Hall/CRC, 3 edition, 2013.
Wei Guo, Molei Tao, and Yongxin Chen. Provable benefit annealed Langevin Monte Carlo for non- log-concave sampling. In The Thirleenth International Conference on Learning Representations, 2025. /tRpstts::/// http://eeeeeeeeeeeeeeeeeooaaaaaa
Wei Guo, Molei Tao, and Yongxin Chen. Complexity analysis of normalizing constant estimation: from jarzynski equality to annealed importance sampling and beyond. In The Fourleenth Interna- fional Conference Learning Representations, 2026a. URL https::/openrer: ew. /
Eorum?id (wwi--id---JAAA 96
Wei Guo, Yuchen Zhu, Xiaochen Du, Juno Nam, Yongxin Chen, Rafael Gomez-Boobarelli, Guan- Horng Liu, Molei Tao, and Jaemoo Choi. Discrete adjoint Schrödinger bridge sampler. arXiv preprint arXiv:2602.082 082200 2026b.
Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2025. URL https: //ww www
gurobi. com
Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo-Enrich, Anuroop Sriram, Daniel S. Levine, Brandon M Wood, Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, Guan- Horng Liu, and Ricky T. Q. Chen. Adjoint sampling: Highly scalable diffusion samplers via adjoint matching. In Forty-second International Conference on Machine Learning, 2025. URL https::/openrevv ew net ftnt//forum?idd- Or
Ye He, Kevin Rojas, and Molei Tao. Zeroth-order sampling methods for non-log distributions: Alleviating metastability by denoising diffusion. Advances in Neural Information Processing Systems, 12771122271111 2024.
12

---

Published as a conference paper at ICLR 2026
Yuchen He and Chihao Zhang. On the query complexity of sampling from non-log- concave distri- butions (extended abstract). In Nika Haghtalab and Ankur Moitra (eds.), Proceedings of Thirty Eighth Conference Learning Theory, volume 291 of Proceedings of Machine Learning Research, pp. 2786- 2787. PMLR, 30 Jun-04 Jul 2025. URL httts::/pprocest http:///pooeeedngg.......
press/ ngs. ml r...m......ss/
v291 2a11/he25a html
Jeremy Heng, Adrian Bishop, George Deligiannidis, and Arnaud Doucet. Controlled sequential monte carlo. The Annals of Statistics, 48(5):2904-2929, 2020.
Byeongho Heo, Song Park, Dongyoon Han, and Sangdoo Yun. Rotary position embedding for vision transformer. In Aleš Leonardis, Elisa Ricci, Stefan Roth, Olga Russakovsky. Torsten Sattler, and Gül Varol (eds.). Computer Vision - ECCV 2024, pp. 289-305, Cham, 2025. Springer Nature Switzerland. ISBN 978-3-031-7268422
Peter Holderrieth, Michael Samuel Albergo, and Tommi Jaakkola. LEAPS: A discrete neural sampler via locally equivariant networks. In Forty-second International Conference on Machine Learning, 2025. URL
formmf d= QAET
Xunpeng Huang, Hanze Dong, Yifan Hao, Yi-An Ma, and Tong Zhang. Reverse diffusion monte carlo. ICLR, 2024.
Leon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Mare Brockschmidr, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pp. 5361-5370. PMLR, 2020.
David Landau and Kurt Binder. A Guide Monte Carlo Simulations in Statistical Physics. Cambridge University Press, 4 edition, 2014.
Krzysztof Latuszyíski, Matthew T Moores, and Timothée Stumpf-Fétizon. MCMC for multi-modal distributions. arXiv preprint arXiv:2501.05900 2025.
Christian Léonard. From the Schrödinger problem to the Monge-Kantorovin problem. Journal of Functional Analysis, 267(4):1879-1190, 2012.
Christian Léonard, Sylvie Rælly, and Jean-Claude Zambrini. Reciprocal processes. A measure- theoretical point of view. Probability Surveys, 2014.
Dong C Liu and Jorge Nocedal. On the limited memory bigs method for large scale optimization. Mathematical programming, 45(1):503-528 1989.
Gang Liu, Jiaxin Xu, Tengfei Luo, and Meng Jiang. Graph diffusion transformers for multi-conditional molecular generation. In The Thirty-eigtth Annual Conference on Neural Information Processing Systems, 2024. URL https: //ttp:://oge http::::///eeeeeeeeeeee www.rrewwwwww..eeeeeeeeeee http://oormmmddfffffff
Guan-Horng Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, and Ricky T. Q. Chen. Adjoint Schrödinger bridge sampler. In The Thirty-ninth Annual Conference Neural Information Processing Systems, 2025. httl:tttss///peeeeeeeee wwtprreviiwwweeeeeeeeee
wnt// ftntt/oorm??dd rMh@B 1hh4c
Jun S. Liu. Monte Carlo Strategies in Scientific Computing. Springer Publishing Company, Incorpo- rated, 2008. ISBN 0387763694.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL
Bkg6RiCqY7
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Amaud Doucet. Continual repeated annealed flow transport monte carlo. In International Conference on Machine Learning, pp. 15196-15219. PMLRR, 2022.
13

---

Published as a conference paper at ICLR 2026
Laurence Illing Midgley, Vincent Stimper. Gregor NC Simm, Bernhard Schölkopf, and José Miguel Heráández-Lobtto Flow annealed importance sampling bootstrap. International Conference on Learning Representations (ICLR), 2023.
Radford M Neal. Annealed importance sampling, Statistics and computing, 11:125-139, 2001.
Kirill Neklyudov, Daniel Severo, and Alireza Makhzani. Action matching: A variational method for learning stochastic dynamics from samples. In International Conference Machine Learning (ICML), 2023.
Nikolas Nüsken and Lorenz Richter. Solving high-dimensions hamilton jucooi--jbbbbbbbbbleee pdes using neural networks: perspectives from the theory of controlled diffusions and measures path space. Partial differential equations and applications, 2(4):48, 2021.
Bernt Øksendal. Stochastic differential equations. In Stochastic Differential Equations, pp. 65 84. Springer, 2003.
Lars Onsager. Crystal statistics. A two-dimensional model with an order-disorder transition. Phys. Rev, 65:117-149 Feb 1944. doi: 10. 10.1105/hhyRev.655117 URL http:tt//iikk...p.. aps.:iiikk..s.................m. do1/10.1103/Payya..
617
Jingyang Ou, Shen Nic. Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. In The Thirteenth International Conference on Learning Representations, 2025a. URL https:/ https://openreview net forumrid= P 8Tanm]
Zijing Ou, Ruixiang ZHANG, and Yingzhen Li. Discrete neural flow samplers with locally equivariant transformer. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025b. URL https: /ttp//peeevveeeeeeeeee/ee net / forum? id
Wk 6k650kms 3T
Neal Parikh and Stephen Boyd. Proximal algorithms. Found. Trends Optim., 1(3):127-339 January 2014. ISSN 2167-3888. doi: 10.1561/2400000000, URL 2400000003
William Peebles and Saining Xie. Scalable diffusion models with transformers. In 2023 IEEF/CVF International Conference on Computer Vision (ICCV), pp. 4172-4182, 2023. doi: 10.1
ICCV510000.023.000387
Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Particle denoising diffusion sampler. arXiv preprint arXiv:2402.06330, 2024.
Yinuo Ren, Haoxuan Chen, Grant M. Rotskoff, and Lexing Ying. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via stochastic integral framework. In The Thirteenth International Conference on Learning Representations, 2025a. URL
https: //tpnrev//peerevieww net. forum 6awxwQ8182
Yinuo Ren, Haoxuan Chen, Yuchen Zhu, Wei Guo, Yongxin Chen, Grant M. Rotskoff, Molei Tao, and Lexing Ying. Fast solvers for discrete diffusion models: Theory and applications of high-order algorithms. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025b. htthttts::// http://///pppeeeeeeeeeeeeeeee
id=0uk116600
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. International conference on machine learning, pp. 1530-1538, PMLR, 2015.
Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations, 2024.
Chris G Rogers and David Williams. Diffussions, Markov processes and martingales: Volume 2, Itō calculus, volume 2. Cambridge university press, 2000.
14

---

Published as a conference paper at ICLR 2026
Sebastian Sanokowski, Sepp Hochreiter, and Sebastian Lehner. A diffusion model framework for unsupervised neural combinatorial optimization. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine ngs
Learning Research, pp. 033466443367 PMLR, 21-27 Jul 2024. htt.thttss///rreeee mlr. h.l.pp.ees/v23//// http:///2255//aooooosk22222/// html
Sehastian Sanokowski, Wilhelm Franz Berghammer, Haoyu Peter Wang, Martin Ennemoser, Sepp Hochreiter, and Sebastian Lehner. Scalable discrete diffusion samplers: Combinatorial optimization and statistical physics. In The Thirteenth International Conference on Learning Representations, 2025. URL /ttttp:://operreveewwe net/ wtrss//foummmii pewppupppdd.ppggppp
Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International Conference on Machine Learning (ICML), 2021.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
Oswin So. Brian Karrer, Chuchu Fan, Ricky T. Q. Chen, and Guan- Liu. Discrete adjoint matching. In The Fourteenth International Conference on Learning Representations, 2026. URL
httts:://pperrev// net Teree/foounne forum? dit.mmmddddxxXxxxxexe
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna-
fional Conference on Learning Representations (ICLR), 2021a.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021b.
Haoran Sun, Etash Kumar Guha, and Hanjun Dai. Annealed training for combinatorial optimization on graphs. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop), 2022. URL httphhtss///ppeeeiiiieeeee wtt.://www.n...eoenneeeeeeed
net www.eoorurio.rr... wwtorrmmmd=fffoo030 TKu
Robert H. Swendsen and Sheng Wang. Replica Monte Carlo simulation of spin-glasses. Phys. httpttts::/ https://link
Rev. Lett, 57:2607-2609,, Nov 1986. doi: 10. 10.1103/PhysReeeeettt..... aps. http:///aaaaa....o.../o//eeeeeta 57. 2607
Robert H. Swendsen and Jiang Wang. Nonuniversal critical dynamics in monte carlo simulations.
Phys. Rev. Lett. 58:86-888 Jan 1987. doi: 10.1103/hhyRReweetts... URL htttps://inm 58.888.8
aps. wwtp////oooooooo/o////// http://o.....o//yyeeeeeeeee
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers & distillation through attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 10347-10357. PMLR, 18-24 Jul 2021. URL htt.t::t///// //poccee ngs.ml: www.sr.pressseee 139/touvro22 html
George Uhlenbeck and Leonard S Ornstein. On the theory of the brownian motion. Physical review, 36(5):823, 1930.
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In International Conference on Learning Representations (ICLR), 2023.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets varia- tional inference: Controlled monte carlo diffusions. In International Conference on Learning Representations (ICLR), 2024.
15

---

Published as a conference paper at ICLR 2026
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic normalizing flows. Advances in neural information processing systems, 33:5933-59444 5944, 2020.
Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations,
volume 43, Springer Science & Business Media, 1999.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling. In International Conference on Learning Representations (ICLR), 2022.
Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Yang Yuan, Quanquan Gu, and Andrew C Yao. On the design of KL-regularized policy gradient algorithms for LLM reasoning. The Fourteenth International Conference on Learning Representations, 2026. URL htttps://operreviees
het//forum2dd=ee06000
Yuchen Zhu, Wei Guo, Jaemoo Choi, Guan-Horng Liu, Yongxin Chen, and Molei Tao. MDNS: Masked diffusion neural sampler via stochastic optimal control. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025, |tLhtt /ttttp:://opereevvee http/foouumidxxxHHSBB
16

---

Published as a conference paper at ICLR 2026
A RELATED WORKS
A.1 CONTINUOUS DIFFUSION SAMPLERS
SOC-based Diffusion Samplers Formulating diffusion- sampling as a stochastic optimal control (SOC) problem has emerged as an active line of research. The Path Integral Sampler (PIS) (hang & Chen 2022) and Denoising Diffusion Samplers (DDS) (Vargas et al., (023) adopt memoryless reference dynamics and optimize the control either through direct gradient updates or adjoint-based training. LV-PIS (Richter& & Berner, (024) introduces a log-variance (LV) formulation, whose optimal solution coincides with the SOC optimal control. LV-based methods decouple the rollout path measure from the training objective, enabling flexible trajectory generation and avoiding the need to store full pathwise gradients or explicit adjoints; however, the loss still aggregates along entire trajectories. Adjoint Matching (AM) (Domingo-Entee Enrich et al. 2025) takes orthogonal approach by replacing full adjoint solves with a lean-adjoint estimator tailored to the SOC structure, thereby improving computational efficiency. Building on AM, the Adjoint Sampler (AS) (Havens) et al. 2025) incorporates a reciprocity (optimality) condition on path measures to further enhance training efficacy. Two recent extensions broaden this framework: ASBS (Liu al.. 2025) generalizes to non- non-memoryless references inspired by the Schrödinger bridge problem, while NAAS (Choi et al.
2025 employs annealed reference dynamics to provide stronger guidance during learning.
CE-based Diffusion Samplers Cross-entropy (CE) approaches replace forward- KL or relative- entropy training with a reverse-KL. projecto,,,ooecctonnn, typically of the form, which reduces to a weighted negative log-likelihood on trajectories. In the diffusion setting, this often appears as a denoising or score-matching style objective with importance weights to correct for mismatch between the rollout law and the target. [Phillips et al. (2024) utilize weighted denoising cross-entropy (WDCE) objective derived from denoising score matching. Their sampler combines importance weighting with diffusion-
learning loss to bridge from simple reference to the target. Blessing et al. (2025) couple CE training with Adjoint Matching (AM), and interpret the resulting criterion within the stochastic optimal control matching (SOCM) viewpoint (Domingo-Enriee Enrich et al., 2024). They introduce a trust-region mechanism that can be cast as a proximal step or an explicit KL bound, yielding more stable CE iterations. This work is concurrent with ours. In contrast, PDNS utilize proximal algorithm with WDCE objective, which can be applied uniformly to continuous time and discrete-tima diffusions, and is designed to be memory- while retaining the benefits of denoising or score-matching objectives.
Sequential Monte Carlo and Annealed Diffusion Samplers A long line of work studies Markoy chain Monte Carlo (MCMC) and Sequential Monte Carlo (SMC) methods that traverse a sequence of intermediate (often annealed) distributions connecting a tractable prior to a complex target (Chopin 2002 Moral et al.) 2006; Guo et al. 2025). Annealed importance sampling (AIS) provides the correction mechanism via importance weights (Neal) 200T Guo et al. 2026a). Early non-neural approaches learned or adapted transition kernels within this pipeline (Wu et al. 2020 Geffner & Domke, 20211 Heng et al.] 22020). With the advent of learnable proposals, normalizing flows (Dinh et al. 2014 Rezende & Mohamed, 2015) Chen al., (018)) have been used parameterize
SMC/MCMC kernels and to refine proposals via variational objectives (Arbel et al., 2021] Matthews ct al.
2022). These methods typically improve sample efficiency but still rely on importance weights
10 ensure correct marginal sampling.
More recently, diffusion-based annealed samplers that leverage score-based generative modeling (Song et al. 20215) have been proposed, including both non-learning-base approaches with rigorous analyses (c.g.. (Huang et al., 22024) He al.) 2024)), and learning-basee approaches (e.g.. (Vargas ct al. 2024; Akhound-Sadeea et al. 2024j). These methods use annealed dynamics derived from denoising diffusion processes and combine importance weighting with learning objective such as score matching or LV to mitigate mismatch between simulated trajectories and the target. Sequential Control for Langevin Dynamics (SCLD) (Chen et al.) 2025) improves training efficiency by off- policy optimization over annealed dynamics, although importance weighting is still required to achieve accurate marginals. A complementary direction augments diffusion models with MCMC components for sampling Boltzmann distributions (De Bortoli et al. 2024 [Akhound-Sadep et al. 2024) or adopts alternative training criteria such as action matching (Albergo & Vanden- Vanden-Eijnden
2025 Neklyudov al. 2023). In practice, these approaches often inherit significant computational
17

---

Published as a conference paper at ICLR 2026
cost from evaluating energies form importance weights, motivating schemes that reduce reuse weighting while retaining correctness.
A.2 DISCRETE DIFFUSION SAMPLERS
The study of neural samplers for continuous distributions has greatly inspired the construction of similar model structures and algorithms for discrete distributions. Holderrieth et al. (2025) proposed LEAPS that leverages the discrete version of Jarzynski equality and learns the transport via locally equivariant network, and Ou et al. 20256) follows this line of study and proposed discrete neural flow samplers (DNFS) with different loss formulation. The work Zhu al. (2025) solved the sampling problem through the SOC perspective similar to Zhang & Chen (022). and proposed samplers based on masked and uniform discrete diffusion. There is also line of research that leverages discrete diffusion for solving combinatorial optimization problems via sampling, e.g., (aanokossi et al. (2024) 2025). More recently, So et al. (2026); Guo al. (2026b) generalized the adjoint sampler (Havens et al.) (225) to discrete state spaces and build samplers trained from discrete versions of adjoint matching.
B COMPLETE ALGORITHMS FOR PROXIMAL DIFFUSION NEURAL SAMPLER
In this section, we present the complete algorithms for PDNS. We first detail the practical components; the replay buffer and the ---chedulerr. We then describe several instantiations of PDNS. One class employs priority-based resampling using computed weights in (15). The other variant uniformly samples from the buffer and uses importance weights the loss functional (see (3)). Moreover, we introduce two versions of computing weights and (10).
Replay Buffer In the beginning of each stage k, we sample set of trajectories (X(1N p or pp-11 Then, we compute the weight
1i=1
w = (x(()))
di (X(1))
dpref
NETT
OR c (1-))_ dipref
dipk (x(()), (17)
for each i 10. .N. To simplify the notation, denote (+) wh w depending the local optimization problem we choose. We store terminal state X T and corresponding weight ( into buffer, i.e., B = ((x:()) w(()
Adaptive Scheduler We provide a concrete implementation of the adaptive scheduler. Given a KL
trust-region radius > O, we aim to choose the largest such that satisfies KL(P% S e. With the samples in the replay buffer B
((x())
w(())
) the KL divergence can be estimated N
)===
via empirical measures, POR- ~ N Σ 8 8xc,, pk.
Zi=1 why 8xto, as follows:
dik
N
(L lIpk' - log
a.a (a)d1 (x)
(18)
(i() (1) is the weight for ()) computed through 17). and here we emphasize its dependence on 7). then select 7h by solving the one-dimensional leaat-squares to the trust russt-region constraint e:
N
= argmin
c+ N
Σ
log N (ia
(19)
1=1
In practice, we optimize this objective with L-BFGS optimizer (Liu Nocedal 1989). Note that the computation is negligible compared to model updates.
Resampling-based & Weight-based Algorithms
At stage k, let the buffer be "())
== ((x((()
Define normalized weights (l0) := w+( - w(), Then the empirical approximation of IP (or FF') ) is
P =
Σ () x(0
i=1
18

---

Published as a conference paper at ICLR 2026
We instantiate the CE loss in two equivalent ways: Weight-based:
N
(0) Σ
(
log Pe (x
(20)
Resampling-basedt draw indices I Ii...., IN (ategorical ,...,, ((N)) and set X ((1.), Then use the unweighted loss
N
(0)
1 N
Σ
log P (x
(21)
We refer these as the weight-basse based and resampling-basee algorithms, respectively, For both continuous and discrete diffuuion models, we can replace the log-likelihood in above two equations by matching losses. Let the score matching (Song et al., (2021b) bridge matching (Shi et al., 2023) loss be denote as L: X' + IR. Then, the corresponding WDCE variants of weight- and resample-base methods are provided in Algorithms and respectively.
Algorithm 2 Proximal Diffusion Neural Sampler (PDNS) vl
Require:
Model parameters 0, target distribution x e -BB reference path measure prof initial path measure 10% Buffer B with buffer size N, batch size B, number of inner iterations L. for in 2...... L do
2: Sample a buffer B == {(xx))
I by using an predefinecd/ddaptie scheduler.
3: for step in 1, ..., L do
4: Randomly select B number of indices I C NJ..
5: Compute the loss Eier ( ( and update
6: end for
7: Set Tp + Tpo
8: end for
Algorithm 3 Proximal Diffusion Neural Sampler (PDNS) v2
Require: Model parameters 0, target distribution x e - reference path measure pref initial path measure 10% Buffer B with buffer size N, batch size B. number of inner iterations L. for in 2....... do
2: Sample a buffer = {xx)) ( by using an predefined/adaptive scheduler.
3: Resample N numbers of tuple in the buffer B proportional to weights {w(1) }, i.e.,
J Cat(fw())),
B = (x hes-
4: for step in 1, 2, do
5: Randomly select B number of samples {X )== in B.
6: Compute the loss 1/2 C ((x(2) and update 0.
7: end for
8: Set Tp0c <- Tpo.
9:
end for
Table 7: Comparison between weight-based (Alg.
and resampling- (Alg. method.
Target
DW-4
IJ--13
Ising, 8, J 1, 9low 0.6
Metric
W/ (L)
E(-) W/ (4)
W/ (4)
E(-) W/2 (4)
ESS (t)
E(-) W/ (L)
Adaptive Training Steps (4)
weight
0.50
0.32
1.52
3.96
0.984
1.052
3080
Resampling
0.51
0.21
1.57
1.01
0.972
0.844
4339
Comparison between Weight- and Resample-based Algorithm We compare two instantiation of
method; weight- and resample-based methods. As shown in Tab. 7 reweighting and resampling
19

---

Published as a conference paper at ICLR 2026
achieve comparable performance in terms of the 2- Wasseerstein distance W2. However, resampling consistently yields better performance on the energy energy-weighter Wasserstein metric E(-) W2 in all tasks, across all tasks, albeit the cost of slower convergence when using the adaptive scheduler.
C
THEORY OF PROXIMAL CROSS ENTROPY AND THE PROOFS
C.1
PROOFS
Proof of (3).
Proof. As Pd Prof μa, by the chain rule of KL divergence, we have
[x)(x) r(X)) (I.(P")
= - ((x)) + E (x(x)) KL (0,77) (a.77oo(- (o)) Po (0.771((-1X66)
p (o.T (X(( (o.T) Xp) -r(X)) log pret
Ep(Xo) Erg (x.r
(0,T (0,7) Xo)
pPo (0,27j0 (X (0,77
((,))) Ego (a
log
Pref
[0,27]0 (X (o,r7] (0,T) X (orr)()))))))
Therefore, we have
P ((o771o(X(02211. prrof (o.T) (X (Xo)e
Let the normalizing constant the right-hand-side be
(X)
Exper (Xoo.T) (Xx)) Epput (xr) (xxx)
which implies
IPt (o (X pref (0.77 (X (X Xoo
P (X) 1prof (x))_(X))-V((0))
Thus, by integrating out (()))), have
Xr) = Po, (Xo, Xx), (xxx)))))
(22)
Due to the memoryless condition,
VolXo) (x)) (XT) E., e-BE
7, BE Therefore, we have which is the normalizing constant of x -BE
(X,, Xx) = ((X))u(XT) 9E(Xr) v(xr) Z = ((Xo)T(XT))
which completes the proof.
Remark. If the memoryless assumption of pref does not hold, then from 22 have
(+7)))= / paatt (Xo, X (xr)
f P (XolXT) Vo(Xo)
pref (Xx)))(Xx))
(Xa)
which is intractable and typically not equal to n, i.e., the optimization problem (3) would not lead to path measure that ends at our target distribution и. We refer readers to Domingo-Enrich et al. (2025) for further discussion on the memoryless assumption.
Proof of Prop. 3.1
20

---

Published as a conference paper at ICLR 2026
Proof. 1. Similar to the proof of 3 have
(x) r(Xr) - KI[(PT)ppeef + 7/6
(X) + log
(I (X))
- (x) -r(Xx) + log pref (X)
(p0x-1 (X))))1))))
(F (X)) tnnl/n
(x)e(x) log peel (X (X)er(Xr) (Pok-) (x)
(X))
p
Me
Ege log (P (P (w/))///111 (1+mm) (10r (1(((t)))
?]
1 + KL
(F | + const.
2. The first claim is easy to prove by induction. To obtain the second claim, it suffices to note the relation IP* (X) oc prof (X)ee(( where r = log z + const.
C.2 ABSTRACT FORM TO CONTINUOUS DIFFUSION MODELS
C.2.1 SOC FORMULATION FOR SAMPLING PROBLEM
Sampling Problem Suppose (2;) x exp( exp(-8E(x)) for E X' C CR. Suppose that the reference
dynamics is given follows:
dx, = (X,) o, Xo μ..
(23)
Let pprel be the path measure induced by 23 Let r(t) - BE(() - log v(x), where = Pref
Our goal is to obtain a sample æ from the target distribution n, i.c.
dm
x π
e (1phpr
(24)
SOC-based Difffusion Samplers that the random variables Xo and X1 generated from the reference process 23 are independent, i.e.
Assume that the reference dynamics are memoryless, meaning
0,7 (Xo, X
Prr (Xr) (Xr).
(25)
Given this setting, now consider the following Stochastic Optimal Control (SOC) problem:
min E x up /0 (x,) - - r(Xr)
(26)
s.t. dX, = (b, (X.) + on (X,)) dt + o.dWe Xo
Let Ta be the path measure induced by the controlled dynamics and let u* be the optimal control. Then, the optimal path measure P : pa" can be analyzed as follows (Domingo-Enne Enrich et al. 2025 Liu et 2025):
(X7) / Por (Xo, Xr)e (Xx)) too Xo)
tooo) (Xo) 0o(x) log / ( dy. (27)
Due to the memoryless condition, up = const there by
(++)) = 1 ((X))u(XX)) (xx) dXo =xX)=x(X7))000 u(XT)o x x(XT).
(28)
Thus, The terminal distribution matches our target distribution, i.e., PP= n. Note that we call the objective in (66 as an relative entropy (RE) loss.
Cross-Entropy (CE)
Cross-ente Cross-entropy counterpart of 26 be defined as follows:
min KL(PP
(29)
21

---

Published as a conference paper at ICLR 2026
Now, decompose this eross-entropy objective into the tractable form:
dlp
KL(PP* IPP)) = xP
log
(X)
= ExxP= [log (ogP"(()] +C
(30)
dlP
cIP*(()
= Rx-FV
(X) log Pa (X) +C,
(31)
where is arbitrarily given control. Note that need auxiliary control sample trajectory since we do not know the optimal path measure P*. Instead, we need to compute Radon-Nikodym
(RN) derivative:
dl
dF
dlPref
memoryless X
1lx (1x vx.,) -dW, + r(xr)
dlpref (X)
dPu (X)
exp
dle
(32)
C.2.2 WEIGHTTD DENOISING CROSS-ENTROPY (WDCE)
Reciprocal Property
The reciprocal property Léonard
2012
Léonard et al. 2014
Chen et al.
2021
Shi et al. 2024) implies that
(*()) = (PT) (X(XT) = Pūr (Xo. XT) (X
(33)
When the reference dynamics memoryless,
(Xo, XT) & (((o))((X))e x x (Xx))=(XT),
(34)
which implies that I is memoryless, i.e., the joint marginal distribution (Xo, (Xx))((x)) is independent. Combining 33 and (34,. we obtain
X1)
(*() p(Xo) (XT 10, Xo, (o,
(35)
WDCE
The following SDE induces the optimal path measure F
dX; (b,( + of Ep-
log Pr ( (Xr x.)) dt + o,dW, Xo
(36)
using Doob's fi-transform theory (Rogers & Williams, 2000). Note that this is generalization of score matching (Anderson, 1982 Song et 2021b).
Then, our CE objective (29) can be extended
denoising CE as follows:
[ 2
argmin KLP argmin E Ex
2 (l,) a, V log Prje (X, ( xx)
(37)
Finally, we obtain the weighted denoising cross entropy optimization problem:
"² argmin Ex
(lu((X)) or Vog (Xr
(38)
argmin Ex Ex E Xx [lr (X+) o. V. log (X
(39)
2
u
dl (X)
(- (lx) V.V log (Xx
argmin Expare E Pxp
(40)
dpv
"
Therefore, combining we can run WDCE by the following iterative process:
1. Sample N trajectories { ~ P". J==
2. Compute weights ( j by (32 for each corresponding trajectory (X(O)N
i==1-
3. Resample {x(() j by following categorical distribution:
{ ~ Cat (( (x where (1)
(41)
ΣN
N w(
4. Update the control ug := u through score matching loss with resampled data.
Remark. Particle Denosing Diffusion Sampler (PDDS) (Phillips et al. 2024) is one of the sampling method which leverages this Importance weighted CE method.
22

---

Published as a conference paper at ICLR 2026
C.2.3 THEORIES ON PROXIMAL DIFFUSION NEURAL SAMPLERS
In this section, we introduce the theory for PDNS written in SDE formulations.
Proximal Relative Entropy Suppose we have the current control w*=1 and the corresponding path measure [pk-1, Then, consider the following iterative method:
u = argmin Ex~e" [-r(XT)] + (F + KL (L(P"IIRR-)
u
7/k
(42)
J.' 24 (X,)11 +
(X,) u (XiII dt - (ee-((X))
arggmnn Ex~P"
277
Note that we denote (42) proximal RE (Zhang 2026
Proximal Cross Entropy Solving the Proximal RE problem in 42 is computationally expensive, as requires simulating trajectories under the current (online) policy and solving numerical optimization problem, often via the adjoint system (Yong & Zhou 1999. Nüsken & Richter 2021
Domingo-Enrinn et al. 22025] Havens al. 2025). To mitigate this cost, we reformulate the Proximal RE objective as a cross-entropy problem by (11 and (12)
(r((T)) dr
KL(P (aPt) x Ex Ex
dp (x)
log (X) (X)+C.
(43)
dlpeef
(L(P"1|P)) x Exxpe-1
((l-Ae)r(Xx) dl (x) log po (X) + C,
(44)
Now, let uh 0x 1 (w) (resp. -
be the control that induces the path measure 1p (resp. pak-1). Using the Girsanov theorem (Kkenddal 2003; Särkkä & Solin 2019), i.c.
(x:) -
dpret dp
8 lu
dx + at 5 (X.) dW,
exp
(45)
Precisely, the weights in
435 and 44
can be computed as follows:
(r((T)) d
dPpb_. (x)
+ + (X,) dW, - ((80))) (46)
(x.)l |
k-
cxp
+ 7k
2
and
(1-A)))(X7) dereff
dIpA-1 (X)
10 (Xx)) + af -x,) .dW, - -(1 101--0111))))) -dV::(1>>))(xx))]
exp
(47)
Proximal WDCE
Similar the derivation of (88 can easily derive
Jo" 2 [1( (X+) V
6 EE x
log pref (XT)
(48)
By leveraging reciprocal property, i.e., poi (X) = (Tx))) analyze the PCE objectives:
IPref |0.7
X Xo, Xr). we can further
[ Iluf(()))) log
o argmin
E Xp=m.Xy-~PPR
E Peet 10. (1Xo.
Pref Pyof
(X,) (x)))
(49)
By applying denosing score (bridge) matching objective (Song et al. 202 Shi et al. 2024). we
obtain proximal WDCE objective:
23

---

Published as a conference paper at ICLR 2026
F(Pr;poi)
E ()
1 1// (Xx)) - oz V log pT] (XT Xx)1
(50)
P Xx)
Nompp X
dipoz
(x))llu((Xx)) - as V log 2 (
(Xr|X:))1
(51)
(0,T) Xoru
fo (1)
d
Xp
C.2.4 PROXIMAL WDCE IN PRACTICE
Desampling-based Algorithm We instantiate Proximal WDCE via 50). To obtain 50 First, to obtain a sample X P required in (3), we apply importance resampling:
1. Draw N trajectories {x(1) fi-1 from po
2. Compute weights
(Xx) d dp (x)
26 76+T
8 (X)
(46
(52)
3. resample (x(O) li=1 by following categorical distribution:
{ N E Wiy (X())
(X)(1) "
N
(8r) ~ Cat
(X(())
10 N
li=1
(53)
i=1
Weighting-based Algorithm
Alternatively, using (1). draw Xyp PT-1 and and incor-
0% porate Wrik directly in the objective, i.e., optimize weighted loss where each sample contributes proportionally to importance weight rather than being resampled.
C.3 ABSTRACT FORM TO DISCRETE DIFFUSION MODELS
C.3.1 SOC FORMULATION FOR SAMPLING PROBLEM
Sampling Problem
We suppose the target distribution is supported on Xo f1, Nj and let the mask-augmented state space be = (x) x -8B supported on Xn-
Let the target distribution be
The reference path measure pref has a generator Qref defined by
20
x' M, m n M:
N
Qi (r. v)
2' M) x =
if otherwise.
7: [0,7] R. be any function satisfying f/" 7())dI oo. By construction, IPgef μ Pmask the delta distribution on the fully masked sequence, and pref Punit is the uniform distribution on Ng. Let = F3..
Consider the following SOC problem:
min E Xxppo
T ["
Σ
Q
log
Oi Qiee
O₂ + giet)
(Xt: r(XT)
s.t. X (Xt)xEjo,7) is CTMC on X' with generator Q", Xo Pmaas:
Let the generator Q" be parameterized by
(1))) (t)ffi M]
if M. x M;
Qx
f = y if otherwise,
24

---

Published as a conference paper at ICLR 2026
where so(z) is d x N matrix whose each row is probability vector. The optimal value of be proved to have the following form (Zhu et al. 2025):
(0( Pr (X' X~x
n n = x xull), if = M.
C.3.2 WEIGHTED DENOISING CROSS-ENTROPY
The WDCE loss follows:
KT. E Ex-ro BA-Uni((0.1)
Σ
log so (T)d.x4
dizd=
E xaplees (or(xx) dlpref dpm
1
E [n-Unif(O.1)
Σ
log so (0()))),
A
d
KT.(PP*|1P0)) = Exare EX-Uni((0.1)
Ā (()T)
Σ d
log 80 (0080)())))) (3)
x
E Xapee (ell-A))(Xx))
dP (x)
deref
Ex-Unt((..))
x
log 80 so()),,
=M
where (+x()))) means independently masking each entry of x with probability A. The weights can be computed by
dipref
1/N
dp3, (X) = exp
expp
Σ
log
S6211 (X. (X)))) x()))((()
8:X, #X.
where assume the jump time happens at the i(t)-th entry, and the total number of jumps is the sequence length d. The design choice resampling- based and weighting- algorithms can be done similarly as discussed in the continuous part above.
D EXPERIMENTAL DETAILS AND FURTHER RESULTS FOR LEARNING CONTINUOUS TARGET DISTRIBUTIONS
D.1 FURTHER EXPLANATION ON THE TARGET DISTRIBUTIONS
In this section, we describe the benchmarks used our continuous-energy experiments. We evaluate seven targets: four synthetic distributions (MW-54, Funnel, GMM40, MoS) and three atomistic potentials (DW4, LJ13, LJ55). For the synthetic benchmarks we follow Choi et al. (2025): Chen et al. (2025). For the atomistic systems we follow Havens al. (2025); Liu et al. (2025). For convenience. let d be data dimension. Unless noted otherwise, we use the hyperparameters in the cited sources.
MW54 We use 5-dimension many-well potential. Precisely, we consider the energy function defined as follows:
E(x) = > - -)2 i=1
(54)
where (Z1, R2), d d= 5, and 4. This leads to 2d 32 wells in total.
Funnel We use a funnel-shaped distribution with d = 10. Precisely, the unnormalized density is defined as follows:
N(z;00,0²) x N (10)
(55)
where (I1, ..... va) and 3.
25

---

Published as a conference paper at ICLR 2026
GMM40 Let 40 be the number of modes for the Gaussian mixture. The target distribution the equally weighted mixture
1 m
x (x)
Σ (X), Ti = He ~ Ua(-400,00) m
(56)
where (a b) is the uniform distribution over the hypercube (a, bjd. This benchmark demonstrates mode exploration in high dimensions.
MoS Following Blessing et al. (2024) consider a heavy-tailed analogue with 50 and m 10:
(④) = Σ (X), (T(x) t2 - (p)), #4 ~ Ua(-10,
(57)
where l2 denotes the multivariate Student's Ł-distribution with two degrees of freedom and identity scale. The heavier tails make this target notably more challenging for samplers.
DW4 This potential was originally proposed in Köhler et al. (2020) and used in Akhound-Sadegh et al. (2024); Havens et al. 2025). is a double-well potential with particles in two dimensions, hence total d = 8 state vector. We denote state = [Ir 22: 403 v4a IRd with E RR.. The energy function reads
E(x) = exp
Σ ドマ
(a(d;) - do) + b(dey - do)2 c(d; - do)*)
(58)
2T
where di = || 2yll2 is the Euclidean distance between particles and j. We follow the configuration with a 0, -. 0.9, do and temperature 7
LJ-13 and LJ-55 The Lennard-Jones (LJ) potentials are classical intermolecular potentials com- monly used in physics to model atomic interactions. Consider n particles in 3-dimension space, state [zi;...; E 3 with positions I; E R The suffix in "LJ-13" or "LJ-55" indicates the particle count (n = 13 or n = 55). We use the following unnormalized energy:
127
E(x)
Σ iej
I'm
I'
Σ
1 - C(x)11
(59)
2T
diy
2
where di ||, 22 is the pairwise distance and C(x) Σ x; is the center of mass.
Di=1
Following prior work (Havens et al. 2025 Liu et al. 2025), We use the parameter values Imm = 1, c 1, 0.5, and Note that the - and D systems correspond to d = 39 and d = 165, respectively.
D.2 IMPLEMENTATION DETAILS
D.2.1 MEMORYLESS DYNAMICS WILD
In this section, we first specify the reference dynamics used in practice. Then, we derive closed-form form expressions for (i) the reference bridge law pref "Ju,r" which yields the conditional distribution needed to
sample Xr given (X,,Xi in (3) and (ii) the score V log T/e appearing in 13 These formulas enable exact sampling from the reference bridge and analytic evaluation of the score.
Orastein-Uhlenbed (OU) Process
Throughout the continuous target benchmarks, we OU
process (Uhlenbeck & Ornstein]
1930) which is defined as follows:
dX, Ot x,dt /o,dWW Xo μ NN(0001)
(60)
Se and a. o Va,. In this case, every marginal of the reference measure follows the stationary Gaussian distribution, i.e., pref == W(0,02)) for L |0, T). Thus, μ == which allows us to use a tractable terminal reward r(x) = (BE(2) log v(x). The conditional probability can be written as follows:
Note that 6x(x) -
pref Pro (-1x) =N (e (e-2 -- ands x, 0² -
- bb5
a.de)1).
(61)
26

---

Published as a conference paper at ICLR 2026
Note that when e $ J a,ds = 0, then the joint of the marginal distributions Por is independent, i.e. prof memoryless. In practice, we set
(1 111-tyo m) + to
(62)
where (aimin (max) a given hyperparameter. To ensure that the reference dynamics (almost) memoryless, we set O'maxx >> 1.
Closed Form Conditional Path Measure for OU Process Suppose the reference dynamics
given 60). Then the conditional probability paref *tJo,r is written as follows:
Pref e/0,7 (,1 XT) =N =N( B. (1 ())
Xo + Cx(1 - -B))
XT, 0² (1 BB) (i) (1-B)XI-CE))
1
(63)
B1
1 B1
1 B1.
where
1]t [ ogds,
(-1 B. exp
1 L' a,
( asds)
C. exp -
Closed Form of Conditional Score Function for OU Process function V., In ppref is written follows: T/
Moreover, the conditional score
me log Ti T (Xr(X,) = -C,(C,,, - --)
(64)
02(1 - (t)
D.2.2 INITIALIZATION OF THE CONTROL
The typical choice for initial control would be u = 0, i.e., ppo pref However, non-informative reference dynamics, i.e. u. (x) 0, it takes long time achieve reasonable performance with our proximal method when the potential is in the high-dimensional sparse domain. Hence, for the use the informative initial control u (or path measure (o6) instead of starting from trivial control. To obtain informative sample XT at the first stage, we use annealed Albergo & Vanden- 2025)
continuous target benchmark,
dynamics (Neal 2001 Matthews et al.) 2022)
Guo al., 2025
defined as follows:
04 VV,(X,) + /a,dW, dx, = -
Xo ~
(65)
We set
Vx(x) (1 - V)V log u(() + tv log (x) (1 V log u(x) (x))
which is a typical annealed dynamics used in Albergo & Vanden- Vanden-Eijnden (2025); Chen et al. (2025;
Choi et al. 2025). At the first stage, we train WDCE objective with the samples {X- (3)
J/V obtained
Trom the annealed dynamics. (We use uniform weights for the initial stage.) For energy with steep landscapes, clip the energy gradient norm with predefined (VE) which is treated hyperparameter in Tab.
D.2.3 HYPERPARAMETEE SETTINGS AND METRIC
Model Backbone For the MW54, Funnel, GMM40, and MoS experiments, we adopt the model architecture Choi et al. (2025) ; Liu et al. (025). For the DW4, LJ13, and LJ55 benchmarks, we follow the architecture in Liu et (2025). which employs an Equivariant Graph Neural Network CCGNN:Satorra, et al. 2021).
Hyperparameters The hyperparameters for all experiments are provided in Tab. g The number of stages is iteration number of outer loop (line 1 in Alg. T and the number of inner epochs is the number of epochs used for inner loop (see line Alg. We use Adam optimizer with no weight decay for all benchmarks.
Evaluation Metrics For Tab. 2 we compute MMD and entropic OT (Sinkhorn, & 10-3) using the implementation of Blessing et al. (2024); baseline results are taken from Blessing et al.) (2024); Chen et al. (2025): Akhound-Sadegh et al. (2024). For Table Tab. 3 DW/LJ tasks where rigid- motion and particle-permutation symmetries are relevant, we additionally include symmetry-aware geometric Wi following Akhound Akhound-Sadegh et al. (2024) with the alignment heuristic of Köhler et al. (2020). All metrics 2.000 generated and 2,000 reference samples. Note that the LJ-55 experiment requires approximately 80 hours of training on a single H100 GPU, with a GPU memory usage of around 14 GB. The total training time is comparable to ASBS (Liu al. 2025
27

---

Published as a conference paper at ICLR 2026
Table Hyperparameters for continuous target distribution benchmarks.
Cluss
Hyperparameters
MW54
Funnel
GMM40
MoS
DW4
LJ13
LJSS
Dimension
5
10
50
50
8
39
165
2.0
3.0
50.0
15.0
2.0
2.0
2.0
(Ominn (mmax)
(0.1, 10)
(0.1, 10)
(0.1.5)
(0.1,10)
(0.1.
(0.1. 10)
(0.1, 10)
SDE
Emax
100.0
100.0
100.0
100.0
100.0
100.0
100.0
NFE
200
200
1000
1000
1000
1000
1000
1.0
1.0
0.1
0.1
0.1
1.0
1.0
Bufler & Scheduler
[8]
10a
10a
10
10
10
10
10*
Stages
20
20
10
10
30
50
20
Training Hyps
# Inner Epochs
50
200
5000
1000
1000
200
500
Batch Size
500
500
2000
2000
500
500
500
Ir (g)
10-4
10-4
10-1
10-4
10-4
10-4
10-4
Optimizer & EMA
(.,0.9)
(.,).9)
(0,0.9)
(0,0.9)
(0,0.9)
(0.0.9)
(0,0.9)
EMA decay
0.999
0.999
0.999
0.999
0.999
0.999
0.999
Type
MLP
MLP
MLP
MLP
EGNN
EGNN
EGNN
Model
# Layers
4
Channels
256
256
2048
2048
128
128
128
t= 1.00
t=1.00
t=1.00
t-1.00
Target
10
10
-10
-10
-0
-20
-20
-20
-10
-20
7/6
0.1
[=1.00 11.00
[=
1-1.0
1=1.00
Target
20
20
0
0-
0-
-20
-70
-20
-20
-20
(b)
0.5
Target
70
20
20
20
+
0
-10
-20
-20
-20
-20
-20
-10
20
10
-30
-10
20
(c) = 1.0
Figure 5: Ablation studies on fixed 7k for stages on MoS benchmark. We fix " constant for all stages and visualize the first four stages (A: 1, 3, 4; left to right). Larger "k. i.e. weak proximal regularization, leads to rapid mode collapse whereas smaller 7x preserves multimodal coverage. These results are consistent with the analysis Sec. 3.1
Tk
to
I+m
D.3 EFFECT OF PROXIMAL POINT METHOD
In this section, we study how the proximal term influences training and sampling quality. We first compare proximal WDCE against the non-proximal WDCE baseline under identical settings. We then perform ablations over the proximal step size }m and its scheduling policy, reporting metrics across epochs (e.g., Sinkhorn and MMD) quantify convergence and mode coverage.
28

---

Published as a conference paper at ICLR 2026
E=0.01
E=0.1
E=1.0
y=0.1
y=0.5
y=1.0
2500
0.6
2000
0.5
Sinkhom
0.4
1500
MMD
0.3
1000
0.2
500
0.1
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Training Epochs
Training Epochs
Figure 6: Ablation on proximal step size 7 and the choice of the scheduler on MoS. We evaluate Sinkhorn (4) and MMD (4) across training epochs for multiple choices of proximal step size aik and scheduling policy. 7 denotes "Yk: == It0s over stage k. The legend entry "y = const" denotes runs with a fixed 7x for all stages k. Note that 0 < 7he S 1; larger 7h weakens the proximal effect and approaches the non-proximal WDCE variant, whereas smaller "ik imposes stronger regularization.
Vi
8=1.0
y=0.1
y=1.0
1.0
35
30
0.6
20
0.2
10
5
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Training Epochs
Training Epochs
Figure 7: Ablation on proximal step size 7/s and the choice of the scheduler on LJ-13. We monitor and energy 2-Wasserstein distance (E(-)W((L) across training epochs for multiple choices of proximal step size 77 and scheduling policy. y denotes ? := I+ over stage k. The legend entry "r const" denotes runs with fixed "/ for all stages k. Note that 0 < 7 S 1; larger " weakens the proximal effect and approaches the non proximal WDCE variant, whereas smaller % imposes stronger regularization.
We
Baseline Proximal We study the step size by using fixed schedule, i.e., constant 7] across stages. Define 7 == I+m which appears as the tempering exponent on the path weights in (8). The original WDCE is recovered at % (7 oc). We visualize generated samples for the first four stages (k 1, 3, 4) at 7 {0.1, 0.5, 1.0). As shown Fig. when Y is large (0.5 1.0), the sampler quickly collapses onto partial modes in the initial stages. In contrast, = 0.1 preserves all modes across stages, indicating that stronger proximal regularization early in training helps prevent mode collapse and yields steadier progression toward the target.
We
Effect of Now we analyze the effect of c. In using adaptive scheduler, e the most important hyperparameter since it determines proximal step size 7)s for each stage k. Decreasing tightens the trust region (smaller 7)k, smaller 7x.), which tempers weights, reduces gradient variance, and improves stability; however, it results the slower convergence. Increasing relaxes the constraint (larger mek Th Rs 1), approaching the WDCE baseline with faster but potentially unstable steps (c.g., weight collapse).
On the MoS benchmark (Fig. 6 a very small radius (c= 0.01) is robust but converges too slowly to reach the best scores within the allotted epochs. A moderate radius (e 0.1) attains steady, monotone improvements, while large radius (e = 1) converges quickly at first but exhibits mild degradation in the final epochs. Overall, the range (..1 < offers favorable stability-speed trade-off.
29

---

Published as a conference paper at ICLR 2026
Effect of Adaptive Scheduling As shown in Fig. under the adaptive rule, the selected Y Tk tends to increase over stages (Fig. 7 small, conservative steps are chosen when the sampler is far from the target, larger steps become feasible as the model's support aligns with the target and more trajectories carry useful signal. This aligns to intuition. In early stages, only few samples lie near the target, producing a highly skewed distribution of pathwise weights in which small set of trajectories dominates the loss; smaller yı (or 7ik) tempers these weights, stabilizes gradients, and prevents premature mode concentration. On LJ-13, the adaptive scheduler consistently outperforms fixed-y schedule.
1+m
Groundruth
PDNS tous!
Gronndouth
PONS (oun))
415
(a) MW4
(b) Funnel
Grouddruth
PDNG
Groundruth
FORS tours)
AUGA
(c) GMM40
(d) MoS
Figure 8: 2D kernel density estimates of ground-truth generated samples, projected onto first two coordinates, for four synthetic benchmarks: (a) MW54, (b) Funnel, (c) GMM40, (d) MoS. The samples shown red dots.
DW4
LJ13
PDNS Grannd Thuth
0.06
Lround ruth
PONS (ours)
Normalized Densty
Normalized Density
0.05
0.03
0.10
0.02
0.05
0.01
0.00
22
-76
-24
-18
-12
70
66
50
40
-30
-00
Energy E(X)
Energy Elx)
Figure 9: Energy histograms for DW-4 and LJ-13, PDNS produces energy distributions that closely match those of the ground-truth samples.
D.4 ADDITIONAL QUALITATIVE RESULTS
Synthetic Energy We present qualitative results of our learned continuous diffusion samplers. First, we visualize 2D kernel density estimation (KDE) plot for synthetic benchmarks (MW54.Funnel GMM40, MoS) projected onto first two coordinates. We use 10000 samples to estimate KDE. The samplers cover all modes on MW-54, GMM40, and MoS, and capture the steep, narrow basin in Funnel. We further compare energy histograms under the true potentials for DW-4 and LJ-13 (Fig. 9): the distributions of E(x) for PDNS closely match those of the ground- -ruth samples, providing qualitative support for the quantitative results in Tab.
30

---

Published as a conference paper at ICLR 2026
Reference samples
Generated samples
Generated samples
10
Donkty
107
Figure 10: Ramachandran plot for AD between ground-trum -tuuh and PDNS.
Figure 11: Ramachandran plot for Tpo
Alanine Dipeptide the AD task, the Ramachandran plot in Fig. shows that model misses a few modes with extremely low probability mass. This mode collapse originates from the initial reference samples produced by the initial path measure IPo, which already suffers from mode collapse (see Fig. 11 Because the proximal updates refine po in a progressive manner, modes absent in IPo may not recovered during training. Note that Fig. 11 plots all samples with energy Exx) < 1000 obtained from 10t terminal samples drawn from F highlighting the missing low-probabilic regions. We leave the exploration of improved initialization strategies for constructing a better 1e0 as future work.
E EXPERIMENTAL DETAILS AND FURTHER RESULTS FOR LEARNING DISCRETE TARGET DISTRIBUTIONS
E FURTHER EXPLANATION ON THE TARGET DISTRIBUTIONS
Lattice Ising and Potts Models We consider Ising and Potts models defined on square lattice with I, sites per dimension, = (1.... I.12. The state space is Xg = [+1]A for the Ising model and No {1, .... 4) for the Potts model with spin states. Given an interaction parameter J R, the potential of any configuration æ E No is given by
Ising: E(x) -JJ Σ x x'x'; Potts: E(x) = -J -J-- Le ANj
(66)
indj
where i ~ j means that i,j E A are adjacent on the lattice. For simplicity, we impose periodic boundary conditions in both the horizontal and vertical directions. With inverse temperature 3 > the target distribution is the Gibbs measure x(x)= x 9E(a).
Combinatorial Optimization problem is defined as E(x) = E 1-9,0 where a; 2x; 1.
The energy of the target distribution for solving the maximum cut
Note that and æ are now graph-dependent, but we omit this in the notation for simplicity.
E.2 IMPLEMENTATION DETAILS
E.2.1 MODEL BACKBONE
Lattice Ising and Potts Models Following the implementation in Zhu al. (2025). we use vision transformers (VIT, Dosovitskiy et al. (2021) as the backbone for the discrete diffusion model. Specifically, we adopt the DeiT (Data-efficient image Transformers) framework (Touvron et al. 2021) with 2-dimensional rotary position embedding (Heo et al. 2025). which effectively captures the 2-dimensional spatial structure. For learning both the Ising model with L and the Potts model with L. 16 and a 4, use 4 blocks with 128-dimensional embedding space and 4 attention heads, and the whole model contains 829k parameters, which larger than the model used in Zhu al. (2025) due more challenging target distributions here.
*We do consider the external field for simplicity.
31

---

Published as a conference paper at ICLR 2026
Combinatorial Optimization We leverage variant of diffusion transformer (DiT. Peebles & Xie (2023) ) known as the Graph DiT (Liu et al. 2024), with 64-dimmesionall hidden space, 6 blocks and 8 heads, and the total number of parameters is 483k.
Suppose the graphs in the data set have at most D nodes. The score network so takes the following
three components as input: (1) x X {0, 1,2 M) (2) ((G E Dx the feature of the graph G obtained from adjacency spectral embedding, and (3) v() E {0,1)D an indicator of whether a node is valid in graph or just dummy (padding) node. 4 is constructed as follows: suppose
RDxD be the adjacency matrix of G, whose eigenvalue decomposition is written as A(G) = U UAT Then given the dimension of graph feature E E [1, D). (() == A (i.. E RDXE, It is easy to verify that adding padding nodes does not change the embedding, except for appending lines of zeros. Through all of our experiments, we choose E = D/2 by default.
the graph G, after padding dummy nodes, has D nodes. Let A(G) E R
E.2.2 TRAINING HYPERPARAMETERS
Lattice Ising and Potts Models For these distributions, we choose the batch size as 256, and use the AdamW optimizer (Loshchilov Hutter 20T9) with a constant learning rate of 0.001. We always use exponential moving average (EMA) to stabilize the training, with decay rate of 0.9999. In computing the loss, we use 4 replicates per buffered sample for the Ising model and 8 for the Potts model due the memory constraint, and resample the buffer every 10 gradient updates with the buffer size B = 512. All experiments are trained on an NVIDIA H100 80GB HBM3.
For learning both models Beritical, we use 25 outer loops, each with 4000 steps of gradient update. In the first 20 outer loops, the A decreases linearly from 1 to O; for the last five outer loops, À is set to 0 (i.e., the step size Tis= oo so that we refine the final target distribution). To learn the models Stow, we start from the best checkpoint at Periiical and directly use infinite step size.
Combinatorial Optimization For the maximum cut problem, we choose the inverse temperature of the target distribution as 3 = 5, and use three outer loops with ' = 0.8, A2 = 0.4, À = O. For each outer loop in Alg. [ we train the model for 20 inner epochs. In each epoch, we first randomly sample 32 graphs from the training set, and for each graph, we generate 16 samples along with their weights as buffer. Then, we iterate over the samples in the buffer one by one for 10 rounds. Each time, the samples that correspond to the same graph along with their weights are used to compute the loss. All experiments are trained on an NVIDIA RTX A6000.
E.2.3 GENERATING BASELINE AND GROUND-TRUTI SAMPLES
Lattice Ising and Potts Models For learning-basee baseline, we mainly compare with LEAPS (Holderrieth et al. 2025). We use the same model backbone as introduced in Holderrritt et al. (2025). For a fair comparison, we adjust the hyperparameters to ensure the models have comparable number of parameters above: we choose kernel sizes [3, 5, 7, 9, 11, 13, 15] and 32 channels. For learning Ising models, the number of input channels is 1, making the total number of parameters 919k, while for learning Potts models with q 4, the number of input channels is 20, making the total number of parameters 928k. For each inverse temperature, we train LEAPS for up to 100k steps for each temperature, which is comparable to ours.
We also include two MCMC baselines. First, we use the Metropolis-Hastiee (MH) algorithm with single-site updates. For both Ising and Potts models, and among all temperatures, we use batch size 1024, and warm up the algorithm with 2 = 1048576 burn-in iterations. After that, we collect the samples every 2 65536 steps to ensure sufficient mixing, and collect for 64 rounds to gather total of 216 65536 samples. Moreover, we run the SW algorithm on all target distributions to generate examples accurately distributed as the ground truth. We use batch size of 256, and warm up the algorithm with 2 8192 burn-in iterations. After that, we collect samples every 128 steps to ensure sufficient mixing of the chain, and collect for 256 rounds to gather a total of 216 65536 samples.
For evaluation, we draw 212 = 4096 samples from each trained model along with their weights to compute the ESS (see Sec. E for explanations for the computation of ESS). We also use the same number of samples from the MH algorithm. To obtain a precise estimate of the ground truth, we use all 65536 samples from SW compute the magnetization and 2-point correlation.
32

---

Published as a conference paper at ICLR 2026
We refer readers to Zhu et al. (2025 App. D.3 and E.2) for the definitions of the evaluation metrics (magnetization and 2- 2-point correlation) for Ising and Potts models, respectively. In particular, Fig.
plots the average 2-point correlation 1 (C (k; k + +r + C (k, k r)) therein.
Combinatorial Optimization For each type graph, we generate 1024 graphs as the training set, and another 32 graphs as the test set. We solve the problems by Gurobi (Gurobi Optimization, LLC 2025 and take the output as the ground truth.
E.3 ADDITIONAL RESULTS
Monitoring Training Procedure of PDNS In Fig. 12 we demonstrate an example of the PNDS
learning procedure for the Potts model with q = 4 states on 16 16 lattice, at Peritical = 1.0986.
In the k-th outer loop, we fit the target path measure IP*
We now illustrate the computation of the effective sample size (ESS). For i..dd trajectories X (1:M) ~ dp* (X(()) computed through (15) and (66 and dF
pd associate ( with local weight W/ normalize them such that Σソ wy = 1. Associate X () with a global weight w/ x
(X oc
dFc
apsed
dFo (X computed through (16), and also normalize them such that Σ
25=1 w/ 1.
M
M w²) Then, the local ESS w.r.t. Pk is defined by (ME
and the global ESS w.r.t P is defined
M/ myy ) Both values are in M 1]. with higher value typically indicating better fit
by M jj=1 between two path measures.
Essential Sample Size (ESS)
Global Local
8.8
Path EL Divergene 8
0.6
0.6
0.4
15
0.2
Eval, global Tran
Eval: local
10
0.0
20000
40000
GD000 80000
100000
20000
40000 60000
80000 100000
20000 40000 coco0 8o3co 100000 0.
Training Epochs
Training Epochs
Training Epochs
Figure 12: An example of PNDS training for Potts model with 4 states on a 16 x 16 lattice, at Peritical 1.0986. Left: the schedule of ' 's. We 25 outer loops with 4000 steps per outer loop. In the initial 20 outer loops, we choose step size 77s 's such that A decrease linearly from to 0. The last outer loops use πæ 30, Å and evaluation during training. Right: the KL divergence between path measures during training. In the middle and right plots, "global" means with respect pa and "local" means with respect to pp
refine the final target. Middle: the ESS for training
Ablation Study Scheduler Throughout this part, the target distribution consider is the lattice Ising model with I. = 8, J 1, and Stow 0.6, which shares the same physical property as the Ising mode on larger lattices presented in Sec. 5.2
In Fig. I3 we test several predefined schedules for 7Js and dg: three of them are with linearly K uses the same number of training epochs, and the remaining 7/6 and the number stages K is chosen such that Ag is around 0.01
decreasing d and each stage three are with constant 7/ = I+m (the final 7x is set to be one). We fix the same number of training epochs (10000). We observe that all schedulers tested are able learn the correct target distribution, while schedulers with constant Th shows faster convergence in terms of the metrics due to its faster decrease in the corresponding X at initial stages.
In Fig. 14 we test the effect of on the adaptive schedules, where we recall that the next target pak" is chosen such that KLPP?- ||[* < e. We choose the values of e ranging from 0.05 to 5. Here, for each subproblem & 1, local ESS w.r.t. ppk 1 (defined at the previous part) is 2 0.95. so, we adaptively select the next T through solving the optimization problem (19). We found that this adaptive scheduler is quite robust
at least 100 gradient updates to the model, and test whether the
33

---

Published as a conference paper at ICLR 2026
A= 17. x= == 4-=1-p K=
A=11-t
K=
40
y.1.11 K= 40
y=02. K=
v K=
1.0
0.8
80
W2
E5S
60
0.2
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Training Epochs
Training Epochs
Figure 13: Ablation the choice of the prefixed schedulers learning lattice Ising model with L = 8, J ===== 1, and 3low = 0.6. Left: the energy 2-Wasserstein distance E(() W2 between the generated and ground-truth samples; Right: the effective sample size (ESS) of the samples from nd with respect to the optimal path measure IPa
to the choice e, and in all choices tested, the total numbers of training steps are smaller than the naively chosen total number of training steps in the prefixed schedulers above.
= 0.05
4- 0.1
x-02
5=5
1.0
122
1.6
102
58
0.5
52
21-
0.2
2500
4020
6000
2003
10003
2000
4002
6002
8000
20005
2000
4000
6000
5000
10030
Training Epochs
Training Epocts
Training Epechs
Figure 14: Ablation on the choice the adaptive schedulers on learning lattice Ising model with T= 8, J = 1, and 3Low 0.6. Left: the adaptive selected X ss Middle: the energy 2- Wasserstein distance (E-) W2 between the generated and ground-trum tauut samples; Right: the effective sample size (ESS) of the samples from goº with respect to the optimal path measure IP*.
Additional Qualitative Results For lattice Ising and Potts models, we visualize the generated samples under each target distribution in Figs. to T8 where we observe that our PNDS method learns the correct target distribution. We also compare the energy histograms in Figs. 19 and 20 The distributions E(x) (defined in (66 for samples from PDNS closely match those of the ground-truth samples generated by the SW algorithm, indicating that PDNS is capable of learning correct target distribution. We use in total 4096 samples from both methods each distribution.
34

---

Published as a conference paper at ICLR 2026
(a) PDNS (ours)
(b) LEAPS
(c) MH
(d) Ground truth (SW)
Figure 15: Visualization of non-cherry-pich samples from the learned 16 16 Potts model with = 1, q = 4, and Beritical 1.0986. (a) MDNS. (b) LEAPS. (c) MH. (d) Ground truth (simulated
with SW algorithm).
(a) PDNS (ours)
(b) LEAPS
(c) MH
(d) Ground truth (SW)
Figure 16: Visualization of non-cherryppicked samples from the learned 16 16 Potts model with == 1, q = 4, and Blow = 1.3. (a) MDNS. (b) LEAPS. (c) MH. (d) Ground truth (simulated with
SW algorithm).
(a) PDNS (ours)
(b) LEAPS
(c) MH
(d) Ground truth (SW)
Figure 17: Visualization of non-cherry-picke samples from the learned 24 x 24 Ising model with J=1 and Beritical 0.4407. (a) MDNS. (b) LEAPS. (c) MH. (d) Ground truth (simulated with SW algorithm).
(a) PDNS (ours)
(b) LEAPS
(c) MH
(d) Ground truth (SW)
Figure 18: Visualization of non-cherryppiced samples from the learned 24 x 24 Ising model with = and Brow = 0.6. (a) MDNS. (b) LEAPS. (c) MH. (d) Ground truth (simulated with SW algorithm).
35

---

Published as a conference paper at ICLR 2026
Bertical 0.4407
Blow 0.6
PDNS (ours)
PDNS (ours)
0.007
(SW)
(SW)
0.025
Normalized Dersity
0.006
0.020
Normalized
0.004
0.015
0.003
0.010
0.001
0.000
0.000
-1000
-90
-800
-700
-600
1150
-1225
-110
-1075
-1050
-1025
-1000
Energy Elx)
Energy E60
Figure 19: Energy histograms for Ising models with I. 24 and J
Bertical 1.0986
Biow 1.3
0.05
0.014
PONS (ours)
PDNS (ours)
(SW)
(SW)
0.032
Normalized Dersby
Nomalized Dersty
0.030
0.03
0.008
0.006
0.004
0.01
0.002
0.000
0.00
-500
-450
-400
-350
-300
-510
-500
-990
480
-470
-460
-450
-440
Energy Elx)
Energy Elx)
Figure 20: Energy histograms for Potts models with L 16, 1, and g
36


Reviews:
[{"review_id": "tV3behGl4n", "reviewer": "Reviewer_ptdZ", "summary": "This paper addresses general-purpose sampling by simulating diffusion inference trajectories in both continuous and discrete settings. A central challenge is that fitting the noisy target distribution with a neural network often leads to mode collapse and training instability. To mitigate these issues, the authors augment standard SOC objectives (e.g., CE, WDCE) with a KL-regularized local objective, yielding a proximal point method in path space. They instantiate this “path-measure proximal point” framework (PDNS) for both continuous SDEs and discrete CTMC diffusion samplers, demonstrating its generality.", "questions": "1. Did authors consider using the Wasserstein two distance as a regularizer to replace the KL divergence between $p^{\\theta}$ and $p^{\\theta_{k-1}}$, by coupling the JKO scheme, the proximal problem formulation may be more intuitive.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The proximal-in-path-measure perspective is clean and broadly applicable across continuous SDE and discrete CTMC settings. In particular, the proximal WDCE naturally tempers importance weights, directly addressing mode collapse and high-variance gradients. \n2. The paper presents the SOC, CE, WDCE, proximal formulations, and Girsanov/CTMC Radon–Nikodym weights with clarity. Practical guidance (e.g., OU reference bridges, conditional score formulas, and step-size schedulers) facilitates reproducibility.\n3. Experiments span synthetic multimodal and high-dimensional mixtures, challenging particle systems (LJ-13/55), discrete Ising/Potts models near criticality, and combinatorial max-cut. PDNS achieves state-of-the-art or near-state-of-the-art results on multiple tasks, measured by domain-relevant metrics (Sinkhorn/MMD, $W_2$, energy $W_2$, magnetization, two-point correlations).", "weaknesses": "1. Proximal WDCE introduces overhead from replay buffers, weight computation (Girsanov/CTMC), and resampling. While results are strong, the paper does not quantify compute and memory costs relative to baselines (e.g., wall-clock, NFE, memory footprint), especially on the largest tasks (LJ-55, Potts). Resource-normalized comparisons would clarify efficiency. \n2. Many derivations rely on independence between initial and terminal distributions (Eq.~(2)). The framework assumes (nearly) memoryless references to obtain closed forms and reciprocal sampling, which may be unavailable for some targets. Although non-memoryless extensions are mentioned in related work, PDNS is not analyzed under such references. \n3. The proximal step size $\\eta_k$ (or $\\lambda_k$) critically trades off stability and speed. Despite an adaptive KL-based scheduler, the paper lacks systematic sensitivity studies and practical heuristics for choosing $\\epsilon$ across tasks and scales. Additional ablations on $\\eta_k$, buffer size, and resampling vs.\\ weighting would aid practitioners.\n\nIf the authors resolve my concerns, I am willing to raise my rating.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "s6ada3IuVN", "reviewer": "Reviewer_Kemc", "summary": "Proximal Diffusion Neural Sampler proposes a method an stochastic optimal control (SOC) sampler that learns a target distribution based on having an energy function alone. It is an extension of the current developments in this area, specifically by applying the proximal point method on the space of path measures. They provide a method that works for both discrete and continuous cases. The end result is an algorithm where one can do a variable amount of optimization problems that empirically (and theoretically) converge to the target distribution without mode collapse. Essentially, a computationally tractable method for \"forward-KL\" solving with SOC.\n\nIn general, the contribution is strong and fairly easy to follow given the complexity (although I suggest improvements below). It seems applicable with fairly good results.", "questions": "# results\n- can you please plot the energy histogram for the particle problems? It tells us much more than W2 and E() W2.\n- can you please work on notation in the methods you are proposing a bit? I know they're all path measures...\n- I know it's a lot to ask, but is there any molecular task you can apply this to? Alanine Dipeptide for example? Strong contender ASBS misses important modes, if your method gets them that would be great!\n- Can you better quantify how much excess training cost is required to solve the k proximal problems? Does it really slow everything down by a lot?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "# overall\n- problem setting is well identified and known to be a practical issue\n- background is well explained and covers both discrete and continuous cases\n- In general, the key results from a hard topic are well presented... (The extensions from the paper need some notation work, in my opinion).\n- The method introduces a scheme of iterates to help exploration, including a proposed scheduler that dynamically decides the number of iterations based on a bounded KL term. This is nice and well motivated.\n\n# results\n- results on non-particle continuous models are very good.\n- results seem fairly good on particle models\n- results are strong on the discrete models", "weaknesses": "# clarity in actual objective\n- From lines 298 - 318 the paper does not make it extremely clear how the actual objective is computed. There are a number of identities used, in particular the right part of (15) that do not seem clearly motived by other equations in the text.\n- Similarly, about clarity, You use so many version of $\\mathbb{P}$ it becomes difficult to compare them. In particular the most confusing are the non-parametric ones: P^k, P^k*. P^k is a non-optimal, non-parametric path measure? It's not obvious why we even need this thing or how it is defined. Maybe it would be worth introducing alternative notation of P to help keep everything straight.\n- ... Similarly, (16) refers to a ratio of path measure derivatives that does not appear in the discrete objective on 322-323. Can you please write it out?\n\n# cost\n- can you better explain the cost of training your model versus others? Is it k-times as expensive?\n\n# unclear about some parts of method\n- ASBS, for example, is extremely hyperparameter sensitive. How about your method?\n- You use a memoryless path, correct? How does this interact with your optimization problem?\n- Can you use a non-memoryless distribution, e.g. like ASBS does?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "C6dZb67mOq", "reviewer": "Reviewer_prWn", "summary": "This paper proposed Proximal Diffusion Neural Sampler (PDNS) to learn a neural sampler for sampling from unnormalised density functions.   Both continuous and discrete formulation is proposed and verified in this paper.\nThe paper argus that the reason behind model collpasing is that the sampling task is too far from the prior and hence PDNS breaks down the entire problem with proximal point method, progressively driving the neural sampler towards the target.\nThis proposed method enriches the family of neural samplers.", "questions": "See my weakness", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and easy to follow. It is complete and well structured. It covers both discrete and continuous domains and presents results for each. The proposed algorithm is interesting and offers insights into mode collapse. The empirical evaluation indicates that the approach is promising.", "weaknesses": "While the paper is motivated by mitigating mode collapse, the experimental section does not convincingly demonstrate this. Several benchmarks (e.g., LEAPS) do not appear to exhibit severe collapse. Do you have any explanation to this?\n\n\nOn line 780: \n> Blessing et al.  (2025) couple CE training... In contrast,PDNS utilize\n proximal algorithm with WDCE bjective,which can be applied uniformly to continuous-time\n and discrete-time diffusions,and is designed to be memory-efficient while retaining the benefits of\n denoising or score-matching objectives. \n\nThe mentioned algorithm has quite similar formulation and has strong connections. Blessing et al.  (2025)  works on continuous. But to my knowledge, it should also be able to work on discrete domain. Is this the main difference between PDNS and their approaches? Or is my understanding wrong?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "BptwPKe5wh", "reviewer": "Reviewer_h5wm", "summary": "The work proposes the “Proximal Diffusion Neural Sampler”. The authors consider the setting of sampling from distribution with given unnormalized densities on both discrete and continuous state spaces. Rather than solving the SOC/sampling problem in one step, they propose an iterative procedure where an intermediate set of path measures are approximated, starting with a simple path measure and a final target path measure. It is shown that these intermediate path measure correspond to a geometric annealing in path space (proposition 3.1). Every step of the optimization problem is solved with proximal WDCE, a variant of WDCE designed for the proximal setup of this work.", "questions": "- Line 72: “path space, which contains all functions from [0, T] to X”. A path measure is only defined for a subset of functions usually (there is a measurable space, etc.)\n- The ESS for the LEAPS algorithm reported here is significantly lower than report in their work. Comparing the setting, you have slightly changed the settings (slightly different grid size) or q=4 instead of q=3 for the Potts model. Why did you change this setting? It would be more compelling if you compare to exactly the same parameters than their work.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Novel idea that is well-motivated with a technically good presentation.\n- A nice and simple unifying framework of existing neural sampler methods. They consider both discrete and continuous setting in one (instead of writing two separate papers, one general principle is presented).\n- Great illustration of mode collapse (section 3.1) motivating the work.\n- Experiments show improvements that illustrate the theoretical innovations.", "weaknesses": "- For a less-experienced reader, the paper would potentially be hard to follow. I recommend including more sentences explaining intuitions and motivations at intermediate paragraphs.\n- The experiments are only on simple distributions. While this unfortunately common in this line of literature, a more complex setting of higher-dimensional distributions would have been more convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)