On scalable and efficient training of diffusion samplers
Authors: Minkyu Kim
Kiyoung Seong
Dongyeop Woo
Sungsoo Ahn
Minsu Kim
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion sampler
Generative Flow Network
Primary area: probabilistic_methods
TL;DR: We propose a scalable and sample-efficient framework for training diffusion samplers by integrating classical sampling methods, suitable for practical applications like molecular conformer generation.

Abstract: We address the challenge of training diffusion models to sample from unnormalized energy distributions in the absence of data, the so-called diffusion samplers. Although these approaches have shown promise, they struggle to scale in more demanding scenarios where energy evaluations are expensive and the sampling space is high-dimensional. To address this limitation, we propose a scalable and sample-efficient framework that properly harmonizes the powerful classical sampling method and the diffusion sampler. Specifically, we utilize Monte Carlo Markov chain (MCMC) samplers with a novelty-based auxiliary energy as a Searcher to collect off-policy samples, using an auxiliary energy function to compensate for exploring modes the diffusion sampler rarely visits. These off-policy samples are then combined with on-policy data to train the diffusion sampler, thereby expanding its coverage of the energy landscape. Furthermore, we identify primacy bias, i.e., the preference of samplers for early experience during training, as the main cause of mode collapse during training, and introduce a periodic re-initialization trick to resolve this issue. Our method significantly improves sample efficiency on standard benchmarks for diffusion samplers and also excels at higher-dimensional problems and real-world molecular conformer generation.

On scalable and efficient training of diffusion samplers
Minkyu Kim
Kiyoung Seong"
Dongyeop Woo'
Sungsoo Ahn'
Minsu Kim 1.2
Korea Advanced Institute of Science and Technology (KAIST)
2 Mila Quebec Al Institute
Abstract
We address the challenge of training diffusion models to sample from unnormalized energy distributions in the absence of data, the so-called diffusion samplers. Although these approaches have shown promise, they struggle to scale in more demanding scenarios where energy evaluations space is digensiinnii..l. To address this limitation, we propose a scalable and samplc-efficient framework that properly harmonizes the powerful classical sampling method and the diffusion sampler. Specifically, we utilize Monte Carlo Markov chain (MCMC) samplers with novelty- based auxiliary energy Searcher to collect off-policy samples, using an auxiliary energy function to com- pensate for exploring modes the diffusion sampler rarely visits. These off- off-policy samples are then combined with on-policy data to train the diffusion sampler, identify
expensive and the sampling
thereby expanding its coverage of the energy landscape. Furthermore, primacy bias, i.e., the preference of samplers for early experience during training, as the main cause of mode collapse during training, and introduce a periodic re-initialization trick to resolve this issue. Our method significantly improves sample efficiency on standard benchmarks for diffusion samplers and also excels at nigher-dimensione problems and wolld molecular conformer generation.
Introduction
Inference in unnormalized densities a central challenge in machine learning, underlying probabilis- tic deep learning [19, 25] and many scientific applications [33, 9]. Traditionally, Markov chain Monte Carlo (MCMC) methods have been used, most prominently Metropolis- adjusted Langevin algorithms (MALA) [37] and Hamiltonian Monte Carlo (HMC) [15], but they incur repeated avergy-gradient evaluations per sample. Amortized inference instead trains deep generative models map noise to samples, enabling evaluation- iorluattioo-freeeeee generation at test time and promising orders-of-magnitude speedups once the model trained.
Researchers have recently focused on diffusion samplers, which parameterize continuous- time diffusion processes with neural networks, approach inspired by successes in high-dimensional settings like image and text generation. The leading methods include flow-annealed importance sampling bootstrap (FAB) [30], generative flow networks (GFlowNets) [3], denoising diffusion samplers (DDS) [43], controlled Monte Carlo diffusion (CMCD) [44], and iterative denoising energy matching (iDEM) [1]. Because samples from the target distribution are unavailable, these samplers iterate between: (1) sample from the neural diffusion model, (2) query the energy. and (3) update the model to better match the target distribution.
Despite their promise, diffusion-based samplers struggle in high dimensions. Early in training, the neural proposal effectively random and is not aligned with the energy landscape, leading to sample-inefficient exploration. This is in contrast with the classic training-fiee samplers, e.g., MALA, which leverage gradient information to steer proposals toward low-energy modes from the start.
Equal contribution, correspondence to: fminkyu-kim, kiyoung. seang))kattn ar.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Techniques to improve the sample efficiency of diffusion samplers, like replay buffers [3] and local energy-guided refinements [21], yield only marginal gains and fail overcome the poor quality of the initial diffusion samples. Indeed, He et al. [18] recently showed that nearly all effective neural samplers rely on Langevin parametrization, i.e., incorporating energy gradients at inference, which erodes the primary efficiency benefit of amortized sampling.
Moreover, diffusion samplers are prone to mode collapse: training on their own outputs leads to overfitting to dominant modes, and the model "locks in" prematurely. Reinforcement learning exploration bonuses [36] can broaden coverage, but at the cost of biasing the sampler's target distribution. Local perturbations [21] help, but require many expensive iterations in large state spaces.
Contribution. We propose search-guided diffusion samplers (SGDS), simple yet powerful frame- work that enables scalable and unbiased training of diffusion samplers in high-dimensional problems. training-free Markov-chain "Searcher" explores the target density augmented with an explicit exploration reward to discover underexplored modes. The diffusion "Learner" then distills these trajectories through the trajectory balance objective [28], preserving theoretical guarantees while incorporating exploration.
At a high level, our SGDS operates in two stages. Stage 1: the Searcher collects informative samples from the target (optionally with exploration incentives) to overcome the random initialization of the Learner. The Learne is trained off- off-policy via trajectory balance self-generated trajectories, rapidly improving sample efficiency. Stage 2: the Searcher employs random network distillation (RND) bonuses [10] to probe modes the Learner has yet covered; the Learner then ingests these enriched trajectories using trajectory balance with weight re-initialization to counter primacy bias [32].
mixture of Searcher- and
We show that SGDS, despite its simplicity, produces substantial gains over baseline diffusion samplers across benchmarks: classical Gaussian mixtures and the Manywell task; particle simulation problems like LJ-13 and LJ-55; and real-world molecules, Alanine Di-, Tri-, and Tetra-peptide. Our method significantly improves sample efficiency and scalability, marking practical path towards high- dimensional diffusion-bart based inference.
2 Preliminaries
2.1
Diffusion samplers as controlled neural SDES
Let &: Rd - R be energy function defining unnormalized density R(x) exp(-E(x)). Sampling from the corresponding Boltzmann distribution (x) = R(x)/Z, with partition function Z z= R(x) dx, can be formulated as controlling the stochastic differential equation (SDE)
dx, = ox r) dr + g(x,, dw,,
Xg ~ t4. 1 E [0. 1].
where w, is standard d-dimensional Brownian motion, ug is the drift function parameterized by e.g., neural networks. and g is the diffusion function. The goal is to choose 0 such that the terminal distribution of induced by Equation (1) matches the target, i.e., p (x) x R(x)).
Euler-Maruyama discretization. With 7' uniform steps of size As == 1/T, the SDE Equation (1) is discretized via the Euler Maruyaaa scheme
X1+Ai Xy, (x(,,, 1) Ar s(x,, IVVM
2,
42)
which defines Gaussian forward kernels ( x2; Analogously, one defines reference backward kernels (x Ar |xy). Common choices for P8 include Brownian motion dx, = B(1) di, for variance-exploding (VE) processes, the time-reversed Omstein-Uhlenheed (OU) kernel dx, -(()xx, -2B(() die, for variance- preserving (VP) processes, and the Brownian bridge dx, #r + odw, where W, is time-reversed Brownian motion.
The forward and backward policies for the complete trajectory (Xo XAr by Pr(r; 0) and (A(T] x1), repectively, are defined as compositions of these kernels across discrete time steps:
x1), denoted
T-1
T-1 [
PF(T:0) = [
Pr (x (isrii), Pn(T xn) =
Pa (x x

---

Algorithm 1 Training search-guidee diffusion samplers (SGDS)
1: Quulfer 0; fix random target net fend; initialize predictor fo and Learner (PF(T; 0). log Zo) outer rounds
for r = lo....Noound do
3:
// /earcchr:: gradient-guided MCMC
[(xx) (x)m fo(x))12 r= I,
E(x),
4:
E(x)
5:
Obtain (x) (1) Malii and log 2 by running Mchain parallel MCMC for Mier steps on ε(x) ii=l
Qbutfer - Qbuffer U (x]] ,E(x] (1)
7:
// Learner: 1 inner iterations (even iterations: on-policy, odd iterations: off-policy) for p- I.....Id. if mod 2 then
9:
on-policy
10:
Sample (Tx ]
k=1 (F(T:
11:
x {x] from rk)
12:
else
off-policy
13:
Sample X = (x1 P(- Qbuffer)
14:
Generate [T/] PB(T (xx)
15:
end
16:
LTB = log R(x (
7a (r
17:
0 Minimize LTB)
diffusion sampler update
18:
Minimize
[ Exxex lo (x1))1
RND predictor update
19:
end
20:
Re-initialize Pp(: but retain log 7
Periodic partial re-initialization
21:
end for
Stochastic control of neural SDEs. Diffusion models typically minimize the forward Kull- back-Leible (KL) divergence
DKL (P xx) Ptrgett PF(T: (F(T))))((()))
which presupposes abundant samples from X1 ~ Ptaget- When such data are unavailable, e.g., in scientific domains, one may instead minimize the reverse KI. divergence
DKL (F:FFTTO) po(x)) PB(T x1) (x))))((1))))
using samples from x1 ~ Pr. Notable methods that optimize this objective include the path-integral sampler (PIS) [47], which employs a VE Brownian-moote motion reference process, and denoising diffusion samplers (DDS) [43], which use VP OU reference process.
2.2 Continuous GFlowNet objective for diffusion samplers
Following Sendera et al. [39], Muler-Maruyamm samplers can be interpreted continuous generative flow networks (GFlowNets) [24]. GFlowNets [3. are off-policy reinforcement -eamming algorithms for sequential decision making samplers. Treating the initial state the forward policy PF acts as an agent that sequentially constructs trajectory т. The trajectory balance (TB) criterion [28] guarantees that the density induced by PF matches the target distribution:
point mass at the origin,
Ze (F(T:)) R(x)) PB(T x1).
VT,
where Za is a learnable scalar that approximates the unknown partition function Z. Existing GFlowNettbased samplers [46, 39] often adopt Brownian-briene bridge kernels for PB.
Applying the TB condition to sub-trajectory of yields the sub-trajectory balance objective [27, 35, 46]. While this variant can improve credit assignment, it estimates marginal densities at intermediate states with higher bias compared to the global TB estimates [39].
Off-policy property of GFlowNeetbaseed diffusion samplers. In contrast to KL-based objectives such as PIS or DDS, using on-poll on-policy training, GFlowNet objectives can be optimized with off-policy trajectories drawn from any proposal distribution with full support. This flexibility enables richer exploration strategies-noisy roll-outs [24], replay buffers, and MCMC-based local search [39]- that crucial for efficient sampling from multimodal distributions.

---

3 Method
3.1 Search-guided diffusion samplers (SGDS): overall framework
In this section, we describe the overall framework of the search-guided diffusion samplers (SGDS). Our SGDS combines the strengths of off-policy training from GFlowNet diffusion samplers with the exploratory power of gradient-guidee MCMC. We follow the setting of Sendera al. [39] for modeling GFlowNet-based diffusion samplers. Each round alternates between two roles:
Searcher (gradient-informett MCMC). The Searcher uses gradient information V log (xx to efficiently generate representative samples from the target distribution. These samples populate replay buffer and simultaneously provide an estimate of the log partition function, log Ex- ploration is guided by an intrinsic reward from random network distillation (RND) [10], which identifies underexplored modes using form of self-supervised learning.
Learner (diffusion sampler). Learner, a neural diffusion sampler, is trained by minimizing trajectory balance loss [24]. blending (i) on-policy trajectories generated from its current policy and (ii) off-policy trajectories replayed from the buffer. Periodic re-initialization of the Learner mitigates primacy bias, enhancing sample efficiency.
This round repeats until the Learner alone generates high-quality samples. For simple targets, training may converge within a single round, while complex targets typically benefit from multiple rounds.
The SGDS tackles two critical challenges in existing diffusion sampling approaches:
Scalability. In high-dimensional spaces, diffusion samplers frequently miss low-energy modes, as their generated samples rarely visit unexplored modes. The Searcher, operating as parallel gradient-informed chains, rapidly identifies these modes. Although the samples collected from the Searcher are biased, the trajectory balance objective enables unbiased training of the Learner.
Sample efficiency. Each expensive gradient evaluation is amortized across multiple Learner updates through off-policy replay. The RND-driven intrinsic rewards direct the Searcher to- wards under-explored areas, maximizing the informativeness of samples. Periodic Learner re-initialization prevents overfitting to initial samples and maintains replay buffer diversity. Collec- tively, these components significantly enhance the efficiency of gradient computations.
Algorithmic details for each component follow in subsequent sections and Algorithm 1.
3.2 Searcher
The Searcher identifies low-energy modes using parallel gradient-guided Markov chains. Meth- ods such as annealed importance sampling (AIS) [31], Metropolis-adjuste Langevin algorithms (MALA) [37], molecular dynamics (MD) are suitable candidates. These methods generate samples by transporting prior samples in the direction of the target density (or its tempered density) via several Markov chains. We use AIS and MALA for synthetic energy functions, and MD for all-atom systems.
In the initial step of the algorithm, we run Mehain parallel chains, estimating log 2 which is explained in Appendix A. The Searcher then stores the collected samples in replay buffer and passes the estimated log 2 to the Learner model. In subsequent rounds, we incorporate exploration uncertainty from the Learner via intrinsic rewards for exploration, modifying the Searcher's 's energy landscape as:
(5) Here, (imuriniic (x) highlights underexplored modes based on previous Learner experiences, and the gradient is used in the drift function of SDEs. Adding repulsive term for exploration resembles the core idea of metadynamics, which biases sampling away from the modes that have already been well captured.
E(x) = E(x) log Fintriosic (x).
Random network distillation (RND). To efficiently guide exploration, we employ RND [10] to quantify state novelty, steering the Searcher towards underexplored consists of a fixed, randomly initialized network fxx and a trainable predictor network fxx Φ) trained by minimizing:
[RND(x) IIl((x) (x: 1|2-
(6)
and, for the Searcher in the next round, we utilize this loss as the intrinsic reward given by:
Tintrinsic (x) = cxp( 6)

---

High prediction errors indicate novel states. RND training uses replay buffer samples and online trajectories, assigning high novelty underexplored modes.
3.3 Learner
With the replay buffer initialized by Searcher's samples, the Learner minimizes the trajectory bal- ance objective through iterative training, combining online and replay trajectories. The training incorporates:
log ZoP=(T:0) R(x)) PB(T xx)
Lfff-policy (0) = ET - Douterr) 2
(8)
(opr(T; Lon-poticy (0) (T)))( log R(x1) PB(T (xx)
(9)
Here P(xi Douffer) denotes rank-based sampling distribution [42] that assigns higher probability to lower energy samples stored in the buffer, focusing replay on promising modes.
We leverage both on-policy and off- off-policy training signals from online trajectories and replayed samples, with replay ratio y determining the frequency of replay updates (default: y = 1).
Re-initialization. Learner re-initialization mitigates primacy bias commonly observed in reinforce- ment learning scenarios. Primacy bias [32] refers the model's tendency to rely excessively on early experiences, being trapped in low-reward or biased samples generated at initial stages, thereby hindering the discovery of high-reward samples and underexplored modes. Periodically re-initializing the Learner model Pp (-10) alleviates this bias by resetting parameters strongly influenced by early samples, allowing faster adaptation to recent, higher-quality experiences. Crucially, we retain the previously learned log Zo parameter and the replay buffer, preserving the accumulated knowledge while allowing the network recalibrate based on updated experiences.
Related works
Classical samplers. Classical sampling approaches primarily rely on MCMC methods. This includes gradient-based algorithms like MALA [37] and HMC [15]. Annealing- based techniques, such as AIS [31] and SMC [12], introduce intermediate distributions to gradually approximate complex targets, mitigating mode collapse. While these MCMC- methods enable sampling from the complex unnormalized density, they require long trajectories and extensive energy evaluations.
Neural amortized inference. Neural amortized inference methods aim to bypass costly MCMC by training neural samplers that generate approximate samples in one or few forward passes. Diffusion- based neural samplers learn stochastic differential equations parameterized by neural networks to map simple priors to complex targets [47, 43], and GFlowNets train stochastic policies whose marginal visitation probabilities match unnormalized density [3, 13]. Boltzmann Generators (BG) is another line of works to amortize inference, such as molecular dynamics simulation. BG utilizes normalizing flows trained simulated data to sample from the Boltzmann distribution and estimate density, enabling statistical reweighting for unbiased estimates [34, 14, 30, 23, 41).
Diffusion-based neural samplers. Diffusion- samplers aim to sample from unnormalized target distributions in data-free settings. Several approaches [47, 43, 44, 5] formulate the sampling objective via KL divergence in path measure space. Akhound-Sadegh et al. [1] further introduces off-policy training via replay buffers. Recent works [8. 11] also explore controllable dynamies, offering improved exploration in complex energy landscapes. While these methods often improve mode coverage by learning reverse-tine time dynamics, they remain computationally intensive, hindering scalability in high-dimensional settings.
Generative Flow Networks. GFlowNets was originally introduced by Bengio et al. [3] and Bengio et al. [4] on discrete spaces where the probability of each outcome is proportional to a given reward sig- nal. Subsequent extensions have connected GFlowNets to continuous space [24], enabling sampling from unnormalized densities in high high-dimensional spaces [13, 29]. Recent work has also explored enhancements to off-policy training strategies [39] and incorporated local search mechanisms [21] al- lowing GFlowNets to more effectively navigate continuous energy landscapes. Additionally, adaptive reward design has emerged as a promising direction for improving mode coverage during training [22], especially in tasks that require structured exploration sparse supervision.
5

---

Table 1: ELBO, EUBO, their gap, and energy calls across high-dimensional Manywell distributions. We use MALA as the local search algorithm. We consume 6M energy calls per searcher (12M total for 2 rounds) and 8M energy calls for the leamer. Bold indicates the best performance per metric, and indicates large absolute values of metrics.
Manswell (df = 64)
Many (d/ (188)
Method
EUBO
BLBO
Energy
BUBO
ELBO
Energy
PIS+LP
500.57 0.37
347.48 0.26
46.91 0.55
130M
601.01 = 0.94
697 107752. 0.49
96.31 a 0.71
130M
TR-1.P
306.47
351.98
0.51
612.45 0.65
706-73 2.59
94.28 3.00
300M
FL SubTTLBB
0.77
352.22 0.62
46,08 0.26
330M
609.55
709.96 2.10
99.61 1.83
330M
TR-ISHP
2.66
199. 34 1.02
26.68 3.57
500M
592.52 2.25
693 65 1.40
101.81 3.62
120M
TB-Expl+LP
93.37
TB-EXpH-LSLL
300.10 1.05
344.85 0:41
44,75 1.39
320M
591.47 0.36
694.93 0.54 0.54
103.45
0.88
320M
321.87
1704.91
1159
48.53
516.30 30 49.67
TR
6.00
853,94 43.35
544,35
20.85
100M
637.01 2.14
1423,35
292.15
766.35
290,46
100M
TB-LS
357.40 4.36
42.91 9.15
290M
573.13 73.49
738.07 10.77
164.95 62.71
290M
TH-EXpiLLS
265.99 95.39
361.00 16.58
41.46
=
698.24 2.81
108.74 10.07
290M
GAPN [36]
$20.88 0.36
573.68 + 29.02
252 80 30.87
100M
100M
1P
353 3.48
462.6
730.90
277.52
DEM [1]
268.99 1,22
414.18 1.06
145.20 1.60
300M
191.28 a 2,94
817. .2 3.22
323.04 5.69
300M
614.41
70.35
(a) Ground Truth
(b) SGDS
(c) PIS
(d) SubTB+LP (c) TB+ExpI+LS
(f) iDEM
Figure Mode coverage comparison using 2D projections of 10,000 samples on Manywell- 128..
Connection to previous works. Using gradient-guides MCMC for improving exploration in off- policy diffusion samplers is not new. Lemos et al. [26] employed gradient-guided MCMC populate replay buffers for GFlowNet diffusion sampler training. Sendera et al. [39] applied parallel MALA initialized from diffusion sampler states, similar to discrete local search GFlowNet methods [21]. Our approach extends the multiple-round algorithm of Lemos et al. [26]. incorporating RI techniques to boost efficiency. It can be viewed as a deeper but shorter-cycle alternative to Sendera et al. [39], whose frequent diffusion-bel based re-initializations overly depend sampler performance (see comparison with TB LS at Table 1, Table 2, and Figure 4a).
Leveraging Learner uncertainty to guide exploration aligns with active learning and related GFlowNet approaches [36, 22]. Following generative augmented flow network (GAFN) [36], direct injection of intrinsic reward was effective, similar to our idea (see comparison with GAFN at Table 1). While Kim al. [22] introduced additional neural samplers called adaptive teachers (AT) as Searchers to covers high loss region it is highly unstable in large scale due to Searcher's s adversarial behavior with non-stationary objective, where our method efficiently employs MCMC-based exploration without additional neural network (see comparison with AT L. at Table 1).
Experiments
In this section.' 1 our primary goal is to demonstrate the performance and efficiency of our proposed framework through several experiments. Specifically, we aim to showcase the sample efficiency and scalability of our method, as well as to validate the effectiveness of the various training strategies we introduced. We focus on presenting results on high-dimensional tasks. In all the experiments, we use four different random seeds and average the results of each run. We provide details of experimental settings in Appendix A.4, Table 4, and Table 5, and additional results Appendix
5.1 Main results
Settings. In this work, compare the performance of proposed framework against baselines on multiple benchmark tasks, including 40GMM, Manywell-32/64/122 LJ-13, and LJ-55. We
Source code: https://githu/: http:///ggggtabbbooo/aaaaa/////s

---

121
120
8008
Method
ELBO1
EUBO J
EU9O ELBO
102
AIS+MLE
244.85 + 10.34
840.45 :: 11.59
Fine-tuning
554.17 107.70
707.07 8.49
82
5225
Re-init w/o log Z w/o RND-searcher
583 63,.6532 + 16.69
691.37 = 1.40
608.01 4.18
SGDS
614.41 =: 3.44
684.76 1.30
Calls (x100)
Energy Calls [x10")
(a) Manywell-128
(b) LJ-55
(c) Ablation results on Manywell-
Figure 2: Trade-off between (UBB--FLBO gap and energy calls in Manywell-1288 (left) and IJ-55 (middle). The results of ablation study on Manywell- 1-2 (right) show the performance of AIS using the same total energy calls with MLE amortizing, taking 2 rounds with fine-tuning instead of re-initialization, and using the Searcher with no RND rewards. All methods use 20M energy calls.
Table 2: The gap between ELBO and EUBO, (equivariant) position Wasserstein-2, energy histogram Wasserstein-2, and energy calls across Lennard-Jones potential. We denote EUBO as the EUBO metrics calculated by the reference samples provided by [1], which are not exact samples from the target distribution. For iDEM, we reproduce the results with the hyperparameter setting which can be found in Appendix A.4. The results of Adjoint sampling are obtained from [17]. Bold indicates the best performance, indicates large absolute values of metrics, indicates inaccessible value from their papers, and "Div." " indicates divergent training.
LJ-13 39)
IJJ55 (d 165)
Method
BUBO ELDO
WW5
E(-WW 1
Energy calls EUBO ELBO
W: L
E(-W5 J
Energy calls
2.04 0.32
370K
4.28 0.05
40.63 13.68
12.53
1.71
10.99 3.114
370K
Div.
Div.
Div.
45K
TR-ExpIeLLS
0.52
2.02 0.04
18.69 4.49
3M
12.46 + 0.19
123.02 + 21.72
IM
IDEM [1]
112,37 + 7.63
4.27
39.92 3.24
300M
17.17 0.32
210.87 + 6.26
120M
Adjoint Sampling [17]
2.40 1.25
IM
58.04 20.98
IM
SODY
1.53 0.25
1.56
0.19
370K
33.01
4.25
32.17 12.90
45K
evaluate methods using three metrics: the Evidence Lower Bound (ELBO), Evidence Upper Bound (EUBO) [7], and the EUBO- ELBO gap. A smaller gap between ELBO and EUBO indicates more accurate approximation of the target distribution.
For fair comparison on the number of energy calls, we train the methods until convergence of ELBO and EUBO. To determine convergence, we evaluate based the moving average of the metrics over the 10 consecutive evaluations, where we evaluate the model every 100 training steps. If a method does not converge within the maximum number of epochs, report the metrics at the final step.
Baselines. For the manywell potential, the basclines are primarily selected based on their strong performance demonstrated in prior work [39], as well as their methodological relevance [36, 22] or their different framework [1]. Specifically, iDEM [1] utilizes trajectories of length T = 1, 000 for SDE integration, whereas other baselines, including PIS [47]. TB [28], AT [22], and GAFN [36], employ shorter diffusion trajectories (T 100) with distinct optimization objectives. We further evaluate enhanced variants of these methods incorporating LP, such PIS+LP, TB+LP, and FL-SubTB+LP along with exploration-enhame [39]. We describe the details of the abbreviations related baselines in Appendix A. For the LJ potentials, we compare against Adjoint Sampling [17], as well as iDEM, PIS, TB, and its variant TB+EXpI+LS. We omit LP-based methods due their divergent training.
(+Expl) or local search (+LS) variants introduced by Sendera et al.
Results. As shown in Table and Table 2, our proposed framework consistently achieves superior performance across all high-dimensionall tasks: Manywell-64, Manywell-128, and LJ-55. Especially, our method demonstrates the best trade-off between performance and efficiency of energy call.
In Figure 1, one can observe that our method better captures the modes in 11n when compared to the baselines. As illustrated in Figure 2a and Figure 2b, even increasing the energy budget of baselines does not allow them to surpass the performance of our proposed approach. Also, as shown in Figure 3, our framework generates high quality samples with low energy. Furthermore, for the LJ-55 potential, the distribution of interatomic distances is similar to the ground truth
7

---

2,000
Bebernes
3005
M
5G26
PIS
3.22
EIN
T
2012
1.02
-00
-50
-40
-00
-20
-303
-323
-270
-220
-150
-105
Energy [RaT]
Energy [KaT]
Interatomic distance
LJ-1 energy histogram
(b) LJ-55 energy histogram
(c) LJ-55 interatomic distance
Figure 3: Histograms for LJ-13/55 energy densities and LJ-55 interatomic distances.
distribution. Additionally, our method obtains competitive results with significantly fewer samples in lower-dimensional tasks such as 40GMM, Manywell-32 (see Appendix B.1), and LJ-13 (see Table 2).
5.2 Ablation study
MCMC sampler with the same budget. In method, consume energy calls during both Searcher sampling and Learner training. To evaluate the efficiency of the Learner's on/off-policy mixing training scheme, conduct controlled comparison where the total energy call budget (20M) entirely allocated to AIS on Manywell- 128. We rum 200 chains on the trajectories with 7 = 10, 000 for AIS. As a result, even though high-reward samples were collected using AIS with much longer trajectories, the MILE Learner failed to perform amortized learning as shown in Figure 2c.
Periodic re-initialization and pre-trained flow. We perform an ablation study to evaluate two design components of our method when proceeding to the next training round: (1) re-initiaalizin the Learner model, and (2) retaining the pre-trained log Z parameter from the previous round. Specifically, to assess the benefit of re- re-initialization mitigating primacy bias, we compare our method against fine-tuning baseline where the second-round Learner continues training from the first-round model weights without re-initialization. To isolate the effect of retaining the estimated log value, we compare against variant where the log Z parameter is also re-initialized at the start of the second round. As shown in Figure 2c, our full method outperforms both ablation variants, confirming that re-initialization is beneficial for mitigating primacy bias, and that employing the log Z parameter across rounds leads to better training stability and performance.
Novelty-based reward in Searcher sampling. We assess the effectiveness of incorporating the novelty-based intrinsic reward derived by RND [10] into the Searcher sampling process in later training rounds. In our framework, starting from the second round, the Searcher sampler drives prior samples in the direction of the target distribution and exploration signal derived from a previously trained RND module, which prioritizes underexplored modes by the Learner sampler. These dynamics guide the Searcher to focus sampling efforts on modes that remain novel and close to the target distribution across rounds. As shown in Figure 2c, at the end of round 2, RND--uugeented approach yields a smaller EUBO-ELB gap compared to a way of repeating the same Searcher sampling without exploration. These results demonstrate that using intrinsic rewards to adaptively bias Searcher sampling toward novel modes improves overall distributional coverage across rounds.
5.3 Application to molecular conformer generation.
We also consider three real-world systems, Alanine Di-, Tri-, and Tetra-peptide, consisting of 23, 33, 43 atoms in vacuum at a temperature of 300K. While some previous works show promising results in sampling conformation of Alanine Dipeptide, they rely on low-dimensional descriptors such as rotatable torsion angles [45]. Solving these three peptides at all-atom resolutions remains challenge for existing diffusion-ben based neural samplers.
Settings. To accurately evaluate molecular energies, we employ TorchANI [16], a PyTorch imple- mentation of ANI deep learning potentials trained on quantum- mechanicca reference data. For the Searcher, we run four parallel 55ps Langevin dynamics simulations under the TorchANI potential. In the first round, simulations are performed at 600 K to efficiently sample slow degrees of freedom; in the second round, we use 300 K to capture faster motions and collect high-reward samples. The Learner and RND models use the E(3) (3)ivvvaarri graph neural network (EGNN) architecture [38] based on atomic coordinates. We provide details in Appendix A.4.

---

Reference
SGDS
MLE
SG0S MLS
TU-Expl-LS
Disapiate
P1S
TB- L
Topopsiate
(a) Ramachandran plots
(b) 3D visualization of sampled conformations
Figure 4: Qualitative results of methods in three peptides. (a) Ramachandran plot of Alanine Dipeptide with two backbone torsion angles (ø. 4/) and (b) 3D visualization of generated conformations.
Baselines. We compare our method with meximum--lkkllbeen estimation (MLE) as well as diffusion- based neural samplers: PIS, TB, and TB+ExpI+LS. We generate reference ensemble using 100 simulation at 300K. In MLE, the likelihood of forward path distribution is
ns Langevin-dynanen maximized over the samples from backward path distribution. To match the number of energy calls of MLE with our method, we collect 2.5 times more samples than our method, using the same searcher as our method without RND. TB+ExpI+LS utilizes same langevin dynamics for local search algorithm. We omit the LP methods since they have large absolute values of EUBO and ELBO. We exclude comparison with Volokhova er [45], as they consider only rotatable torsion angles, and with Midgley et al. [30], which employs a discrete normalizing flow on internal coordinates, whereas our method utilizes a diffusion model in atomic coordinates.
Results. As shown in Table 3, our method out- performs diffusion- based neural samplers and MLE on three peptides. As illustrated in Fig- ure 4, both our method and MLE capture the free energy landscape and generate physically plau- sible conformations by leveraging high-fidelity samples from the Langevin dynamics Searcher. By contrast, PIS, TB, and TB+ExpI+LS fail to reconstruct the target free-energy surface to produce realistic structures, since their forward policies insufficiently explore the complex land- scape. We note that the Langevin dynamics used in the Searcher yields higher-quality samples than those obtained by local searches of TB+EXpI+LS from forward-poll policy outputs. Furthermore, our approach refines the biased samples from the high- temperature Searcher through unbiased TB objective, improving performance compared MLE.
Table 3: EUBO-ELBO gaps on peptides in units of 10 Bold indicates the best performance. in- dicates large absolute values of metrics, and "Div." indicates divergent training.
Method
Dipeptido
Tripeptide
Tetrapeptide
(IM)
84.65 a 9.93 Div.
35.23 2.43
Div.
TB-ExplALS (3M)
18.86
6,41 0.34
Div.
MLE (IM)
17.41 0.14
6.11 0.02
29.07 0.56
SGDS (LM)
17.11 0.16
5.60 a 0.03
26.60 0.21
Conclusion
We have proposed a scalable and sample-efficient sampling framework SGDS that integrates an MCMC Searcher with diffusion Learner. By leveraging high-quality samples from replay buffers and training the Learner model via on/off-policy TB objectives, our method effectively bridges classical sampling with neural amortization. The inclusion of novelty-based intrinsic rewards by RND further enhances the exploration of the Searcher, enabling informed guidance to underexplored modes throughout multiple rounds.
Our work opened promising directions for integrating learning-based amortization with classical sampling, particularly for tasks where both diversity and precision are crucial. Future extensions include designing multi-agent search systems that leverage classical sampling methods for coopera- tive strategic exploration in sigh-dimensional spaces and developing advanced off- off-policy learning schemes, such as adaptive filtering strategies for the replay buffer.

---

Acknowledgements
This work was partly supported by Institute for Information & communications Technology Planning & (valuation(IITP) grant funded by the Korea governmen((MSST) (RS-2019-I1190075, Artificial In- telligence Graduate School Support Program(KAISTT National Research Foundation of Korea(NRF) grant funded by the Ministry of Science and ICT(MSIT) (No. RS-2022-NR0721
GRDC(Global
Research Development Center) Cooperative Hub Program through the National Research Foundation of Korea(NRF)) grant funded by the Ministry of Science and (MST)) (No. RS-2024-004361663 and the Institute of Information & Communications Technology Planning & (valuation(IITIP) grant funded by the Korea government(MSIT) RRS-2025-02304900 AI Star Fellowship(KAISST) Minsu Kim was supported by KAIST Jang Yeong Sil Fellowship.
References
[1] Tara Akhound-Sadegh, Jarrid Rector-Brooks 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. In International Conference on Machine Learning (ICML), 2024.
[2] Michael S Albergo and Eric Vanden- Nets: A non-equilibrium transport sampler. arXiv preprint arXiv:2410.02711, 2024.
[3] Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
[4] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, 24(210):1555 2023.
[5] Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion- based generative modeling. Transactions on Machine Learning Research (TMLR), 2024. ISSN 2835-8856.
[6] Julius Berner, Lorenz Richter, Marcin Sendera, Jarrid Rector-Brooks, and Nikolay Malkin. From discrete-time policies cintinuous-time diffusion samplers: Asymptotic equivalences and faster training. arXiv preprint arXiv:2501.06148. 2025.
[7] Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Beyond ELBOs: A large-scale evaluation of variational methods for sampling. In International Conference on Machine Learning (ICML), 2024.
[8] Denis Blessing, Julius Berner, Lorenz Richter, and Gerhard Neumann. Underdamped diffusion bridges with applications to sampling. In International Conference Learning Representations (ICLR), 2025.
[9] Ignasi Buch, Toni Giorgino, and Gianni De Fabritus. Complete reconstruction of enzyme- inhibitor binding process by molecular dynamics simulations. Proceedings of the National Academy of Sciences, 10182)) 2011.
[10] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations (ICLR), 2019.
[11] Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anandkumar. Sequential controlled langevin diffusions. International Conference Learning Representations (ICLR), 2025.
[12] Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):41 2006.
[13] Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien. Ste- fan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. Conference Uncertainty in Artificial Intelligence (UAI), 2022.
10

---

[14] Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and boltzmann generators. Physical Review Research, 4(4):L042005, 2022.
[15] Simon Duane, A.D. Kennedy, Brian J. Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physics Letters B, 195(2):216-22) 1987.
[16] Xiang Gao, Farhad Ramezanghorbani, Olexandr Isayev, Justin S Smith, and Adrian E Roitberg, Torchani: a free and open source pytorch- based deep learning implementation of the ani neural network potentials. Journal of chemical information and modeling, (17):3400-3415, 2020.
[17] Aaron J Havens, Benjamin Kurt Miller, Bing Yan, Carles Comingo-Enriche 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 International Conference on Learning Representations (ICLR), 2025.
[18] Jiajun He, Yuanqi Du, Francisco Vargas, Dinghuai Zhang, Shreyas Padhy, RuiKang Ou Yang, Carla Gomes, and José Miguel Hernández-LLoatto No trick, no treat: Pursuits and challenges towards simulation-free training of neural samplers. arXiv preprint arXiv:2502. 06685, 2025.
[19] Geoffrey E. Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.
[20] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning (ICML), 2022.
[21] Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search GFlowNets. International Conference on Learning Representations (ICLR), 2024.
[22] Minsu Kim, Sanghyeok Choi, Taeyoung Yun, Emmanuel Bengio, Leo Feng, Jarrid Rector- Brooks, Sungsoo Ahn, Jinkyoo Park, Nikolay Malkin, and Yoshua Bengio. Adaptive teachers for amortized samplers. In International Conference on Learning Representations (ICLR), 2025.
[23] Leon Klein and Frank Noé. Transferable boltzmann generators. Neural Information Processing Systems (NeurIPS), 2024.
[24] Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández- Garcia, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of continuous generative flow networks. International Conference on Machine Learning (ICML), 2023.
[25] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, et al. A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
[26] Pablo Lemos, Nikolay Malkin, Will Handley, Yoshua Bengio, Yashar Hezaveh, and Laurence Perreault-Levasseur Improving gradient- guided nested sampling for posterior inference. In International Conference on Machine Learning (ICML), 2024.
[27] Kanika Madan, Jarrid Rector-Brooke Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2022.
[28] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. Neural Information Processing Systems (NeurIPS), 2022.
[29] Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GelowNets and variational inference, International Conference on Learning Representations (ICLR), 2023.

---

[30] Laurence Illing Midgley, Vincent Stimper. Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernaneez-Lobtto. Flow annealed importance sampling bootstrap. In International Conference on Learning Representations (ICLR), 2023.
[31]
Radford M Neal. Annealed importance sampling. Statistics and computing, 11:125-139, 2001.
[32]
Exgenii Nikishin, Max Schwarzer, Pierluea D'Oro, Pierre-Lue Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning, In International Conference on Machine Learning (ICML). 2022.
[33] Frank Noe, Christof Schütte, Eric Vanden-Eijnden, Lothar Reich, and Thomas R Weikl. Con- structing the equilibrium ensemble of folding pathways from short eeuiiiiiuiiim simulations. Proceedings of the National Academy of Sciences, (96(45):19011-1900 2009.
[34] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 1147, 2019.
[35] Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GelowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023.
[36] Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative augmented flow networks. In International Conference on Learning Representations (ICLR), 2023.
[37] Peter Rossky, Jimmic D Doll, and Harold L Friedman. Brownian dynamics as smart monte carlo simulation. The Journal of Chemical Physics, 69(10):4628-45)) 1978.
[38] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International Conference on Machine Learning (ICML), pages 9323-9332. PMLR, 2021.
[39] Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector- Brooks, Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off-policy training of diffusion samplers. Neural Information Processing Systems (NeurIPS), 2024.
Rector-Brooks,
[40] Yuji Sugita and Yuko Okamoto. Replica-exchange molecular dynamics method for protein folding. Chemical physics letters, 314(1-2):141-153 1999.
[41] Charlie B. Tan, Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. Scalable equilibrium sampling with sequential boltzmann generators. In Frontiers in Probabilistic Inference: Learning meets Sampling, 2025.
[42] Austin Tripp, Erik Daxberger, and José Miguel Hernández-Lobato. Sample-efficient optimiza- tion in the latent space of deep generative models via weighted retraining. Neural Information Processing Systems (NeurIPS), 2020.
[43] Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. Interna- tional Conference on Learning Representations (ICLR), 2023.
[44] Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational inference: Controlled Monte Carlo diffusions. International Conference Learning Representations (ICLR), 2024.
[45] Alexandra Volokhova, Michal Koziarski, Alex Hernández-Garrda, Cheng-Hao Liu, Santiago Miret, Pablo Lemos, Luca Thiede, Zichao Yan, Alán Aspuru-Guzik, and Yoshua Bengio. Towards equilibrium molecular conformation generation with gflownets. Digital Discovery, (5):1038-1047, 2024.
[46] Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusson generative flow samplers: Improving learning signals through partial trajectory optimization. International Conference on Learning Representations (ICLR), 2024.
[47] Qinsheng Zhang and Yongxin Chen. Path integral sampler: a stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022.
12

---

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

---

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

---

Answer: [Yes]
Justification:
We provide the code to reproduce all of our experimental results in Section
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
ionninsssiooooooo onooliccyyy olicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //wpps.ccpppuiii
cc/ wwwpccccpccccceeee.ceeee..eccele http:pppbpipdp//////eeep httpa//gggideeeeeeuuubiiiiiii http///ooSuuiiiseeeeeee corbbbmiisioni ieiiionooolily for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See the experiment sections Section 5 and references to appendix material Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All experimental tables include standard deviation and indicate significance of the best metric.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
15

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See experiment sections and references to appendix material.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the ines?
NeurIPS Code of Ethics https: cttpss//aarrrpp..cpppp.cc/c
Answer: [Yes]
Justification:
We believe there are no violations of the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
The paper studies ML problem with immediate societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

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

---

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

---

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

---

A
Experiment details
Code is available at htttps:/ggtte httph//....///iakyyy/////////000
And the reference samples can be downloaded from
A.1
Description of notations
PIS (Path Integral Sampler) [47], TB (Trajectory Balanced) [28], and FL-SubTB (Forward-Looking SubTB) [39] denote the types of objectives. GAFN (Generative Augmented Flow Networks) [36] GFlowNet-
method that directly injects intrinsic rewards into TB loss. AT (Adaptive additional neural sampler (the Teacher) trained to focus high-loss regions of the Student. LP (Langevin Parametrization) [39] refers to a drift construction technique where the model's 's drift term is combined with the gradient of the target energy function, helping the trained dynamics to better follow the target energy landscape. LS (Local Search) [39] denotes a refinement step where short Markov chain guided by the target energy gradient is applied from the final state of the generated trajectory to improve sample quality. Expl (Exploration) [39] represents a noise-scheduling technique applied during trajectory generation, in which additional stochasticity is injected more in the early training phase to promote broad exploration and gradually reduced over time for exploitation later.
Teacher) [22] introduces
A.2 MCMC samplers for Searcher
Annealed importance sampling (AIS). Annealed importance sampling (AIS) [31] is an MCMC sampling method for estimating the partition functions of target distributions. AIS bridges between an toay--oppaaplee initial distribution xx(x) and target distribution πy (x) through a sequence of
intermediate distributions (x) (x) } where T the length of trajectory or chain. Each intermediate
distribution N, (x) typically has the form:
2i (x) no(x)² ( 0= Bo B1
B1= 1,
(10)
where {Br} } is predefined annealing schedule, and Br + in framework. AIS generates samples through an MCMC transition kernel at each intermediate distribution with the following SDE simulation:
dx, = v (x;) v2aw,,
(11)
where V ll = (1-B, B, log no(x,) + B,V log (x;, is the score function of the annealed distribution (unnormalized). Then it accumulates importance weights given by:
27 (X1-1)
WAIS
(12)
(x
[=1
and the expectation of these weights provides an unbiased estimator of the partition function ratio between πT (x) and (x()):
N w
Zr
(i)
(13)
Zo
i=1
where (i is the importance weight computed for the i-th AIS trajectory, and N is the total number of trajectories. We compute the unbiased estimation of the log scale of the partition function for Manywell experiments by
log 2, log
W
(14)
where log Zn because the initial distribution is Gaussian in our framework.
Metropolis-Adjustel Langevin Algorithm (MALA). The Metropolis-Adjuster Langevin Algorithm (MALA) is an MCMC method that uses the gradient of the energy function to generate samples from
20

---

target distribution π(x). MALA starts by sampling initial states xo To(x)), where To(-) (-) is some proposed initial distribution (in most cases, NO(, II). It then iteratively proceeds transition from to Xy+1 by simulating the following Langevin dynamics:
dx, dV v2dW,
(15)
Here, is the current state at time 1, W, denotes the standard Brownian motion, and (xx) is the energy function of target distribution 7xx,, i.e. -V((x) log (x;)
The proposed sample X,+1 is then accepted or rejected according to the Metropolis Hasttings accep- tance probability:
a min (1. (xx+11)((x |
(16)
(xx+)((xx+1 [x,)
where q( 1-) denotes the Gaussian transition density induced by the Langevin proposal:
Xy+1 = X VE(x)) V2 z N)..
(17)
The step size Ar is a key factor influencing the quality of sampling. For all tasks, we utilize the scheduling of step size, by comparison between the current acceptance rate and the target acceptance rate (57.4%). We use MALA as Searcher on 40GMM, LJ-13, and LJ-55.
Also, since MALA trajectory forms Markov chain, consecutive samples still correlated and therefore { j=1 are not strictly i.i.d. To reduce the most severe correlations we discard the first Moum-in iterations as burn-in and use all subsequent states directly. We then compute rough estimate
1 Σ exp (-8(x:))
log 2 = log
(18)
where this estimator is biased since xi ~ x ideally and E, [exp (-E(x)) = z xx) < Z. Despite the bias, the estimation can provide numerically reasonable heuristic value for the initialization of the Learner's log Zg.
Underdamped Langevin dynamics. For MCMC Searchers of three peptides, we adopt underdamped Langevin dynamics as our molecular dynamics (MD). This framework combines deterministic forces with stochastic fluctuations, which is essential for accurately capturing thermal motion and inertial effects of the molecules. The resulting dynamics are governed by the following system of stochastic differential equations:
dx, dr, dv, = - VV di - yv, dr + 2 M-1 dw,
(19)
Here, x, the position at time 1, V, is the velocity, M is the mass matrix (symmetric positive definite), E(x) is the potential energy function, and V&(x,) is its gradient with respect to position, i.e., the negative force. The parameter y is the friction coefficient, KB is the Boltzmann constant, T is the absolute temperature, and W, denotes standard Brownian motion.
For peptides, we use underdamped Langevin dynamics as MD with high (emp)rature(600KK. We use Fuler-Maruyama integration to discretize the Langevin dynamics. As in MALA, we compute log 2 for the initialization of log Ze in Learner, using Equation (18).
A.3 Metries
In this subsection, formally define the evaluation metrics used to assess the Learner's quality. All metrics are derived from the same importance-weigg formulation based on the target partition function.
begin with the exact log partition function log Z, which can be written using forward- path importance sampling. Let (xo..tr.....x..) denote sample trajectory drawn from the forward
21

---

policy (T()), and let R(x)) be the reward associated with the final state x1. Then, the partition function can be expressed as
P( (x))
log Z = log (r-))(r)
(20)
PP(T)
where PH(T xx) is the backward policy conditioned on the final state.
Since directly optimizing this quantity is intractable, we use two surrogate bounds. The first is the evidence lower Bound (ELBO), defined as
R(x))PB(T x1)
ELBO By-PF(T)
log
(21)
PF(T)
By Jensen's 's inequality, ELBO is always a lower bound on the true log Z. It is commonly used as training objective and can reflect how well the forward policy Pp concentrates on high-reward trajectories. However, ELBO can be misleading in practice. A high ELBO does not necessarily imply that all important modes are captured, the forward policy may collapse to small subset of modes while still achieving high reward [7].
To address this limitation, we also evaluate the evidence upper Bound (EUBO), which flips the sampling distribution:
R(x)) (x1) |
EUBO Pr-P log
(22)
PF(T)
Unlike ELBO, EUBO acts as diagnostic metric. It is an upper bound of log Z and penalizes missing probability mass. EUBO is driven to penalize missing probability mass and therefore exposes mode- collapse that ELBO may hide [7]. And then, true log Z is consequently bounded by two bounds, i.e., ELBO & log EUBO.
smaller gap between the two bounds yields tighter estimate of log Z, making this gap a useful indicator of the Learner's sampling quality.
Table 4: Searcher configurations of SGDS
Benchmark
40GMM
Manywell 32
Manywell 64
Manywell 128
LJ-13
LJ-55
Peptides
Type
MALA
AIS
MALA
MALA
MD
# of Chains
300
60K
60K
60K
16
1
Chain length
4K
100
100
100
4K
10K
110K
Burn-in
2K
2K
4K
10K
init. step size
le-3
le-3
1e-3
1e-3
le-5
1e-5
0.5fs
Table 5: Learner configurations of SGDS
Benchmark
40GMM
Manywell Manywell Manywell
IJ-13
LJ-55
Peptides
Brownian bridge std (o))
10.0
1.0
1.0
1.0
0.2
0.2
0.2
Buffer size
600k
60k
60k
60k
50K
10K
S00K
Batch size
300
300
300
300
Architecture
MLP
MLP
MLP
MLP
EGNN
EGNN
EGNN
hidden dim
256
256
256
256
64
128
of layers
2
2
RND weight
100
100
100
100
1
10
A.4 Experimental setup
We reproduce iDEM [1] with modified hyperparameter settings ensure comparable number of energy calls. Specifically, we adjust the number of MC samples for score estimation and the number of iterative loops. Furthermore, we exclude the additional refinement steps originally applied the LJ-55 potential in iDEM to maintain consistency across all evaluated methods.
For the diffusion-based neural samplers synthetic benchmarks, follow the setup of [39].
Gaussian mixture model with 40 modes (40GMM), Training proceeds in one or two rounds. Our framework achieves competitive performance against basclines even with only single round, and
22

---

shows marginal improvement with a second round. We use MALA as the Searcher, running 300 parallel chains of length 4K, discarding the first 2K steps as burn-in. We maintain a target acceptance rate of 57.4% through step size scheduling, resulting in total of 2.4M energy evaluations. We use the Gaussian prior with a standard deviation of 21.0 for MALA.
All methods adopt the PIS architecture [47, with a joint network consisting of a two-layer MLP with 256 hidden dimensions. The RND network consists of three layers in the predictor network and the target network, with 256 hidden dimensions. We adopt Brownian bridges as the backward process, with a Brownian motion coefficient of 10.0. We run 25K epochs in both the first round and the second round.
Manywell distributions. We proceed with one two rounds for training Manywell distributions. We use AIS as the Searcher, running 60K parallel chains (3K chains * 20 iterations) of length 100, only taking the final step samples. We use the Gaussian prior with a standard deviation of 1.0.
All methods adopt the PIS architecture [9], with a joint network consisting of a two-layer MLP with 256 hidden dimensions. The RND network consists of three layers in the predictor network and the target network, with 256 hidden dimensions. We adopt Brownian bridges as the backward process, with Brownian motion coefficient of 1.0. We run 25K epochs in the first round and 30K in the second round.
Lennard-Jones (L.)) potentials. Training proceeds in two rounds. We use MALA as the Searcher for two rounds: IJ-13 we run 16 parallel chains of length 4K corresponding to 64K energy evaluations, discarding the first 2K steps as burn-in and retaining 57. 7. accepted samples among remaining 32K samples; in LJ-55 we run a single chain of length 10K corresponding to 10K energy evaluations, discarding the first 4K steps and retaining 57.4% accepted samples among remaining 6K samples. We use the Gaussian prior with a standard deviation of 1.75 for MALA.
All methods utilize five EGNN layers with 64 hidden dimensions, Following [20, 23], we design an E(3)eequivariant generative model initialized from a Dirac delta at the origin, using a mean- free forward transition kernel in inference. The RND network comprises three layers in the predictor network and two in the target network. We adopt Brownian bridges as the backward process for diffusion-basse neural samplers, with a Brownian motion coefficient of 0.2. For LJ-13, we run 5K epochs in the first round and 10K in the second round; for LJ-55, 10K and then 20K epochs.
Specifically, we note that the reported performance of the iDEM on Table 2 differs from the original paper [1] due to adjustments, except Tmax and mmin of the noise scheduling, made avoid significant discrepancies in energy call usage compared to our method. We reduce the EGNN hidden dimension to 64 and the batch size to 8, and limit the total number of training epochs, including both inner and outer loops, to 15K accordingly. And while the latest iDEM codebase employs 10 steps of Langevin dynamics refinement before evaluation, particularly for LJ-55, we omit this step for fair comparison and instead set the number of samples for MC estimation to IK. While iDEM reports a lower bound of log Z computed via importance sampling with its learned proposal density (xx given by OT-CFM model, we omit this result in our tables. We compute the lower bound based on trajectory-level estimators without training auxiliary models, i.e., CFM. Thus, our reported values are not directly comparable to those from iDEM.
Additionally, in LJ-55, we maximize the log-likelihood of the forward path distribution under the backward process for the first 5K epochs of each round, discretizing backward paths from Brownian bridges initialized with empirical samples collected by Searchers. We also use randomized time scheduling introduced in [6] for our method. We train PIS at learning rate of le 4, TB at a learning rate of 2e- 4. and SGDS at learning rate of Se 4. We use 4 and 32 batch sizes for all methods except PIS in LJ-13 and LJ-55, respectively. For PIS, we halve these sizes due to the memory limitation required by the forward SDE computational graph.
Peptides. We perform two rounds of search using under-damped Langevin dynamics. In each round, we four parallel simulations of 55 ps each, with time step of 0.5 fs, requiring 440K energy evaluations. We discard the first 5 ps of each trajectory as burn-in, then collect 400K samples. Each simulation starts from the same initial position drawn from Dirac delta distribution, with all initial velocities set to zero. We integrate equations of motion using the Maler-Maruyama integrator, set the friction coefficient y = 1, and use temperature 600K for the first round Searcher and 7 300K for the second round Searcher.
23

---

Similar to LJ potentials, all models utilize five EGNN layers with 128 hidden dimensions. We use Dirac delta prior distribution at the origin and mean-free forward transition kernel guarantee equiviauur of the marginal density in inference. The Learner network comprises five EGNN layers, while the predictor network and target network in the RND framework contain three and two layers, respectively. As in LJ potentials, we use the Brownian motion coefficient of 0.2. We run 10K epochs in the first round and 20K epochs in the second. As in LJ-55, we maximize the log-likelihood for the first 5K epochs each round. We also utilize randomized time scheduling for our method. We train PIS at learning rate of le and all other methods at 5e 4. We use a batch size for all methods except PIS, which uses an 8 batch size due the memory limitation required by the forward SDE computational graph.
In inference time, we follow [23]. We first align the topology of generated samples with the target bond graph since the architecture and machine learning potential have degree of freedom in atom ordering. We first match the bond graphs of generated samples with given bond graph of interest and then correct the chirality of the generated sample to fit the target molecular configuration. The generated sample is rejected if the bond graph is not isomorphic to the target bond graph.
A.5 Task details
40-Component Gaussian Mixture Model (40GMM). The 40-com componmnnntt Gaussian Mixture Model (GMM) consists of mixture distribution of 40 Gaussian components, each characterized by distinct mean vector μt, The energy function for the GMM is defined as:
n Σ n
8(x) = log
where n 40, the weight of each i-th Gaussian component is the same, and Hi (2)) is the probability density function of the multivariate Gaussian distribution.
WanyWell distributions. The Manywell potential describes high-dimensional energy landscape containing multiple wells (local minima), each representing stable states with distinct energy levels. The energy function of Manywell distribution is given by:
11
(x))=
Σ
(x
6x2
2 x2h11 2x2))+ +C,
where и d/2 is the number of wells, and d is the dimensionality of the landscape. Adjusting the dimensionality 2n allows varying the number of wells and complexity, creating tasks like Manywell-32,, Manywell-64 and Manywell- 128..
Lennard-Jones (IJ)) potentials. The Lennard-Jones potential models the interactions between particles. The energy function is defined as:
12
6
((x) 2x
Σ 1siJJEN
2
Ir: remll?
(23)
rij
where and are parameters defining the depth of the potential well and the energy factor, respectively. Tiy = ||x; xxy| represents the Euclidean distance between particles and j. or is the characteristic distance at which the potential between two particles vanishes, often interpreted as the van der Waals radius. In our experiments, we set all parameters to 1.0, i.c., K = o = = 1.0. Adjusting the number of particles creates tasks such LJ-13 and LJ-55, increasing the complexity of the particle interactions and resulting in a rugged energy landscape.
TorchANI potential for peptides. We leverage TorchANI [16], PyTorch implementation of ANI deep-learning potentials trained on quantun-mmehanical reference data, to accurately calculate molecular energies. It provides transferable machine learning potential trained on organic molecules for efficient energy and force evaluation with accuracy comparable to density-function theory (DFT). In particular, TorchANI excels modeling small peptides.
24

---

Table 6: ELBO, EUBO, their gap, and Energy calls 40GMM and Manywell-32
(i/ 2)
Manywell (df
32)
Method
ELBOT
EUBO
Gap
Energy
BUBO
Gap
Energy calls
PIS+LP
2.42 0.07
3,75 0.22
300M
160.83
180,49 4.76
19.66 4.78
300M
TB+LP
0.03
0.04
0.87 0.03
160M
195.89
8.14
34.37
8.15
300M
FL-SubTB+LP
-..36
0.58 0.08
0.94 0.07
260M
160.74
215.93
4
55.19 4.52
350M
TB+LS+LP
0..38
0.32 = 0.02
0.69 = 0.02
320M
162.95 # 0.08
166.30 0.11
3.35 0.314
320M
TB-EspleLPP
-0.37+0.01
0.32
0.69
300M
160.76 0.13
14,90 166.25
TRe
-0.37+0.01
0.34 0.02
0.71 0.02
320M
162.97 0.06
3.28
320M
PIS
-2.03+0.22
55.48 10.71
57.50 9.02
100M
159.71
1,70
333.79 3.99
174.08 4.33
100M
TR
-1.35
99.04 6.01
100.40 5.67
100M
160.58
0.87
4.9928
166.52
278.70
166.49
100M
TB+LS
0.05
0.83 0.46
1.21 =: 0.38
290M
163.12 0.10
166.05 a 0.12
2,93 a.
290M
TB+ExpH+LS
-0.38
0.58 0.34
0.96 a 0.34
290M
160.87 1. 3.31
168.27 1.49
7,40
290M
GAFN
282/40
121.38 2.02
IDEM
12:75 0.41
3.67
14.89 3.70
300M
142.23 0.40
211.56
69.33 2.56
300M
SGDS fround 1)
0.01
0.33 0.02
0.73 a 0.02
6M
162.49 a 0.05
166.60 0.01
4,11 0.05
9M
SGDS fround 2)
0.40
0.33 = 0.05
0.73 = 0.05
12M
162.63 + 0.01
166. 166.488-2 0.03
3.85 0.03
20M
(a) Ground Truth
(b) SGDS
(c) PIS+LP
(d) TB+ExpI.+LS
(c) iDEM
Figure 5: Mode coverage comparison 40GMM.
B Additional experimental results
B.1 Low-dimensional standard benchmarks
Baselines and settings. We benchmark our framework on two standard low-dimensional tasks: 40GMM and Manywell-32. Consistent with the high-dimensional experiments, report ELBO, EUBO, their gap, and the number of energy calls required during training. We employ the same baseline methods and trajectory configurations (including trajectory length and training objectives) as in the high-dimensional settings. We provide detailed configurations, including diffusion scales for each task, in Table 5.
Results. As demonstrated Table 6, our method achieves competitive performance on lower- dimensional standard tasks, producing EUBO and ELBO metrics comparable to the strongest base- lines, while using significantly fewer energy calls. On the 40GMM task, despite some baselines reporting strong ELBO and EUBO scores, they notably fail to capture the mode located at the bottom-right corner (see Figure 5). In contrast, our framework reliably identifies all modes without sacrificing performance metries. We report both the first-round and second-round performances of our method, showing that our method attains robust performance on low-dimensional tasks even in the first round, with slight but consistent improvement observed in the second round.
B.2 Debiasing of Learner from MCMC Searcher
To address potential biases inherent in MCMC sampling due to finite-length chains, our framework incorporates both off-policy TB training using samples from the Searcher and on-policy TB training. This design choice aims to mitigate biases arising from the Searcher samples alone by enabling the Learner model adjust toward the target distribution.
To evaluate whether the Learner effectively debiases the samples collected by the Searcher, we compare kernel density estimations (KDE) of samples obtained by the AIS Searcher with those generated by the on/off on/off-policy TB Learner on Manywell distributions. Figure 6 illustrates these KDE comparisons across dimensions 32, 64, and 128.
Due to the varying mode masses assigned Manywell distributions, even when AIS successfully covers all modes with limited budgets, struggles to preciscly capture the relative mode masses.
25

---

(a) Ground Truth
(b) AIS
(c) SGDS (d = 32)
(d) SGDS (d = 64)
(c) SGDS (d 128)
Figure 6: KDE figures of AIS(T 100), ours, and true samples Manywell-32/64/122
Table SUBO-ELBO gap at 20% of the second training for different RND weight values.
RND weight
100
10a
102
103
Manywell-128
549.54
542.60
538.79
570.49
LJ-13
3.12
2.67
3.90
4.77
LJ-55
34.76
36.72
44.78
51.03
ALDP
17,531.95
17,381.04
17,417.75
17,410.98
In contrast, the KDE of the samples generated by the Learner aligns more closely with the true density, effectively reflecting the relative importance of different modes. This result highlights the effectiveness of combining on- and off-policy training to achieve better density approximation than relying solely on finite-budget AIS samples.
B.3 RND Weight Calibration
To evaluate the sensitivity of our method to the choice of the RND weight, we conducted additional experiments across approximately four different RND weight values per task. The results, summarized Table Table 7, show that the performance largely robust this hyperparameter. The calibration of the RND weight is straightforward:
run 20% of the second-round training estimate performance and select a reasonable value. As shown in the ELBO-EUBO gap measured at this stage, reported for Manywell-1228 LJ potentials, and peptides, the variation across different RND weights is minor, indicating that extensive hyperparameter tuning unnecessary.
Overall, these results suggest that the method performs consistently across wide range of RND weights, with negligible degradation in stability or performance. Although a more systematic tuning strategy could he explored in future work, the current approach provides reliable results with minimal calibration effort.
B.4 Consistency with parallel tempering MD
We further validate the consistency our method by replacing the Searcher with parallel tempering MD [40]. Parallel tempering (or replica exchange) MD is an enhanced sampling method that runs
26

---

multiple independent simulations (or replicas) at different temperatures and periodically attempts to exchange them, allowing -ewteempertture simulations to overcome energy barriers and escape local minima. Also, we align the number of energy evaluations by collecting additional data for MLE, ensuring a fair comparison. Table 8 shows the ELBO, EUBO, and their gap, demonstrating that our method achieves better metrics than MLE. This result indicates that our approach consistently works even when combined with advanced sampling techniques such as parallel tempering.
Table Comparison main metrics between parallel tempering MLE (forward KL) training and method with parallel tempering Searcher. We align the number of energy calls by collecting more data for MLE.
Method
ELBO [x10']
EUBO [x10'] |
Gap [x10'] 1
MLE
520.71 + 0.02
538.03 + 0.00
17.32 0.02
SGDS
521.03 + 0.01
538.02 + 0.00
16.99 0.01
C Limitations
While our framework demonstrates strong empirical performance, several limitations remain.
First, the effectiveness of intrinsic rewards from RND depends on careful tuning of the novelty scale parameter a. Poorly calibrated can overly emphasize exploration, producing noisy or irrelevant samples, or conversely yield overly conservative exploration. This could mitigated by employing adaptive strategies that dynamically adjust a during sampling based on diversity metrics or exploration progress signals.
Additionally, the quality of samples provided by the Searcher sets fundamental exploration limit. the Searcher fails to adequately explore challenging modes, the Learner will inevitably inherit these limitations, particularly in barrier energy landscapes, Introducing enhanced exploration strategies, such as parallel tempering or more advanced proposal schemes like HMC, could improve coverage of harplto-sample modes.
27


Reviews:
[{"review_id": "20kGgQL6Th", "reviewer": "Reviewer_Kvbn", "summary": "This paper addresses the problem of sampling from an unnormalized energy with learned diffusion samplers. They combine an MCMC-based “searcher” with intrinsic novelty reward and a combination of offline and online diffusion sampler training using the trajectory balance loss. By combining these strategies, they are able to show improved results compared to previous diffusion sampler approaches (iDEM, PIS, TB) on several classical energies and a small peptide neural-network energy function.", "questions": "-Since the ground-truth is also generated by some MCMC variant, why wouldn’t we just do flow/score matching onto the MCMC samples and achieve better performance? How good do these proposal samples from MCMC need to be to effectively learn a diffusion sampler with the TB loss?\n\n- What is the run-time compared to other algorithms? About how much time is spent on MCMC during each “searcher” phase? A run-time comparison would be helpful, especially for the case of Alanine dipeptide which uses the neural network energy function.\n\n\n- How is the ELBO evaluated for diffusion samplers? It’s my understanding that for samplers like PIS, you use the path-weights via the Radon-Nikodyn derivative, which is not exactly the model-likelihood of the sample x_1.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- This paper combines MCMC sampling and offline diffusion sampler training in a novel way that yields good results on classical energy benchmarks which previous diffusion samplers have struggled with. \n\n- The incorporation of an intrinsic reward in the energy for exploration is a novel heuristic for diffusion samplers (despite concerns below).\n\n- This work presents many empirical evaluations of combining a number of interesting heuristics that seem to address mode collapse issues common with diffusion samplers.\n\nWeaknesses:\n\n- The abstract mentions addressing the limitation of “expensive energy functions”, where Manywell and Lennard Jones are very cheap to evaluate. It seems like the MCMC-based searcher might not scale well for expensive energies, and I suspect the majority of the run-time would be spent on this phase. It would be good to see neural network energy examples other than Alanine-dipeptide and their run-times (see [1] for results on larger peptides).\n\n- I’m not very convinced by the heuristic “searcher” exploration reward as a contribution. The mean squared error to this random function at some point x will depend on things like the Lipschitz constant. It could be a good heuristic for novelty in the beginning of training, but won’t prevent that mode from being lost once it has been visited. In the ablation on Manywell 128, “w/o RND-searcher” performs only marginally worse than with RND.\n\n[1] Tan et al 2025 “Scalable Equilibrium Sampling with Sequential Boltzmann Generators”", "was_revised": true, "final_justification": "The authors have provided additional experiments and clarified most of my concerns about the comparison to MCMC+flow-matching. The results are interesting and I like the idea of adding some MCMC strategy in the inner-loop. I do wish that the RND exploration loss provided a more principled understanding of addressing mode exploration / coverage, as it is somewhat heuristic from my understanding. That is why I lean towards weak accept rather than accept.", "rebuttal": "Dear Reviewer Kvbn,\n\nThank you for your thoughtful and encouraging review. We appreciate your comments on the novelty of the methodological incorporation and several evaluations of our work. Below, we address your points in detail.\n\n### **W1: It seems like the MCMC-based searcher might not scale well for expensive energies. It would be good to see neural network energy examples other than Alanine Dipeptide and their run-times.**\n\nTo resolve your concern, we provide additional results on Alanine Tripeptide and Tetrapeptide considered in Tan et al.[1].\n\nTable 1 shows that the searcher is not the bottleneck for scaling to larger systems. Table 2 shows that our method still outperforms MLE (or forward KL) training on samples produced by the searcher. All model architectures and training settings for the tripeptide and tetrapeptide follow those used for the dipeptide, with the only change being batch sizes of 8 for the tripeptide and 4 for the tetrapeptide (versus 16 for the dipeptide). Upon further inspection, we observed that our model successfully generates low‑energy conformations, and we plan to include 3D visualizations of these structures in the future manuscript.\n\n###### **Table 1. Runtime (minutes) of neural network energy for a single round using NVIDIA GeForce RTX 3090.**\n\n| Task | Searcher | Learner | Total |\n| - | - | - | - |\n| Alanine Dipeptide | 128 | 1181 | 1309 |\n| Tripeptide | 139 | 2497 | 2636 |\n| Tetrapeptide | 157 | 3858 | 4015 |\n\n###### **Table 2. Performance comparison with MLE on Alanine Tripeptide with the same energy calls (1M).**\n\n| Tripeptide | Gap $[\\times 10^3]\\downarrow$ | ELBO $[\\times 10^3]\\uparrow$ | EUBO $[\\times 10^3]\\downarrow$ |\n| - | - | - | - |\n| MLE | $6.117 \\pm 0.019$ | $858.947 \\pm 0.015$ | $865.064 \\pm 0.002$ |\n| **Ours** | $\\mathbf{5.600 \\pm 0.030}$ | $\\mathbf{859.455 \\pm 0.031}$ | $\\mathbf{865.055 \\pm 0.005}$ |\n\n###### **Table 3. Performance comparison with MLE on Alanine Tetrapeptide with the same energy calls (1M).**\n| Tetrapeptide | Gap $[\\times 10^3]\\downarrow$ | ELBO $[\\times 10^3]\\uparrow$ | EUBO $[\\times 10^3]\\downarrow$ |\n| - | - | - | - |\n| MLE | $29.072 \\pm  0.560$ | $1,115.880 \\pm  0.560$ | $1,144.953 \\pm  0.002$ |\n| **Ours** | $\\mathbf{26.600 \\pm 0.214}$ | $\\mathbf{1,118.342 \\pm 0.213}$ | $\\mathbf{1,144.942 \\pm 0.013}$ |\n\n---\n\n### **W2-1: The Searcher exploration reward with RND won't prevent the previously visited mode from being lost once it has been visited.**\n\nIt seems that you are concerned about the following case: when the RND method assigns too small exploration reward for a previously visited mode, the Learner loses this mode during training. \n\nOur algorithm avoids this issue for several reasons. First, in our training, the RND reward decays slowly even for previously visited modes, so that the Learner properly learns the mode before the RND reward becomes negligible. Next, we use a top-k priority buffer to keep the low energy states, hence the Learner does not lose the mode. Finally, even if the RND reward becomes negligible, the Searcher will have a relatively high probability of re-visiting the mode due to its low energy.\n\n### **W2-2: In the ablation on Manywell 128, “w/o RND-searcher” performs only marginally worse than with RND.**\n\nThe improvement is significant since it is in log scale. In the ablation on Manywell 128, the RND searcher improves ELBO by $6.4$ of ELBO, which actually represents $\\exp(6.4) \\approx 403.43$ times better performance in lower bounding the marginal probability of the Learner generating the true sample.\n\n---\n\n### **Q1-1: Since the ground-truth is also generated by some MCMC variant, why wouldn’t we just do flow/score matching onto the MCMC samples and achieve better performance?**\n\nOur algorithm outperforms flow/score matching on MCMC samples since our online learning collects informative samples in a sample-efficient manner with a good exploitation-exploration trade-off. That is, once the Learner is sufficiently trained, the Learner generates low-energy samples that are hard to collect from short MCMC trajectories (exploitation), while the Searcher is still necessary to collect samples from unseen modes (exploration).\n\nWe already verified this in Figure 2-\\(c\\), where TB loss outperforms MCMC + MLE equivalent to score matching on the MCMC samples. Also, we show the debiasing effect of TB on Manywell distributions in Figure 6 in Appendix B.2.\n\n### **Q1-2: How good do these proposal samples from MCMC need to be to effectively learn a diffusion sampler with the TB loss?**\n\nFor Manywell 128 case, 300 parallel chains for Annealed Importance Sampling (AIS) with length 100 are enough to help the model learn well. For LJ-55, a single chain for Metropolis-Adjusted Langevin Algorithm (MALA) with length 10K is enough to learn well. For Alanine Dipeptide, four parallel chains of high-temperature molecular dynamics (MD) with length 110K are enough to capture diverse modes and train the learner well.\n\n---\n\n### **Q2: A run-time comparison would be helpful, especially for the case of Alanine Dipeptide, which uses the neural network energy function.**\n\nWe present the table of runtime of the searcher and the learner in Alanine Dipeptide. Table 1 shows that the run-time of Searcher is marginal compared to that of Learner. \n\n###### **Table 4. Runtime (minutes) of neural network energy using NVIDIA GeForce RTX 3090. Note that the first round searcher requires computing only neural network energy (51 minutes), but the second round searcher utilizes the RND network additionally (128 minutes).**\n\n| Alanine Dipeptide | Run-time |\n| - | - |\n| MLE | 1,583 |\n| **Ours** | 1,850 (searcher 179 + learner 1671) |\n\n---\n\n### **Q3: How is the ELBO evaluated for diffusion samplers? It’s my understanding that for samplers like PIS, you use the path-weights via the Radon-Nikodym derivative, which is not exactly the model-likelihood of the sample $x_1$.**\n\nWe evaluate the ELBO by the following steps:\n\n* Compute $\\log P_F(x_{t+\\Delta t}|x_t ; \\theta)$ along the forward path with the drift $u_\\theta(x_t,t)$ and the noise schedule $\\sigma$, and $\\log P_B(x_t|x_{t+\\Delta t})$ along the Brownian bridge for $t\\in[0,1]$.\n* Average $\\log\\frac{R(x_1)\\sum P_B(x_t|x_{t+\\Delta t})}{\\sum P_F(x_{t+\\Delta t}|x_t ; \\theta)}$ along the trajectories generated by the Learner.\n\nWe follow prior works that use the ELBO as a performance metric [2-4] as a lower bound to the model likelihood of samples.\n\n---\n\n[1] Tan et al., \"Scalable Equilibrium Sampling with Sequential Boltzmann Generators\", ICML 2025.\n\n[2] Sendera et al., \"Improved off-policy training of diffusion samplers\", NeurIPS 2024.\n\n[3] Kim et al., \"Adaptive teachers for amortized samplers\", ICLR 2025.\n\n[4] Vargas et al., \"Transport meets variational inference: Controlled Monte Carlo diffusions\", ICLR 2024.\n\n---\nWe hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification."}, {"review_id": "RlZOmMXgpI", "reviewer": "Reviewer_DVcd", "summary": "In this paper, the author's improve diffusion sampling--the application off diffusion models to sampling from a distribution only with access to its unnormalised density function--specifically targetting settings where computing the gradient of the log density is expensive and in high-dimensional settings where exploration is key. They propose SGDS which combines off-policy training with biased MCMC algorithms that searches for representative examples.", "questions": "- Does the paper's limited consideration of real-world datasets reflect what is standard in the literature?\n- In the limitations, it is discussed that calibration of alpha is important. How simple is this callibration?", "limitations": "Yes, but the limitations are in the appendix. I would encourage the authors to find space for this in the main body", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper is well-written and provides a good overview of the literature and the tools that they borrow in their algorithm.\n- Diffusion samplers are an important and recent class of algorithms and papers targetting practical improvements to these algorithms, like this paper, are important.\n- The toy example experiments indicate a clear benefit of SGDS over existing algorithms.\n\nWeaknesses:\n- I am not overly familiar with the ALDP example but I think it would be good to compare with some papers that target this task (it seems that the comparison is mainly against generic sampling or diffusion sampling algorithms). For example, if the baseline from Midgley et al. was included, this would help.\n- The contributions made in this paper appear to be the use of two existing methods to the setting of diffusion samplers. Due to my lack of knowledge of the literature, it is not immediately clear to me how much of a novel contribution this is.", "was_revised": true, "final_justification": "I do not know a great deal about developments in this area but the paper is well-written and has convinced me that they have made a contribution to an important field.", "rebuttal": "Dear Reviewer DVcd,\n\nThank you for your thoughtful and encouraging review. We greatly appreciate your recognition of the clarity of our writing, the importance of our field, and the significance of our experimental results. Below, we address your comments in detail.\n\n### **W1: It would be good to compare with some papers that target this task, e.g., the baseline from Midgley et al. [1].**\n\nThank you for the suggestion. We did not include FAB [1] in our comparison because, as noted in [6], FAB cannot amortize sampling and relies on annealed importance sampling, which demands many energy evaluations. Instead, we compare against Adjoint Sampling method from [6], which also targets small peptides like Alanine Dipeptide. As shown in Table 1, when the model is trained solely on Alanine Dipeptide, our method outperforms Adjoint Sampling.\n\n###### **Table 1. Comparison of Wasserstein-2 distance of positions and energy histogram on Alanine Dipeptide.**\n| ALDP | $\\mathcal{W}_2 \\downarrow$ | $E(\\cdot)$ $\\mathcal{W}_2 \\downarrow$  |\n| - | - | - |\n| Adjoint Sampling | 33.29 | 599.74 | \n| SGDS (ours) | $\\mathbf{32.90}$ | $\\mathbf{487.54}$ |\n\n---\n\n### **W2: Unclear novelty in combining two existing methods.**\n\nWe resonate with your concern, but we also believe that our method is not just a combination of two methods. \nThe diffusion‑sampling field sits at the crossroads of diffusion models, probabilistic inference, stochastic control, and reinforcement learning, so progress often comes from *properly hybridizing* techniques that each community describes in its own language.\n\n* Prior work has imported existing **stochastic‑control** formulations into diffusion models [1],\n* utilized existing **replay buffer** auxiliary [6],\n* applied existing **local credit‑assignment** tricks [7] from RL, and\n* mapped existing **search algorithms** onto diffusion sampling [8].\n\nOur contribution follows this line by combining sample-efficient RL techniques with MCMC inside a diffusion sampler, which is not previously explored and therefore a meaningful addition to the literature.\n\n---\n\n### **Q1: Does the paper's limited consideration of real-world datasets reflect what is standard in the literature?**\n\nYes. In fact, our work expands the standard benchmarks [7-9], i.e., GMM, Lennard-Jones, and Manywell, by additionally including Alanine Dipeptide experiments. In the future, we expect researchers to further the algorithms and apply them to real-world problems, e.g., fine-tuning biomolecular foundation models like BioEmu [10] with energy function evaluations. \n\n---\n\n### **Q2: In the limitations, it is discussed that calibration of alpha is important. How simple is the calibration of RND weight?**\n\nWe simply choose from four possible RND weights (1, 10, 100, 1000) based on the performance at 20% of the second round, while the performance has a dependency on the RND weight.\n\nTo alleviate your concern, we conducted experiments across approximately four different RND weight values per task to ensure stable training and meaningful results in Table 2.\n\n###### **Table 2. ELBO-EUBO gap at 20% of the second round for Manywell-128, LJ potentials, and ALDP.**\n\n| RND weight | $10^{0}$ | $10^{1}$ | $10^{2}$ | $10^{3}$ |\n| - | - | - | - | - |\n| Manywell-128 | 549.54 | 542.60 | **538.79** | 570.49 |\n| LJ-13 | 3.12 | **2.67** | 3.90 | 4.77 |\n| LJ-55 | **34.76** | 36.72 | 44.78 | 51.03 |\n| ALDP | 17,531.95 | **17,381.04** | 17,417.75 | 17 410.98 |\n\nWe found that reasonable values could be selected with minimal tuning, though a more systematic approach could be explored in future work.\n\n---\n\n[1] Midgley et al., \"Flow annealed importance sampling bootstrap\", ICLR 2023.\n\n[2] Noe et al., \"Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning\", Science, 2019.\n\n[3] Jing et al., \"Torsional diffusion for molecular conformer generation\", NeurIPS 2022.\n\n[4] Hassan et al., \"Et-flow: Equivariant flow-matching for molecular conformer generation\", NeurIPS 2024.\n\n[5] Havens et al., \"Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching\", ICML 2025.\n\n[6] Akhound-Sadegh et al., \"Iterated Denoising Energy Matching for Sampling from Boltzmann Densities\", ICML 2024.\n\n[7] Zhang et al., \"Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization\", ICLR 2024.\n\n[8] Sendera et al., \"Improved off-policy training of diffusion samplers\", NeurIPS 2024.\n\n[9] Zhang et al., \"Path Integral Sampler: a stochastic control approach for sampling\", ICLR 2022.\n\n[10] Lewis et al., \"Scalable emulation of protein equilibrium ensembles with generative deep learning\", Science 2025.\n\n---\n\nWe hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification."}, {"review_id": "ur9zSQkSFF", "reviewer": "Reviewer_fmfd", "summary": "The paper proposes search-guided diffusions samplers (SGDS) to sample from unnormalized target densities.\nThe method trains a diffusion sampler via trajectory balance from two sources: Forward trajectories from the sampler (on-policy), and backward trajectories starting from points obtained from an MCMC searcher (off-policy).\nThe method is evaluated on several synthetic datasets (40GMM, Manywell), as well as atomic systems (LJ, Alanine Dipeptide) and achieves superior accuracy and mode coverage compared to prior work.", "questions": "# Questions\n\n- Why is there such a large discrepancy between TB+LS and SGDS? When running TB+LS with different hyperparams (fewer local searches with more depth) would you effectively recover SGDS (less minor improvements such as weight re-init and RND)\n- How are the MCMC chains for the searcher initialized? Would it make sense to draw the initial states from from forward trajectories?\n- SGDS appears to be exceptionally efficient with regards to the number of energy evaluations. Why does it require so much fewer evaluations (for example compared to TB+LS)? Is it correct that SGDS only consumes energy evaluations during search (1 per chain and MCMC step) and on-policy steps (B evals per opt step), but none during off-policy steps?\n\n# Minor points\n\n- Eq (3): The indices of the product in $P_B$ seem to be off by one\n- Equations after line 84 and 86: Shouldn't $P_F(\\tau, \\theta)$ also be conditioned on $x_0$ (analogous to $P_B$ being conditioned on $x_1$)? If I understand correctly you later choose $\\mu_0$ to be a dirac delta at the origin but at this point (and Eq. 3) it seems unintuitive.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# Strengths\n\n- Comprehensive evaluation with strong performance of SGDS across experiments\n- Ablation study of proposed improvements\n\n# Weaknesses\n\nThis appears to be a collection of various existing \"tricks\" (longer searches, weight re-initialization after rounds, RND-based penalization of already discovered modes) yielding overall a strong paper. I would primarily encourage to more clearly discuss the intuition behind these improvements, in particular when comparing to similar existing variants.\n\n- The method seems to be primarily an incremental improvement on existing work, in particular Sendera et al (which the authors cite and compare against approriately). For single round training (where re-initialization and random network distillation do not apply) the only key difference appears to be that SGDS runs fewer/longer local searches vs. the frequent/short local searches of Sendera et al. It would be helpful to qualitatively discuss these differences to highly related methods more thoroughly.\n- While the manuscript gives a good overview of the field, its heavy use of abbreviations makes reading key tables such as Tab. 1 hard to read for readers not intimately familiar with the field. Either naming the abbreviations in the table caption or (perhaps better) a short glossary in the supplement would be very helpful. As an example (+Expl) is only defined via a reference to Sendera et al, which in turn itself does not contain an easy-to-find definition.", "was_revised": true, "final_justification": "I keep my initial assessment that this is a solid paper, which outperforms existing sampling techniques.\nMy initial concerns were primarily around readability of the paper and qualitative comparison against existing techniques: The authors have promised to address both in the final paper.", "rebuttal": "Dear Reviewer fmfd,\n\nThank you for your thoughtful and encouraging review. We appreciate your recognition of the clarity of our experiments and evaluations. Below, we address your comments in detail.\n\n### **W1: The method seems to be primarily an incremental improvement on existing work, in particular, Sendera et al. It would be helpful to qualitatively discuss the differences.**\n\nOur contribution is a conceptual shift in a mixture of training schemes compared to prior works. While existing diffusion samplers rely on full online learning [1,2] or use a buffer dependent on the current model's performance [3,4], SGDS utilizes both offline guidance from an independent MCMC-based searcher and online learning from the diffusion model itself. This is a meaningful improvement in the context of the training structure.\n\n---\n\n### **W2: It is difficult to read the tables because of their heavy use of abbreviations.**\n\nTo resolve your concern, we plan to incorporate the following clarifications for Table 1 in the future manuscript as an extension of the description written in Section 5.1.\n\n* **PIS (Path Integral Sampler)** [1], **TB (Trajectory Balanced)** [5], and **FL-SubTB (Forward-Looking SubTB)** [6] denote the types of objectives. **GAFN (Generative Augmented Flow Networks)** [7] is a GFlowNet-based method that directly injects intrinsic rewards into TB loss. **AT (Adaptive Teacher)** [8] is a framework that introduces an additional neural sampler (the Teacher) trained to focus on high‑loss regions of the Student.\n\n* **LP (Langevin Parametrization)** [4] refers to a drift construction technique where the model’s drift term is combined with the gradient of the target energy function, helping the trained dynamics to better follow the target energy landscape.\n\n* **LS (Local Search)** [4] denotes a refinement step where a short Markov chain, guided by the target energy gradient, is applied from the final state of the generated trajectory to improve sample quality.\n\n* **Expl (Exploration)** [4] represents a noise-scheduling technique applied during trajectory generation, in which additional stochasticity is injected more in the early training phase to promote broad exploration and gradually reduced over time for exploitation later.\n\n---\n\n### **Q1-1: Why is there such a large discrepancy between TB+LS and SGDS?** \n\nThe differences of RND reward, initial distribution, and search length make a large discrepancy between TB+LS and SGDS. \n\nWhile TB+LS relies on only Langevin dynamics with the target energy function, SGDS utilizes RND reward together for sample-efficient exploration, which is one of the key components in our framework.\n\nAnd TB+LS always seeds each walk from the current diffusion sample [4]; if that sampler is low‑quality in training, exploration is constrained, while SGDS does not depend on the current training states.\n\nAlso, TB+LS search is too short to reach high‑reward regions on larger, sharper tasks. SGDS extends each walk, letting trajectories climb farther before termination.\n\nThese three factors synergistically contribute to the substantial performance gap between SGDS and TB+LS.\n\n\n### **Q1-2: When running TB+LS with different hyperparams, would you effectively recover SGDS (except for re-init and RND reward)?**\n\nNo, they share some motivations, yet the methods are fundamentally different in where to begin the search. TB+LS starts search from the samples generated by the current Learner; if that sampler is low‑quality in training, exploration is constrained, while SGDS does not depend on the current training states.\n\nAdditionally, we conducted an experiment on Manywell-128 to address your question of how TB+LS performs with a double-fewer freuquency but double deeper chains.\n\n###### **Table 1. ELBO-EUBO gap of SGDS and TB+LS with different hyperparameters for local search on Manywell-128.**\n\n| Energy | Gap $\\downarrow$ | LS cycle | Chain length | Energy calls |\n| - | - | - | - | - |\n| TB+LS | 164.95 | 100 | 200 | 290M |\n| TB+LS | 96.45 | 200 | 400 | 290M |\n| SGDS | **70.35** | - | - | 20M |\n\nFewer and deeper local search with TB could make TB+LS better, but SGDS is still superior in performance and sample efficiency.\n\n---\n\n### **Q2-1: How are the MCMC chains for the searcher initialized?**\n\nWe initialize the MCMC chains for the searcher with a Gaussian prior in MALA and AIS, and relaxed positions starting from a fixed position in pdb file in MD. \n\n### **Q2-2: Would it make sense to draw the initial states from forward trajectories?**\n\nDrawing the initial states from forward trajectories does not make sense in our algorithm. We initialize the MCMC chains with a Gaussian prior. Seeding the search with forward trajectories reduces exploration and biases it toward the learner’s current distribution—especially problematic when the learner is immature. This choice helps explain why SGDS outperforms TB + LS, which draws its initial states from the diffusion sampler’s forward policy.\n\n---\n\n### **Q3-1: Why does SGDS require fewer evaluations (for example, compared to TB+LS)?**\n\nSGDS requires fewer energy evaluations than other methods since it consumes energy calls for collecting the Searcher samples, and it recycles the samples and rewards for off-policy training. For example, TB+LS methods waste too many samples for searching high-reward regions as they initialize the local search start point too frequently using the Learner. \n\n### **Q3-2: Is it correct that SGDS only consumes energy evaluations during search (1 per chain and MCMC step) and on-policy steps (B evals per opt step), but none during off-policy steps?**\n\nYes.\n\n---\n\n### **Minor1: Eq (3): The indices of the product in $P\\_B$ seem to be off by one.**\n\n### **Minor2 : Equations after line 84 and 86: Shouldn't $P\\_F(\\tau;\\theta)$ also be conditioned on $x\\_0$ (analogous to $P\\_B$ being conditioned on $x\\_1$)?**\n\nWe thank you for pointing out the minor typos in the equations. We will correct $\\sum_{i=0}^{T-1}P_B(x_{(i-1) \\Delta t} | x_{i \\Delta t})$ to $\\sum\\_{i=1}^{T}P_B(x\\_{(i-1) \\Delta t} | x\\_{i \\Delta t})$, and $P\\_F(\\tau;\\theta)\\mu\\_0(x\\_0)$ to $P\\_F(\\tau|x\\_0;\\theta)\\mu\\_0(x\\_0)$ in the future manuscript.\n\n---\n\n[1] Zhang et al., \"Path Integral Sampler: a stochastic control approach for sampling\", ICLR 2022.\n\n[2] Vargas et al., \"Transport meets variational inference: Controlled Monte Carlo diffusions\", ICLR 2024.\n\n[3] Akhound-Sadegh et al., \"Iterated Denoising Energy Matching for Sampling from Boltzmann Densities\", ICML 2024.\n\n[4] Sendera et al., \"Improved off-policy training of diffusion samplers\", NeurIPS 2024.\n\n[5] Malkin et al., \"Trajectory balance: Improved credit assignment in gflownets\", NeurIPS 2022.\n\n[6] Zhang et al., \"Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization\", ICLR 2024.\n\n[7] Pan et al., \"Generative augmented flow networks\", ICLR 2023.\n\n[8] Kim et al., \"Adaptive teachers for amortized samplers\", ICLR 2025.\n\n---\n\nWe hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification."}, {"review_id": "Lr2VghXfzG", "reviewer": "Reviewer_HM7Y", "summary": "The paper proposes a scalable sample-efficient method for generating samples from an unnormalized density with increased exploration capabilities. To authors propose to collect off-policy samples using MCMC using the target energy with an exploration bonus that is computed using random network distillation. This data is then used to learn a diffusion model with off-policy TB loss and combined with on-policy data collected by the diffusion model. Additionally, the authors suggest to periodically reset the weights of the networks to prevent primacy bias. The method is evaluated on several toy and many-body problems.", "questions": "- What does LP stand for?\n- What is the benefit of this intrinsic motivation term, compared to just running MCMC on a higher target temperature? The latter also should encourage exploration.\n- Did the authors try the Log-variance loss instead of TB?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strenghts\n\n- The paper tackles the very important problem of improving exploration in diffusion-based sampling.\n- The paper uses mode-covering metrics such as EUBO.\n- The approach seems very sample efficient.\n\nWeaknesses\n\n- The RND weight varies between several orders of magnitude between different tasks (Table 5), suggesting the need for an intensive hyperparameter search per task to obtain a good exploration behavior. An ablation study would help to understand the sensitivity of the method for different values of the RND weight. However, the reviewer acknowledges that this was already included as a limitation of the method.\n- I currently doubt that the method outperforms parallel tempering / replica exchange MCMC / MD + forward KL training. It is well known that MD can easily obtain ground truth samples from ALDP. However, the performance of SGDS on ALDP doesn’t seem convincing to me due to significant differences in the Ramachandran histogram.\n- The ablation study on ManyWell shows that the RND searcher does not significantly boost performance.\n- A comparision to Adjoint sampling seems missing, which has shown to be very scalable and sample efficient as well.\n\n[a] No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers", "was_revised": true, "final_justification": "I remain somewhat skeptical about the novelty and practical impact of the work, but I appreciate the authors' efforts and will increase my rating to 4.", "rebuttal": "Dear Reviewer HM7Y,\n\nThank you for your thoughtful and encouraging review. We appreciate your recognition of the contributions and evaluations of our work. Below, we address your comments in detail.\n\n### **W1: The RND weight varies between several orders of magnitude for different tasks. An ablation study would help to understand the sensitivity of the method for different RND weights.**\n\nTo resolve your concern, we provide experimental results that verify that intensive hyperparameter search is not needed, and the results from different RND weights are not sensitive.\n\nTaking these factors into account, we conducted experiments across approximately four different RND weight values per task to ensure stable training and meaningful results. The calibration of the RND weight $\\alpha$ is important but simple: run 20% of the second round to measure the performance.\n\n###### **Table 1. ELBO-EUBO gap at 20% of the second round for Manywell-128, LJ potentials, and ALDP.**\n\n| RND weight | $10^{0}$ | $10^{1}$ | $10^{2}$ | $10^{3}$ |\n| - | - | - | - | - |\n| Manywell-128 | 549.54 | 542.60 | **538.79** | 570.49 |\n| LJ-13 | 3.12 | **2.67** | 3.90 | 4.77 |\n| LJ-55 | **34.76** | 36.72 | 44.78 | 51.03 |\n| ALDP | 17,531.95 | **17,381.04** | 17,417.75 | 17 410.98|\n \nWe found that reasonable values could be selected with minimal tuning, though a more systematic approach could be explored in future work.\n\n---\n\n### **W2: I currently doubt that the method outperforms forward KL training combined with one of parallel tempering, replica exchange MCMC, or MD.**\n\nTo dispel your doubt, we show that our method outperforms MLE (or forward KL) training on samples from parallel tempering (or replica exchange) MD in Table 2.\n\n###### **Table 2. ELBO-EUBO gap comparison between parallel tempering + MLE (or forward KL) training and ours with parallel tempering searcher. We align the number of energy calls by collecting more data for MLE.**\n\n| Method | ELBO $[\\times 10^{3}] \\uparrow$ | EUBO $[\\times 10^{3}] \\downarrow$ | Gap $[\\times 10^{3}]\\downarrow$ |\n| - | - | - | - |\n| MLE | $520.71\\pm0.02$ | $538.03\\pm0.00$ | $17.32\\pm0.02$ |\n| Ours | $\\mathbf{521.03\\pm0.01}$ | $\\mathbf{538.02\\pm0.00}$ | $\\mathbf{16.99\\pm0.01}$ |\n\n---\n\n### **W3: The ablation study on ManyWell shows that the RND searcher does not significantly boost performance.**\n\nWe believe the RND searcher significantly boosts the performance since the reported values are in log scale. The presence of the RND searcher results in an improvement of $6.4$ of ELBO in the ablation study, which actually represents $\\exp(6.4) \\approx 403.43$ times better performance in lower bounding the marginal probability of the Learner generating the true sample.\n\n### **W4: A comparison to Adjoint Sampling seems missing, which has shown to be very scalable and sample efficient as well.**\n\nWe reproduced Adjoint Sampling algorithm and conducted the experiments to compare SGDS and Adjoint Sampling on LJ-13 and LJ-55 potentials. The results show that SGDS outperforms Adjoint Sampling in the Wasserstein distances of positions and energy histograms.\n\n###### **Table 3. Comparison of Equivariant position Wasserstein-2 and energy histogram Wasserstein-2 between SGDS and Adjoint Sampling on LJ-13. The parameters of LJ-13 energy function are set to be the same as the settings in our experiments.**\n\n| LJ13 | $\\mathcal{W}_2 \\downarrow$ | $E(\\cdot)$ $\\mathcal{W}_2 \\downarrow$  | # of energy calls |\n| - | - | - | - |\n| Adjoint Sampling (reproduced) | 1.95 | 10.76 | 1M |\n| Adjoint Sampling (reported) | 1.67 | 2.40 | 1M |\n| SGDS (ours) | **1.63** | **0.7589** | 370K |\n\n###### **Table 4. Comparison of Equivariant position Wasserstein-2 and energy histogram Wasserstein-2 between SGDS and Adjoint Sampling on LJ-55.**\n\n| LJ55 | $\\mathcal{W}_2 \\downarrow$ | $E(\\cdot)$ $\\mathcal{W}_2 \\downarrow$  | # of energy calls |\n| - | - | - | - |\n| Adjoint Sampling (reported) | 4.50 | 58.04 | 1M |\n| SGDS (ours) | **4.17** | **23.0422** | 45K |\n\nWe reproduced Adjoint Sampling by closely following the public codes and variable settings, but the training on LJ-55 exploded. So we instead compare our method against the value reported in Adjoint Sampling paper for LJ-55 task.\n\n---\n\n### **Q1: What does LP stand for?**\n\nLP stands for **Langevin Parameterization**—a neural‑network design where the input includes the energy gradient as mentioned in Section 1, and the network outputs a scaled and biased version of this gradient, thus embedding a Langevin‑style inductive bias [2].\n\n---\n\n### **Q2: What is the benefit of this intrinsic motivation term, compared to just running MCMC on a higher target temperature?**\n\nThe intrinsic-motivation term benefits from the ability to incorporate previously visited samples to better guide the exploration process.\n\n---\n\n### **Q3: Did the authors try the Log-variance loss instead of TB?**\n\nYes, we did. To provide a clean and fair comparison for your question, we conducted experiments using four different seeds across 40GMM and Manywell-32/64,/128. Table 5 is the experiment results of comparing the EUBO-ELBO gap between TB and Log-variance for the Learner's objective: \n\n###### **Table 5. Comparison of EUBO-ELBO gap between TB and Log-variance on 40GMM, Manywell-32/64/128.**\n\n| Energy | 40GMM | Manywell-32 | Manywell-64 | Manywell-128 |\n| - | - | - | - | - |\n| Log-variance | 24.79 | 74.41 | 95.95 | 392.46 | \n| TB (ours) | **0.73** | **3.85** | **16.26** | **70.35** |\n\nHere, we believe that TB outperforms Log-variance since it estimates $\\log Z$ using samples from the Searcher, while Log-variance estimates $\\log Z$ on-the-fly using a smaller number of samples collected from the Learner.\n\n---\n\n[1] Sugita et al., \"Replica-exchange molecular dynamics method for protein folding\", Chemical physics letters 1999.\n\n[2] Sendera et al., \"Improved off-policy training of diffusion samplers\", NeurIPS 2024.\n\n---\n\nWe hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification."}]

Author rebuttal:

Meta-review: The paper introduces a novel diffusion sampler. The algorithm is a combination of an MCMC searcher with intrinsic novelty reward and offline and online diffusion sampler training using the trajectory balance loss.
The paper is well written and it combines ideas from sample-efficient RL with MCMC. The method is showcased in several benchmarks in the paper and during the rebuttal phase, as response to reviewers' questions.
Some of the concerns raised are about the conceptual novelty, runtime, and the practicality on real world tasks. The paper however focuses on tasks that are standard in the neural sampler literature, and does a through comparison with state of the art methods, showing consistent improvements.
Decision: Accept (poster)