Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities
Authors: Tara Akhound-Sadegh
Jungyoon Lee
Joey Bose
Valentin De Bortoli
Arnaud Doucet
Michael M. Bronstein
Dominique Beaini
Siamak Ravanbakhsh
Kirill Neklyudov
Alexander Tong
Venue: NeurIPS 2025 spotlight ()
Keywords: Boltzmann Sampling
Diffusion Models
Molecular Conformations
Sequential Monte Carlo
Primary area: machine_learning_for_sciences
TL;DR: We propose a diffusion-based approach to sampling from Boltzmann densities based on temperature annealing.

Abstract: Sampling efficiently from a target unnormalized probability density remains a core challenge, with relevance across countless high-impact scientific applications. A promising approach towards this challenge is the design of amortized samplers that borrow key ideas, such as probability path design, from state-of-the-art generative diffusion models. However, all existing diffusion-based samplers remain unable to draw samples from distributions at the scale of even simple molecular systems. In this paper, we propose Progressive Inference-Time Annealing (PITA) a novel framework to learn diffusion-based samplers that combines two complementary interpolation techniques: I.) Annealing of the Boltzmann distribution and II.) Diffusion smoothing. PITA trains a sequence of diffusion models from high to low temperatures by sequentially training each model at progressively higher temperatures, leveraging engineered easy access to samples of the temperature-annealed target density. In the subsequent step, PITA enables simulating the trained diffusion model to *procure training samples at a lower temperature* for the next diffusion model through inference-time annealing using a novel Feynman-Kac PDE combined with Sequential Monte Carlo. Empirically, PITA enables, for the first time, equilibrium sampling of $N$-body particle systems, Alanine Dipeptide, and tripeptides in Cartesian coordinates with dramatically lower energy function evaluations.

Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities
Tara Akdound-Sadegh" 1,;; Jungyoon Lee Avishek Joey Bose Valentin De Bortoli",
Arnaud Doucet Michael M. Bronstein 4.6, Dominique Beaini?.3.2 Siamak Ravanbakhsh
Kirill Neklyudov 3,2,8t, Alexander Tong9.3.2t
McGill University, Miila Quebec AI Institute, 'Université de Montréal, +University of Oxford, GGoogle DeepMind, AITHYRA, VValence Labs, Institut Courtois
Abstract
Sampling efficiently from a target unnormalized probability density remains a core challenge, with relevance across countless high-impact scientific applications. A promising approach towards this challenge is the design of amortized samplers that borrow key ideas, such as probability path design, from state-of state-of-the-art generative diffusion models. However, all existing Liffusion-basee samplers remain unable to draw samples from distributions at the scale of even simple molecular systems. In this paper, we propose PROGRESSIVE INFERENCE- - ANNEALING (PITA) novel framework to learn diffusion- based samplers that combines two comple- mentary interpolation techniques: I.) Annealing of the Boltzmann distribution and II.) Diffusion smoothing. PITA trains a sequence of diffusion models from high to low temperatures by sequentially training each model at progressively lower temperatures, leveraging enginecred easy access to samples of the temperature- annealed target density. In the subsequent step, PITA enables simulating the trained diffusion model to procure training samples at lower temperature for the next diffusion model through inference-timg annealing using a novel Feynman-Ke Kac PDE combined with Sequential Monte Carlo. Empirically, PITA enables, for the first time, equilibrium sampling of N-body particle systems, Alanine Dipeptide, and Tripeptide in Cartesian coordinates with dramatically fewer energy function evaluations. Code available at: https: // //gi thub. temp/ttrak///aa aaita.
Introduction
The problem of sampling from an unnormalized target probability distribution arises in numerous areas of natural sciences, including computational biology, chemistry, physics, and materials science (Frenkel and Smit, 2023; Liu, 2001; Ohno et al., 2018; Stoltz et al., 2010; Noé et al., 2019). In many of these high-impact scientific settings, the problem's 's complexity stems from operating in molecular systems where the unnormalized target (Boltzmann) distribution at a low temperature of interest is governed by highly non-convex and non-smooth energy function, under which there is limited to no available data (Hénin et al., 2022). As result, the sampling problem necessitates solving an equally hard exploration problem: finding the modes-in correct proportion- of the target distribution.
To address the general sampling problem, extensive research has been dedicated to Markov chain Monte Carlo methods (MCMC), Sequential Monte Carlo (SMC), and, particularly in physical systems, Molecular Dynamics (MD) (Leimkuhler and Matthews, 2015). To enhance scalability, Monte Carlo approaches often employ an interpolating sequence of probability distributions that
'Equal contribution 'Equal advising. Correspondence to tara. aaaraaathhomdaadegg@0ii... quebec,,
k. www.eedoo@gaail.coom com, atongCai at
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Inference time annealing
Noising samples
Di Diffussin inference
Temperature ameea Ling
T
7=
Diffusion
Figure 1: Illustration of the proposed PITA framework combining two complementary processes: temperature annealing of the target Boltzmann density and the diffusion process applied the collected samples. Annealed inference allows for decreasing the temperature (increasing (3) of trained diffusion model, thus generating samples from the annealed target. These samples can reused training a lower-temperature diffusion model.
transitions from easily sampled reference distribution to the desired target distribution via annealing/tempering strategies. This powerful concept underlies methods such as parallel tempering (Swendsen and Wang, 1986), Annealed Importance Sampling (Jarzynski, 1997; Neal, 2001), and SMC samplers (Del Moral et al., 2006). MD, conversely, involves integrating equations of motion using finely discretized time steps. Despite their effectiveness, both classes of methods possess inherent limitations that complicate their application to practical systems of interest: annealing modifies the masses of distribution modes depending their widths (a phenomenon often referred to as mass teleportation (Woodard et al., 2009)), while MD requires computationally expensive time discretization on the order of femtoseconds to simulate millisecond-scale phenomena.
Diffusion-based samplers are an alternative and emergent class of sampling techniques (Zhang and Chen, 2022; Vargas et al., 2023; Akhound-Sadegh et al., 2024; Berner et al., 2024; Blessing al., 2024; Havens et al., 2025) exploiting modern developments sample complex multi-modal distributions by leveraging prescribed interpolating probability path. However, instead relying on annealing, these samplers utilize a noising mechanism which theoretically enjoys favorable mixing properties compared to annealing (Chen et al., 2023).
generative modeling. They
Diffusion-based based samplers, despite their appeal, have not yet proven effective for even small molecular systems in Cartesian coordinates. This is primarily because of the absence of training data to accurately approximate the logarithmic gradients of the marginal densities, i.e. the Stein scores challenge distinct from generative modeling settings. Additionally, standard training objectives, such as reverse Kullback-Leibles are mode dropping and often yield too high-variance score estimates for stable training (Blessing et al., 2024). Crucially, current diffusion-based samplers require too many energy function evaluations for training. Indeed, when normalized by the number energy evaluations, carefully tuned MCMC methods with parallel tempering empirically competitive Liffusion-basse samplers (He et al., 2025).
with, not superior to, state-of-the-e
Main Contributionss In this paper, we introduce PROGRESSIVE INFERENCE- TIME ANNEALING (PITA), novel framework for training diffusion models to sample from Boltzmann distributions. PITA leverages two complementary interpolation techniques to significantly enhance training scalability: temperature annealing (increasing the system's 's temperature) and interpolation along conventional diffusion model's probability path. This combination is motivated by learning framework designed exploit their distinct advantages: temperature annealing mixes modes by lowering high-energy barriers, while diffusion paths avoid mass teleportation.

---

Annealing the target distribution transforms the challenging sampling problem into an easier one by removing high-energy barriers and flattening it. This crucial step enables the cheap collection of an initial high-temperature dataset via classical MCMC, which in turn facilitates the efficient training of an initial diffusion model. Subsequently, we define a novel Feynman- Kac PDE that, when combined with MMC-based resampling, allows us to simulate the trained diffusion model's 's inference process to produce asymptotically unbiased samples at lower temperature. This effectively allows us to train the next diffusion model, enabling the progressive and stable training of sequence of diffusion models until the target distribution is reached, as illustrated Fig. 1.
We test the empirical performance of PITA on standard N-body particle systems and short peptides in Alanine Dipeptide and tripeptides. Empirically PITA not only achieves state-of-the performance in all these benchmarks but is the first diffusion- in their native Cartesian coordinates. More importantly, we demonstrate that progressing down our designed ladder of diffusion models leads to significantly lower energy evaluations compared to MD, which a step towards realizing the promise of amortized samplers for accelerating equilibrium sampling.
sampler that scales to our considered peptides
Background
2.1 Diffusion models
A diffusion process defines interpolating path between easy-to-sample reference density, such as a multi-variate Normal, and a desired target distribution (x). When samples from the target distribution are available, it is possible to generate samples from intermediate marginals of the diffusion process (x) through the following Gaussian convolution:
Pr (+(x) = N((x,,,,, = dy (x ary, ())
As a result, this means that samples from Pr (x) can be generated as I+ a,y + O,E, where y ~ (yy) and NO, 1). Selecting specific schedules for a, and o, one can ensure the following boundary conditions. For 7 = 0, a 1, = 0 and (==)( = x (x)) i.e. the marginal matches the target distribution. For r = a, 0, a- and p+=1( = (0, 11), i.e. the marginal matches the standard multivariate normal distribution.
Importantly, despite the simplicity of sample generation, the evaluation of density pr px(x) is not straightforward. and one has to use deep learning models to approximate either scores V log (x)
or marginal densities pr (xx). Furthermore, the model density Pr (2) or scores V Prg))( (2x) can be
used to generate new samples from (Ix) using the reverse-time SDE. In particular, the marginals
introduced in Eq. (1) describe the marginal densities of the forward-ting (Ornsten--Uhlenbed) SDE
① bog Orr
θ 20
dx,
dr
log dw 7, Ir=o ~ mxx), Ør
(2)
ar
where W,, is the standard Wiener process, and marginals follow the Fokker-Planck PDE. After inverting the time variable in this PDE, i.e. t T, the time-evolution of marginals pr(x) is
Op. (x) (V, (V.)(())(aa))) Øl
2 Ap.(2) = (oo())=-(V,,((ee---- V
G (Voopp(()))
(3)
which shows that one can sample from the marginals {pr(=))eefoo.) via simulating the following SDE
(-axt
Gi (1 (1++))Vlogp((e di +GW 2
dr:
While the marginals are correct for any & > O, there are two important special cases: for &, = the equation becomes the reverse-SDE with the same path-measure as Fq. (2), and, for & = 0, the SDE becomes an ODE, In practice, this SDE simulated using either the density model exp(-U, Pz(x) (Du et al., 2023) or the score model (z; 0) se V log Po(r) (Song et al., 2021).
x
3

---

2.2 Annealing
Annealing defines a family of "simpler" problems when we have access to the unnormalized density by interpolating or scaling the target log-density (negative energy). Formally speaking, given the unnormalized density x (z), the annealed density is defined
{ Z3
/
T (z)
ZB :a=] dx T())3
where 3 is the inverse temperature parameter (i.e., 3 1/TT) controlling the smoothness of the target density. Thus, for high temperature 7 > (hence, 3 < 1), the target density becomes smoother and easier to explore via MCMC algorithms. Importantly, getting the unnormalized density for x (²()) (x) can be simply achieved by raising the unnormalized density π (2) to the power |3.
2.3 Feynman-Kac Formula
The Feynman-Kaa Partial Differential Equation (PDE) is generalization of the Fokker-Planck PDE and is defined as follows
Op)(2) (V. (V.))(())))) + Ge p(() -Bpp ) 2
(6)
Ot
where the first term corresponds to the probability mass transport along the vector field et (œ), the second term corresponds the stochastic moves of samples according the Wiener process W and the last term is responsible for rewcighting the samples according to coordinate dependent weighting function ge (z). For any test-function ((x), the Feynman- formula relates its expected value to the expectation over the SDE trajectories 24, i.e.
E (=)p((()
1 Zr
E elee de gi(x)) ((())) where dr. (x GdW, zo Po(x), (7)
and Zr is a normalization constant independent æ. In practice, the exponential term is computed as "weight" we of the corresponding sample x; and can be integrated in parallel with the simulation,
dr. e4 (x.)) + GdW,, dloog we (xx))tt, initialized as To po()) log up 0. (8)
Finally, one can estimate the normalization constant Zr by considering (zx) 1 in Eq. (7) and get biased but consistent Seff-Normalized Importance Sampling (SNIS) estimator (Liu, 2001), i.e.
1 B eele gx( p(xT))
Ieeo dt g((z))p(aT)
wh
((TT)
(9)
Zr
Eelo dt (1().)
where (x w/) are the samples from the simulation of the SDE in Eq. (8).
3 Progressive Inference-Time Annealing
In this section, we combine diffusion and annealing processes into an efficient learning algorithm for sampling from the target density x (xx). To design this method, we build on the fact that diffusion and annealing are complementary ways to simplify or "smoothen" the target distribution (see Fig. 1).
Namely, for the high-temperature version of the target distribution " (I), assume having samples from (x)) (1:) and learn the density model of the marginals defined by the diffusion process (see Section 3.2). For instance, this can be done by running MCMC chains that face little challenge mixing in high temperatures. For the given density model of the diffusion process, perform annealing of all the marginals and generate samples from the lower temperature target T (x)) Bi41 < 3, (see Section 3.1). We detail every step of our method in the following subsections.
3.1 Inference-Time Annealing
In this section, we discuss the inference-time annealing process, which allows us to modify the trained diffusion model to generate samples from the lower temperature target density. Namely, for a diffusion process with marginals (x)) and the end-point Pt- (x) x (x), we assume having two models:
score model s, (z;00 0) ze V Po(r) and energy based model Uz(2; 1)) se - -log pc(x) const

---

with parameters & and respectively. Given the score and the energy models trained to sample from a higher temperature density n (z), we define a new sequence of marginals that correspond the Boltzmann density of the energy model but with a lower temperature
(x(2:) x 2xp(--A-1//..
(2) x exp( 1--=1111 n)) = (=(=))134411/8 (10)
The following proposition derives the Feynman- Kac PDE that describes the evolutvol of the marginals q, (z) and allows for importance sampling via the Feynman-Kaa formula.
Proposition 1. [Inference-tine Annealing] Annealed density of the energy-based model qe(a)) x exp(-?U((I;n) matches the marginal densities of the following SDE
0 + 27 (S((xz) 7Ex (EVUU(((
)at dt + GeV/Ed, 20 ~ -t=0(()) (11)
dx;- =
2
S 22 (V. si(x)) - (U)(() -00 (( 2
OU1( dt. (12)
d loo Wt =
Öt
where (z) is any vector field, az St. & are analogous to parameters from Eq. (4), and the sample weights wt correspond the SNIS estimator of the Feynman- Kac formula Eq. (9).
See Appendix A.I for the proof. Intuitively, this result defines an importance sampling scheme, where Eq. (11) generates samples from the proposal distribution and Eq. (12) integrates the density ratio between the sampled proposal and the target density 9: (z). Different choices of the vector field st sxxx) and the noise schedule & yield different proposal distributions. Theoretically, one can choose different parameters a; S. as well, but below we argue for setting them according to Eq. (2).
The dynamics in Proposition are not unique, and there exists a continuous family of PDFs that We motivate this dynamics minimizing the variance of the weights for the case when there is no annealing (r = 1). Indeed, if the trained EBMs and score models approximate the diffusion process perfectly, then, for ? 1, the weights become constant, so SNIS equally weights all the samples; thus, eliminating the need to resample at all. We formalize this result in the following proposition (see Appendix A. 1.1 for the proof).
follow the marginals (xx): exp(-?U,(x;
Proposition 2. [Convergence to Diffusion] For 7
1 and perfect models st(x) -U, n) = V log Pr (z), the variance of the weights in Proposition becomes zero.
In the case of unbounded support of the target distribution, c.g. su)p(>(x)) = R4, increasing the temperature might cause numerical instabilities. Indeed, π ( x Uniform(Rd) is not normalizable. To avoid this potential issue, in Appendix A.2, we consider geometric averaging between some simple prior and the target densities, c.g. (10,
Integrating the dynamics from Propositions and 3 we generate set of weighted samples ((of_1> w ) k=1 that converge to the samples from (t=1(x) when K + co. In practice, this density is defined the Boltzmann distribution of the corresponding energy model, i.e. (x)) x exp(-3++1
K
Ut=1 (2; 1))), which approximates (x)+ but does not necessarily match exactly. We discuss several possible ways to bridge this gap between the density model and the target density in Appendix B.
3.2 Training using PITA
The proposed algorithm consists of interleaving the inference-time annealing (described in the previous two subsections) and model training on the newly generated data from the annealed target distribution, which we describe here. Throughout this stage we assume availability of samples from 8 which were previously generated at the sampling stage?.
For the target distribution x0++122) (z), we define the diffusion process with the marginals ((()) obtained convolution of the samples from the target gB+4 (x) with the Gaussian A (a,x,, of1). o?1). To
For the very first iteration of our algorithm, we assume that there exist such & that samples from T can be simply collected by conventional Monte Carlo algorithms.
5

---

Algorithm 1 Training for single temperature 1/3r41
Require:
samples æ from x 9111
for training iterations do sample (n( ~ ([Pmaan: sld)
add noise Iy ~ (2 on -1 0f -11
Denoising Score Matching(O) = Vo E EE x()) De(x,; 0)1
Target Score Matching(8) Vo E, | V log (2) + x (crr;: I(t 2 tuneest)
EBM Distillation(1) = V, En,,, x(1) of (2 7)) + - O)l
Energy Pinning(n) v, (- (1: (z) - 3/+1 log (x))|²
0 < FirstOrdeerppiimizerr (0, Score Matching ()), Target Score(O))
7 f FirsstOrderOptimmizerrl Energy Matching Energy Pinning(1))
end for
return trained parameters 0",77
learn the score function st(;; 0) s V pz), we follow the standard practice and parameterize the denoising model Dr(x;; 0) = of st (x);; 0) + Ie. which we learn via the denoising score matching (DSM) objective (Ho et al., 2020), i.e.
Denoising Score Matching(O)
a= Dr(s,; 0)1?²
(13)
where the expectation is taken w.r.t samples from the annealed target 8*)50(x), (x), noised samples 2, ~ (xx 01-12, _f))) (). and time parameter sampled with log(1 - t) ~ (N Pata)
largely following Karras et al. (2022).
However, the DSM objective is not sufficient for training a good score model close to the target distribution (7 = 1 / = 0) due to the high variance of the estimator. Indeed, for r = - / = 0, it has no information about the target distribution. Target Score Matching (De Bortoli et al., 2024)
overcomes this issue by explicitly incorporating the score of the target unnormalized density into the objective, which is as follows
Target Score Matching(O) E. |
(1) + x (o)
(x(= I[t 2 fmeesh)
(14)
where the expectation is taken w.r.t. the same random variables as in Eq. (13), but the time variable is restricted to the interval [fuhresh: because the variance of the objective estimator grows with the noise scale (De Bortoli et al., 2024). Notably, for larger noise scales, the Denoising Score Matching objective results in a stable training dynamics; thus, these objectives complement each other and result in stable training dynamics across the entire time interval.
To train the energy model (2;;;) n)) which plays the central role in the inference-time annealing (see Section 3.1), we follow Thornton et al. (2025) and distill the learned score model to the parametric energy model via the following regression loss (w.I.t. 77, i..
(15) where, the expectation is taken w.r.t. the same random variables as in Eq. (13). Note that, in contrast to the denoising score matching loss in Eq. (13), the "target" in Eq. (15) does not depend on x, which means that its variance for the same 2, is zero, stabilizing the training of the energy based model U,(
EBM Distillation(1)) = R,, (0)]oo? (-V Va(e; m)) + x2 D,(a,;
Finally, to use all the supervision signal available in the problem, we use the target unnormalized density x 9 as the regression target for the end-point energy-based model Uz=1 (z;; n)), and introduce the following loss
Energy Pinning(n) = ( (--
Bi+1 log (x))
(16)
Notably, this loss allows for fixing the gauge present due to the shift invariance of the energy-based model (V;((();))) V- (U,((z);) + const)). In practice, we observe that this loss significantly stabilizes the training and improves the final performance. We present the pseudo-code for the training loop in Algorithm train a diffusion model at temperature 1/8 In practice, we find that sequential training of models demonstrates the best performance. Furthermore, we initialize the model for the next temperature 1/3:+1 with the parameters of the trained model for the temperature 1/3.
where we simultaneously optimize all the introduced loss functions to

---

4 Related work
Diffusion- Sampling. A variety of amortized samplers that use properties of diffusion models have recently been proposed in the literature. Simulation-base approaches that also exploit the fast mode-mixing of diffusion models include Berner et al. (2024); Vargas et al. (2023); Zhang and Chen (2022); Richter et al. (2024); Vargas et al. (2024), which exploit diffusion processes for fast mode mixing. Conversely, simulation-free methods like iDEM (Akhound-Sader) Föllmer drift (Huang et al., 2025), and TSM (De Bortoli et al., 2024) offer more scalable approaches to learning the score but suffer from inefficient and high variance score estimates far from the data. Finally, new diffusion bridges have also risen prominence with underdamped dynamics (Blessing et al., 2024), known modes (Noble et al., 2025), and bridges with SMC (Chen et al., 2025).
et al., 2024), SB with
Inference-time Resampling. The inference-tine annealing scheme proposed in Proposition 1 connects several recently proposed methods. Namely, for & = O, it closely matches the importance sampling of the continuous normalizing flows proposed in (Köhler et al., 2020). Indeed, Eq. (11) becomes a probability flow ODE, and Eq. (12) becomes an integration of the log-density the target density can be defined either as a linear interpolation of log-densities point as the target density. Furthermore, Proposition is an application of the Feynman Kac formula to annealing and non non-equilibrium dynamics simultaneously. Indeed, for Y = 1, this proposition becomes the result proposed in Vaikuntanathan and Jarzynski (2008); Albergo and Vanden-Fijnden (2025); whereas, for (x)) VU (1;; n) = V log pt (x)), this proposition becomes the result from (Skreta et al., 2025). In practice, however, these equalities are not satisfied because we use learned models for the vector field se (xx) = (xx;;; 0) and the energy-based model (2; 17). In concurrent work, Rissanen et al. (2025) also propose an annealed sampling scheme using diffusion models. They use a different inference-time annealing procedure which restricts importance re-sampling to the final timestep, whereas in this work, we benefit from annealed importance sampling over diffusion time.
where
only in the final
Boltzmann Generators. Noé et al. (2019) proposed training a probabilistic model and resampling the generated samples according to the target Boltzmann density via importance sampling. Various probabilistic models have been used, e.g. continuous normalizing flows (Chen et al., 2018) and the flow matching objective (Lipman et al., 2023) also allow for efficient training and resampling under the Boltzmann generators framework (Köhler et al., 2020; Klein et al., 2023). Boltzmann Generators can also be combined with Annealed Importance Sampling, which enhances their scalability (Tan al., 2025a). However, as we demonstrate empirically, the straightforward resampling with target density of a different temperature (Schopmans and Friederich, 2025) results in high variance of importance weights. Thus, one has to deviate from the Boltzmann Generators framework to perform the inference-time annealing.
5 Experiments
We evaluate PITA on molecular conformation sampling tasks including a toy Lennard-Jones system of 13 particles (LJ-13) and Alanine peptide systems of varying sizes (Alanine Dipeptide and Tripeptide) in Cartesian coordinate space. Throughout, we assume access to a short MCMC chain run at high temperature. Note that we do not require these chains to be well mixed, but only require them to cover the modes, a much less stringent requirement (See Appendix D). For metrics, we use sample-based metrics such as 2- D-ssssesttein distance on Ramachandran coordinates (T-W2, and energy distribution (E-W1, &-W2), to assess mode coverage and precision, respectively. We also compare the KL divergence between the Ramachandran plots the ground-truth MD samples and the generated samples (RAM-KL)), well as Wasserstein distances on the first two TICA lagged Independent Component Analysis) coordinates of ground- truth and generated samples. Finally, we report the computational expense of all methods using the total number of energy evaluation calls.
Baselines. We compare PITA with three different baselines: Temperature Annealed Boltzmann Generators (TA-BG, Schopmans and Friederich (2025)), normalizing flow model trained on molecular simulation (MD) data collected at the target temperature (MD-NF), diffusion model trained on molecular simulation (MD) data collected at the target temperature (MD-Diff) and importance sampling using continuous normalizing flows (Köhler et al., 2020). For the LJ-13 dataset, we additionally compare the performance of the model with two oate-offthe-art diffusion- based sampling algorithms: namely, iDEM Aakhond-Sadeeg et al., 2024) and adjoint sampling (Havens et al., 2025); however, none of the current diffusion- based approaches are able to achieve competitive
7

---

KD Samples
Semples
PITA
1.4
TA-06 (TarFlce)
TA-80 (TarFlow)
Norinalized Density
8.12
Normallzer 1zed Density
1.2
8.10
1.0
6.95
8.8
8.06
0.6
0.64
0.4
0.02
-60
-40
-20
20
49
68
U/11(x)
Interatomic Distance
Figure 2: LJ-13 sampling task. We compare distribution of the interatomic distances and energy of the particles in the MCMC dataset (ground-truth), samples generated using a PITA model, and TA-BG progressively trained from high temperature to sample from the target distribution.
performance on the small protein tasks, we only compare to the TA-BG baseline for those datasets. Additionally, for the ALDP experiment, we compare against two other annealing strategies for the diffusion-only approach: FKC and Score Scaling. FKC refers the annealing scheme in Skreta et al. (2025), and score scaling is essentially scaling the score with the corresponding temperature factor Y. Further experimental details of the baselines, as well as additional baselines, are provided in the Appendices E and I.
Architecture.
training PITA, we use EGNN (Satorras et al., 2021) as the model backbone for LJ-13, and DiT (Peebles and Xie, 2023) for Alanine Dipeptide and Alanine Tripeptide. Our training follows sequential temperature schedule, proceeding from high to low temperatures. After training at given temperature for fixed number of epochs, we generate samples at the next lower temperature and continue training at that temperature. For LJ-13, we train a single diffusion model conditioning it on B and using the data across all previously seen temperatures. For molecular conformation sampling tasks, we adopt a fine-tuning approach, where at each temperature step, the model is trained only on the newly generated samples corresponding to the current temperature without revisiting earlier ones. For the TA-BG baseline, we train TarFlow (Zhai et al., 2025) with adaptations suited to molecular data for all three systems, and for MD-Diff, we use the same DiT architecture we used in PITA. We parameterize the energy network using the parameterization in Thornton et al. (2025) and use the preconditioning (c,, Cout, Can. cς) of Karras et al. (2022) for both energy and score networks. Further training details and hyperparameters are provided in Appendix I.
Hyperparameters. Proposition 1 allows for many choices of the vector field (z In practice, we set proportional to the score model st x sz 6) and try several scaling coefficients (see Appendix G). Finally, one can easily add the time-dependent schedule Y, by adding the extra term to the weights, i.e. dl log w; = Eq. (12) - (x 1))0n// ( /idl which we study in Appendix G. We use the noise schedule from Karras et al. (2022). where, for all experiments, we set dmax = 80 and 7. For LJ-13 and molecular conformation sampling tasks, we use mnin = 0.05, 0.01, respectively.
5.1 Main results
LJ-13. We first consider a Lennard- Jones (LJ) system of 13 particles to demonstrate the effectiveness of train- ing a sampler at a high temperature (Tz- = 4), followed by annealing to a lower temperature (I's = 1). As shown in Table we compare the per- formance of PITA with existing base- lines. A visual comparison to TA-BG is provided in Fig. We evaluate each method using the 2- Wassersstni distance computed over interatomic distance distributions, energy distributions, and 3D geometric coordinates. We omit the Distance-W, metric for Adjoint Sampling, as its results could not be reproduced and no code is available at this time; the reported Energy- and KeeometrccWW, values are taken directly from the original paper (Havens al., 2025). To ensure consistency, we exclude samples with energy above 1000 across all methods; this notably impacts TA-BG, removing
Table I: IJ-13 sampling task. The starting temperature is Tr 4, annealed to Tg = 1.
Algorithm
Distance-Ws W/
Energy-Wa
Geometric-W9 J
iDEM
0.127
30.78 24.46
1.61 0.01
Adjoint Sampling
2.40 # 1.25
1.67 # 0.01
TA-BG (TarFlow)
1.21 = 0.02
61.47 = 0.12
4.16 = 0.01
PITA (Ours)
0.04 = 0.00
2.26 + 0.21
1.65 # 0.00
8

---

ND
0.86
FITA ND-D1ff
8.67
(w/o relaxation)
No Normalized Density
Norma Normalized corzalized Density
9.05
MD-D1ff
0.06
TA-BG (w/o relaxation)
TA-BG
8.05
0.84
0.83
8.03
0.82
8.82
0.81
8.01
8. 0.00
68
-40
-28
20
48
68
88
108
0.00
-150
-180
-50
8
58
188
Ut=1(x)
Ux=1(x)
(a) Alamine Dipeptide
(b) Alamine Tripeptide
Figure 3: Molecular conformation sampling tasks. We compare the energy distribution of the ground-truth MD dataset and the samples generated using different models at 300K. We use 30k samples for the plots.
Table 2:
Performance of methods for the ALDP sampling task. The starting temperature is Tz. 1200 K.
annealed to target 7's- 300 K. Metries are calculated over 10k samples and standard deviations over seeds.
Rama-KL
Tica-W
Tea-W;
Baergy-W 4
EnergyW/2
T-W;
#Eaergy Bvals
PITA
4.773 0 0.460
0.112 + 0.006
0.379 = 0.028
1,530 + 0.068
1.615 = 0.053
0.270 + 0.023
5 10"
MD-DiI
0.072
0.113
0.004
3.627 0.023
3.704 0.026
0.001
11'
MD.NF
13,533 1 0.024
0,138 1 0.003
0.586 1 0.003
0,551
1 0.062
1.198 1 0.069
0.403 1 0.045
6 10"
TA-BG
14.993
0.002
0.219 0.013
83.457
0.070
86.176 0.104
0.012
5x 10"
FKC
14.392 + 0,909
0,217 + 0.000
0.649 0.001
11.281
+ 0.025
11,466 = 0.027
2.120 0.024
6 10
Score Scaling
4.588 1 0.467
0.183 + 0.002
0.608 0.008
10.282 + 0.020
10:460 = 0.019
0.550 1: 0.036
5 107
Table 3: Performance of methods for the AL3 sampling task. The starting temperature is Tr 1200 K. annealed to target Ts= = 300K. Metries are calculated over 10k samples and standard deviations over 3 sceds.
Rama- KL.
Twa-W:;
Tic->11/
Energy Wi
Energy-93
T-);; #Eosrgy Evals
PITA
1,209- = 0.144
0.272 : 0.017
0.52
0055
2.567
0.108
2,892 :: 0.107
0.521 0.006
10"
PITA (wlo relaxation)
8,535 = 0.251
0.405 1: 0.014
0,999
1 0.013
86.270 1: 0.294
87.695 2: 0.294
0.651 1 0.013
5 10"
MD-Diff
9.662 0.085
0.059 0,006
0.426 0.010
7.13600.130
7,599 + 0.137
0.424 + 0.011
10"
TA-BG
2.078
2,088
D182 + 0.001
0.454
0.00
4.782=
4,863 +: 0.082
0.347
0.014
10"
TA-BG (wlo relaxation)
14,98- = 0.009
0,321 1 0.001
0.618 0.000
173.042 0.771
178.558
1: 0.732
1,300 1 0.001
5 107
approximately 60% of its samples. Even under this filtering, PITA consistently outperforms TA-BG and other baselines trained directly at the target temperature.
Alanine Dipeptide. We apply PITA the task of sampling Alanine Dipeptide a target temperature of Ts = 300 K, given initial samples at a higher temperature of T,. 1200 K. We use annealing steps of 1200 K, 755.95 K. 555.52 K, 300 K. These temperatures correspond to a subset of the temperatures from Schopmans and Friederich (2025), as PITA does not require as many annealing steps to achieve competitive performance. We also analyze the performance of the model, taking larger annealing steps in Appendix G. As shown in Table 2, PITA consistently outperforms both the diffusion- based baseline and TA-BG across all evaluation metrics, achieving a partieularly large margin in melaty- metrics. We further present TICA plots of the generated samples at the target temperature in Fig. 4. PITA successfully recovers the essential slow collective dynamical modes of the system, which baseline methods fail to capture. Additionally, we find that while TA-BG performs reasonably well at earlier stages of training at higher temperatures, its performance deteriorates as temperature decreases. Such decline is likely due to the increasing difficulty in generating high-quality proposals as the temperature decreases, which is crucial in the importance sampling used for subsequent training stages, as well as the difficulty in importance sampling between large temperature gaps (Appendix D.2). Additional details on training dynamics across all temperatures for PITA and TA-BG can be found in Appendix H.
Alanine Tripeptide. We further evaluate the performance of PITA on a larger molecular system, Alanine Tripeptide (AL3). We employ temperature annealing schedule with intermediate steps at 1200 K, 755.95 K, 555.52 K, 408.2 24 K, 300 K. As shown in Fig. 5, PITA again successfully recovers the essential dynamical modes of the system, indicating its capability of generating samples that align with the dominant kinetic features of the underlying dynamics. practice, we also observe that performing a short additional MD refinement at the target temperature (300 K) after generation

---

MD (ground-truth)
PITA
ND-Diff
TA-BG
Figure 4: TICA plots for Alanine Dipeptide (ALDP) at 300K obtained from different methods using 30k samples. Each panel shows the free energy landscape along the top two TICA components which capture the dominant slow transitions the system.
(ground-turth)
(wto relavation)
ND-Difr
relasation)
Figure 5: TICA plots for Alanine Tripeptide (AL3) 300K obtained from different methods using 30k samples.
further improves the physical plausibility and smoothness the trajectories. This leads to more accurate estimates of the free energy landscape, for both PITA and TA-. In Table 3, provide quantitative analysis of the performance of the models. Notably, despite resulting a better mode coverage, PITA performs worse than baselines according Tica- Wi and fica-Wa, which suggest that it does not fully recover the correct relative weights the modes.
Conclusion
In this paper, we propose PITA a new framework to train diffusion-basse samplers by introducing two mechanisms of interpolating target Boltzmann density by changing the temperature and defining a diffusion noising process. We demonstrated that PITA allows the progressive training of sequence of diffusion models that go from high temperature. where ground truth data is simple to collect, the lower temperature target temperature. Using PITA we demonstrated equilibrium sampling of N-body particle systems, and, for the first time, equilibrium sampling of alanine dipeptide and tripeptide in Cartesian coordinates. Importantly, we demonstrate PITA requires drastically fewer energy evaluations than existing diffusion samplers.
We believe that PITA represents step forward in the scalability of diffusion-based samplers and opens up ripe avenues for future work including improved training strategies and regimes for energy- based models that are in service of the PITA framework. Natural directions for future work include automatically determining the optimal temperature jump when instantiating our Feynman- Kac PDE generate asymptotically unbiased samples at lower temperatures or transferable sampling (Klein and Noé, 2024; Tan et al., 2025b).
importance sampling scheme, has define density model the generated samples. For this, PITA relies on training an additional energy-based model, which notoriously challenging task. Furthermore, simultaneous training and inference of both the score model and the energy-based model introduces additional computational and memory requirements.
Linitations. To obtain consistent estimator
10

---

Acknowledgments
This project was partially sponsored by Google through the Google Mila projects pro- gram. The authors acknowledge funding from UNIQUE, CIFAR, NSERC, Intel, and Sam- sung. The research was enabled in part by computational resources provided by the Digi- tal Research Alliance of Canada / ht//aallaancccnnnco ca), Mila fattps://mila / and NVIDIA. KN was supported by IVADO and Institut Courtois. AJB is partially sup- ported by an NSERC Post-doc fellowship. This research is partially supported by the EPSRC Turing AI World- World-Leading Research Fellowship No. EP/X040062/1 and EPSRC Al Hub No. EP/Y028872/1.
References
Akhound-Sadeee T., Rector-Brooks, J., Bose, A. J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y., Malkin, N., and Tong, A. (2024). Iterated denoising energy matching for sampling from Boltzmann densities. In International Conference on Machine Learning. 2,7
Albergo, M. S. and Vanden-Bijnden, E. (2025). Nets: A non-equilibrium transport sampler. In International Conference on Machine Learning.
Berner, J., Richter, L., and Ullrich, K. (2024). An optimal control perspective on diffusion-based generative modeling. Transactions Machine Learning Research. 7
Blessing, D., Jia, X., Esslinger, J., Vargas, F, and Neumann, G. (2024). Beyond ELBOs: A large-scale evaluation of variational methods for sampling. In International Conference Machine Learning. 2,7
Chen, J., Richter, L., Berner, J., Blessing, D., Neumann, G., and Anandkumar, A. (2025). Sequential controlled Langevin diffusions. International Conference on Learning Representations. 7
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. (2018). Neural ordinary differential equations. In Neural Information Processing Systems. 7
Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., and Zhang, A. R. (2023). Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference on Learning Representations.
De Bortoli, V., Hutchinson, M., Wirnsberger, P., and Doucet, A. (2024). Target score matching. arXiv preprint F-XXi:2402.00667 6,
Del Moral, P., Doucet, A., and Jasra, A. (2006). Sequential Monte Carlo samplers. Journal of the Royal Statistical Society: Series B, 68(3):411-436. 2, 33
Du, Y., Durkan, C., Strudel, R., Tenenbaum, J. B., Dieleman, S., Fergus, R., Sohl-Dickstein, J., Doucet, A., and Grathwohl, W. S. (2023). Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. In International Conference on Machine Learning.
Frenkel, D. and Smit, B. (2023). Understanding Molecular Simulation: from Algorithmms to Applica- tions. Elsevier. 1
Havens, A., Miller, B. K., Yan, B., Domingo-Enrich, C., Sriram, A., Wood, B., Levine, D., Hu, B., Amos, B., Karrer, B., Fu, X., Liu, G.-H., and Chen, R. (2025). Adjoint sampling: Highly scalable diffusion samplers via adjoint matching. In International Conference on Machine Learning. 2, 7,
He, J., Du Y., Vargas, F., Zhang, D.. Padhy, S., Yu R., Gomes, C., and Hernández-Lobasd M. (2025). No trick, no treat: Pursuits and challenges towards simulation-free training of neural samplers. arXiv preprint 2xiv:2550200688.
Hénin, J. Lelièvre, T., Shirts, M. R., Valsson, O., and Delemotte, L. (2022). Enhanced sampling methods for molecular dynamics simulations. Living Journal of Computational Molecular Science, 4(1). 1
11

---

Ho. J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. In Neural Information Processing Systems.
Hoffman, M. D. and Gelman, A. (2014). The No-U-Tum sampler: Adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research. 32
Huang, J., Jiao, Y., Kang. L., Liao, X., Liu, J., and Liu, Y. (2025). Hharödinger-Föllmee sampler: sampling without ergodicity. IEEE Transactions on Information Theory. 7
Jarzynski, C. (1997). Nonequilibrium equality for free energy differences. Physical Review Letters, 78(14):2690-299) 296-236666696... 2, 33
Karras, T., Aittala, M., Aila, T., and Laine, S. (2022). Elucidating the design space of diffusion-based generative models. In Neural Information Processing Systems. 6, 8, 24, 30
Klein, L., Krämer, A., and Noé, F. (2023). Equivariant flow matching. In Neural Information Processing Systems.
Klein, L. and Noé, F. (2024). Transferable Boltzmann generators. In Neural Information Processing Systems. 10
Köhler, J., Klein, L., and Noć, F. (2020), Equivariant flows: exact likelihood generative learning for symmetric densities. International Conference on Machine Learning. 7.
Leimkuhler, B. and Matthews, C. (2015). Molecular Dynamics: With Deterministic and Stochastic Numerical Methods. Springer.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. (2023). Flow matching for generative modeling. In International Conference on Learning Representations.
Liu, J. S. (2001). Monte Carlo Strategies in Scientific Computing. Springer. 1,4
Neal, R. M. (2001). Annealed importance sampling. Statistics and Computing, 11(2):125-139. 2,
Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. (2023). Action matching: Learning stochastic dynamics from samples. In International Conference on Machine Learning. 24
Noble, M., Grenioux, L., Gabrié, M., and Durmus, A. O. (2025). Learned reference-based diffusion sampling for multi-modal distributions. In International Conference on Learning Representations.
Noé, F., Olsson, S., Köhler, J., and Wu, H. (2019). Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 1,7
Ohno, K., Esfarjani, K., and Kawazoe, Y. (2018). Computational Materials Science: From Ab Initio to Monte Carlo Methods. Springer. 1
Peebles, W. and Xie, S. (2023). Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference Computer Vision. 8, 24, 30
Richter, L., Berner, J., and Liu, G.-H. (2024). Improved sampling via learned diffusions. International Conference Learning Representations. 7
Rissanen, S., Yu R., He, J., Chen, W., Heinonen, M., Solin, A., and Hernández-Lobato, J. M. (2025). Progressive tempering sampler with diffusion. In International Conference Machine Learning. 7, 25
Satorras, V. G., Hoogeboom, E., and Welling, M. (2021). (n) equivariant graph neural networks. In International Conference on Machine Learning. 8, 30
Schopmans, and Friederich, (2025). Temperature-annelled Boltzmann generators. In Interna- tional Conference on Machine Learning. 7. 25, 32
12

---

Skreta, M., Schound-Sadege T., Ohanesian, V., Bondesan, R., Aspuru-Guzioe A., Doucet, A.. Brekelmans, R., Tong, A., and Neklyudov. K. (2025). Feynman Kac correctors in diffusion: Annealing, guidance, and product of experts. In International Conference on Machine Learning. 7,
-ac
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations.
Stoltz, G., Rousset, M., and Lelièvre, T. (2010). Free Energy Computations: A Mathematical Perspective. World Scientific. I
Swendsen, H. and Wang, J.-S. (1986). Replica Monte Carlo simulation of spin-glasses. Physical Review Letters, 57(21):2607. 2, 33
Tan, C. B., Bose, A. J., Lin, C., Klein, L., Bronstein, M. M., and Tong, A. (2025a). Scalable equilibrium sampling with sequential Boltzmann generators. In International Conference on Machine Learning. 7
Tan, C. B., Hassan, M., Klein, L., Syed, S., Beaini, D., Bronstein, M. M., Tong, A., and Neklyudov, K. (2025b). Amortized sampling with transferable normalizing flows. In Neural Information Processing Systems. 10
Thornton, J., Bethune, L.., Zhang, R., Bradley, A., Nakkiran, P., and Zhai, S. (2025). Composition and control with distilled energy diffusion models and sequential Monte Carlo. In International Conference on Artificial Intelligence and Statistics. 8, 24
Vaikuntanathan, S. and Jarzynski. C. (2008). Escorted free energy simulations: Improving conver- gence by reducing dissipation. Physical Review Letters, 100(19):190001.. 7
Vargas, F., Grathwohl, W. S., and Doucet, A. (2023). Denoising diffusion samplers. In International Conference on Learning Representations
Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. (2024). Transport meets variational inference: Controlled Monte Carlo diffusions. In International Conference Learning Representations. 7
Woodard, D., Schmidler, S., and Huber, M. (2009). Sufficient conditions for torpid mixing of parallel and simulated tempering. Electronic Journal of Probability. 2
Zhai, S., Zhang, R., Nakkiran, P., Berthelot, D., Gu, J., Zheng. H., Chen, T., Bautista, M. A., Jaitly, N., and Susskind, (2025). Normalizing flows are capable generative models. In International Conference on Machine Learning. 8,32
Zhang, Q. and Chen, Y. (2022). Path integral sampler: stochastic control approach for sampling. In International Conference on Learning Representations. 2,
13

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA]-
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1 (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd a proper justification is given (e.g., "error bars expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
the checklist one of the factors in their evaluation.
reported because it would be too computationally
"[No] " or "[NA] is not grounds for rejection. While the questions are phrased eeknowledge that the true answer is often more nuanced, so please just use your best, judgment and write justification elaborate, All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can found.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
We provide a new theory for training diffusion-based samplers demonstrated in our propositions. The experiments are done on N-body particle systems and short peptides in Cartesian coordinates claimed. We also provide the number of energy evaluations highlighting the efficiency of PITA.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: The conclusion includes dedicated paragraph limitations. These include the challenges of training an Energy based model and determining the temperature jumps during inference annealing, among other things.
Guidelines:
14

---

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might 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 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: [Yes]
Justification: The paper includes full set of assumptions and completed proofs for each theoretical result the appendix.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper includes experiments on public datasets. Moreover, we include anonymized code with runnable commands in the supplementary to aid in reproducing our results. We also outline exact experimental setup and hyperparameters in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
15

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way 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 a large language model), releasing of a model checkpoint, or 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) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All our experiments are on public datasets and energies. We provide scripts to reproduce our experimental findings.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a 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: http://ccppllcgguuicieeee
/CodeSubmis ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
16

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting//etails
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: We include precise experimental setup for all settings in our appendix. We also provide hyperparameters and how they were chosen in these details along with scripts to the code. which is provided in the supplement.
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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We include error quantification in all of our reported metrics. More precisely, we report error bars within one standard deviation of the mean.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, 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).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-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 enrors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We run our experiments on H100 GPUs and the exact run time, energy function query cost, and overall run time outlined in the experimental setup.
Guidelines:
17

---

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 the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: h
Answer: [Yes]
Justification: All authors aware of the code of ethics and abide by it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: The paper makes a fundamental advance in sampling from a methodological perspective. While the future potential for accelerating equilibrium sampling poses great excitement, the current scale of experiments are too far removed to have any significant real-world impact on application domains like drug discovery.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate 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 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
18

---

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: [No]
This paper introduces a new diffusion-basse sampler for which we believe there is no current risk of being misused and result safe guards are needed.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but 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: [NA]
Justification: We only use public assets and only rely on self created code.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, phcorswithcode. thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper only uses public assets which freely available.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl 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 is used.
19

---

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not conduct crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: This paper does not involve any crowdsourcing research or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: This paper does involve LLMs as any important component, original, or non-standard component for this project.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non standtadd components.
Please refer to LLM policy (https: / //wnnnurip..../ cc/Confeerncee///0000 for what should or should not be described.
20

---

A Proofs
A.1 Inference-Tinee Annealing
Proposition [Inference-time Annealing] Annealed density of the energy-based model (e(x)) & exp(->Ur(I;)) matches the marginal densities of the following SDE
2 (s((2)) ( (El (E.VU.((E))) )dt &.VV,,W,, 20 ()=o(xx) (11) dæ,- -
S2 St (vi)) - 7 VUl(a a,2 s0i)) OU, (U,(Z,i1)
d log w
dt (12)
Öt
where (z) is any vector field, at St. & are analogous to parameters from Eq. (4), and the sample weights correspond to the SNIS estimator of the Teynman-Kac formula in Eg. (9).
Proof. For the Energy- -Based Model U2(2; m)) we denote the corresponding Boltzmann density as q: (z) x exp(->U,,2e n)), where y is the target inverse temperature. Taking the time-derivative of ge (x(x) we get the following equation
(17)
(x)
(t) (-7
(-1 aU1(z;n)
aU.(2; m)
(((())
(18)
at
0t
which can be simulated by reweighting the samples from (==0(x) = (0, according to the following weights
(x1)) = exp(-y.J dt aU.(I;; n)/0t)
(19)
(20) exp(-> for dt (U.(m:))/00
Although this reasoning is theoretically justified, in practice, the variance of this importance-weighted estimate (or resampled distribution) is prohibitively large. That is why one has to introduce additional terms that move samples around tracing the original diffusion process. Namely, we consider the following PDE
(g)(z) Øt
s((x)) ) + +1((x) (-> OU.(z:1) + + (=((()(--axx+
OU,(;;
Ōt
Øt
$2
= V. (-xx + s (x))
- (x))) + (x)
(20)
2
(x) = 7 (VU,(z; 7), ax2) - (U,(:;)) (x)) - (V,
2
(21)
2 (z))
(U;n);))
(22)
2
01
where the term (V, arr)) does not depend on x and cancels out when in the reweighting term. Fur- thermore, can introduce the noise term by adding and subtracting the score ( log (x). i.c.
0q((x)
v. V. -ax + G s. 2
- G
(vllae))) + & &
7E,
Aq((2)+
(23)
Ø1.
+ ((((() (x))
(24)
(x(x)= > (ui((:))) -a + b s(()))
22 (V, -
aU,(x:n)
(25)
Ot
which can be simulated as
57
dx. -, +
dt CiV EidWe,
(26)
2
d d W. ( VU((Zri)) -a + 52
(x()
)
+ g 2
s( - 00, (Ue(1 m)
dt.
(27)
2
Ot
21

---

Proposition 2. [Convergence Diffusion] For ? 1 and perfect models se(x) -VUi(z;)) = V log (z)), the variance of the weights in Proposition becomes zero.
Proof.
Indeed, for 7 = 1, Eq. (25) becomes
S (x(x)) = (U;(;; - ax ((((a)
St (v,)(()))
(U((2;n)
(28)
Øt
VUi( au - ?
G (((())
aU.(x;
(((2))
ay V. V
+ +da
2
2
ōt
where d the dimensionality of the state-space. For sx(2)
vU,(2;
(or), have
& log (x
i log Ze
(x(x)) (c(x)
V. PP(E))
v log
da
(29)
2
Ot
Öt
2 v logp(())
(1) Ø1
Ölog 74
(x) ax
+ da
(30)
(r(
=0 duc to Eq. (3)
where Ze fd (xr). The term in the brackets equals zero due Eq. (3) since the ground true marginals (2) are defined as the marginals of the diffusion process. Hence, we have
i log Zr (x(x)
da, constant of x,
(31)
Ōt
which becomes zero after the normalization g=(x) (x)) which concludes the proof.
A.2 Inference-Time Geometric Averaging
For the diffusion process with marginals pr(x)) and the target distribution Pt= (x) x
N(x 00 ())=8))(=)8) we assume having the energy model U,(2; m) and the score model sx; 0).
Then, for the following density
(x) x exp (ssp-1
(-Bun BB++ U+++U(((;)))) B4+1- B Bi
log ((x10, + di-)))
(32)
B,
we have (x= N(x 0, 11)1--11110) ())
To sample from this density, we derive another SDE that performs inference-time geometric averaging. Analogously to Proposition for y 1 and perfectly trained models, the weights become constant, and this SDE yields the reverse-timg diffusion SDE.
Proposition 3. [Inference-time Geometric Averaging] For the geometric averaging of the energy-based model qe (e(x) & exp((1 (),-U 1)) + + (xgx((()) 0, a))) the weighted samples from (x) can be collected by ranning the following SDE
dx. acti + (1 (1-))))1 Gt 2
G 2
(s((x)) (.V(((e))))) - - 10
1 1+
It + GiVE,,
d log W (-(1-7)VU,(0;; -1 - D a, + (1 (1-7)
1
2
8 (xx)
T )
(33)
2
& (V, (x))) + (1-7) aU,( + +11
,o, 0t
7 o
|
1l
(34)
0t
where (x)) any vector field. Finally, unweighted samples from (2) can be sampled using SNIS from Eq. (9).
Proof. For the Energy-Bass Based Model (x; we denote the corresponding geometric averaged density as
(x) x exp((1 - -- log
(35)
22

---

where 7 is the target inverse temperature. Taking the Lime-dervative of (z( we get the following equation
@q((x) 0t = (1(2)()((x)) - (x)
(36)
(1-7))0000)) aU,(,;n) 1
0o, ||=]||
(x = (1)
1 of
(37)
Ot
Ot
Assuming that the change of the density is close the trained diffusion process, introduce the drift-term corresponding to the score of the marginals
@q((z) # V. (7.)(()((aa ax (1 7) S sc(2)
1 -
(1((1)(((x)) - (x))
fictitions term
(e(x) =1 (1 7)) 0U (U((;;; 1))
00: 0t
+
||||||
(38)
Ot
Moving the positive term to the weights and interpreting the negative term as the continuity equation, we get
@q((1)
1
(1 st -7
(g((() - (x)))
01
ge(x)
( (-11- (1) -7 0, - (1 7) 52 (c(1)
1
- )
7
1
(39)
+ +11 & (V, (x)) + (1 - 7) OU.(1; n)
1
dox
+ 7
||
(40)
06
0t
Finally, we introduce the noise term by adding the drift
&1 2 voo 2
= S (- ((1-))VU((E;))) - -
(41)
Thus, we get (1)(
-- -
(V.u(0)) -1At
(cal0))aatt0( (1 (1- 22 (s((x) - (VU,((55)) - 1
))
o
1 fe
x
Ön
Ee St Aq((z) + (z(( E (x)))
(42)
2
(x( = (-11 - VU)(())) - 30² -x + (1 - 7) 2?
(z())
-
x +
(43)
2
+ +11 (1-7) ? (V, s(()) + (1 7) (E:(::))
do
+
||0|||
(44)
2
0t
o
0t
The corresponding SDE is
dx. = CxT1 + (1 -
4 (x(()))
- &. -
(U((s)))) - 7
1 +
E1
C
2e + &.VV E,d.,
2
dlog ((11 ((ll-7)V46 (1(1->))U, VU
1
- -, - ap2t + (1 7)
(Ire
xe )
(45)
2
o?
+ (1 - -22
22
V s. + (1 - 7) OU
1
do:
||
(46)
2
Øt.
0
01.
B Bridging the Gap at the End-Point
Integrating the dynamics from Propositions and we generate a set of weighted samples {(ff_ W k that converge to the samples from (z) when K + cc. In Section 3.1
23

---

we assume that this density is defined as the Boltzmann distribution of the corresponding energy model, i.. qt-1 (x) & exp(-3;+1 (z=1(z: n)), which approximates T but does not neces- sarily match it exactly. Here we describe two possible ways to bridge the gap between the density model and the target density.
The first way to sample from 8(1) is via Self-Normalized Importance Sampling (SNIS). The integrated weights = edo de (m)) represent the density ratio between the distribution of the samples =²=1 and the density of the integrated PDE (see discussion in Section 2.3). Correspondingly; to sample from (2)) we have to take into account the density ratio π (2))++) a new estimator, i.e.
//==1(z)) and obtain
(()) =()))
x E edo de ge(ss) (T)) (p(x))
E (x x (x)
(47)
(1(x)
21 (1(21)
B (()))
Σ
कt = el ( / (x p(1)), k= w1
(48)
where the new weights w are obtained from the old ones w by multiplication with the corresponding density ratio. Note, that the following empirical distribution approximates the target density π )
8(x W
x1))
(49)
The alternative to importance sampling with the density model proposal is the gradual interpolation between the density model and the target during the integration. In particular, one can satisfy the boundary conditions by defining smooth interpolant between the boundary densities pr-o (0, I), Pt=1 B and the annealed density model follows
() t1
DogN(0,1) 1) (1-1-1)
(x))) exp
Bi+1 log 1og)(x)-
(50)
4) 1 -1-11111
1-1)
Be+ Uc(a:n)
(51)
Bi
where (x)+ = max 22, 0 < t1 < t2 < are the hyperparameters that define switch times between models, and i the smoothness parameter. Thus, we guarantee that (x=) x π (1-) However, in practice, we found that this interpolation technique results in high variance of importance weights.
C Network Parameterization and Preconditioning
We condition our score network s and our energy network Ug based on findings in EDM (Karras et al., 2022), use energy parameterization based on Neklyudov et al. (2023) and Thornton al. (2025), and include a new pre-conditioning on 3. All of our networks are based on backbone (o f, 1) (IRd x [0, oo) x oo)) + Id a flexible network architecture based on a diffusion transformer (DiT) backbone (Peebles and Xie, 2023). Specifically, parameterize denoiser network De as:
D0 t, f) = (1 ((C (t) 1)x, + BCout (C (t)x2, Cnoise (t))
(52)
which allows us to define our score network 89 as
so so( = Do t. 33) xt
(53)
We pre-condition the energy as
CcCokip
EeCoue(t
t,8) = 8
| 2
Cin (t)o? (x; (, Cnosse (t)))
(54)
207
24

---

D Molecular Dynamics Analysis
D.1 MD mode mixing across temperatures
In Fig. 6 and Fig. 7, we analyze the mixing behaviour of MD simulations for ALDP and AL3 across various annealing temperatures. Specifically, we examine simulations consisting of 50 million sleps-matching the quantity of MD data used for training PITA at 1200K. As the temperature decreases, the sampling quality deteriorates: the chains exhibit poorer mixing and fail to explore significant regions of the configuration space, missing major modes of the distribution. This is shown both in Ramachandran and TICA plots, as well as the trace plots of the internal angle ó and the second TICA axis. More specifically, for ALDP, mode 5.8%, 3.0%, 1.0% and 0% of the time at temperatures 1200K, 755,95K, 555.52K and 300K, respectively. For AI.3, this happens
that the chain switches out of the main
rate of 12.7%, 9.2%, 5.8% and 0%.
This motivates training at a higher temperature then annealing to a lower temperature as is done in PITA. As we are able to take advantage of relatively quick mode mixing at higher temperatures and the ability of inference time annealing to recover samples from a lower temperature.
1200K
755.95K
555. 52K
300K
n/2
-02
m/2
-02
102
-1/
n/2
(a) Visualization of the Ramachandran plots of the MD chain time, the lines are colored from purple yellow over million MD steps. The gray plot shows the Ramachandran plot of chain with billion MD steps.
1200K
755 05E
555 52K
SANK
11 5 11 12
17
(05)
(n6]
En6
(b) Trace plot for the internal angle Φb. The blue lines show the value of á across the 50-million-sepp MD chain. The orange line indicates when the chain switches out of the main modes.
Figure 6: Analysis of the mixing of the MD chains for ALDP, for 50 million-step MD simulation, across the annealing temperatures.
D.2 Temperature Annealing ESS
To show the difficulty of directly performing importance sampling directly across large temperature jumps in coordinate space using methods similar to Schopmans and Friederich (2025); Rissanen et al. (2025) we calculate the ESS assuming these models could learn the source temperature perfectly. To perform this experiment we take 33k samples from the test chain and calculate the (normalized) effective sample size of these points when importance sampled of this experiment are presented in Fig. 8 for ALDP and Fig. 9 for AL3. Here we report the Log normalized effective sample size for each temperature jump considered in this paper. We compute the log normalized ESS using Kish's 's formula normalized by the number of samples as:
lower target energy. The results
EN
2
log ESS (SS(((;))))))) ( log
(55)
N ii W
Where w; is the importance weight of the ith sample. We note that log ESS of -5 would mean that direct importance sampling would require 000 times many points at the higher temperature as
25

---

1280K
755.95K
S2K
380K
(a) AL3 TICA plots of ground-truth MD samples across different temperatures, using MD chains of 50 million steps. For all temperatures, the TICA axes are matched to those of 300K.
TATE
1200 (02]
T2m2
(19)
Tamc
(b) Trace plot second TICA axis of ground-trutt MD samples. The blue lines show value second TICA axis across the 50-millionstep MD chain. The orange line indicates when the chain switches out of the main mode.
Figure 7: Analysis of the mixing of the MD chains for AL3, for a 50 million-step MD simulation, across the annealing temperatures.
Log ESS
1200
-2
Target Temperatur
755,80
-4
-6
555-55
-3:11
-247
0.20
-6
9.35
-10
1200
735
Source Temperature
Figure 8: Log effective sample size (ESS) values for importance sampling from different source and target lemperatures for ALDP. The ESS values are computed on 33000 samples from the MD chain. The low values indicate that purely relying importance sampling is not sufficient sample from lower temperature targets.
are needed at the lower temperature, This makes importance sampling computationally infeasible for these temperature jumps, even with perfectly learned models.
This motivates the guided approach taken in PITA, where we are able guide the samples preferen- tially towards the target temperature, avoiding the problem of low ESS between these temperature jumps.
E
Additional Baselines
For the ALDP experiment, we compare against two additional baselines of annealed MCMC with sequential Monte Carlo (SMC) and parallel tempering (PT) strategies. For SMC, we take 10 annealing temperatures between 1200 and 300 K using geometric schedule, taking 166 steps per temperature with 30k particles. For PT, we use the same 4 annealing temperatures as PITA, running 25 MD steps between particle exchanges for total of 50k iterations. As can be seen in Table 4, PITA performs better overall in capturing the distribution given the same budget of energy evaluation. In Table 5,
26

---

Log
6:93
Target Temper cilperature
535.55
1:00
2.51
0:33
10.22
18.00
5.13
0.00
1250
755.95
$95.52
408
350
Source
Temperature
Figure 9: Log effective sample size (ESS) values for importance sampling from different source and target temperatures for AL33 The ESS values are computed 33000 samples from the MD chain. The low values indicate that purely relying on importance sampling is not sufficient to sample from lower temperature targets.
Table 4: Performance of the SMC and PT baselines ALDP
Rama-Ki KL.
Tica-Wi >
Tica-W2 L
Energy-W,
Energy-Wa2
T-W2
MEnergy Evals
PITA
4.773 0.460
0.112 + 0.006
0.379 + 0.028
1.530 + 0.068
1.615 + 0.053
0.270 + 0.023
5 x 10F
PT
7,306
1.077
0.625 0.010
0.895 + 0.016
4.632 4: 0.015
4.689 4: 0.014
0.911 + 0.004
5 x 10"
SMC
5.935 0.228
0.372 0.006
0,425 + 0.969
1.002 0.072
0.874 0.016
5 x 10"
also compare the wall-clock time for inference on 30K samples for a trained PITA model as well as classical baselines, SMC and PT. The training time for PTTA ALDP is 564.3 minutes. We note that PITA is roughly an order of magnitude faster than PT and two orders of magnitude faster than SMC after amortization. Therefore, the wall time breakeven point is around 120k samples SMC and 300k samples vs PT for PITA.
Table Inference time for 30k samples ALDP
Time (min)
PITA
4.7
SMC
139.2
PT
59.5
Additional Results
In Table 6, report the Effective Sample Size (ESS) during PITA inference when importance weights are accumulated without SMC-based resampling on the ALDP dataset. The values are computed using 30000 samples for each of the temperature annealing steps. Note that these results correspond to a purely importance sampling-based setting (AIS) and thus differ from the full SMC formulation employed in PITA. The relatively low ESS values with AIS underscore the necessity of resampling for maintaining particle diversity and numerical stability. Additionally, we visualize the ESS as function of the integration time during inference in Fig. 10. The results indicate that most of the ESS variation occurs at the early stages f < 0.5 (i.e., at higher noise levels), after which the ESS stabilizes. These findings suggest possible avenues for future work, such exploring drift modifications aimed at reducing the variance of the importance weights.
G Ablation Studies
To evaluate the impact of our design choices, we perform series of ablation studies examining: (1) the effect of annealing to 300K using different temperature jump sizes, (2) the choice of the Yt schedule, and (3) the role of resampling and different loss components, well as comparing
27

---

Table 6: Log effective sample size (ESS) values without resampling for each temperature annealing step calculated across 30k samples on ALDP
Tz to Ig
ESS (x10
T200K to 755.95K
2.77
755,95 95K 555.52K
8.66
555.52K to 300K
3.79
Log Effective Sample Size (LogESS)
18
18
10
18"
18
0.2
0.3
0.4
8.5
8.6
0.7
9.8
0.9
Integration time
Figure 10: Log effective sample size (ESS) over integration time during inference.
out method with simple classifier-free guidance approach. The temperature jump size ablation is performed for both ALDP and AL3, while the remaining studies are conducted on ALDP.
Temperature Jump Sizes in Annealing. To evaluate the effectiveness of the progressive annealing schedule, we compare the performance models where the system is annealed from different starting temperatures to 300K. For ALDP, skipping intermediate temperatures has the most pronounced impact on energy distribution metries, as shown Table 7a and Fig. 12a. In the case of AL3, we find that sequential training is essential for reliably capturing all modes at the lower temperature, as illustrated Fig. 1 lb. However, energy metrics degrade, likely due to small deviations in the sampled buffers at each annealing step, which accumulate over time. Accurately capturing the energy distribution in the sampled buffers at each intermediate temperature appears to be more difficult for AL3, which may contribute to the observed degradation. Nonetheless, capturing the correct modes remains key priority, as modes lost during training are difficult to recover later. In contrast, mild deviations in the energy distribution can often be corrected through short MD relaxation steps, as demonstrated in Section 5.1.
Tt schedule. We analyze the effects of using different schedules for time-dependent % during inference. In particular, we annneal from 555.52K to 300.0K using constant schedule, a linear schedule which linearly increases from 7 = to = 1.85, and sigmoid schedule again, increasing from ? y 1.85. Table 8 shows that the linear schedule generally performs best across the different evaluation metrics. It achieves the lowest values the TICA metrics, while showing comparable performance to the sigmoid schedule on the energy-based metrics.
Resampling, Energy Pinning Loss, and Classifier Free Guidance. To evaluate the impact of resampling on sample quality, we perform inference from 556K 300K without applying resampling. Across all metrics, we observe that resampling consistently improves performance. To assess the roles of the energy pinning loss and classifier-free guidance, we retrain a model for the 556K to 300K transition with each component removed. Omitting the energy pinning loss results in slight improvement in TICA metrics but leads a noticeable decline in energy metrics, indicating that the loss plays an important role in maintaining accurate energy distributions. Finally, we train diffusion model on the data generated from PITA 555.52K, and anneal to 300K simply by scaling the score by (similar to classifier-free free guidance approaches). This approach shows mixed results, offering no consistent improvement over MDf- (which is directly trained on samples at 300K) and performing below the level of PITA in all metrics.
28

---

Table Effect different starting temperatures on anncaling performance for ALDP and AL3, evaluated at the final temperature of 300K.
(a) ALDP.
Tica-Wi 4
Tica- W2 +
Energy- W1 |
Energy-Wy
T-W2
Tz to Ts
1200K to 300K
0.100
0.004
0.297 + 0.019
6.438
0.024
6.531 + 0.021
0.301 0.023
755.95K to 300K
0.180
0.002
0.611 + 0.003
5.639
0.072
5.683 + 0.070
0.358 0.018
555.52K to 300K
0.121 + 0.004
0.404 + 0.019
1.541 + 0.009
1.619 + 0.010
0.270 + 0.023
(b) AL3.
Tica- Wi +
Tica-Wa TeaaWW/ +
Energy
Energy-Wz
T-W2
Ty to Ts
1200K 300K
0.291 + 0.005
0.558 + 0.003
0.521 + 0.122
0,597 + 0.110
1.351 + 0.014
755.95 555.95K 300K
0.234 0.009
0.663 0.019
17.147 + 0.105
17,429 0.107
0.751 0.006
555.52K 300K
0.158
0.004
0.329 0.025
40.222 0.198
40,978 0.208
0.621 0.038
Samples
1200K
756K
555K to 300K
1008
1028
1238
1208
(a) Alanine Dipeptide
MD Samples
1280K to 300K
756K to 300K
555K to 380K
1039
120
(b) Alanine Tripeptide
Figure 11: TICA plot of ALDP and AL3 samples obtained via annealing from various starting temperatures 300K.
Table 8: y schedule ablation on ALDP.
Tica-W;
Tica- Sea-WW2.1 4
Energy- Wi
Energy-Wa, W2 4
T-W2 4
? Schedule
Constant
0.115 + 0.005
0.389 0.014
1.485 + 0.100
1.584 t 0.095
0.258 + 0.030
Linear
0.095 + 0.009
0.243 + 0.048
1.453 + 0.099
1.555 - 0.099
0.275 + 0.058
Sigmoid
0.113 + 0.009
0.339 0.046
1.443 + 0.087
1.550 = 0.087
0.345 + 0.027
Table 9: Additional Ablation Results on ALDP.
Tica-Wi $
Tica-Wi 1
Energy-Wi
Energy- 4
T-W2 J
PITA
0.121 0.004
0.404 + 0.019
1.541
0.009
1.619 0.010
0.270 0.023
w/o resampling
0.140 1 0.007
0.452 + 0.027
1.606
0.094
1.676 1 0.075
0.363 = 0.023
w/o energy pinning loss
0.098 1: 0.012
0.291 + 0.065
4,709 4: 0.091
4.722 + 0.090
0.219 = 0.021
MD-Diff + CFG
0.137 + 0.007
0.446 + 0.029
8.106 + 0.025
8.190 + 0.026
0.383 = 0.042
29

---

Sanpics 1200K 380K
NO Samples 1200K
No Normallza Lized Density
388K
Normalized Density
8.85
388K
9.86
0.04
8.89
0.07
9.62
0.00
-50
-40
-30
-20
0.00
-100
-150
-140
-120
-109
-00
Ux-1(x)
Uz-1(x)
(a) Alanine Dipeptide
(b) Alanine Tripeptide
Figure 12: Energy distributions of ALDP and AL.3 samples obtained via annealing from various starting temperatures to 300K. The model is trained on ground-truth MD samples at 1200K, and on annealed intermediate samples at 755.95K and 555.52K - the latter being our default training setting.
Table 10:
Metrics across temperatures
Temperature
Model
Tica-Wi
Tica-Wa
Energy-W,
Energy-Wa |
T-W; |
PITA
0.024 + 0.004
0.125
0.021
2.855 + 0.083
2.886
0.079
0.134
0.013
755.95K
TA-BG
0.040 0.002
0.178
0.008
2.065 = 0.044
2.140
0.042
0.355
0.007
PITA
0.141 + 0.001
0.836 + 0.004
1.420 = 0.030
1,430 + 0.027
0.219 + 0.015
555.52K
TA-BG
0.337 + 0.009
0.967 + 0.013
48.486 + 0.042
56.897 + 0.059
1.135 + 0.014
PITA
0.112 + 0.006
0.379 + 0.028
1,530
0.068
1.155
0.053
0.270 + 0.023
300K
TA-BG
0.219 + 0.013
0.685 + 0.034
83,457 0.070
86. + 0.104
0.979 + 0.012
H Training Dynamics Across Temperatures
In this section, we analyze the performance of the models (PITA and TA-BG) across different temperatures during annealing toward the target temperature on ALDP. Table 10 presents quantitative metrics, demonstrating that PITA consistently achieves lower discrepancies across all temperatures. Additionally, Figure 13 shows the Ramachandran plots temperatures, further illustrating the ability of the model to generate physically realistic samples that capture the cunperature-dependdee conformational landscape each step of the annealing process. TA-BG demonstrates reasonable performance at 755.95K when initialized with ground-trutt samples, reflecting its ability to model nigh-teeperattee distributions under ideal conditions. However, its performance deteriorates when transitioning to lower temperatures using recursively generated samples for importance sampling, indicated by the mode collapse in the Ramachandran plots, where the conformational diversity sharply diminishes.
I Additional Experimental Details
L.1 Parameterization Details
PITA. For IJ-13, we use equal loss weights for energy pinning, denoising score matching, and EBM distillation. We use the noise schedule of Karras et al. (2022), with the following parameters: Juin 0.05, maxx = 80 and p = 7. The model uses EGNN (Satorras et al., 2021) with approximately 90k parameters, consisting of three layers and a hidden dimension of 32. For ALDP and AL3, the energy pinning, denoising score matching, and EBM distillation components of the loss are weighted equally at 1.0, with an additional target score matching loss weighted 0.01. We use the same noise schedule as the LJ-13 experiment, using a smaller min of 0.01. We use DiT (Peebles and Xie, 2023) comprising six layers and six attention heads, with a hidden size of 192 and a total of roughly 12 million parameters. All models are trained with a learning rate of 1 x 10-3 without any weight decay. For ALDDP and AL.3, we use Exponential Moving Average (EMA) with a decay rate of 0.999, updating every gradient step.
30

---

755. 95K
555. 52K
300K
MD Samples
-112
9
R/2
11
17
-7/2
n/2
-1
-n/2
0
n/2
PITA Samples
-71/2
1122
-77
m1/
n//
-1
11/2
n/2
TA-BG BG Samples
-712
11/2
17
17
-71/2
0
n/2
Figure 13: Ramachandran plots Alanine Dipeptide (ALDP) obtained from different temperatures using 30k samples, We compare the samples from PITA and TA-BG with the ground truth MD samples.
31

---

MD-Diff. We train the diffusion model on the MD trajectories generated at the target temperature. This serves strong baseline, since we have direct access the ground-truth samples, unlike PITA and TA-BG. To ensure controlled comparison, the length of the MD chain used to train the diffusion model chosen such that the total number of energy evaluations matches the computational budget used to train PITA over all annealing steps. We provide further analysis on the mixing properties of different lengths MD chains at low and high temperatures in Appendix D. We use min value of 0.005, while keeping the rest of the model hyperparameters the same as PITA.
TA-BG. In Schopmans and Friederich (2025). TA-BG trains a normalizing flow by minimizing the reverse Kulbback-Leibler (KL) divergence at high temperature and progressively refining the model via importance sampling as the temperature annealed toward the target distribution. We carefully adapt their training pipeline to ensure consistent and fair comparison. Specifically, we initialize the training with ground-truth MD data rather than learned high-temperature samples, represent molecular configurations in Cartesian rather than internal coordinates, and use the same temperature annealing schedule as PITA. We use TarFlow (Zhai et al., 2025), configured with four meta blocks, each containing four attention layers and hidden size of 256, resulting in approximately 12 million param- eters. We use a learning rate of 1 x 10-4 and employ 60,000 samples at the end of training for each temperature to compute the importance weights used in generating the buffer for the next temperature.
MD-NF. Similarly to MD-Diff, we train the normalizing flow model on the MD trajectories generated at the target temperature. We use TarFlow with the same model hyperparameters that we used for training TA-BG.
Score Scaling. We use a Score Scaling baseline which is a simple (but biased) modification to the score attempt sample from different temperature. Specifically, given score function (xa))) where we would normally sample with the SDE
1 &r dx. (-aret + G
8x sx(e)) di + Si WiV,,,W,, To (t=o(())
(56)
2
we instead sample with
dx: = (-OLDs + 762
1 +
s((0)) ) dt Se G.V,,,W, so ~ (==
(57)
2
1.2 Metries
We evaluate model performance using both sample-based metrics and metrics that assess energy distributions. To compare energy distributions between generated samples and ground-truth molecular dynamics (MD) samples, we compute the ID I-Wasserstein and 2-Wasserstein distances on the energy histograms. For sample-based evaluation, we measure the 2D wrapped 2- 2-Wasserstein distance of the internal dihedral angles, & and & (denoted as T-W,). Additionally, we calculate the 2D 1- Wassersstei and 2- Wasserstein distances between the first two TICA axes of the ground-truth and generated samples.
I,3 MD Parameters
LJ-13 Parameters. The Lennard-Jones (LJ) potential is intermolecular potential that models inter- actions of non-bonding particles. The energy is a function of the interatomic distance of the particles:
12
EL
(z)
2T Σ
Tm deg
I'm
(58)
di
where the distance between two particles i and j is dij = ||:;; - x. 25 and rm c, and Cose are physical constants. As in Köhler et al. (2020), also add a harmonic potential the energy so that gLJ-system =gL (xr + Cosce Eosc (x) This harmonic potential is given by:
gost (x)
Σ || ICOM
(59)
where #COM is the center of mass of the system. We use I'm = =,,c= the IJ--3 dataset, we draw MCMC chains using the Co-U-Turn-Samplee (NUTS) (Hoffman and Gelman, 2014)
2.0 and c= 1.0. For
32

---

Alanine Parameters. For MD data on ALDP and AL,3, we run two chains one for training and one for test. We use the same simulation parameters for both. For training data we sample shorter chains more frequently (every 100 md steps). To conserve disk space for long test chains, we save every 10k steps. Further parameters can be found Table 11 and Table 12.
Table OpenMM simulation parameters.
Force field
amber-14 14
Integration time step
is
Friction coefficient
0.3 ps
Temperature
300
Nonbonded method
CutoffflonPerioiic
Nonbonded cutoff
2 2m
Integrator
Langeviinhiddllelnteggaaao
Table 12: Training and evaluation dataset parameters.
Train
Test
Burn-in period
50 ps
50 ps
Sampling interval
0.1 ps
10 ps
Simulation time
50 ns
ss
Pseudocode
In this section provide Python pseudocode for easy of understanding and reimplementation.
def resampled_ inldddinfeeencce (x0, т, acore_model, energy_model, gamma,
zeta, xi):
T
for t 1inapace T+1) [: 11:
# Define variables at score_model (xt, t)
energy_model (xt, t)
grad_Ut grad (ut, xt)
dUt_ dt = (ut, t)
# Equation (11)
drift = (-a(t) xt) 2eta / (st gamma * (t)
grad_Us Ut)
12
diffusion zeta(t) (t)) randn like (xt)
13
xt += drift * dt + diffusion * sqrt (dt)
# Equation (12) aA div (xt) zeta(t) 2/ 2 gamma grad_Ut ((-a(t) 2
15
xt) zt)
st) gamma dU_dt
16
At dA dt # Resample xt proportional to At with quasi monte carlo xt resample (xt, At)
19
return
Listing I: Python implementation of resampled inference.
K Extended Related Work
Annealed Importance Sampling. In the context of AIS (Jarzynski, 1997; Neal, 2001), SMC samplers (Del Moral al., 2006) and parallel tempering (Swendsen and Wang, 1986), our method reduces the number of energy evaluations by learning the models of intermediate marginals. Indeed, when the buffer of samples from the current temperature is sampled, training of the diffusion model does not require new energy evaluations (note that the gradients for target score matching can be
33

---

cached). Thus, the only time we need to evaluate the energies is for the importance sampling at the final step of the inference-time annealing and for the collection of samples via MCMC at high temperature. Obviously, for sampling from the target density x (xr), the trained diffusion model, unlike AIS, allows producing uncorrelated samples without restarting the chain from the prior distribution.
34


Reviews:
[{"review_id": "fPjJu9QbpV", "reviewer": "Reviewer_JqJ5", "summary": "This paper introduces Progressive Inference-Time Annealing (PITA), a generative framework that progressively trains a cascade of diffusion models to sample from low-temperature Boltzmann distributions. The core idea is to combine two interpolation strategies: (1) temperature annealing, which flattens energy barriers and enables better mode mixing; and (2) diffusion smoothing, which enables efficient resampling at inference time via Feynman–Kac-weighted reverse SDEs. By alternating between training a diffusion model at a given temperature and then resampling to a colder temperature using learned score and energy models, PITA can produce samples that align closely with the Boltzmann target at low temperatures.", "questions": "1. Since the method relies on both a score model and an energy model at each temperature level, is the computational and training cost substantially higher than standard diffusion models or Boltzmann generators?\n2. The proposed method appears to be reminiscent of annealed Langevin dynamics and predictor-corrector samplers discussed in [9], which also gradually anneal the target distribution by adjusting the noise scale. Could the authors discuss the relationship and trade-offs? What are the advantages of your approach?\n3. In [10], the authors proposed an energy model with adaptive temperature control by reparameterizing the energy function, enabling built-in temperature adjustment in the generative model. Could the authors comment on this approach and discuss the advantages of PITA compared to such adaptive-temperature energy models?\n\n**References:**\n\n[9] Score-based generative modeling through stochastic differential equations, ICLR 2021\n\n[10] Learning energy-based models by cooperative diffusion recovery likelihood, ICLR 2024", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\nThe paper is well-written and presents a clean formulation of its algorithmic ideas. The integration of Feynman–Kac theory into diffusion inference is interesting and, to my knowledge, novel in the generative modeling literature. Moreover, the approach achieves significant gains in sample quality and coverage compared to prior Boltzmann generator and diffusion methods, especially in physically meaningful molecular domains.\n\n**Weaknesses**\n\n1. The core novelty of the two interpolants in this work (temperature annealing and diffusion smoothing) is not sufficiently distinguished from existing literature. It seems that both techniques have been widely studied in the MCMC, diffusion-based modeling or energy-based modeling literature.\n2. Annealing has long been used in MCMC [1] and has been applied to diffusion models [2–4].\n3. The idea of diffusion smoothing looks similar to annealed importance sampling [5] and was later extended to deep energy-based models [6–8].\n4. From my understanding, the method reads like a combination of known components, and the authors may wish to clarify why this framework yields fundamentally new insights or performance advantages beyond prior annealed sampling or diffusion-EBM hybrids.\n5. The scope of the evaluation is limited. While the authors demonstrate promising results on small physical systems like LJ-13 and alanine peptides, it is unclear whether this method generalizes to broader settings (e.g., image or text generation, or large protein folding datasets). The method is currently tailored for low-dimensional molecular systems where accurate energy evaluations are feasible.\n\n\n**References:**\n\n[1] Bayesian learning via stochastic gradient Langevin dynamics, ICML 2011\n\n[2] Generative modeling by estimating gradients of the data distribution, NeurIPS 2019\n\n[3] Denoising diffusion probabilistic models, NeurIPS 2020\n\n[4] Improved techniques for training score-based generative models, 2020\n\n[5] Annealed importance sampling, Statistics and computing, 2001\n\n[6] Learning energy-based models by diffusion recovery likelihood, ICLR 2021\n\n[7] A Langevin-like sampler for discrete distributions, ICML 2022\n\n[8] Score-based diffusion meets annealed importance sampling, NeurIPS 2022", "was_revised": true, "final_justification": "Thank you for your comprehensive rebuttal. Your clarifications on the core novelty and the re-contextualization of your evaluation's scope have been effective. The new experiments provide evidence supporting your claims and successfully distinguish your work from prior work. You have addressed my concerns. I have raised my score to 4.", "rebuttal": "We thank the reviewer for their valuable feedback and the time spent. We appreciate that the reviewer finds the paper to be \"well-written\" and \"the integration of Feynman–Kac theory into diffusion inference to be interesting and novel\". In what follows, we address the potential misunderstandings and concerns raised.\n\n## Distinguishing the novelty of the two interpolants\nWe acknowledge the reviewer's comment that the core novelty of our two introduced interpolants may not have been sufficiently distinguished from existing literature. We now clarify this important point further.\n\n> The core novelty of the two interpolants in this work [...] is not sufficiently distinguished from existing literature.\n\n> From my understanding, the method reads like a combination of known components, and the authors may wish to clarify why this framework yields fundamentally new insights or performance advantages beyond prior annealed sampling or diffusion-EBM hybrids.\n\nThis framework yields performance advantages beyond prior annealed sampling or diffusion-EBM hybrids by being the first model to perform consistent simultaneous annealing in both the diffusion and temperature axes. \n\nOne of the core novelties of this work is the Sequential Monte Carlo (SMC) scheme for the inference-time annealing of diffusion models, which enjoys both the efficiency of expressive diffusion models and the consistency of SMC estimators. **This combination of diffusion and temperature annealing in a consistent manner is what distinguishes PITA from the existing literature**. In particular, compared to the previous literature, this procedure has the following merits:\n1. Unlike Langevin Monte Carlo [1], which converges only in the limit of an infinite number of steps, it provides a consistent estimator w.r.t. the number of simulated particles. This fact allows us to curate new lower temperature (consistent) datasets to train new diffusion models---a fact that is not possible with simple Langevin Monte Carlo without incurring bias.\n2. The critical advantage of PITA is that by combining both interpolation strategies, we achieve more stable and performative training of the final diffusion sampler. This combination, to the best of our knowledge, is novel and in particular benefits from our inference time SMC scheme that is tailor-made for diffusion models. As we stress, in the main paper, no existing diffusion sampler is able to tackle the molecular energy systems we consider in Cartesian coordinates. This is because training such samplers without data is extremely challenging and numerically unstable due to the poor estimation of the Stein score, as well as finding the modes in a complex energy landscape to estimate the score in the first place.\n\n> Annealing has long been used in MCMC [1] and has been applied to diffusion models [2–4].\n\nThe reviewer makes an astute observation that Annealing has been applied to diffusion models. **In new experiments, we demonstrate that PITA performs better than Annealing in MCMC, and existing annealed diffusion models**.\n\n1. We first add two baselines of annealed MCMC with sequential Monte Carlo and parallel tempering (PT) strategies.  As it can be seen in Table R1 below, PITA performs better overall in capturing the distribution given the same budget of energy evaluations\n\n**Table R1**: Performance of the SMC and PT baselines on ALDP. For SMC, we take 10 annealing temperatures between 1200K and 300K using a geometric schedule, taking 166 steps per temperature with 30k particles. For PT, we use the same 4 annealing temperatures as PITA, running 25 MD steps between particle exchanges for a total of 50k iterations.\n\n| Method| Rama-KL ↓| Tica-W₁ ↓| Tica-W₂ ↓| Energy-W₁ ↓| Energy-W₂ ↓| T-W₂ ↓| #Energy Evals|\n|-|-|-|-|-|-|-|-|\n| PITA |**4.773±0.460**|**0.112±0.006**|**0.379±0.028**|1.530±0.068|1.615±0.053|**0.270±0.023**|5×10⁷|\n| PT|7.306±1.077|0.625±0.010|0.895±0.016|4.652±0.015|4.689±0.014|0.911±0.004| 5×10⁷|\n| SMC|5.935±0.228|0.372±0.006|0.425±0.003|**0.969±0.078**|**1.002±0.072**| 0.874±0.016|5×10⁷|\n\n2. Second, we add a baseline of \"score scaling\", which is the typical strategy in classifier-free guidance approaches to approximately anneal diffusion models [Table R2 Reviewer Ldmh]. PITA substantially outperforms this method. We note that PITA is consistent where score scaling is not, which explains PITA's superior performance. Score scaling does not produce correct marginals both theoretically and empirically, as also demonstrated in the recent Feynmann-Kac Correctors paper [11], Section 5.2.\n\n> The idea of diffusion smoothing looks similar to annealed importance sampling [5] and was later extended to deep energy-based models [6–8].\n\nCompared to Annealed Importance Sampling (AIS) [5], it follows the diffusion path raised to the power $\\beta$ (i.e. $p_t(x)^\\beta$) instead of the geometric averaging (i.e. $\\pi(x)^{t\\beta}p_0(x)^{(1-t)}$). The use of diffusion paths is crucial for our setting as it reduces the variance of importance weights.\n\nIn addition, we note that previous works on deep energy-based models [6-8] combine the standard inference with Langevin dynamics. This means that the inference procedure converges only in the limit of an infinite number of iterations. This is in stark contrast with the proposed consistent estimator used in PITA, which, as described above, allows for the consistent curation of lower temperature datasets.\n\n## Scope of Evaluation\nWe value the reviewer's concern that the scope of evaluation may initially appear limited. However, we would like to politely push back against this characterization. We begin by recalling that our paper is concerned with sampling from an unnormalized density---a problem domain that is well known to be significantly harder than conventional generative modelling. This is because, unlike generative modelling, we are afforded **no samples from the target distribution**, and yet we must draw i.i.d. samples from it. Due to this challenge, a majority of work has focused on smaller molecular systems in bespoke internal coordinates, which are of smaller dimensionality and are provided through exact domain knowledge for a given system. In contrast, energy functions in Cartesian coordinates are more challenging as they represent the **full dimensionality of the molecular system** and, as a result, are much larger. As a result, we wish to again reiterate this subtle but important point, existing diffusion samplers have failed to even scale to Alanine Tripeptide (AL3) with 99 dimensions, which our method, PITA, effectively tackles. Consequently, we argue that such an empirical demonstration goes beyond the current benchmarks in the literature.\n\nWe thank the reviewer again for their comment, and we will further update the manuscript to highlight the challenge of sampling in Cartesian coordinates.\n\n## Questions \n> Since the method relies on both a score model and an energy model at each temperature level, is the computational and training cost substantially higher than standard diffusion models or Boltzmann generators?\n\nTraining a score model and an energy model incurs higher cost than a standard diffusion model, which only has to train a score model. Training the energy model is around 2x the number of flops as the score model for 3x the cost. However, since training the energy model has substantial benefits performance-wise over only training a score model, this additional cost is a necessary price to pay. \n\n> The proposed method appears to be reminiscent of annealed Langevin dynamics and predictor-corrector samplers discussed in [9][...]. Could the authors discuss the relationship and trade-offs? What are the advantages of your approach?\n\nFirst, we would like to address the potential confusion about the usage of the term \"annealing\". Although it is sometimes used for different noise scales of the target distribution (as in [9]), in this paper, we use annealing exclusively for temperature annealing (i.e., changing the power of the marginal density $p_t(x)^\\beta$). Note that changing the temperature $p_t(x)^\\beta$ and adding the noise $p_t(x) = \\int dy\\; \\mathcal{N}(x|y,\\sigma_t^2) * p_0(y)$ are completely different operations resulting in different marginals and requiring different models for their simulation. We will add the corresponding discussion to the text to clarify this terminology.\n\nSecond, the Predictor-Corrector scheme proposed in [9] samples from the correct marginals only in the limit of an infinite number of corrector steps (Langevin dynamics), while our proposed inference-time annealing uses a consistent SMC estimator in finite steps by incorporating the importance weights.\n\n> In [10], the authors proposed an energy model with adaptive temperature control by reparameterizing the energy function [...]. Could the authors comment on this approach and discuss the advantages of PITA compared to such adaptive-temperature energy models?\n\nWe thank the reviewer for their question. The work in question [10] does not consider sampling from the temperature annealed marginals. Instead, at inference, it runs samples from the marginals with different noise levels, which is similar to the Predictor-Corrector scheme, and cannot be used directly for consistent sampling.\n\n## Concluding remarks\nWe thank the reviewer again for their feedback. We hope that our response and additional experiments clarify all the great points raised in the review and help the reviewer understand the core novelty of our contributions and caliber of our empirical evaluation better. If the reviewer is satisfied with all of our responses, we very politely encourage the reviewer to potentially upgrade their score if they deem it appropriate. We are also more than happy to answer any follow-up questions the reviewer may have.\n\n## Additional References\n\n[11] Skreta, Marta, et al. \"Feynman-kac correctors in diffusion: Annealing, guidance, and product of experts.\" ICML (2025)."}, {"review_id": "OzGOhou3DK", "reviewer": "Reviewer_Ldmh", "summary": "This paper proposes a novel method that combines annealing and diffusion-based samplers to sample from unnormalized target distributions. The key idea is to begin at a high-temperature distribution, where mixing is easier for standard MCMC methods, and progressively anneal toward the target distribution. At each annealing step, the approach alternates between two components: (i) learning a diffusion-based sampler via time-dependent score and energy models at that temperature, and (ii) generating samples using a Feynman-Kac PDE framework for the next, lower temperature. \n\nIn empirical study, the method is evaluated on a series of synthetic particle systems. The proposed method generally produces qualitatively and quantitatively better results than existing baselines. On the most challenging task, Alanine Tripeptide, the method achieves good mode coverage but underperforms on certain evaluation metrics.", "questions": "- If I understand correctly, the method involves training separate models at multiple temperature levels. Are these models trained independently, or is it amortized across temperatures? It would also be helpful to clarify how this overhead compares to other methods in terms of computational efficiency and scalability.\n\n- I wonder why the authors chose to use two separate models for energy function and score (or equivalently, denoiser)? In principle the score function can be obtained by taking the gradient of the energy function. Is this separation primarily to avoid the computational and memory burden of backpropagation? \n\n- Have the authors conducted an ablation study on the four types of losses used in Algorithm 1? Are all of them necessary? Do they receive equal weight in the final objective? Are they trained under the same optimizer settings?\n\n- Have the authors considered a baseline where a diffusion model is trained at a high temperature, and then sampling is performed using the Feynman-Kac PDE framework at the target (lower) temperature?\n\n- How do you tune the hyper-parameters for all methods? \n\n- Do you have any guidance in setting the annealing schedules for new tasks?\n\n- Do diffusion models trained at different temperatures require different optimization settings or hyperparameter tuning? If yes, do you have practical guidance on that?\n\n-This paper [1] also explores diffusion-based sampling with access to the target-mode information and may serve as a useful point of comparison:\n\n[1] LEARNED REFERENCE-BASED DIFFUSION SAMPLING\nFOR MULTI-MODAL DISTRIBUTIONS\nMaxence Noble∗, Louis Grenioux∗, Marylou Gabrie & Alain Oliviero Durmus\n\nMinor and writing-related: \n\n-  In this expression \"U_t = -log p_t(x) + const\", is the constant independent of t?\n\n- Eq 10: Can you explain the right equation? What does the \"\\approx\" mean and why it holds?\n\n- Algorithm 1: the loss function names in the optimizer lines are inconsistent with those in previous lines\n\n- Algorithm 1: It would be helpful to make the algorithm self-contained by specifying inputs such D_t and U_t. \n\n- Table 2 and 3: does the reported #energy evals for the proposed method  include the evaluations used in the initial MCMC sampling phase? Could the authors provide the MCMC configuration details (e.g., number of steps, acceptance rates), and are the final results sensitive to these settings? \n\n- Table 3: what is \"PITA\" vs \"PITA (w/o relaxation)\"?\n\n- On page 5 line 164: \"one subsection\" instead of \"two\"?", "limitations": "One limitation acknowledged by the authors is the need to train both score and energy networks, which increases the training complexity \n\nAdditional imitations include\n\n- The need to train these models at several temperature levels, further raising the training cost. \n\n- Lack of comparison to some of the classical sampling approaches.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The proposed method is intuitive and well-motivated, leveraging a progressive annealing path and coarse simulation data (by MCMC) to address challenges in existing diffusion-based methods.\n\n- The proposed method empirically scales to systems where prior diffusion-based approaches struggle.\n\nWeakness:\n\n- Some parts of the writing could be improved—there are inconsistencies and instances of imprecise or vague explanations (see specific comments). \n\n- The paper lacks empirical comparisons to classical baselines such as replica exchange or sequential Monte Carlo (with temperature annealing), making it difficult to assess the overall significance of the proposed approach. \n\n- The proposed approach can be less flexible as it requires training diffusion models at different temperature, and relies on specifying the initial temperature and annealing schedule.", "was_revised": true, "final_justification": "Most of my questions are resolved. \n\nWhile PITA incurs additional training overhead, it provides faster inference and has the potential to build transferrable samplers across systems.", "rebuttal": "We thank the reviewer for the detailed and constructive feedback. We appreciate the positive assessment of the proposed method. Namely, the fact that the reviewer finds it to be \"intuitive\", \"well-motivated\", and \"addressing existing challenges in diffusion-based methods\". In what follows, we address the concerns raised and answer their questions grouped by theme.\n\n\n## Empirical comparison with classical baselines \n\nWe acknowledge the reviewer concern that our paper could benefit from the addition of classical baselines(SMC and Replica Exchange/Parallel Tempering). We agree with this feedback and have included the suggested baselines in this rebuttal, as shown in Table R1 in our response to Reviewer JqJ5. \n\n## Computational overhead and efficiency of training multiple diffusion models\n\nWe appreciate the reviewer’s comment that the discussion of the computational overhead and efficiency of training separate models amortized across temperatures may not have been sufficiently clear in the manuscript. We now attempt to add increasing clarity to these points. We begin by first noting that in practice we find that sequential fine-tuning of the model demonstrates the best performance---a point thats outlined also in lines 199-201 of the main paper. More precisely, for the subsequent training of a diffusion model at a lower temperature $1/\\beta_{i+1}$ we fine-tune the model trained at the previous temperature $1/\\beta_{i}$. Notably, this strategy allows for reducing the number of required parameters (compared to amortization over temperatures) and the number of training iterations (compared to training from scratch for every temperature). \n\nFurthermore, we wish to highlight that by cost we include both the training time as well as the number of target energy function evaluations needed, which is an important cost associated to real world Boltzmann distributions for molecular systems. We will clarify these computational efficiency gains and their overhead in the updated draft of the main paper. \n \n ## Questions\n \n> I wonder why the authors chose to use two separate models for energy function and score (or equivalently, denoiser)? In principle the score function can be obtained by taking the gradient of the energy function. Is this separation primarily to avoid the computational and memory burden of backpropagation?\n\nThe reviewer makes an interesting observation. Indeed, the score function can be obtained as a gradient of the energy function. There are three reasons why we use the a separate score model.\n1. As noted by the reviewer the computational burden is reduced, and in particular this scheme allows us to avoid the computation of the laplacian of the energy-based model for computing the weights.\n2. Training separate score and energy models is the current SOTA in energy model training following [1]. We get the score function for free given our training scheme.\n3. The score function in general produces better samples than the gradient of the energy function, therefore is a better proposal distribution in practice.\n\nWe think this is a cool and important component of PITA and will add additional discussion in the text.\n\n> Have the authors conducted an ablation study on the four types of losses used in Algorithm 1? Are all of them necessary? Do they receive equal weight in the final objective? Are they trained under the same optimizer settings?\n\nThe four losses in Algorithm 1 consist of two groups: 2 losses for training the score model and 2 losses for training the EBM. For training the score model when both training samples and the target unnormalized density are available, we follow the best practice in the field so far [2]. For training the energy-based model, we follow [1] and distill it from the score model. The novel loss introduced here is the \"energy pinning\" loss, which we ablate in Appendix E (see Table 6 for comparison). We hope this now sufficiently clarifies and the reviewers question.\n\n> Have the authors considered a baseline where a diffusion model is trained at a high temperature, and then sampling is performed using the Feynman-Kac PDE framework at the target (lower) temperature?\n\nWe have updated our results below to include \n1. Training a diffusion model and only rescaling the score to anneal from 1200K to 300K (ScoreScaling),\n2. Using the Feynman-Kac Correctors [5] framework to anneal from 1200K to 300K (FKC).\n\nAs can be seen in the table below, neither of these frameworks can successfully capture the lower temperature distribution on this task. \n\n**Table R2: Performance of methods for the ALDP sampling task.**  The starting temperature is $T_L$ = 1200 K, annealed to target $T_S$ = 300 K. Metrics are calculated over 10k samples and standard deviations over 3 seeds.\n\n|Method|Rama-KL|Tica-W₁↓|Tica-W₂↓|Energy-W₁↓|Energy-W₂↓|T-W₂↓|#EnergyEvals|\n|------|-------|--------|--------|-----------|-----------|------|-------------|\n|PITA|4.773±0.460|**0.112±0.006**|**0.379±0.028**|1.530±0.068|1.615±0.053|**0.270±0.023**|5×10⁷|\n|MD-Diff|**1.308±0.072**|0.113±0.001|0.579±0.004|3.627±0.023|3.704±0.026|0.310±0.001|5×10⁷|\n|MD-NF|13.533±0.024|0.138±0.003|0.586±0.003|**0.551±0.062**|**1.198±0.069**|0.403±0.045|5×10⁷|\n|TA-BG|14.993±0.002|0.219±0.013|0.685±0.034|83.457±0.070|86.176±0.104|0.979±0.012|5×10⁷|\n|FKC|14.392±0.909|0.217±0.000|0.649±0.001|11.281±0.025|11.466±0.027|2.120±0.024|5×10⁷|\n|ScoreScaling|4.588±0.467|0.183±0.002|0.608±0.008|10.282±0.020|10.460±0.019|0.550±0.036|5×10⁷|\n\n> How do you tune the hyper-parameters for all methods?\n\nFor all methods we start with the parameters from SOTA methods in the literature then perform coordinate-wise grid searches over the most relevant parameters to our task using the best performing setting on a validation set at temperature 1200k. For loss weights we simply approximately balanced the magnitude of the losses on the training set, e.g. we used a weight of 0.01 for the target score matching loss to ensure it is on the same scale as the other losses.\n\n> Do you have any guidance in setting the annealing schedules for new tasks?\n\nFor simplicity, we follow the standard practice in the literature and consider the geometric annealing schedule [3]. This, however, can be potentially improved by incorporating modern advanced schedule optimization techniques [4]. We leave exploration of annealing schedules, including dynamic or learned ones, as interesting directions for future work. In theory, the standard practice of annealing to match a target ESS should also work in our setting.\n\n> Do diffusion models trained at different temperatures require different optimization settings or hyperparameter tuning? If yes, do you have practical guidance on that?\n\nWe find that our approach performs consistently well across different temperatures without additional hyperparameter tuning. Specifically, we keep the training settings (e.g., learning rate, number of iterations) constant across all temperature levels.\n\n> This paper also explores diffusion-based sampling with access to the target-mode information and may serve as a useful point of comparison: LEARNED REFERENCE-BASED DIFFUSION SAMPLING FOR MULTI-MODAL DISTRIBUTIONS Maxence Noble∗, Louis Grenioux∗, Marylou Gabrie & Alain Oliviero Durmus\n\nWe thank the reviewer for pointing out the relevant piece of literature. We have added a discussion of this work into the related work section.\n\n## Minor and writing-related comments\nWe thank the reviewer for the detailed minor writing comments.\n\n> In this expression \"U_t = -log p_t(x) + const\",  is the constant independent of t?\n\nYes, that is the log normalization constant.\n\n> Eq 10: Can you explain the right equation? What does the \"\\approx\" mean and why it holds?\n\nIt is approximate as $U_t$ is a learned model and $\\pi$ is the true distribution.\n\n> Table 2 and 3: does the reported #energy evals for the proposed method include the evaluations used in the initial MCMC sampling phase? Could the authors provide the MCMC configuration details (e.g., number of steps, acceptance rates), and are the final results sensitive to these settings?\n\nYes, the reported number of energy evaluations includes those in the initial sampling phase. Note that we collected samples using MCMC for LJ-13 and Molecular Dynamics for the peptide experiments. These configurations are provided in Appendix G.3. Since we use a fairly converged chain, the final results are not very sensitive to these configuration details.\n\n\n## Concluding remarks\nWe thank the reviewer again for their detailed feedback that helped us improve our work. We very politely invite the reviewer to increase their score if they feel our rebuttal was successful in answering all their raised concerns. We also remain available for further discussion and for addressing any lingering doubts. Please do let us know.\n\n\n## References\n[1] Thornton, James, Louis Béthune, Ruixiang Zhang, Arwen Bradley, Preetum Nakkiran, and Shuangfei Zhai. \"Composition and control with distilled energy diffusion models and sequential monte carlo.\" arXiv preprint arXiv:2502.12786 (2025). \n\n[2] De Bortoli, Valentin, Michael Hutchinson, Peter Wirnsberger, and Arnaud Doucet. \"Target score matching.\" arXiv preprint arXiv:2402.08667 (2024). \n\n[3] Del Moral, Pierre, Arnaud Doucet, and Ajay Jasra. \"Sequential monte carlo samplers.\" Journal of the Royal Statistical Society Series B: Statistical Methodology 68, no. 3 (2006): 411-436. \n\n[4] Syed, Saifuddin, Alexandre Bouchard-Côté, Kevin Chern, and Arnaud Doucet. \"Optimised annealed sequential Monte Carlo samplers.\" arXiv preprint arXiv:2408.12057 (2024).\n\n[5] Skreta, Marta, et al. \"Feynman-kac correctors in diffusion: Annealing, guidance, and product of experts.\" ICML (2025).\n\n[7] Karras, T., et al. \"Elucidating the design space of diffusion-based generative models.\" NeurIPS (2022)."}, {"review_id": "G3KpwkqFqM", "reviewer": "Reviewer_9xad", "summary": "The paper proposes a scheme for training a diffusion sampler of a molecular Boltzmann distribution given access to the energy function and a set of high-temperature samples. In the scheme, diffusion models trained at the kth temperature are used to importance sample the target distribution at the (k+1)th temperature. These samples are used to train the next diffusion model. The process repeats until the desired temperature is reached.", "questions": "See above.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nThe core idea proposed in the paper - that of bootstrapping a sampler by learning a sequence of distributions (using diffusion models) that anneal towards the target - is novel and conceptually solid. The formulation is rigorous and the connections with existing and prior literature are laid out well. The paper is overall well written and the supplementary experiments are detailed.\n\n**Concerns or questions**\n\n* The paper makes repeated claims about using less energy evaluations than previous methods, but these comparisons are not shown (on alanine peptides). Indeed, rather few prior methods are mentioned - standard diffusion sample baselines (DDS, iDEM, etc), while conceivably worse, are not shown in the results.\n\n* As in persistent in this field, evaluations are made in settings of questionable practical utility - the small system size conceals problems that may make applications to larger systems impossible. In particular (1) poor coverage of dominant modes at high temperatures, e.g. phase transitions; (2) increasing number of intermediate temperatures required for effective importance sampling; (3) divergence computations to calculate the particle weights (4) distillation of the energy based models.\n\n* Also as in persistent in this field, comparisons are made only on the basis of energy evaluations rather than wall-clock time.\n\n* Why is it necessary to distill the score model into an energy based model? It should be possible to importance sample the target temperature with path RNDs computed only using the SDE drift terms. I would have guessed the weights at intermediate times, enabled by an explicit U_t, allow for frequent resampling using SMC - but the authors don't seem to do this.\n\n* Authors should at least mention the vast body of prior literature on temperature replica exchange simulations.\n\n* Authors notably omit ESS metrics in final comparisons, despite the key role of importance sampling in the training method.\n\n* The development of a robust diffusion-based importance sampler of molecular systems in Cartesian coordinates is of independent interest and has not been reported in prior comparable work (iDEM, adjoint sampling). Authors should elaborate on the details and performance of the importance sampling weights in inference time annealing \n\n* What is the meaning of the number of energy evaluations for the MD-diff baseline?\n\n* To best contextualize the performance of method, baselines of normalizing flows trained on MD trajectories (i.e., SBG) should be shown.", "was_revised": true, "final_justification": "The authors have presented an approach to a longstanding problem that meaningfully advances current capabilities. Concerns about clarity were resolved in the rebuttal stage.", "rebuttal": "We thank the reviewer for their time and effort in assessing our paper. We are pleased to hear that the reviewer finds the key ideas in PITA of learning a sequence of diffusion models through annealing is \"novel and conceptually solid\". Moreover, we appreciate that the reviewer finds our formulation \"rigorous\" in conjunction with prior literature, and the overall paper \"well written\" with detailed supplemental experiments. We now turn to the main points of clarification raised by the reviewer below.\n\n## Cost of energy evaluations in comparison to standard diffusion sampler baselines\n\nWe acknowledge the reviewer's point that the high cost of energy function evaluations for standard diffusion sampler baselines, such as iDEM and DDS, is not provided in the current manuscript. We first note that the expensive energy evaluation cost of standard diffusion samplers like iDEM has been previously reported in prior works like \"Adjoint Sampling\" [1], Table 1, and in concurrent work \"Adjoint Schrodinger Bridge Sampler\" [2], Fig. 4.\n\nMoreover, the scalability of past diffusion samplers to molecular systems, in the **Cartesian coordinates** setting, remains challenging, if not downright infeasible. Since iDEM is simulation-free and more scalable than DDS, we bolster such folklore knowledge with private correspondence with the iDEM authors, who told us that iDEM failed to scale to ALDP even with substantial effort. We further attempted to reproduce these statements and trained iDEM on ALDP ourselves in Vacuum instead of implicit solvent, and observed that iDEM is unable to successfully sample in this easier setting with most modes missing and a poor energy distribution. We will update the manuscript with a discussion on this point and will provide new Rama plots for iDEM in the appendix. \n\n## Evaluation on small peptides \n\nThe reviewer makes an astute observation that the molecular systems considered in this paper are small peptides, and real-world application domains may demand conformation sampling at a much larger scale than considered here. We agree with the reviewer on this point. However, we would like to definitively clarify that prior to our work, PITA, **no other generative sampler** has been demonstrated in **Cartesian coordinates** for even ALDP, much less AL3. In addition, there is a burgeoning literature of combining generative models with classical MCMC techniques in an effort to \"learn to sample\" complex Boltzmann distributions. While preliminary, these methods promise the full benefit of deep learning, such as amortization of cost, faster than MD sampling of new points, and transferability to new, unseen molecular systems within a reasonable family. Consequently, in this literature, we position PITA as the current most scalable diffusion sampler, with expectations that future work---given the rapid progress of generative models for sampling---will enable scaling to more meaningful systems of interest. Thus, we believe PITA represents an inflection point wherein future work on generative samplers will be tested beyond even simpler synthetic energy functions such as 2D GMMs and Lennard-Jones potentials. \n\nWe thank the reviewer for raising this point. We will include a discussion about the current limitations of diffusion samplers due to current scaling challenges in the conclusion.\n\n## Further Questions\n\n> Also as in persistent in this field, comparisons are made only on the basis of energy evaluations rather than wall-clock time.\n\nWe acknowledge the reviewer's comment regarding computational cost being tied to energy evaluation rather than simply wall clock time. In our Boltzmann sampling, the wall-clock time heavily depends on the quality of the energy approximation. For instance, using an energy function that takes into account quantum effects is orders of magnitude more expensive. Consequently, we report the number of energy evaluations in line with standard practice in the field.\n\n> Why is it necessary to distill the score model into an energy-based model? It should be possible to importance sample the target temperature with path RNDs computed only using the SDE drift terms. I would have guessed the weights at intermediate times, enabled by an explicit U_t, allow for frequent resampling using SMC - but the authors don't seem to do this.\n\nWe appreciate the reviewer's nuanced comments. Parameterizing and learning the energy-based model is necessary for the importance sampling with the target density $\\pi^\\beta(x)$ after the inference-time annealing. Indeed, note that the inference-time annealing samples from the density $\\propto\\exp(-\\gamma U_t(x;\\eta))$, which does not perfectly match the target density. We discuss this in lines 156-161 and Appendix B.\n\nFurthermore, we perform an additional experiment using this method,  which we label as \"FKC\" in Table R2 of our response to Reviewer Ldmh. We note that this method is only correct if the learned scores perfectly align with $U_t$, which is difficult to achieve in practice.\n\n> Authors should at least mention the vast body of prior literature on temperature replica exchange simulations.\n\nThank you for your comment. We will include further discussion on temperature replica exchange methods and their relation to our work on PITA.\n\n> Authors notably omit ESS metrics in final comparisons, despite the key role of importance sampling in the training method.\n\nThank you for the note. ESS is not applicable here because our method uses SMC and resamples at every step. While ESS could be considered in settings without resampling, the performance differs drastically with and without resampling in our task. Therefore, we believe that reporting ESS for the AIS method would not accurately reflect the performance of our method as a whole, and we instead rely on other metrics. We will add a note to clarify this in the text.\n\n> The development of a robust diffusion-based importance sampler of molecular systems in Cartesian coordinates is of independent interest and has not been reported in prior comparable work (iDEM, adjoint sampling). Authors should elaborate on the details and performance of the importance sampling weights in inference time annealing.\n\nWe thank the reviewer for this thoughtful suggestion and for recognizing the broader interest in this direction. We agree and will definitely elaborate on these points in the updated draft. \n\n> What is the meaning of the number of energy evaluations for the MD-diff baseline?\n\nFor the MD-diff baseline, we run a Molecular Dynamics (MD) simulation with the predefined budget of energy evaluations to collect the training data at the target temperature and then train a diffusion model on the collected samples. The number of energy evaluations corresponds to the budget for the MD simulation.\n\n> To best contextualize the performance of the method, baselines of normalizing flows trained on MD trajectories (i.e., SBG) should be shown.\n\nWe provide results of the normalizing flow trained on the MD trajectory data used for our MD-Diff baseline, which we call \"MD-NF\" [Table R2 Reviewer Ldmh]. We note that while this model can capture the energy distribution well, it fails to capture all the modes, as is evident from the TICA and Ramachandran (Rama) plots that will be included in the updated manuscript.\n\n## Concluding remarks\n\nWe thank the reviewer again for valuable comments and detailed feedback. We hope that our rebuttal addresses all of their concerns and encourages the reviewer to potentially consider updating their score if they so deem it. We are also available to answer any further questions that arise.\n\n## References\n\n[1] Havens, Aaron, et al. \"Adjoint sampling: Highly scalable diffusion samplers via adjoint matching.\" ICML (2025).\n\n[2] Liu, Guan-Horng, et al. \"Adjoint Schrodinger Bridge Sampler.\" arXiv preprint arXiv:2506.22565 (2025)."}, {"review_id": "rgPKRSMaAa", "reviewer": "Reviewer_SiSU", "summary": "The goal is to draw samples from a distribution knowing its unnormalized density $\\pi(x)$ — the authors design a way to additionally get approximate samples to make the sampling easier. The idea is to sequentially draw samples from tempered distributions $\\pi(x)^{\\beta}$, increasing the inverse temperature $\\beta$ from $0$ (excluded) to $1$, as do classical Sequential Monte Carlo (SMC) algorithms. \n\n**Repeat** until $\\beta = 1$:\n\n1. Increase $\\beta$.\n\n2. Use the unnormalized density $\\pi(x)^\\beta$ to get approximate samples. \n\n3. Use the unnormalized density $\\pi(x)^\\beta$ and approximate samples to get better samples. To do so, train a diffusion model to learn and simulate energies of $\\pi_t(x)^{\\beta}$ where $t$ indexes the diffusion. \n\nSMC algorithms typically do Steps 1 and 2. The authors' algorithm distinguishes itself by adding Step 3. \n\nMoreover, the authors' implementation of Step 2 is also more sophisticated than SMC. The first time, step 2 is achieved by running an MCMC: it should converge relatively fast given that the multi-modality is flattened away. The rest of the time, step 2 is achieved by running an SMC (or FK-PDE) on the diffusion path at the previous temperature, raised to a power.", "questions": "Could the authors address the different points in the weaknesses section above?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strengths\n\nThe authors' method is creative and interesting and seems to be improve over previous models on biological problems. It addresses many current questions:\n\n- how to combine probability paths\n\nThe diffusion path avoids mass teleportation: it is usually simulated from samples. \nThe temperature path preserves the mode locations: it is usually simulated from an unnormalized density. \n\n- how to obtain \"good enough\" samples from an unnormalized density in order to train a diffusion model to get \"better\" samples\n\nThis is an important challenge in the literature of diffusion samplers, as explained for example in *He et al., No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers*.\n\n- how to use both combine samples *and* an unnormalized density to draw samples\n\n## Weaknesses\n\n**Dataset benchmarks**. It is encouraging to see sampling methods being used to make progress on real-world datasets. However, given the relevance and centrality of the authors' method to the sampling literature, it would be useful to include at least one experiment on a synthetic dataset such as a mixture of 50 Gaussians in 2 dimensions. This is because the sampling community as a whole does not have technical expertise in biological problems to judge the experiments section, because other competing methods are often benchmarked on a toy dataset, and also because it provides intuition on the performance of the method. \n\n**Comparison with SMC in the experiments**. The most natural competitor to the author's PITA algorithm is a basic SMC that follows a temperature-annealing path. I appreciate the fact that some of the baselines (e.g. Temperature Annealed Boltzmann Generators) are very similar to SMC, but a vanilla SMC with a Langevin Kernel and a classical schedule for choosing the intermediate temperatures would be a convincing baseline in the experiments. \n\n**Ablation study.** \n\nQ1. For inference-time annealing, is it really helpful to do run an SMC through the diffusion path from the previous temperature, raised to a certain power? Can we not just directly reweigh the samples from $\\pi^\\beta$ to $\\pi^{\\beta + \\Delta \\beta}$ and then learn a diffusion model?\n\nQ2. What is the difference in performance if PITA was run only with inference-time annealing and no model training?\n\n**Comparison with SMC in the text**. Again, the most natural counterpart to the author's PITA algorithm is a basic SMC that follows a temperature-annealing path. It is also an algorithm most people in sampling are familiar path. Discussing how PITA and vanilla SMC differ would be very helpful to understanding PITA, as in the summary of this review. \n\n**Lack of clarity**. I am quite familair with the topic and it took me many reads to understand what the algorithm what doing. Some things that could be more clearly explained:\n- at a given temperature $\\beta$, there are two probability paths that are being used: first, the diffusion path from the previous temperature, raised to a new power (\"`inference-time annealing\"), and then the diffusion path at the new temperature (\"model training\"). \n- why model training is useful in addition to inference-time annealing? \n- Figure 1 is unclear to me. The two taxes of diffusion and temperature are very helpful. But the arrows are very confusing. I did not understand from the arrows that two probability paths are being used at a given temperature: the old diffusion path raised to a power (\"inference-time annealing\") and the new diffusion path (\"model training\"). \n\n**Unclear statements**. \n\nLines 175-177: \"the DSM objective is not sufficient for training a good score model close to the target distribution [...] it has no information about the target distribution\". This is not very clear: near the target distribution, the score model in Eq 13 is trained using samples $x_t$ that are almost from the target distribution. How is this compatible with the statement: \"it has no information about the target distribution\"?\n\nLines 61-62: \"temperature annealing mixes modes by lowering high energy barriers, while diffusion paths avoid mass teleportation.\" This is a strong statement which would benefit from citing a few references. While it may be clear to experts, it is not obvious to any reader. For example, if one has not done the derivations in closed-form of the diffusion path from a Gaussian to a mixture of Gaussians, it is not obvious that mass is not teleported.", "was_revised": true, "final_justification": "I have raised my score to 5. \n\nThe reviewers had addressed my concerns with additional experiments. I have also read through other reviews and am satisfied with the clarifications the authors have provided.", "rebuttal": "We thank the reviewer for the detailed and constructive feedback. We greatly appreciate the positive evaluation of our work. In particular, the reviewer finds the proposed method \"creative\", \"interesting\", \"addressing many current questions\" and addressing \"important challenges\". In what follows, we address the concerns raised and provide additional evaluation requested by the reviewer.\n\n## Dataset benchmarks\n> **Dataset benchmarks.** It would be useful to include at least one experiment on a synthetic dataset such as a mixture of 50 Gaussians in 2 dimensions.\n\nWe thank the reviewer for their suggestion. We now include in this rebuttal the comparison on the synthetic Mixture of Gaussians dataset, training at $T_L=4$ and annealing to $T_S=1$. The comparison is presented in the table below, where we take the results for baselines from [1] and add the performance of PITA. Notably, PITA significantly outperforms all the considered baselines.\n\n**Table R3: Performance of methods on the mixture of 40 Gaussians in 2 dimensions.**\n|Method|Distance-W₂|Energy-W₂|Total Variance|\n|-------|------------|----------|--------------|\n|DDS|15.04±2.97|305.13±186.06|0.96±0.01|\n|PIS|6.58±1.68|79.86±7.79|0.95±0.01|\n|FAB|9.08±1.41|47.60±7.24|0.79±0.07|\n|iDEM|8.21±5.43|60.49±70.12|0.82±0.03|\n|**PITA**|**2.64±0.2**|**0.42±0.04**|**0.67±0.001**|\n\n## Comparison with SMC in the experiments.\n\nWe thank the reviewer for suggesting the inclusion of a vanilla SMC baseline with a Langevin kernel and a classical annealing schedule. We have incorporated this baseline in our experiments, and the results are reported in Table R1 (see our response to Reviewer JqJ5).\n\n> Comparison with SMC in the text. Again, the most natural counterpart to the author's PITA algorithm is a basic SMC that follows a temperature-annealing path. It is also an algorithm most people in sampling are familiar path. Discussing how PITA and vanilla SMC differ would be very helpful to understanding PITA, as in the summary of this review.\n\nThank you for this constructive suggestion! We will add the corresponding discussion in the text.\n\n## Questions about Ablations\n> For inference-time annealing, is it really helpful to do run an SMC through the diffusion path from the previous temperature, raised to a certain power? Can we not just directly reweigh the samples from $\\pi(x)^\\beta$ to $\\pi(x)^{\\beta+ \\Delta \\beta}$ and then learn a diffusion model?\n\nFollowing the reviewer's suggestion, we report the performance of inference-time annealing with importance sampling from $\\pi(x)^{\\beta+ \\Delta \\beta}$ using $\\pi(x)^\\beta$ as a proposal distribution, in Table R2 of our response to Reviewer Ldmh.\n\n> What is the difference in performance if PITA was run only with inference-time annealing and no model training?\n\nThe inference-time annealing procedure heavily relies on the parametric models and is designed to sample from the unnormalized density of the energy-based model, i.e. $q_t(x)\\propto \\exp(-\\gamma U_t(x;\\eta))$. Thus, sampling entirely without training would not produce a meaningful distribution. Although one could perform sampling without training on the samples from the intermediate annealed densities, this approach consistently underperforms (see ablation study in Appendix E, line 858).\n\n## Questions about Clarity of presentation\n> Some things that could be more clearly explained: at a given temperature $\\beta$, there are two probability paths that are being used: first, the diffusion path from the previous temperature, raised to a new power (\"inference-time annealing\"), and then the diffusion path at the new temperature (\"model training\").\n\nWe thank the reviewer for their insightful suggestion. We agree with the reviewer that the proposed phrasing can help clarify the two probability paths used in the PITA algorithm. We will update the corresponding subsections of the paper to better emphasize this distinction.\n\n> why model training is useful in addition to inference-time annealing?\n\nThe inference-time annealing procedure is tightly coupled with the parametric models and is designed to sample from the unnormalized density of the energy-based model, $q_t(x)\\propto \\exp(-\\gamma U_t(x;\\eta))$. Training ensures that the parametric models capture these distributions accurately, which in turn makes the inference-time annealing effective. Jumping too far in temperature with inference-time annealing reduces performance of the method due to increased variance of the importance sampling weights. Model training is useful to \"reset\" these weights at a new lower temperatue. \n\n> Figure 1 is unclear to me. The two taxes of diffusion and temperature are very helpful. But the arrows are very confusing. I did not understand from the arrows that two probability paths are being used at a given temperature: the old diffusion path raised to a power (\"inference-time annealing\") and the new diffusion path (\"model training\").\n\nWe appreciate the reviewer's feedback. In the current version, the diffusion path is represented by the horizontal line, as the standard diffusion does not involve any change in temperature. In contrast, the inference-time annealing path changes the temperature of marginals, which is illustrated by the curved line leading to the marginal $\\pi(x)^\\beta$ at a higher value of $\\beta$. We hope this clarifies the interpretation of Figure 1.\n\n> Lines 175-177: \"the DSM objective is not sufficient for training a good score model close to the target distribution [...] it has no information about the target distribution\". This is not very clear: near the target distribution, the score model in Eq 13 is trained using samples that are almost from the target distribution. How is this compatible with the statement: \"it has no information about the target distribution\"?\n\nAs $t \\to 1$, equation (13) becomes \n\n$\\mathbb{E}_{x_t,x} \\lambda (t) ||x-D_t(x_t;\\theta)||^2$ \n\n$\\to \\mathbb{E}_{x_t,x}\\lambda(t)|| \\sigma_t^2 s_t(x;\\theta)||^2 $\n\nbecause $\\sigma_t \\to 0$ and $x_t \\to x$. This leads to unstable training of the score function $s_{t\\approx 0}(x;\\theta)$ with a finite dataset, as it converges to modeling the score of a uniform mixture of Gaussians with tiny variance around the dataset. This limitation of the DSM objective is well known, and we follow Target Score Matching [2] to address it. We will clarify this point in the text.\n\n> Lines 61-62: \"temperature annealing mixes modes by lowering high energy barriers, while diffusion paths avoid mass teleportation.\" This is a strong statement which would benefit from citing a few references. While it may be clear to experts, it is not obvious to any reader. For example, if one has not done the derivations in closed-form of the diffusion path from a Gaussian to a mixture of Gaussians, it is not obvious that mass is not teleported.\n\nThis is a valuable suggestion. We will add the discussion of the original work on diffusion models with SDEs [3], which proposed the reverse-time SDE and probability-flow ODE, demonstrating that diffusion paths can be simulated without probability mass teleportation. As well as [4] which nicely contrasts temperature and diffusion interpolations visually.\n\n## Concluding remarks\n\nWe thank the reviewer again for the thorough review and valuable feedback which strengthens the clarity and results of our paper through new ablations. We hope that our rebuttal definitely clarifies the all the great points raised in the review, and we politely invite the reviewer to ask any further questions or potentially endorse this paper more strongly if they are satisfied.\n\n\n## References\n[1] OuYang, RuiKang, Bo Qiang, and José Miguel Hernández-Lobato. \"Bnem: A boltzmann sampler based on bootstrapped noised energy matching.\" arXiv preprint arXiv:2409.09787 (2024). \n\n[2] De Bortoli, Valentin, Michael Hutchinson, Peter Wirnsberger, and Arnaud Doucet. \"Target score matching.\" arXiv preprint arXiv:2402.08667 (2024). \n\n[3] Song, Yang, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. \"Score-based generative modeling through stochastic differential equations.\" arXiv preprint arXiv:2011.13456 (2020).\n\n[4] Máté, Bálint, and François Fleuret. \"Learning interpolations between boltzmann densities.\" arXiv preprint TMLR (2023)."}]

Author rebuttal:

Meta-review: The paper present a novel contribution based on good theoretical foundation and with solid empirical results.
Decision: Accept (spotlight)