BoltzNCE: Learning likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation
Authors: Rishal Aggarwal
Jacky Chen
Nicholas Matthew Boffi
David Koes
Venue: NeurIPS 2025 poster ()
Keywords: Noise Contrastive Estimation
Flow Matching
Score Matching
Free Energy Estimation
Boltzmann Distribution
Boltzmann Generators
Primary area: machine_learning_for_sciences
TL;DR: Learning likelihoods instead of calculating expensive jacobians on generative ODEs to calculate free energy differences and obtain the Boltzmann distribution.

Abstract: Efficient sampling from the Boltzmann distribution given its energy function is a key challenge for modeling complex physical systems such as molecules. Boltzmann Generators address this problem by leveraging continuous normalizing flows to transform a simple prior into a distribution that can be reweighted to match the target using sample likelihoods. Despite the elegance of this approach, obtaining these likelihoods requires computing costly Jacobians during integration, which is impractical for large molecular systems. To overcome this difficulty, we train an energy-based model (EBM) to approximate likelihoods using both noise contrastive estimation (NCE) and score matching, which we show outperforms the use of either objective in isolation. On 2D synthetic systems where failure can be easily visualized, NCE improves mode weighting relative to score matching alone. On alanine dipeptide, our method yields free energy profiles and energy distributions that closely match those obtained using exact likelihoods while achieving $100\times$ faster inference. By training on multiple dipeptide systems, we show that our approach also exhibits effective transfer learning, generalizing to new systems at inference time and achieving at least a $6\times$ speedup over standard MD with only a bit of fine-tuning. While many recent efforts in generative modeling have prioritized models with fast sampling, our work demonstrates the design of models with accelerated likelihoods, enabling the application of reweighting schemes that ensure unbiased Boltzmann statistics at scale. Our code is available at https://github.com/RishalAggarwal/BoltzNCE.

BoltzNCE: Learning Likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation
Rishal Aggarwal
Jacky Chen
CMU-Pitt Computational Biology
CMU-Pitt Computational Biology
Dept. of Computational & Systems Biology
Dept. of Computational & Systems Biology University of Pittsburgh Pittsburgh, PA 15260
University of Pittsburgh Pittsburgh, PA 15260
rishal. aggarwal@pitts edu
jackychen@pitt. edu
Nicholas M. Boffi
David Ryan Koes
Machine Learning Dept. Dept. of Mathematical Sciences Carnegic Mellon University Pittsburgh, PA 15213
Dept. of Computational & Systems Biology University of Pittsburgh Pittsburgh, PA 15260 dkoesOpitt.
nboff Gandreeandwww...m.omm cmu edu
Abstract
Efficient sampling from the Boltzmann distribution given its energy function is a key challenge for modeling complex physical systems such molecules. Boltz- mann Generators address this problem by leveraging continuous normalizing flows transform simple prior into a distribution that can reweighted to match the target using sample likelihoods. Despite the elegance of this approach, obtaining these likelihoods requires computing costly Jacobians during integration, which is impractical for large molecular systems. To overcome this difficulty, we train an energy-based model (EBM) to approximate likelihoods using both noise con- trastive estimation (NCE) and score matching, which we show outperforms the use of either objective in isolation. On 2D synthetic systems where failure can be easily visualized, NCE improves mode weighting relative score matching alone. On alanine dipeptide, our method yields free energy profiles and energy distributions that closely match those obtained using exact likelihoods while achiev- ing 100x faster inference. By training on multiple dipeptide systems, we show that approach also exhibits effective transfer learning, generalizing to new systems at inference time and achieving at least a 6x speedup over standard MD with only a bit of fine-tuning. While many recent efforts in generative modeling have prioritized models with fast sampling, our work demonstrates the design of models with accelerated likelihoods, enabling the application of reweighting schemes that ensure unbiased Boltzmann statistics at scale. Our code is available cow.aa/aaaaaaaaaaa///////AACCCC
Introduction
Obtaining the equilibrium distribution of molecular conformations defined by energy function a fundamental problem in the physical sciences [1-3]. The Boltzmann distribution describes the equilibrium probability density and is given by p(z) α cxp( KBT)) where E(x) is the energy of molecular conformer æ, K is the Boltzmann constant, and T is temperature. Sampling from this distribution is particularly difficult for molecular systems due the non-convex nature of the energy
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Noise 4(x)
Generated Distribution (xx)
Predicted Likelihoods (x f(x)
Desired Distribution μ(x) x e-000//aT
Bottzmann Emulator
BoltzNCE (EBM)
Pa (ctrccle) (asquare))
Reweight
M0Other (6
Py (triingle)
Weighe
Figure 1: Overview. BoltzNCE offers an accelcrated alternative to exact model likelihood compu- tation. Samples from a prior are first transformed to distribution of conformers by Boltzmann Emulator, which is easy to sample from but difficult to evaluate likelihoods for. The generated samples are then reweighted with likelihoods estimated by an energy-based model (EBM), which we train to approximate the emulator distribution with hybrid score matching and noise contrastive estimation scheme. This EBM gives access to likelihoods a single function call, enabling us to reweight to the target Boltzmann distribtion up to x faster than exact computation.
landscape, leading to the presence of widespread energy basins, metastability, and slow transitions. Traditional approaches for sampling conformers, such as Markov chain Monte Carlo (MCMC) and molecular dynamics (MD) [4, 5], often get trapped in these energy wells, which necessitates long simulation timescales to produce uncorrelated samples [6]. Consequently, is particularly inefficient to obtain samples from independent metastable states.
In recent years, several generative deep learning methods have been proposed address the molecular sampling problem. One prominent class is Boltzmann Generators (BGs) [7-10], which transform a simple prior distribution (e.g., a multivariate Gaussian) into distribution over molecular conformers that can be reweighted to approximate the Boltzmann distribution. When reweighting is not applied, the model referred to Boltzmann Emulator [8], whose primary aim to efficiently sample metastable states of the molecular ensemble. While often qualitatively reasonable, Boltzmann Emulators alone fail to recover the true Boltzmann distribution and require the reweighting step for exact recovery of the system's 's equilibrium statistics.
To compute the likelihoods of generated samples, BGs are constrained to the class of normalizing flows. While earlier methods built these flows using invertible neural networks [7, 11], more recent approaches prefer using continuous normalizing flows (CNFs) [8, 12] due their enhanced expressitivity and flexibility in model design. Despite these advantages, computing likelihoods for Cen-generated samples requires expensive Jacobian trace evaluations along the integration path [13, 14]. This computational overhead limits their scalability, particularly for large-scale protein systems, In this work, we ask the question:
Can the likelihood of sarge-scale scientific generative models be efficiently amortized to avoid the prohibitive path integral?
Here, as means to learn such a likelihood surrogate, we investigate the use of energy-based models (EBMs). EBMs learn the energy function of synthetic Boltzmann distribution, po(x) o cxp(E) (x)) where Eg is the energy to be learned [15]. Scalable training of EBMs remains major challenge due to the need for sampling from the model distribution, which often requires simulation during training [16, 17J. As result, developing efficient training algorithms for EBMs continues to be an active area of research [15, 18-20].
We adopt noise contrastive estimation (NCE) [21], which trains a classifier to distinguish between samples drawn from the target distribution and those from a carefully chosen noise distribution. The key advantage of this approach is that it circumvents the need to compute intractable normalizing constants [22, 23]. Despite this, NCE can suffer from the density-chasm problem [24, 25J, whereby the optimization landscape becomes flat when the data and noise distributions differ significantly [26]. To address this issue, we introduce an annealing scheme between a simple noise distribution and the data distribution using stochastic interpolants [27]. Annealing mitigates the density chasm problem by introducing intermediate distributions between noise and data, which facilitates the generation of more informative samples. In particular, it ensures that negative samples lie closer to positive samples, improving the effectiveness of NCE optimization [24, 28]. We further enhance the training process by combining an InfoNCE [29] loss with score matching [30] objective defined over the law of the
2

---

stochastic interpolant. Notably. our proposed method for training the EBM is both simulation-free and avoids the computation of normalizing constants, making scalable large systems.
Contributions. On synthetic 2D systems, we show that training with both losses performs signif- icantly better than either individually. For the alanine dipeptide system, our method recovers the correct semi-empirical free energy surface while achieving 100x speedup over exact likelihood calculations. On other dipeptide systems, we further demonstrate the method's 's ability to generalize to previously unseen molecular systems. To summarize, main contributions are:
Training. We develop a scalable, simulation-free framework for training EBMs by combin- ing stochastic interpolants, score matching, and noise contrastive estimation.
Fast likelihoods. We show that learned likelihoods can replace expensive Jacobian compu- tations in the reweighting step, recovering exact Boltzmann statistics.
Empirical validation. We achieve 100x speedup on alanine dipeptide compared exact
likelihoods, and demonstrate 6x speedup over MD on unseen dipeptide systems.
2 Methodological Framework
In this work, we introduce new class of generative models for sampling from a Boltzmann distribu- tion that features accelerated likelihood computations (Figure 1). To this end, we train standard Boltzmann Emulator and an EBM, which each enable efficient sampling and likelihood evaluation, respectively. Given access to an EBM approximating the output distribution of the Boltzmann Emulator, we can evaluate sample likelihoods in just single function call, enabling rapid reweighting for accurate estimation of observables. We train these EBM models with new hybrid approach that we call Bolt As critical components of our approach, we first provide background on stochastic interpolants and flow matching [27, 31], then show how these can be used to build Boltzmann Emulators and finally, we introduce the innovations underlying the BoltzNCE method.
2.1 Background: stochastic interpolants and Boltzmann emulators
stochastic interpolant [2] is a stochastic process that smoothly deforms data from a fixed base distribution Pu into data sampled from the target distribution P1 pp. Under specific choices of the hyperparameters, stochastic interpolants recover standard settings of diffusion models [30, 33], flow matching [31], and rectified flows [34]. They can be used to learn generative models, because they provide access to time-dependent samples along dynamical process that converts samples from the base into samples from the target. Concretely, given samples {af ) with æ] f1 sampled from the target, we may define a stochastic interpolant as the time-dependent process
1 04 + 3e2
Above, a [0, R and 3: [0. + R are continuously differentiable functions satisfying the boundary conditions dg 1, = 0, 9 0. and 3, 1. In the absence of domain knowledge, often take the base distribution to be a standard Gaussian, Po N(O, I).
The probability density e Law(I:) induced by the interpolant coincides with the density of probability flow that pushes samples from Po onto P1.
by (2t),
(xx) = Eli, I. x],
where (x) is given by the conditional expectation of the time derivative of the interpolant at fixed
point in space. The score (xx) = V log (x()) is further given by the conditional expectation [27]:
st (x)) = 02
E
[ro
T xx]
which we will use to train our energy based model as likelihood surrogate.
Flow matching. Given coupling p( x1) between Po and pl, the vector field (2) can be approximated with a neural network via the regression problem
(b) = [I(,(t.) - i.113]
(4)
3

---

where E denotes an expectation over (11) The objective (4) can be further modified so that the model is trained to predict the clean data #1 instead of the vector field [35, 36]. We refer the reader to Appendix A for more details on the endpoint objective.
Once trained, the learned model can be used to generate samples Ty by solving the differential equation
(2z), To Po-
(5)
The log density log á associated with the generated samples 2 can be calculated with the continuous change of variables formula,
!. V V- b(()))t
log (())) = log po
(6)
While (6) gives a formula for the exact likelihood, its computation is expensive due to the appearance of the divergence of the estimated flow b.
Boltzmann Generators. Boltzmann Generators leverage generative models to sample conformers and to compute their likelihoods, which enables reweighting the generated samples to the Boltzmann distribution. In practice, these models can be built using the stochastic interpolant framework described in (1), (2) and (4) given target data generated via molecular dynamics. We can generate unbiased samples from the target distribution by first sampling Φ P1 by solving (5) and (6), and
E((1)) 1/(_(1)) With this weight, then by reweighting with the importance weight w(61) exp( KsT we can also approximate any observable O under the Boltzmann distribution j using self-normalized importance sampling [7, 8, 12]:
EN (( E w((+)
(0)₂ = EziNo [w()))O(+1]] ~
While the likelihood integral (6) has been used in prior implementations of Boltzmann generators [8, 12]. in general its computational expense prevents it from scaling to large molecular systems. In this work, our aim to amortize the associated cost by learning a second energy-based model that can estimate the likelihoods log P1.
2.2 BoltzNCE
Our method designed to calculate free energies and enable the computation of observables via (7) in an efficient and scalable manner. It proceeds in two stages: (i) standard flow training, and (ii) amortization of the likelihood via EBM training. the first stage, we train a Boltzmann emulator on
dataset D of conformers using the stochastic interpolant framework described in Section 1. The trained emulator is then used to generate samples á ~ PI via (5), leading to a dataset of generated conformers D. In the second stage, we train an EBM on D to approximate the energy U of the generated distribution,
(1 = (x))
Z z
exp(U(z))dd.
(8)
The generated samples are then reweighted to the Boltzmann distribution using (8) in (7).
Training the emulator. We train the Boltzmann Emulator using stochastic interpolants as described in Section 2.1. Boltzmann Emulator models are trained using either the vector field (4) or the endpoint objectives. The endpoint parameterization [35, 36] is where the network predicts the clean endpoint In rather than the velocity field directly. The corresponding velocity field for sampling is given by
6,(x)) ( + (š.a, - 3,)2, (t.x))
(9)
where é ()) is the network's predicted endpoint at time given noisy sample æ. A derivation of this velocity field is provided in Appendix A. Since aj O, this velocity field diverges at t = 1; in practice, we integrate from f = O to t = 1- - le 1 to avoid this singularity. While the endpoint parameterization works well generating samples, it leads unstable likelihoods due to the singularity at t=1, drawback that is addressed by the use of EBMs in the BoltzNCE method. Once trained, generate the dataset D by sampling Φ1 P1 via (5).

---

Training the EBM. Rather than learn a single energy function corresponding to PI as in (8), we propose define a second stochastic interpolant from the base to Ộ1,
I. = 04 + B
(10)
and estimate the associated time-dependent energy function Ug for Pe Law(
Zi Zi= / exp(Uc(z))dz.
pc(x) = exp (U((2)) ZZz,
(11)
By construction, we have that م1 = وı, though in general Ae * Pr = Law( Given access to a model Ü. of Ut, we can evaluate log Đı (a:) up to the normalization constant 21 with a single function evaluation Ú1 (z). eliminating the need for (6). We train U using combination of denoising score matching and an InfoNCE objective [29]. The score matching objective leverages (3) to yield
LsM (0) E |as VŪ,(T)) + +xol?
(12)
where the E is over the draws of (f, To, â1) defining I, To define the InfoNCE objective, we first write down the joint distribution over (t,, It) given by p(t, x) (x))p.(z) for t ~ p(t). In practice,
we choose time uniformly, so that p(t)
and p(t,x) = pr (z). Given this joint density, we may
define the conditional distribution p(t[xr) x(txx)//(xx Pr (x)/p(x) where p(x) f p(t, x)
pt )t is the marginal distribution of æ. This conditional distribution describes the probability that a given observation x was sample at time The InfoNCE objective maximizes the conditional likelihood by minimizing its negative log-lik oggllikelhoodd We can write this intractable quantity
] = -E log ( exp(ú (1 p(x)
log 2,) log )dt ] (Ú), (13)
(c(x))
NLL = EE log
eexp(úp(I))
where the expectation is over the draws of 2o, 21) defining I, This expression depends unknown normalization constants 2, and involves an integral over time t'. To avoid explicit normal- ization, we parameterize U, directly approximate Ue(x) log Ze. absorbing the logmonamalization constant as a time-dependent bias. To approximate the integral, InfoNCE leverages Monte Carlo ap- proximation, leading to multi-class classification problem: given sample 7, we aim to distinguish the true time from set of negative times (fx 1
This yields the tractable objective:
LinfONCE -E log Ere(A exp(Uv )
exp(úi(T()
(14)
We train Ó by minimizing weighted combination of both objectives:
LBONANC(()) (Ú) LInGaNCE(
(15)
Since the noise level changes continuously with t, only nearby times t' ~ t have non-negligible conditional likelihood Le). We therefore sample negatives {fx} from a narrow Gaussian centered providing informative contrast for learning fine-grained temporal discrimination. We note that both objective functions are simulation-fie after generation of D. as they only require sampling the interpolant L. While, in principle, the EBM could be trained on the original interpolant 1, (1), training on the generated interpolant 7, (10) preferable because it approximates the emulator's energy function and enhances transferability to new molecular systems where timescale MD data unavailable.
For more details on model training and inference, refer to Appendices B, E.8 and E.
Related Work
Boltzmann Generators. Boltzmann Generators have become an active line of research since their
initial development with invertible neural networks [7], having now been used to sample both molec- ular systems [10, 11, 37-40] and lattice systems [10, 41 41-43). Tan al. [44] introduce a more stable reweighting scheme that takes advantage of Jarzynski's equality attain the equilibrium distribution.
5

---

However, most of these methods have required input through gystem-specifee featurizations such as internal coordinates, hindering transferability. The emergence of CNFs and equivariant neural network architectures has enabled the development of BGs on Cartesian coordinates [8, 12] thereby not heing system specific and enabling transferability. Despite these advancements, transferability has so far only been demonstrated on small systems such dipeptides, primarily due to the computational limitations associated with likelihood evaluation at scale.
Boltzmann Emulators. Boltzmann Emulators, unlike BGs, are designed solely to produce high- quality samples without reweighting to the Boltzmann distribution. Because they are not required to be invertible, they can typically be applied to much larger systems. This flexibility also enables the use of a wider range of generative approaches, including diffusion models. Boltzmann Emulators have been employed to generate peptide ensembles [45], protein conformer distributions [35, 46-49], small molecules [50-52], and coarse-grained protein structures [53, 54]. However, they are inherently limited by the data distribution they were trained on. As a result, they are generally unsuitable for generating unbiased samples from the Boltzmann distribution or for performing free energy calculations independently. In this work, we aim to leverage the strengths of Boltzmann Emulators and bridge the gap between Emulators and Generators using energy-basee models (EBMs). While EBMs may still yield biased estimates, they bring closer to true Boltzmann statistics at a fraction of the computational cost.
Energy-Based Models (EBMs) are particularly appealing in the physical sciences, as they describe density functions manner analogous the Boltzmann distribution. This similarity enables the use of various techniques from statistical physics to compute thermodynamic properties of interest [39, 55]. Despite their promise, training EBMs remains a challenging task. Recent advancements have introduced training objectives inspired by noise contrastive estimation [18, 20, 24, 25, 55, 56], contrastive learning [29, 57], and score matching [15, 58, 59]. Ou Yang et al. [60] have also proposed an cnnergy-matching objective to train neural sampler on the Boltzmann distribution; however, more work needs to be done to make this approach practical for molecules.
Energy-Basse Models.
4 Experiments
In this section, we validate BoltzNCE on several synthetic low-dimensional systems as well as on molecular conformer generation for dipeptides. Our experiments demonstrate three key results: (i) Combining the denoising score matching (12) and InfoNCE (14) objectives to form (15) significantly improves EBM quality compared to using either alone; (ii) BoltzNCE enables mudpoint-parameteriae emulators to function as Boltzmann Generators by decoupling likelihood estimation from the sampling vector field, avoiding the numerical instability that makes Jacobian trace computation intractable; (iii) The learned EBM provides accurate and efficient likelihood estimates, enabling the calculation of Boltzmann statistics at speeds up two orders of magnitude faster than exact likelihood computations.
4.1 Low-dimensional systems
The effectiveness of the score matching and InfoNCE loss functions (12) and (14) tested on low-dimensional synthetic systems. Specifically, we study an 8-mode Gaussian mixture in two dimensions and the two-dimensional checkerboard distribution. The results for forward evaluation of the trained EBMs are shown in Figure 2. We leverage a simple feedforward neural network that takes the point x R² input and outputs a scalar value. KL-divergence values between the predicted and ground truth densities are also reported in Table 1.
Objective Functions
8-Mode Gaussian Mixture
Checkerboard
InfoNCE
0.2395
3.8478
Score Matching
0.2199
0.8384
Score Matching InfoNCE
0.0150
0.1987
Table 1: Low-dimensionall Systems: Quantitative Results. KL-divergence (4) when using different objective functions to train the EBM. For both systems, using our combined objective (15) significantly outperforms using either (12) or (14) individually.

---

True Density
NCE
Score Matching
Score Matching & NCE
Figure 2: Low-dimensions, Systems: Qualitative Results. EBM density learned synthetic two- dimensional systems. (Above) An 8-mode Gaussian mixture. (Below) The checkerboard distribution. In both cases, the density is shown in the leftmost column, and the results obtained with different methods are shown to the right. Using both objectives (right) provides the best performance.
Our results indicate that combining both objective functions leads to significantly improved perfor- mance compared using either alone. Intuitively, the InfoNCE loss appears to act as a regularizer for the score matching objective, helping to improve the relative weighting of different modes in the learned distribution. This effect is not observed when using the score matching loss its own. For subsequent experiments on molecules, both objective functions are used to train the EBM.
Alanine Dipeptide
As complex system, we first study the alanine dipeptide molecule. Our aim here is obtain the equilibrium distribution of the molecule specified by the semi-empirical GFN-xTB force field [61]. Running a simulation with this force field computationally intensive, so we use the same setup as Klein and Noć [8]. Conformers are generated using molecular dynamics with the classical Amber f1995Bildn force field and subsequently relaxed using GFN-xTB. We use dataset variants: unbiased, corresponding to the original distribution, and biased, in which the positive state is oversampled for equal metastable state representation (Figure 6).
We train geometric vector perceptron (GVP)[62, 63] based Boltzmann Emulators using both the vector field objective (4) (GVP-VF) and the endpoint objective (GVP-EP) on the unbiased and biased datasets, and compare their performance to Equivariant Continuous Normalizing Flow (ECNF) models from Klein and Noé [8] trained on the same datasets. Additionally, we use the Graphormer architecture [64] to parameterize the EBM. Further details on data featurization, model architectures, and hyperparameters are provided in Appendices and E.6.
To evaluate Boltzmann generation, els trained on the bi- ased dataset, yield more accurate esti- mates of free energy dif- ferences. Free energy differences between the positive and negative metastable states com- puted, since this transition corresponds a slow dy- namical process (Appendix D.4, Figure 6). Ground-truth free energy values are obtained via umbrella sampling simulations from Klein et al. [12]; further details on umbrella sampling are provided in Appendix E.2. In addition, we compute the energy-based (E-W2) and torsion-based (T-W,) Wasserstein-2 distances between the unbiased dataset and the generated (proposal or reweighted) distributions.
Table 2: Methods Overview Flow matching objectives and likelihood estimation methods for different models tested on Alanine Dipeptide.
mod-
these
Method
FM Objective
Likelihood Estimation
ECNF [8]
Vector Field
Jac-trace integral
GVP Vector field
Vector Field
Jac-trace integral
GVP Endpoint
Endpoint
Jac-trace integral
BoltzNCE Vector field
Vector Field
EBM forward pass
BoltzNCE Endpoint
Endpoint
EBM forward pass
7

---

Ramachandran plot
Energy distrilution
Free energy projection
Fron enuegyy hgg
XD slood axpple
Free / RayT
Boast now ghtel
Cabrik eupias
0.3
BalaDOC
0:33
0.30
Energy /
kgT
Figure 3: BoltzNCE: Qualitative Results. Results for BoltzNCE alanine dipeptide trained on the biased dataset. We use GVP vector field as the Boltzmann Emulator. BoltzNCE successfully captures the energy distribution and the free energy projection. (Left) Ramachandran plot of generated samples. (Middle) Energy histogram along with BoltzNCE reweighting. (Right) Calculated free energy surfaces for the angle yp on the right.
Method
AF/kET
AF Err. Prop. E-W, Rew, E-W2 Prop. T-W2 Rew. T-W2
Inf. (h) Train (h)
Umbrella Sampling
4.20
ECNF [8]
4.09 0.05
0.01 = 0,05
9.37
3.85
ECNF (reproducd)
4.07 0.23
0.03 = 0.23
8.08 + 0,56
0.37 0.02
1.10 0.00
0.59 + 0.00
9.37
3,85
Landpiint
4.89
26.2
GVP Vector Field
4.38 0.67 0.28
= 0.67
7.20
+ 0.13
0.46
0.05
1.09 0.01
0.60 0.00
18.4
4.42
BoltzNCE Endpoiat
4.14
0.04
= 0.94
6,24
+ 0,52
278
+ 0.04
1.12 + 0.01
0.59
0.01
0.16
12.2
BolizNCE Vector Field
4.08 0.13 0.02 1 0.13
7.12
1 0.15
0.27
1 0.02
112 1 0.00 0.57 0.00
0.09
12.2
Table 3: BoltzNCE: Quantitative Results. Dimensionless free energy difference, energy E-W2 and torsion angle T-W2 Wasserstein-2 distances calculated by different Boltzmann Generator and BoltzNCE models. Standard deviations are shown across 5 runs. Free energy difference values for umbrella sampling and ECNF taken from Klein and Noé [8], which consider to be ground truth. BoltzNCE Vector Field provides the best performance/inferencc time tradeoff as compared to all other methods.
The ECNE, GVP-VF, and GVP-EP models estimate likelihoods using the Jacobian trace integral and act Boltzmann Generators. Energy-based models (EBMs) trained samples generated by the GVP-based emulators are also evaluated and referred to as BoltzNCE-VF VF and BoltzNCE-EP. The specific flow-matching objectives and likelihood estimation methods for each model are summarized in Table 2.
GVP models excel as Emulators but fail as Generators due to poor likelihood estimates. The results for the Boltzmann Emulator models trained on the unbiased dataset are given in Appendix F.I. In general, the two GVP models demonstrate better performance than the ECNF. The GVP-EP model performs the best, making a strong candidate for Boltzmann Emulator, however, as shown below, it fails as Boltzmann Generator as its vector field (Eq. 9) diverges as + 1, which we find corrupts its likelihood calculation.
All Boltzmann generation results are summarized in Table 3. Comparing the GVP and ECNF models, find that while GVP models match exceed ECNF performance as emulators, they produce less accurate free energy differences and exhibit higher reweighted &-W2 and T-W, scores. This inaccuracy may stem from unreliable likelihood estimates produced during ODE integration, which requires the divergence of the model with respect its input to be accurate and well-behaved. As result, Boltzmann Generators face additional design constraints to ensure stability and reliability of their likelihood computation. In the following, we show how BoltzNCE resolves this issue by learning likelihoods separately, decoupling emulator quality from likelihood tractability and enabling greater flexibility in the design of the emulator.
BoltzNCE enables fast and accurate likelihood estimation for Emulators. In contrast, the BoltzNCE models yield more accurate estimates of the free energy difference and achieve lower reweighted E-W2 and T-W2 scores, with BoltzNCE-VE providing the best performance. This indicates that the likelihoods predicted by BoltzNCE are generally more reliable than those obtained via the Jacobian trace integral. Representative energy histograms and free energy surfaces along the slowest transition (pp dihedral angle) for the BoltzNCE-VF model are shown in Figure 3. For energy histograms and free energy projections of other methods, refer to Appendix F. As demonstrated in
8

---

Energy Distribution
MD Simulation
0.03
Botzzmann Emalstor
Bolx&CE Reweglted
0.01
KDE
Ap
0.02
0.01
0.013755
0.00
Energy / KgT
Free energy projection
Free energy projection TICA
Free energy / kgg
20
MD
Free energy / knf
MD
Boitarann Emulater
Bolszmam Exculator
BoltaCC Reweighzed
DahuNCE Reweighted
10
-15
0.0
M15 TICO
10
Figure 4: BoltzNCE Results NY dipeptide. BoltzNCE inference results for NY dipeptide (top left) after fine-tuning. Energy distribution (top right). free energy surfaces along the ip angle (bottom left) and the first TICA component (bottom right). BoltzNCE successfully captures the right energy distribution and free energy projections for the dipeptide.
the figure, the BoltzNCE method is able to accurately capture the free energy projection and energy distribution for alanine dipeptide.
We report inference time costs in Table 3, including the time generate and estimate likelihoods for 10c conformers. BoltzNCE provides an overwhelming inference time advantage over the standard Boltzmann Generator by two orders of magnitude while matching accuracy.
To further evaluate the accuracy of EBM likelihoods, we show that it both more accurate and more computationally efficient than the Hutchinson trace estimator (Appendix F.3). We also demonstrate further flexibility in the design of the EBM training algorithm by comparing OT to independent coupling (Appendix F.).. It is important note, however, that BolızNCE has upfront cost of training the EBM associated with it. In principle, this upfront cost could be reduced by training the EBM parallel on the original interpolant (I) and then fine-tuning on the generated interpolant (10).
4.3 Generalizability on dipeptides
Finally, we demonstrate BoltzNCE's ability to generalize to unseen molecules using systems of dipeptides. For this experiment, was originally developed Klein et al. [65]. The training set consists of classical force field MD simulations of 200 dipeptides, each run for about 50 ns. Since this dataset may not have reached convergence, we bias the dataset in a similar manner to alanine dipeptide to ensure equal representation of the modes along the oo angle. For testing, utilize js long simulations of randomly chosen dipeptides not present in the training set.
the same setup and dataset from Klein and Noé [8], which
In this experiment, we benchmark BoltzNCE against independent 1us MD runs and the TBG- ECNF model [8] trained on the biased dataset. We use the BoltzNCE-Vi method as it achieved the best performance the alanine dipeptide system. We also compute time-lagged independent components (TICA) [66] from the test MD simulations and plot the free energy projections along the first component. For details TICA, refer to Appendix E.5.
The inference procedure with BolIZNCE is modified to improve generalizability. Samples generated from the flow-matching model are passed through a conformer matching and chirality checking procedure to check validity generated samples. For more details, refer to Appendices E.3 and E.4. The EBM, on the other hand, is first pretrained using conformers of training peptides generated by the

---

Method
yp AF Error
E-W2
T-W2
Inference Time (h)
MD Baseline
0.18 + 0.22
0.17 = 0.03
0.22 = 0.03
24.04
TBG-ECNF# [8]
0.13 - 0.10
0.36 + 0.12
0.34 = 0.06
123.07
BoltzNCE
0.43 + 0.21
1.08 = 0.65
0.44 = 0.13
4.005
Table 4: Generalizability Results. Generalizability of the BoltzNCE method in comparison to TBG and MD simulations on systems of 7 dipeptides. *Fewer samples (30,000) used due high GPU compute time. BoltzNCE provides a significant time advantage over the other methods while achieving good performance.
www-matching model and then fine-tuned on the dipeptide of interest during inference. In addition, we exclude the top 0.2% of importance weights during reweighting to reduce variance, following the approach introduced in Tan al. [44].
BoltzNCE yields accurate Boltzmann statistics at a fraction of MD/TBG computational cost. Quantitative evaluations across seven dipeptide systems (Table 4) show that BoltzNCE closely repro- duces Boltzmann-weighted energy distributions and free energy surfaces obtained from molecular dynamics (MD), as illustrated for the NY dipeptide in Figure Although the TBG-ECNF method attains the highest accuracy in free energy estimation, it incurs orders-of caddrs-of-magnntude higher compu- tational cost due to its reliance on exact likelihood calculations, thereby serving as an upper bound on BoltzNCE performance. In contrast, BoltzNCE achieves comparable accuracy by approximating likelihoods at substantially lower computational expense.
Visual inspection of all test systems (Figure 10) confirms that this minor performance drop remains acceptable, with BoltzNCE exhibiting excellent agreement with MD-derived energy distributions and free energy landscapes. A small reduction in E-W, scores, primarily driven by a single outlier in the NF dipeptide (Appendix G), does not affect overall fidelity. Collectively, these results demonstrate that BoltzNCE provides an efficient, scalable, and generalizable framework for amortized Boltzmann sampling on unseen molecular systems, maintaining high thermodynamie accuracy fraction of the computational cost.
5 Discussion
In this work, we introduce a novel, scalable, and simulation-free framework for training energy- energy-based models that integrates stochastic interpolants, InfoNCE, and score matching. We show that InfoNCE and score matching act complementarily to enhance model performance. Our approach learns the density of conformers sampled from a Boltzmann Emulator, eliminating the need for costly Jacobian trace calculations and achieving orders-of sonss-of-aaen magnntude speedups. On alanine dipeptide, BoltzNCE can even surpass the accuracy of ODE-based divergence integration. Across multiple dipeptide systems, the method generalizes to unseen molecules with minimal fine-tuning, providing substantial computational savings over conventional molecular dynamics. This framework bridges the gap between Boltzmann Emulators and Generators, removing the dependence on invertible architectures and expensive likelihood computations, while enabling high-fidelity, scalable Boltzmann sampling.
6 Limitations and Future Work
The present work is limited to dipeptide molecular systems. While on ADP the method demonstrates excellent accuracy, the performance drops in the generalizability settings. This limitation could likely be addressed through more extensive exploration of model architectures and hyperparameter optimization, which we have only minimally investigated here. The method also has the potential to be scaled to larger molecular systems; however, the accuracy of the method needs to be further tested in nigher-dimensionee settings.
Training the energy-based model requires applying the score matching loss to its gradients, which increases compute requirements beyond typical levels for neural networks training. Additionally, since the likelihoods estimated by the EBM are approximate, degree of mismatch between the samples and their predicted likelihoods is inevitable.
10

---

Although the current work is limited to molecular setting, we believe the proposed EBM training framework could be broadly applicable in other domains where energy-based models are useful, such as inverse problems in computer vision and robotics.
7 Acknowledgments
We thank Leon Klein for making the code and data for his original Boltzmann Generator and Equivariant Flow Matching methods readily available.
This work is funded through R35GM140753 from the National Institute of General Medical Sciences. The content is solely the responsibility of the authors and does not necessarily represent the official views of the National Institute of General Medical Sciences or the National Institutes of Health.
References
[1] Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-55) 500, 2024.
[2] Haiyang Zheng and Jin Wang. Alphafold3 in drug discovery: A comprehensive assessment of capabilities, limitations, and applications. bioRxiv, pages 2025-04, 2025.
[3] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. BioRxiv, 2022: 500902, 2022.
[4] Benedict Leimkuhler. Molecular dynamics. Encyclopedia Applied and Computational Mathematics, pages 931-940 Springer, 2015.
[5] David W Borhani and David E Shaw. The future of molecular dynamics simulations in drug discovery. Journal of computer-aided molecular design, 26(1): 2012.
[6] John Klepeis, Kresten Sadorff--aassen Ron O Dror, and David Shaw. tinesiiee molecular dynamics simulations of protein structure and function. Current opinion in structural biology, 19(2):120- 2009.
[7] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365 147, 2019.
[8] Leon Klein and Frank Noé. arXiv:2406. 14426, 2024.
Transferable boltzmann generators.
arXiv preprint
[9] Alessandro Coretti, Sebastian Falkner, Jan Weinreich, Christoph Dellago, and Anatole von Lilienfeld. Boltzmann generators and the new frontier of computational sampling in many-body systems. arXiv preprint arXiv:2404.16500 2024.
[10] Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and boltzmann generators. Physical Review Research, 4(4):L04 2022.
[11] Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. Advances in Neural Information Processing Systems, 34:2796-280,, 2021.
[12] Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances Neural Information Processing Systems, 36:59886-59910, 2023.
[13] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[14] Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01366 2018.

---

[15] Yang Song and Diederik P Kingma. How to train your energy energy-based models. arXiv preprint arXiv:2101.03288, 2021.
[16] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in neural information processing systems, 32, 2019.
[17] Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. arXiv preprint arXiv:2012.01366, 2020.
[18] Sumeet Singh. Stephen Tu, and Vikas Sindhwani. Revisiting energy based models as poli- cies: Ranking noise contrastive estimation and interpolating energy models. arXiv preprint arXiv:2309.05880 005803, 2023.
[19] Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. arXiv preprint arXiv:2003. 2020.
[20] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on robot learning. pages 158- 188.168.77 PMLR, 2022.
[21] Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international con- ference on artificial intelligence and statistics, pages 297-304. JMLR Workshop and Conference Proceedings, 2010.
[22] Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation of unnormalized statistical models, with applications to natural image statistics. The journal of machine learning research, 13()):307-361, 2012.
[23] Dyer. Notes on noise contrastive estimation and negative sampling. arxiv. arXiv preprint arXiv: 110 2014.
[24] Benjamin Rhodes, Kai Xu, and Michael U Gutmann. Telescoping density ratio estimation. Advances in neural information processing systems, 33:4905-4916, 2020.
[25] Bingbin Liu, Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. Analyzing and improv- ing the optimization landscape of noise-contrastive estimation. estimaaion arXiv preprint arXiv:21 10. 1271, 2021.
[26] Holden Lee, Chirag Pabbaraju, Anish Sevekari, and Andrej Risteski. Pitfalls of gaussians as a noise distribution in nce. arXiv preprint arXiv:2210.00188 2022.
[27] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
[28] Omar Chehab, Aapo Hyvarinen, and Andrej Risteski. Provable benefits of annealing for estimating normalizing constants: Importance sampling, noise-contrastive estimation, and beyond. Advances in Neural Information Processing Systems, 36:45945-45970, 2023.
[29] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807 1/807.03745 2018
[30] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1. 13456, 2020.
[31] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[32] Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas Boffi, Eric Vindenn and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pages 23-40. Springer, 2024.
[33] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851,, 2020.
12

---

[34] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arX1v:2209.03003,, 2022.
[35] Bowen Jing. Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402.04844 2024.
[36] Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Harmonic self- conditioned flow matching for multi-ligand docking and binding site design. arXiv preprint arXiv:2310.0576, 2023.
[37] Peter Wirnsberger, Andrew J Ballard, George Papamakarios, Stuart Abercrombic, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, and Charles Blundell. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 153(14), 2020.
[38] Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Mernández-Lobato Flow annealed importance sampling bootstrap. arXiv preprins arXiv:2208.0189) 2022.
[39] Xinqiang Ding and Bin Zhang. Deepbar: fast and exact method for binding free energy computation. The journal of physical chemistry letters, 12(10))2509-251: 2021.
[40] Joseph C Kim, David Bloore, Karan Kapoor, Jun Feng, Ming-Hong Hao, and Mengdi Wang. Scalable normalizing flows enable boltzmann generators for macromolecules. arXiv preprint arXiv:2401.04246 2024.
[41] Rasool Ahmad and Wei Cai. Free energy calculation of crystalline solids using normalizing flows. Modelling and Simulation in Materials Science and Engineering, 30(6):065007, 2022.
[42] Kim A Nicoli, Christopher J Anders, Lena Funcke, Tobias Hartung, Karl Jansen, Pan Kessel, Shinichi Nakajima, and Paolo Stornati. Estimation of thermodynamic observables in lattice field theories with deep generative models. Physical review letters, 126(3):032001, 2021.
[43] Maximilian Schebek, Michele Invernizzi, Frank Noé, and Jutta Rogal. Efficient mapping of phase diagrams with conditional boltzmann generators. Machine Learning: Science and Technology, 5(4):045045, 2024.
[44] Charlie B. Tan, Avishek Jocy Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. Scalable equilibrium sampling with sequential boltzmann generators. arXiv:2502.18462 [cs.LG], 2025. URL https: /ttp///axvvv///////a//eee/eaaav/ org/ abs 250.. 18462.
[45] Osama Abdin and Philip M Kim. Pepflow: direct conformational sampling from peptide energy landscapes through hypemetwork-condiiiion diffusion. bioRxiv, pages 2023-06, 2023.
[46] Shuxin Zheng, Jiyan Hc, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng, Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, et al. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence, 6(5):558-500 2024.
[47] Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. Advances in Neural Information Processing Systems, 36:36499-3646622222 2023.
[48] Yan Wang, Lihao Wang, Yuning Shen, Yiqun Wang, Huizhuo Yuan, Yue Wu, and Quanquan Gu. Protein conformation generation via force-guided (3) diffusion models. arxiv preprint 2400:22003.14088 2024.
[49] Sarah Lewis, Tim Hempel, José Jiménez-Luna, Michael Gastegger, Yu Xie, Andrew YK Foong, Victor García Satorras, Osama Abdin, Bastiaan S Veeling, Iryna Zaporozhets, et al. Scalable emulation of protein equilibrium ensembles with generative deep learning. Science, 389(61))): eadv9817, 2025.
[50] Juan Viguera Diez, Sara Romeo Atance, Ola Engkvist, and Simon Olsson. Generation of conformational ensembles of small molecules via surrogate model-assisted molecular dynamics. Machine Learning: Science and Technology, 5(2):025010, 2024.
13

---

[51] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. Advances in neural information processing systems, 35:24240-24253, 2022.
[52] Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei-Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36:549-568, 2023.
[53] Nicholas E Charron, Felix Musil, Andrea Guljas, Yaoyi Chen, Klara Bonneau, Aldo S Pasos- Trejo, Jacopo Venturin, Daria Gusew, Iryna Zaporozhets, Andreas Krämer, et al. Navigating protein landscapes with machine-learned transferable marse-grained model. arXiv preprint arXiv:23 18278, 2023.
[54] Jonas Kohler, Yaoyi Chen, Andreas Kramer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining of molecular dynamics without forces. Journal of Chemical Theory and Computation, 19399992-952,, 2023.
[55] Anish Sevekari, Rishal Aggarwal, Maria Chikina, and David Koes. Accelerating nce conver- gence with adaptive normalizing constant computation. In ICML 2024 Workshop on Structured Probabilistic Inference { Generative Modeling, 2024.
[56] Kristy Choi, Chenlin Meng, Yang Song, and Stefano Ermon. Density ratio estimation via infinitesimal classification. In International Conference on Artificial Intelligence and Statistics, pages 2552-2573. PMLR, 2022.
[57] Hankook Lee, Jongheon Jeong. Sejun Park, and Jinwoo Shin. Guiding energy-based models via contrastive latent variables. arXiv preprint (rXiv:2303.03022 2023.
[58] Shahar Yadin, Noam Elata, and Tomer Michaeli. Classification diffusion models: Revitalizing density ratio estimation. arXiv preprint arXiv:2402. 10095, 2024.
[59] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and memc. In International conference machine learning, pages 8489 -510. PMLR, 2023.
[60] RuiKang Ou Bo Qiang, Zixing Song, and José Miguel Hernández-Lobato. Bnem: A boltz- mann sampler based on bootstrapped noised energy matching. arXiv preprint arXiv:2409.09787, 2024.
[61] Christoph Bannwarth, Scbastian Ehlert, and Stefan Grimme. Gfn2-xtb an accurate and broadly parametrized self-consistent tight-binding quantum chemical method with multipole electrostatics and density-ddeendert dispersion contributions. Journal of chemical theory and computation, 15(3):1652-166 1671, 2019.
[62] Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael Townshend, and Ron Dror. Learn- ing from protein structure with geometric vector perceptrons. arXiv preprint 1rXiv:2009.01411, 2020.
[63] Bowen Jing, Stephan Eismann, Pratham N Soni, and Ron O Dror. Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint arXiv:2106. 03844,, 2021.
[64] Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie Liu. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877-28885 2021.
[65] Leon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Mare Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36:52863-52883, 2023.
[66] John D Chodera and Frank Noé. Markov state models of biomolecular conformational dynamics. Current opinion in structural biology, 25:135-144, 2014.
14

---

[67] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E (n) equivariant graph neural networks. International conference on machine learning, pages
PMLR, 2021.
[68] Ian Dunn and David Ryan Koes. Accelerating inference molecular diffusion models with latent representations of protein structure. ArXiv, pages arXiv-2311, 2024.
[69] Ian Dunn and David Ryan Koes. Mixed continuous and categorical flow matching for 3d de novo molecule generation. arXiv:2404 197399 (q-bio.BM]) 2024. URL WtLhttts://aaae //axxivv org/ abs/2404.0 19739.
[70] Junhan Yang. Zheng Liu, Shitao Xiao, Chaozhuo Li, Defu Lian, Sanjay Agrawal, Amit Singh, Guangzhong Sun, and Xing Xie. Graphformers: Gnn-nested transformers for representation learning on textual graph, 2023, URL https://arxiv. org/ /205 02605.
[71] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boishunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021. URL http: jmlr wwwgp/ppas//////// wvtp://vvvevv22ee .ttll.
[72] Moritz Hoffmann, Martin Scherer, Tim Hempel, Andreas Mardt, Brian de Silva, Brooke E Husic, Stefan Klus, Hao Wu, Nathan Kutz, Steven L. Brunton, et al. Deeptime: a python library for machine learning dynamical models from time series data. Machine Learning: Science and Technology, 3(1)0015009 2021.
[73] Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.0044 2023.
[74] Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Courna- peau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, et al. Scipy 1.0: fundamental algorithms for scientific computing in python. Nature methods, 17(3)2661-72, 2020.
[75] Patrick Kidger, Ricky T. Q. Chen, and Terry Lyons. "hey, that's not an ode": Faster ode adjoints via seminorms. International Conference Machine Learning, 2021.
15

---

NeurIPS Paper Checklist
The checklist is designed to 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]. [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiiaton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors their evaluation. While "[Yes] " is generally preferable to "[No]" it is perfectly acceptable to answer " "[No]" " provided proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA]" is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often
nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract match experimental results section 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations and future work section has been included (section 6)
16

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, authors should discuss possible limitations of their approach to address problems of privacy and faireess
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All the information is present in the experimental section and appendix (section B,C.E.6.E.8). The code and data is also available at
Guidelines:
The answer NA means that the paper does not include experiments.
17

---

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 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:
The code https//github..om/Rii wwtp.////iiiiaaaaaga/////oleaee.
and
data
available
at
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 reproduce the results. See the NeurlPS code and data submission guidelines (https: ctwps.ccpppbllc///eeccc guides/ /CodeSubmis ssi sisoiiiy) 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).
18

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g.. data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Experimental details have been specified in section 4 and in the appendix sections CE.6, E.8.
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: Mean and standard deviations are provided across 5 runs for all metrics (section
4)
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report I-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources have been specified in training and inference time sections in the appendix D.E.8.
Guidelines:
The answer NA means that the paper does not include experiments.
19

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttrps:///eurppss cc/ppcccppcccccccccccccccccccccc httpppubllttticccccceeeeeeeeeee ines?
attps::/naurrps..cpbbbb.
Answer: [Yes]
Justification: The paper conforms with all guidelines and rules of the NeurlPS code of Ethics
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The work is more foundational research and does not have any impact directly society
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out, For example, it is legitimate point that 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 train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
20

---

The paper does not contain any models/methods that could be considered risk
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The authors all the data and benchmark models have been credited through citations
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, pacorswithcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
thcode
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
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: [Yes]
Justification: Reproducible https///ghhhbcoommmmillaaaaaaa//
code
along
with
license is provided at
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
21

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the 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 an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our L.LM policy (https : htt///eeerppp/pp/ wcc/onneeeeeccs//22222 for what should or should not be described.
22

---

A Endpoint Objective
Stochastic interpolants anneal between To (0, and (z)) with:
4. 0:40 + 3
(16)
solving for Zn²
T B,31
20
(17)
a2
I evolves according to the vector field given by the conditional expectation:
(x()) di, (x = x] + 3, [2 1 xc]
(18)
Substituting 17 in 18 we get:
bx(x)= ór(( 3,B [1 To x²)
+ 8, [21 14 2]
(19)
Ot2
bx(x) = oi (dix + (3.00 - 3,) [2 = x))
(20)
Similarly, the model estimate of the vector field given by:
(x, a ( + (Bror - 3 3,)
(21)
Where (t, 4)) the predicted endpoint by the model. The objective is then given by:
LEP 1-"
- bbr)))
(22)
T
Blor Br ($i(t, 10000=2)ppp =m)1p
LEP
Jo
E
(23)
a1
B
EBM training algorithm
Input and sample
(c,x
InfoNCE lass function
Energy based Model 8
exp(Eol (t', xx'))
V (t,
log Σ Lresfu exp (Eg (t+ x!))
1=1
Score Matching loss
(()x))
Negative time z
Figure 5: Energy Based Model training workflow
23

---

A diagrammatic representation of the method used for training the energy- mmeggy-based model is shown in Figure The model takes a sample z and time point as input and outputs predicted energy Eg x).
The gradient of the output with respect to the sample V, Ea( (t) is used for the score matching loss. The same sample also passed with negative time points { and the predicted energies (1', x)) are used along with the previously output energies Eg (t. x) for the InfoNCE loss.
We also provide pseudocode block for training the energy-based model with stochastic interpolants, InfoNCE, and score matching in algorithm block 1.
Algorithm 1: Training EBM with stochastic interpolants, InfoNCE, and score matching
Input:
Energy-Based model 0, samples from prior Xo, generated samples X1, interpolant functions dt, 3r, negative time sampling variance
for epoch + 1 to epoch do
for batch (Zo, 21) in (Xo, X1) do
(Io. - coupling function(zo,ã1)
sample t ~~ ((0,1) 1)
+ a + 3,01
LSM + N E las Vo,(", [.) x
sample t' ot,
LIMOANCE: En= log exp(Ee(1 (er]))((""""))+eexppp (o .!?))
N
exp(Es((" (t))
+ LsM + LInfONCE
0 + Update(0, VoL)
Output: Updated model parameters 0
C Data featurization and Model Architecture
C.1 Data featurization
The data is featurized such that all atom types are defined according to the peptide topology. For transferable models, amino acid identity and positional information included along with atomic features. Molecular structures are represented as fully connected graphs, and both models operate directly on the Cartesian coordinates of the atoms.
C.2 Geometric Vector Perceptrons
Boltzmann Emulators are parameterized with an SE())-equvvarian graph neural network that lever- ages geometric vector perceptrons (GVPs) [62]. Briefly, GVP maintains a set of equivariant vector and scalar features per node that updated in an SE(3))qqivarraaniinvarin manner via graph convolutions. We utilize this architecture as it has been shown to have improved performance over equivariant graph neural networks (EGNNs) [67] in molecular design tasks [68].
For our models, we use a modified version of the GVP which has been shown to increase performance described in Dunn and Koes [69]. The message passing step constructed by applying the GVP message passing operation defined in Jing et al. [63].
(m{2) m m212)) c'M (119) ddy) ve:
(24)
Here m and m (2) are the vector and scalar messages between nodes i, j. hús di are the scalar features, edge features, and radial basis embedding respectively, while represents the coordinates of the node. For the detailed Node Position Update and Node Feature Update operations, refer to Appendix C of Dunn and Koes [69].
24

---

C.3
Graphormer Operations
Our EBMs are implemented using the graphormer [64] architecture, which has demonstrated state- of-the-art performance in molecular property prediction tasks. Graphormers function similarly to standard transformers, with the key difference being the incorporation of an attention bias derived from graph-specific features. In 3D-graphormers, this attention bias is computed by passing Euclidean distance matrix through a multi-layer perceptron (MLP).
Graphformers are neural network architectures where layer-wise GNN components are nested alongside typical transformer blocks [70]. For our EBMs, we follow the implementation of the Graphformer with one minor modification. For the original Graphformer, each attention head is calculated as:
QK vd
head = softmax
+ +B)V
(25)
where is learnable bias matrix. In our implementation, B is calculated by passing the graph's euclidean distance matrix through an MLP.
D
Metrics
D.1 NLL
To calculate the NLL of the holdout conformers, we take (6) and evaluate the ODE in the reverse direction for given sample. This provides the NLL of the sample. NLL values are reported over batches of 10 samples.
D.2 Energy- W2
In order quantify the difference in energy distributions between generated molecules and MD relaxed samples, we calculate the Wasserstein-2 distance between the two distributions. This can be intuitively thought of as the cost of transforming one distribution to another using optimal transport. Mathematically, we solve the optimization process with the loss:
E-W, = inf c(x, y)² dr( v)
f
(26)
where (æx y) represents coupling between two pairs (a, y) and c(2,, y) is the euclidean distance. We use the Python Optimal Transport package in our implementation [71]. E-W,, values are reported over batches of * 10" samples.
D.3 Angle W2
Similar to the E metric, we seek quantify the differences the distributions of dihedral angles generated and those from MD relaxed samples. Here, following the convention defined in Tan et al. [44], we define the optimal transport torsional angle space as:
T-W, = inf c(a, c(a.v)) v)
fola,
dr( v)
(27)
where x (I. y) represents a coupling between two pairs (x, y). The cost metric on torsional space is defined as:
1
c(2, v)
Σ ((x; yo 1%)%)
(28)
where (x. 7
Similar to Energy-W2 calculations, we use the Python Optimal Transport package for implementa- tion [71]. T-W2 values are computed in radians over batches of 10 samples.
25

---

D.4 Free energy difference
We believe the free energy projection of system a relevant baseline for the following two reasons:
It represents a high dimensional integral probability: The equation free energy projection along reaction coordinate is given by:
F(r') ----KBl -KTT -n
(29)
where p(r') is the probability density of observing the system at position r' along the coordinate,
p(r') - 7 / (r)(())-r)e((-ll)))kkk da
(30)
In principle, if we solve the above integral along all points of the reaction coordinate, solve a (D 1) dimensional integral, where x E RDD Thus, this serves as good metric on how well the model matches the ground truth RD distribution.
For dipeptides, the process of going from the negative & angle to the positive &o angle is considered a slow process as there are regions of high energy/low probability in between the two. Therefore, this serves an ideal reaction coordinate study dipeptide systems.
Domain relevance: Applied studies in the bioghysica/liochhemiil/ssss
biology domain work on elucidating the free energy projection along a reaction coordinate. This helps the researchers identify the relative stability of different modes along the coordinate as well as the rate reaction along the coordinate.
Free energy differences are computed between the positive and negative metastable states of the yp dihedral angle. The positive state is defined as the region between 0 and 2, while the negative state encompasses the remaining range. The free energy associated with each state estimated by taking the negative logarithm of the reweighted population count within that state.
The code for calculating the free energy difference is follows:
left = 0. right 2
hist, edges np. histogram(phi, bins=100, density=True, weightsseetthtttteggttt... centers 0.5 (edges [1:] edges[:-1])
centers_par pos (centers > left) & (centers < right)
lree_energ_diiffeee = -np .ppp.logthistt(contrrrpppp] sum ()/
hist "est["cettrs_por sum())
Where phi is numpy array containing the yp angles of the generated dataset (yp (-H, T)) and weights is an array containing the importance weight associated with it.
D.5 Inference times
Inference time for free energy estimation measured over 1 x 106 samples. Specifically, we use batch size of 500 and generate 200 batches of conformers. During sample generation, Boltzmann Generators also compute the Jacobian trace. All run times are recorded on NVIDIA L40 GPUs. Reported values represent the of five independent across individual runs on the seven test-system dipeptides in the generalization setting.
for alanine dipeptide and the average
E
Technical Details
E.1
Dataset Biasing
Since transitioning between the negative and positive yp is the slowest process, with the positive yp state being less probable, follow the convention of Klein and Noć [8], Klein et al. [12] and use
26

---

Experiment
Model Type
Architecture
Parameters
Alanine Dipeptide
Flow Matching
ECNF
147,599
Alanine Dipeptide
Flow Matching
GVP
108,933
Alanine Dipeptide
Energy-Based Model
Graphormer
4,879,949
Dipeptides (2AA)
Flow Matching
ECNF 1,044,239
Dipeptides (2AA)
Flow Matching
GVP 735,109
Dipeptides (2AA) Energy-Basse Model
Graphormer 6,102,153
Table 5: Model size of different neural network architectures used in this work.
version of the dataset with bias to achieve nearly equal density in both states, which helps in obtaining a more accurate estimation free energy. To achieve the biased distribution, weights based on the von Mises distribution, fuMs are incorporated and computed along the yp dihedral angle as
w(p)) r. f.M (1 1 1, = 10) +
(31)
Where is the ratio of positive and negative yp states in the dataset. To achieve dataset biasing, samples are drawn based on this weighted distribution.
E.2 Umbrella sampling
Umbrella sampling is physics- based method used to estimate the free energy profile along reaction coordinate. It involves selecting set of points along the coordinate, performing separate simulations around each point using biasing potential, typically harmonic restraint, and then combining the resulting data to reconstruct an unbiased free energy landscape. The biasing potential keeps the system near the target point while also promoting sampling of regions that are otherwise rarely visited due to energy barriers.
For alanine dipeptide, Klein et. al[]2] ran umbrella sampling simulations with the GFN2-xtb force- field. We utilize the data from the same simulation and treat it as the ground truth value of the free energy projection.
E.3
Conformer matching
For the generalizability experiments, the bonded graph of a generated sample inferred using empirical bond distances and atom types in similar manner to Klein and Noé [8]. The inferred graph then compared with a reference graph of the molecule and the sample discarded the two graphs are not isomorphic.
E.4
Correcting for chirality
Since SE(3) equivariant neural networks are invariant to mirroring, the Emulator models tend to generate samples from both chiral states. To account for this, we fix chirality post-hoc following the convention set by Klein and Noé [8], Klein et al. [12].
E.5 Time-lagged Independent Component Analysis (TICA)
TICA is dimensionality reduction technique introduced for analyzing time-series data. In general, it is used identify directions that maximize the autocorrelation at chosen lag time. Projecting data onto TICA components yields lower dimensional representations that preserve the system's slowest timescales. Similar to Klein and Noé [8], we construct TICA components using the deeptime library [72] at lag time of 0.5 ns.
E.6 Model hyperparameters
Each GVP-Boltzmann Emulator model consists one message-passing GVP layer and one update GVP layer. The alanine dipeptide (ADP) emulators use 5 hidden layers with vector gating, whereas the dipeptide emulators use 9. ADP emulators are configured with 64 scalar features and 16 vector features, while the dipeptide emulators use 128 scalar and 32 vector features.
27

---

The Graphormer- hrsse potential models are instantiated with 366-iimensional node embeddings and matching 256-umit feed-forward layers within each transformer block, with total of 8 layers for ADP and 10 layers for dipeptides. Self-attention is applied with 32 heads over these embeddings, and interatomic distances are encoded using 50 Gaussian basis kernels. The total parameter counts for each model used in this work are reported in Table 5.
E.7 Endpoint training weights
The Endpoint loss function for training the Boltzmann Emulator is given by:
|| Blor - - (61 10) - x1) B 0t
LEP = (C1.)))))
(32)
Note that, the coefficients á. become divergent near we threshold the min and the max value of these coefficients as follows:
31 Therefore, practice,
9,
3,ar 8e 1), 100)
tu min(max(0.005,
(33)
a4
And optimize the following objective:
E'Pmo = 12-22(0011), (Co0))))))))
(to [ (t,, (t) - =11
(34)
E.8
Training protocols
Emulator models for ADP were trained for 1,000 epochs, while those for dipeptides were trained for 12 epochs. Both were trained using the Adam optimizer with learning rate of 0.001 and a batch size of 512. A learning rate scheduler was employed reduce the rate by factor of 2 after 20 consecutive epochs without improvement, down to a minimum of 1e -5, An Exponential Moving Average (EMA) with 0.999 was applied to the models and updated every 10 iterations. For ADP, batches were coupled using mini-batch optimal transport[73], while for dipeptides independent coupling with rotational alignment was employed. Mini-batch optimal transport was computed using the SciPy linear__e sum___mm__ assignnnsn inssenttttt function [74]. All models were trained on NVIDIA L40 GPUs with batch size of 512.
The EBMs in both settings are trained with independent coupling. For ADP, the training set consists of 100,000 conformers generated by the emulator, while for dipeptides the training set includes 50.000 conformers generated by the emulator across the 200 dipeptides in the dataset. The ADP EBM is trained for 1,000 epochs, whereas the dipeptide EBM is trained for 10 epochs.
For ADP, the negative time point is sampled from a Gaussian distribution with standard deviation of 0.025, while for dipeptides it is sampled from Gaussian with standard deviation of 0.0125. Both models are optimized using Adam with a learning rate of 0.001. The learning rate is reduced by half after 30 consecutive evaluations/epochs without improvement, down minimum of le Training is performed with a batch size of 512.
E.9 Interpolant Formulation
We specify the interpolant process following the design choices explored in Ma et al. [32]. The Emulator models are trained with linear interpolants while the energy-basee models use trigonometric interpolants, both of which satisfy the constraints to generate unbiased interpolation process.
Linear Ot =
B =
(355
Trigonometric: a, = cos ( T(),
3: = sin TT))
(36)
Trigonometric interpolants are called general vector preserving interpolants (GVP) in Ma et al. [32]. However, we change the naming of this notation to avoid confusion with geometric vector perceptrons (GVP), which are repeatedly discussed in our paper.
28

---

MD Relaxed
Blased MD Relaxed
knf
Froe euengy Kall
10
Five energy
Figure 6: Alanine Dipeptide System A visualization of the alanine dipeptide system. Cartoon representation of the alanine dipeptide (left) with its rotatable dihedral angles labeled, Ramachandran plots of unbiased (center) and biased (right) datasets. The biased MD upweights the low frequency mode along the yp dihedral.
GVP Vector Field
GVP Endpoint
Energy cietribution
Tire energy
Free entengy
Energy /hgT"
Figure 7: Evaluation of different emulator models. The Ramachandran plots of dihedral angles of both endpoint and vector field models are displayed on the left, and center. We see that the energy distributions (right) of both endpoint and vector field emulators well as previous method, ECNF, deviate from the distribution of the MD data. However, the Endpoint model distribution deviates the least.
E.10 Integration scheme
All models were integrated with the adaptive step size DOPRI5 solver implemented in the Torchdiffeg package [75]. The tolerance atol and rtol values were set to - for alanine dipeptide and le for systems of dipeptides. Vector field model integrals are evaluated from 0 to 1, while endpoint models are evaluated from 0 to -e in order to avoid the numerical instability that occurs with endpoint parametrization at time t 1.
F
Additional Results
F.1
Vector Field Vs Endpoint Objectives
Inference results for the Boltzmann Emulators presented in Table and Figure 7. In this section, we aim to quantify what training objective makes the best emulator and, surprisingly, whether a better emulator will always make better generator.
Method
&-W₂
T-W2
NLL
NLL std
ECNF
6.22
0.12
0.27
0.01
-125.53
0.10
5.09 0.09
GVP Vector Field
4.99 + 0.50
0.27 + 0.02
-125.42 + 0.15
6.92 + 0.62
GVP Endpoint
3.11 0.70
0.26 + 0.02
-92.04 + 3.24
175.12 + 35.51
Table 6: Boltzmann Emulator results. Comparison NLL and W2 metrics of Boltzmann Emulators across 5 runs (#= indicates standard deviation). GVP Endpoint emulator captures the energy and torsional target distribution the best. The ECNF model provides the best NLL values despite having the worst W metric values, indicating likelihood integration errors for the GVP models. This is also demonstrated with the higher intra-run NLL std deviation values for the GVP models.
29

---

Energy distribution
Free energy projection
0.08
MD reloadd samples
Emulator
KnT
Unakeella samplng
Ventor fildd reweigited
BG, Vector_Sold rewngghted
0.00
KDE
Free energy
= c.
Energy KBT
Energy distribution
Free energy projection
0.12
MD relsed sumples
17.5
Unicella somplng
Boltzman Enmiator
BG Eulpoint reweighted
BC, Eudppint reweighted
KDE
12.3
Froe energy
0.66
0.6
5.0
Energy KgT
Energy distribution
Free energy projection
MD relexd sumple
knf
Unakedla wampling
0.00
Boltzmam Emmlator
BolsNCE reweigated.
BoltoCCR enppoint remigghtd
0.5
KDE
Free energy
m x
002
-100
-81
Energy
KgT
Figure 8: Energy histograms and free energy projections with confidence intervals for the GVP- WeVVectoor Field (top), GVP-Endpon (center) and BoltzNCE BoltzNCE-Endpoint (bottom) models.
The energy (E-W;) and torsion angle (T-W2) Wasserstein-2 distances quantify the discrepancy between the distributions of energies and torsional angles of generated conformers and those in the dataset. The results show that while the T-W2 distance remains relatively consistent across all methods, the GVP models capture the dataset's energy distribution better, with the Endpoint model showing the best performance (Figure indicating that is very good Boltzmann Emulator this dataset.
The ECNF and GVP-VF models are comparable on the Negative Log Likelihood (NLL) metric, whereas the GVP-EP model yields the worst values. important to note, however, that the endpoint vector field (Eq. 9) diverges at time-point 1. Consequently, the likelihoods for the GVP-EP model were evaluated starting from a later time point
1 Furthermore, the divergence at can lead to inaccurate likelihood estimates due to instability in the ODE integration. The standard deviation of NLL values within each run is also reported, and the large variance observed for the GVP-EP model further highlights the potential unreliability of its likelihood computations.
F.2 Boltzmann Generator Results
Energy histograms and free energy projections for GVP Vector Field, GVP Endpoint, and BoltzNCE Endpoint methods are shown Figure 8. The free energy values and energy histograms match up best with the BoltzNCE Endpoint method.
F33 EBM in comparison the Hutchinson trace estimator
Likelihoods predicted by the EBM are directly compared the ground truth likelihoods obtained from applying the change of variable equation the flow matching generated samples. Note
30

---

EBM Prediction Jacobian Trace Integral
Hutchinson Estimate vs Jacobian Trace Integral
Jacoblen Trace Integral
Jacabion Traan Integral
EBM Prediction
Hutchinson Estimate
16
Figure 9: Likelihood Calculation scatter plots: log likelihoods estimated by the EBM and Hutchin- son estimator vs ground truth estimates from the continuous change variable equation. EBM predicts likelihoods at high level of accuracy.
Metric
EBM
Hutchinson call
Hutchinson calls
Hutchinson 4 calls
Hutchinson calls
Spearman (1)
0.93
0.87 8 min
0.88
0.88
0.87 23 min
Time (1)
0.5s
28 min
23 min
Table 7: Likelihood estimation results: Correlation between likelihood estimation methods and exact likelihoods. EBM performs best.
that the likelihoods estimated by the EBM only estimate the exact likelihoods up to constant. The Hutchinson trace estimator is also implemented as a comparative benchmark. The Spearman correlation and time required to estimate likelihoods for 10,000 samples is reported in table 7.
The EBM output exhibits strong agreement with the exact likelihoods. Furthermore, the EBM outperforms the Hutchinson estimator, achieving better correlation while being over two orders of magniude faster at inference on 10,000 samples. Interestingly, increasing the number of estimator calls in the Hutchinson method does not improve its correlation. Due to the use of an adaptive step-size ODE solver, 4 and 8 estimator calls are actually faster than 2 calls.
F,4 Coupling Function Benchmark
To evaluate the effect of different coupling functions on EBM training, we compare independent coupling mini-batch OT coupling on ADP. The results are presented in Table The results indicate that both coupling functions can be used to train the EBM and both achieve similar performance.
Method
AF Error
E-W₂
T-W2
Independent Coupling OT Coupling
0.02 0.13
0.27
0.02
0.57
0.00
0.03 + 0.12
0.23 0.04
0.56 + 0.005
Table 8: Coupling Benchmark. Coupling functions test to train the EBM model on ADP. Both cou- pling functions provide similar performance, indicating that the training algorithms are independent of coupling functions.
G Dipeptides generalizability results
Quantitative results on the test systems of dipeptides are reported in Table 9, Representative energy histograms and free energy surfaces for the dipeptides shown in Figure 10. In general, BoltzNCE is able to approximate the right distribution within acceptable error limits at 6x compute time improvement over MD simulations.
31

---

Method
Dipeptide
AF Error
&-W2
T-W2
MD
AC
0.048
0.192
0.213
TBG-ECNF
AC
0.009
0.202
0.290
BoltzNCE
AC
0.356
0.431
0.318
MD
ET 0.069
0.152
0.182
TBG-ECNF
ET
0.187
0,492
0.341
BoltzNCE
ET
0.222
1.329
0.280
MD
GN 0.545
0.122
0.171
TBG-ECNF
GN
0.355
0.198
0.267
BoltzNCE
GN
0.502
1.374
0.296
MD
IM
0.504
0.142
0.269
TBG-ECNF
IM
0.026
0.430
0.362
BoltzNCE
IM
0.688
0.459
0.454
MD
KS
0.044
0.138
0.232
TBG-ECNF
KS
0.133
0.474
0.434
BoltzNCE
KS
0.477
0.419
0.626
MD
NY
0.0003
0.198
0.229
TBG-ECNF
NY
0.034
0.491
0.425
BoltzNCE
NY
0.090
1.293
0.591
MD
NF
0.072
0.221
0.214
TBG-ECNI
NF
0.102
0.275
0.301
BoltzNCE
NF
0.701
2.318
0.536
Table 9: Dipeptides results. Quantitative results of different methods all 7 dipeptide systems. BoltzNCE delivers acceptable performance while offering a substantial time advantage.
32

---

610
Famry Distribttonn
Free eDergY projection
Free merry projection TICA
816
Sockito Tubrevss
Raedito Bktodtt thwrigtad
iin [ReT"
Ditaee Redda
MD
TubokER Boodgfaed
-122
BaccK Booigtt.
Free carcggy
72
Free
010
010
Energy / KgT
405
TCCO
Energ Distribution
Free oneegy projection TICA
8.66
Free energ projection y
930
MD Sanbases
Bdoanas Busase
Fixe coneg kaT
Free enargs kgT
Exekite
Beseit Larvegied
BaiaiCE Rowalial
880
011
0.0
Easrgy knl
TICI
Energy Distribation
eneegy projection yp
Free eneryy projection TICA
8.07
816
Bobonzas Bariase
100 Staaitke
Free energy Kpd
610
010
Eneryy
TICI
Faerpy Distribtton
energy projection yp
Free eaergy projection TICA
8.65
MD Strkktee
Tebresss Tesiov
Frve apeegy R
Frne energy kyj
Bac Bevadaol
Bhwiggted
810
BubsOOC Reenggoo
0.00
Energy kyT
TICO
Energy Distribation
Fine enery projection
Pie energy projection TICA
8.10
Siaaitio BubeXCR Revigtad
Dive eaergy InT
Fiee energy Kal
Handate
BabdkC Brwig
Bateeeet Hovstated
80
010
Energy Kn
TICD
Energy Distribution
Free jrojection TICA
8.66
energy projection
Five coungy kaT
Elce cunrgy Kajd
Telmauir Feckate
814
BolNNC
Balubc Rowaggoo
BelaNCE Bowalial
800
810
810
2000
2.0
100
-05
100
155
120
15
Emergy
knT
T1CO
Figure 10: Qualitative results on dipeptides. Energy histograms and free energy projections along the &p dihedral angle and the first TICA component for test dipeptides. order from top to bottom, the figures represent results the following dipeptides: AC, ET, GN, IM, KS, NE. In all cases, BoltzNCE achieves good approximations of the energy distribution and free energy surfaces.
33


Reviews:
[{"review_id": "mgWE3FXQw1", "reviewer": "Reviewer_jMkc", "summary": "This paper proposes a scalable, simulation-free framework for learning sample likelihoods from energy-based models to accurately model the Boltzmann distribution. By combining score matching and noise contrastive estimation with stochastic interpolants, the method avoids the costly exact likelihood computations that are usually impractical for large molecular systems. This approach enables efficient estimation of free energy differences and recovery of the Boltzmann distribution, delivering comparable accuracy with substantial computational speedup.", "questions": "- Is it possible to just use EBM for sampling, since it already provide a score of the model density.\n- In your dipeptide experiments, how does the model perform if the emulator training dataset are more biased?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\nThe authors combine stochastic interpolants with the Boltzmann Generator and propose a contrastive learning approach to train the energy-based model (EBM), which circumvents the need to compute the intractable integral in the normalizing constant. The method is evaluated on both 2D toy examples and dipeptide data, showing performance improvements over ECNF.\n\n**Weakness**\n- The method requires training an additional EBM, whereas other approaches typically train only a single CNF model.\n- The architecture choices are complex, yet architectural ablations are missing. A discussion is needed to clarify that performance gains are not merely due to increased parameter counts.\n- The endpoint models diverge near t = 0, requiring manual truncation of the integration. This instability reduces the effectiveness of likelihood-based sampling.\n- If the Boltzmann  emulator  is trained on biased data with missing modes, the EBM is unable to recover those missing modes.", "was_revised": true, "final_justification": "The additional experiment with alternative architectures has address my concerns that the performance gain is from methods itsefl", "rebuttal": "We thank the reviewer for their thoughtful review and questions, and we are pleased that they found our approach practical for scaling to larger systems. We would also like to mention that in addition to results on alanine dipeptide, we now include new experiments demonstrating transferability. For transferability we train on a set of 200 dipeptides and infer on 7 dipeptide systems not present in the training set and show accurate free energy estimates and low energy/torus w-2 scores. Visually our energy histograms have nearly perfect overlap and free energy projections have great agreement with the test set. For more details on transferability, refer to our reply to reviewer vcq2.\n\n> The method requires training an additional EBM, whereas other approaches typically train only a single CNF model.\n\nWe acknowledge that our method requires training two models instead of one. However, we believe that this is a key strength, rather than a limitation. Flow matching models are good for sampling but hard for likelihood evaluations. EBMs on the other hand provide easy access to likelihoods but are not great for sampling. By combining the two types of model we leverage the best of both worlds.\n\nMoreover, our method design significantly reduces inference-time costs, as evaluating the likelihood of a sample requires only a forward pass through the EBM. This is beneficial as it reduces the constraints on the emulator/neural sampler as it gets rid of the requirement of invertibility for Boltzmann generation. Furthermore, in transferability scenarios, i.e. inferring on systems not present in the training set, the EBM requires only a bit of fine tuning for significant inference time gains. The two model approach is also advantageous when scaling to higher dimensional systems, where training an additional neural network is less computationally expensive than calculating Jacobian trace integrals.\n\n> The architecture choices are complex, yet architectural ablations are missing. A discussion is needed to clarify that performance gains are not merely due to increased parameter counts.\n\nWe acknowledge that our flow matching architectures are slightly more complex than the ECNF model (to account for multiple vector features), however, the GVP models our study has significantly fewer parameters (108,933) than the benchmark ECNF (147,599). We now report the parameter counts in our updated paper.\n\n> The endpoint models diverge near t = 0, requiring manual truncation of the integration. This instability reduces the effectiveness of likelihood-based sampling.\n\nWe agree with the reviewer that the endpoint model is not well-suited for likelihood-based sampling. However, we do not view this as a limitation of our method. Our approach leverages an EBM to estimate likelihoods, thereby avoiding the computational burden of the continuous change-of-variable equation. A key insight from our method is that introducing a secondary EBM for likelihood estimation relaxes the design constraints on the Emulator models. Since efficient likelihood evaluation is no longer required for Boltzmann generation, the Emulator can be trained using the endpoint objective. We demonstrate this in our experiments on alanine dipeptide and emphasize that this is the central takeaway from our use of the endpoint objective.\n\n> If the Boltzmann emulator is trained on biased data with missing modes, the EBM is unable to recover those missing modes.\n\nWe agree with the reviewer regarding the issue of missing modes; however, this is a limitation common to most generative models and not specific to our method. The emulator needs to be able to sample sufficiently from the modes of the dataset for effective reweighting with the EBM. Notably, the authors of the original ECNF paper[1] also reported needing to bias the data toward underrepresented modes in order to obtain accurate free energy differences. Our experimental design follows the same approach.\n\n> Is it possible to just use EBM for sampling, since it already provide a score of the model density.\n\nThis is an interesting idea worth exploring further. However, we emphasize that the sample distribution must approximately match the likelihoods predicted by the EBM to achieve the Boltzmann. In practice, achieving this with an EBM typically would require running Langevin dynamics for large amounts of time, which may be computationally infeasible.\n\n> In your dipeptide experiments, how does the model perform if the emulator training dataset are more biased?\n\nWe biased our dataset to account for the underrepresented mode along the positive $\\varphi$ angle. This step is necessary to obtain accurate free energy differences method. We adopted the same method for biasing as employed in the original ECNF work[1] to ensure a fair and equitable comparison between the two methods. Again, we acknowledge the need for our Emulator model to be able to cover all the modes, but this is a limitation that is common to most generative models and is discussed in previous work as well[1].\n\n[1]Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36:59886–59910, 2023.\n\nOverall, we thank the reviewer for their thoughtful comments and hope they will consider increasing their score in light of the new experimental results and our responses to their concerns."}, {"review_id": "4Z7ZPUz4N3", "reviewer": "Reviewer_WnHv", "summary": "The paper addresses the accuracy of learned generative models for Boltzmann distributions (Boltzmann generators of short). When these models are not perfect, the estimation of downstream quantities of interest is improved by weighting generated samples with the ratio of true and estimated probabilities. However, the calculation of the estimated probability is expensive for the currently best models (continuous normalizing flows). Therefore, the authors propose to train an additional energy-based model to calculate these probabilities more efficiently. The key innovation is a simulation-free and thus fast training algorithm for the additional network.", "questions": "* Can you add proper derivations of key formula?\n* Can you add a table saying exactly which algorithms are used in which experiment?\n* Can you provide more extensive experiments, or at least demonstrate the agreement between the EBM and the generative model?", "limitations": "The authors acknowledge that experiments restricted to the alanine dipeptide have limited significance. They also mention the limited accuracy of the probability estimates from the EMB, although it is unclear why this has not been investigated in more detail (the item 2. under \"Weaknesses\").", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "The proposed idea of simulation-free training for an auxiliary energy-based model for efficient reweighting is very good (although the idea was in the air, given the field's recent evolution). The experiments demonstrate that the new method achieves accuracy competitive with a continuous normalizing flow at 1/50 ... 1/100 of the inference time, which is highly practical relevant.\n\nHowever, the paper has several major shortcomings:\n\n1. Although there is a lot of theory, the presentation completely lacks derivations of the formulas, nor have sections 3 and 4 any citations. Thus, it is next to impossible to check if the results are correct. In particular, I have concerns with the coupling $C(x_0, x_1)$ between data and latent samples, which must be a joint distribution whose marginals $p(x_0) = \\int C(x_0, x_1) dx_1$ and $q(x_1)=\\int C(x_0, x_1) dx_0$ are the data and latent distribution respectively. While it is known that the objectives eq. (9), (10) and (16) converge for any valid coupling, this is probably not true for the InfoNCE objective (15). I think that the coupling here must be the optimal coupling resulting from Schrödinger bridge matching, as derived in Shi et al. \"Diffusion Schrödinger bridge matching\", NeurIPS 2023. The authors mention briefly (lines 183-184) that they approximate the coupling by minibatch optimal transport via the Hungarian algorithm, but they fail to discuss to what degree this approximation fulfills the requirements for the coupling, nor what these requirements are in the first place. \n\n2. In the experiments, it is hard to understand exactly which sub-algorithms according to which formulas are included or excluded in the different model variants (\"GVP Vector Field\", \"GVP Endpoint\", \"BoltzNCE reweighted\", \"BoltzNCE Vector Field\", and \"BoltzNCE Endpoint\") mentioned in the diagrams and tables. The information can probably be deduced somehow from the text, but should be made obvious in a table or similar to allow for unambiguous replication of the experiments. Moreover, the \"Umbrella Sampling\" method, which apparently serves as a ground truth estimate, is entirely undefined. It is also unclear why reweighting was not included in the experiment in figure 4 and table 1 and how other methods fail to estimate the free energy in figure 5.\n\n3. The experiments are somewhat weak, as they only cover the alanine dipeptide (although table 1 does not even say what the data are). For a NeurIPS publication, I would expect a more extensive set of experiments. Moreover, when samples are generated with one model, but weighted by another, it must be ensured that both models agree on the generative probability. Since the probabilities of the generative model can in principle be calculated via the continuous change-of-variables formula (even if this is too expensive for most practical applications), this comparison can and should be conducted.", "was_revised": true, "final_justification": "I thank the authors for addressing my concerns and would like to encourage them to implement these improvements (especially the ones concerning the theory) as far as possible in the final version. Moving a non-essential section to the appendix to make room for more detailed explanations and derivations of the essential parts is a good idea. I've raised my score to \"Accept\".", "rebuttal": "We thank the reviewer for their feedback and we are happy to see they agree that our method provides orders of magnitude speed up with competitive accuracy. \n\n**1. Experimental Results** \n\n\n> The experiments are somewhat weak, as they only cover the alanine dipeptide\n\n\nWe agree that additional experimental validation would benefit our paper and provide stronger proof that our method works. To this order, we show that our method is **transferable** by training on 200 dipeptides and testing on 7 unseen dipeptides. With a bit of fine tuning on the unseen dipeptide, we achieve good performance while attaining a significant time advantage over MD and benchmark ECNF methods. Visually our energy histograms have nearly perfect overlap and free energy projections have great agreement with the test set. For more details on transferability, refer to our reply to reviewer vcq2. These results are discussed in depth in our reply to reviewer vcq2 (reviewer 3).\n\n>when samples are generated with one model, but weighted by another, it must be ensured that both models agree on the generative probability\n\nWe agree! To showcase this, we calculated Spearman correlation coefficients between the EBM output and exact likelihood as calculated by the Jacobian trace integral for the alanine dipeptide syteme. Our final paper version will include the scatterplots but we report just the quantative value below. \n\nAs shown in Table 1, the EBM output exhibits strong agreement with the flow-matching likelihoods. Moreover, we found that our EBM outperforms the Hutchinson estimator, achieving better correlation while being over two orders of magnitude faster at inference on 10,000 samples. Interestingly, increasing the number of estimator calls in the Hutchinson method does not improve its correlation. Due to the use of an adaptive step-size ODE solver, 4 and 8 estimator calls are actually faster than 2 calls. \n \n| Metric     | EBM      | Hutchinson (1 estimator call) | Hutchinson (2 estimator calls) | Hutchinson (4 estimator calls) | Hutchinson (8 estimator calls) |\n|:-----------|:--------:|:-----------------------------:|:-----------------------------:|:-----------------------------:|:-----------------------------:|\n| Spearman R | **0.93** | 0.87                          | 0.88                          | 0.88                          | 0.87                          |\n| Time       | **0.5 s**| 8 min                         | 28 min                        | 23 min                        | 23 min                        |\n\n**2. Formulas and equations** \n>lacks derivations of the formulas, nor have sections 3 and 4 any citations\n\nThank you for the comments. We have now included multiple new citations throughout the methods and results sections so that the reader can follow where some parts of the method are from. Some of the main papers we cite heavily include:\n\n* For stochastic interpolant equations\n    * Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pages 23–40. Springer, 2024.368\n\n* For Boltzmann Generator formulations and equations \n    * Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural  Information Processing Systems, 36:59886–59910, 2023\n\n* For Endpoint Objective \n    * Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402.04845, 2024\n* For mini-batch OT\n    * Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.00482, 2023.459\n    \nWe have also included citations to other seminal score matching and flow matching works.\n\n**3. Coupling concerns**\n> The coupling here[in the EBM] must be the optimal coupling resulting from Schrödinger bridge matching\n\nThank you for your comment. To test your hypothesis, we trained versions of the EBM model with both OT and independant coupling. We found that the models converged in both cases and we saw only minor changes in model performance. We report the performance of both methods in the table below. With this we conclude that OT coupling may not be necessary for model convergence.\n\n\n\n|   | $\\Delta F$ Error | E W-2 | T W-2 |\n|----------|----------|----------|----------|\n|OT Coupling          |     0.03   ±0.12      |  0.23±0.04       |      0.56±0.005    |\n|Independent Coupling|     0.02 ±0.13   |   0.27±0.02       |       0.57±0.00   |\n\n> The authors mention briefly (lines 183-184) that they approximate the coupling by minibatch optimal transport via the Hungarian algorithm, but they fail to discuss to what degree this approximation fulfills the requirements for the coupling, nor what these requirements are in the first place.\n\nThank you for your comment. As shown in [1], minibatch OT approximates entropic regularize OT coupling of the two distributions. The hungarian algorithm is of similar flavor and is mainly used to for scalability to larger batch sizes. \n\n**4. Generator and Emulator Clarifications**\n>  it is hard to understand exactly which sub-algorithms according to which formulas are included or excluded in the different model variants\n\nWe apologize for the confusion, we added the following descriptions more explicitly in the tables and figures in our newest version in a similar manner as the table below.\n\n| Method                            | FM Objective | Likelihood Estimation |\n| --------------------------------- | ------------ | --------------------- |\n| GVP Vector field                  | Vector Field | Jac-trace integral    |\n| GVP Endpoint                      | Endpoint     | Jac-trace integral    |\n| BoltzNCE Vector field             | Vector Field | EBM forward pass      |\n| BoltzNCE Endpoint                 | Endpoint     | EBM forward pass      |\n\nWe have now also added a description of the umbrella sampling method in our paper.\n\n> It is also unclear why reweighting was not included in the experiment in figure 4 and table 1\n\nFigure 4 and Table 1 aimed to compare the performance of the *emulators* not the *generators*. ie. we wanted to show the impact of the choice of objective functions without the influence of reweighting. Furthermore, our results showed that a better emulator does not neccesarily create a better generator due of the design constraints that generators have to follow. However, we agree that this section may confuse readers and detract from the main point of the paper and have moved this discussion to the appendix. We have, however, now added additional W-2 results, before and after reweighting on the Boltzmann Generators comparison table. See further detail on this in our response to reviewer vcq2.\n\n> How other methods fail to estimate the free energy in figure 5\n\nThank you for the question. We have provided free energy projections for the other methods in the supplementary material. We hypothesize that the other methods have slightly more inaccurate free energy projection predictions due to inaccurate likelihood estimates that can occur due to error propogation during ODE integration. This is especially prevelant when one tries to integrate likelihoods on CNF models trained with the endpoint objective.\n\nOverall we thank the reviewer for their valuable comments, we hope that they see our novel experiments on transferability provides even stronger proof that our method unlocks the potential for a truely boltzmann generator. We also address concerns regarding couplings for the EBM with emperical results and provide additional clarification on methods and equations. With this, we hope the reviewer considers improving their score on our paper.\n\n\n[1] Tong, Alexander, et al. \"Simulation-free schrodinger bridges via score and flow matching.\" arXiv preprint arXiv:2307.03672 (2023)."}, {"review_id": "TEbV8YlPWN", "reviewer": "Reviewer_vcq2", "summary": "The paper introduces BoltzNCE, a scalable and simulation-free method for training energy-based models (EBMs) to estimate likelihoods for reweighting samples from a generative model toward a target distribution. The approach combines noise contrastive estimation and score matching within the stochastic interpolant framework, avoiding the expensive Jacobian trace computations required by normalizing flows. The learned EBM density enables accurate reweighting of generated samples to recover the Boltzmann distribution. The method is validated on toy datasets and the alanine dipeptide molecule, achieving performance comparable to exact-likelihood normalizing flows while offering orders-of-magnitude faster inference.", "questions": "1. Echoing the above concerns, how does the EBM perform when across the metrics when reweighting samples for alanine dipeptide and for larger systems such as the tetrapeptide and hexapeptide variants of alanine (AL4 & AL6), as well as for free energy difference? And, how does the EBM perform when transferring across different systems? \n\n2. It is not immediately clear to me how the statement made on line 54 relates to the density chasm problem. The statement on line 54 is “We tackle this issue by annealing between a simple noise distribution and the data distribution using stochastic interpolants” where the issue here is the density chasm problem mentioned in the previous paragraph. In this case, NCE distinguishes between a sample drawn from the interpolation and a negative sample drawn from a narrow Gaussian centered at that sample. Would the authors please elaborate more on this?\n\n3. Algorithm 1 in appendix B states the EBM uses data samples but the method states the samples at t=0 comes from the generative model. Can the authors clarify this?\n\n4. From what I understand, the framework currently operates with a generative model and EBM being trained separately. I am curious to see how does the generative model perform if the score and vector field are parameterized using \\nabla E_\\theta from the EBM?\n\nShould the authors address my current concerns (specifically, the additional experiments with scalability to AL4 / AL6 and transferability mentioned in the first point), I would consider raising my score.\n\n\nNot questions but typos found:\n- line 146: enery -> energy\n- appendix eq 21: the dots should be \\cdot for t \\cdot x_1 + (1 - t) \\cdot x_0", "limitations": "Yes. No potential negative societal impact is involved.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n1. The paper is clearly written and motivated.\n2. The proposed method (Vector field) achieves comparable performance to normalizing flows on free energy difference estimation while being significantly faster, encouraging the possibility of scaling to larger systems.\n3. The authors provide a simulation-free method to train EBMs through the stochastic interpolant framework.\n\n**Weaknesses**\n1. While the method demonstrates accurate likelihood estimation on alanine dipeptide, evaluating it on larger systems - such as tetrapeptides or hexapeptides - would better showcase its scalability and inference efficiency. Additionally, for alanine dipeptide, in table 1, reporting energy W2, torus W2, etc for both the proposal and the reweighted samples would help support the effectiveness of the EBM.\n2. Although the GVP-based models support the ability to transfer across systems, the authors do not explore transferability of the EBM across systems. Demonstrating the method’s ability to transfer across molecular systems would strengthen its practicality.", "was_revised": true, "final_justification": "The authors have addressed my questions and concerns in the rebuttal, and I appreciate the additional experiments provided.\n\nI will raise my score to a 4, as it better reflects the paper’s strengthened experimental design and results. While my concerns have been addressed, the work does not introduce a fundamentally novel methodology or conceptual insight. As such, I do not find grounds for awarding a 5, but I still recommend acceptance.", "rebuttal": "We thank reviewer vcq2 for their insights and are glad that they find our work well motivated. We are happy to see that the reviewer agrees with us that our EBM method achieves performance comparable to exact-likelihood normalizing flows while offering orders-of-magnitude faster inference.\n   \n\n**2. Additional Experiments**\n\nAs suggested by the reviewers, we conducted additional experiments to evaluate the transferability of our method. Specifically, we train our model on 200 dipeptides and test it on 7 unseen dipeptides, closely following the experimental setup described in [1]. Similar to our ADP experiments, we bias the training set to ensure adequate coverage of all relevant modes along the $\\varphi$ angle.\n\nTo establish a molecular dynamics baseline, we run 1 μs simulations and compute E-W2, T-W2, and $\\Delta F$ estimation errors with respect to the test set provided in [1]. We also train the model from [1] on our biased dataset. However, inference with this model requires several days of computation, so we were unable to obtain results for it within the rebuttal period. We will include these results in the camera-ready version. \n\nDuring inference, we fine-tune our pre-trained EBM using samples generated for each unseen dipeptide. Additionally, we exclude the top 0.2% of importance weights during reweighting, following the approach introduced in [2], which significantly improves the stability of the reweighting process. Our results show that the method successfully transfers across systems while offering a substantial inference-time advantage over MD and ECNF baselines. The performance of our method compared to the MD baseline across the 7 dipeptides is summarized in the table below.\n\n\n| Metric              | Energies W2   | Angles W2    | $\\varphi$ $\\Delta F$ Error  |Inference time|\n|:--------------------|--------------:|-------------:|-------------------:|-------------------:|\n|  MD Baseline | 0.17 ± 0.03   | 0.22 ± 0.03  | 0.18 ± 0.22 | 1 day |       |\n|  BoltzNCE | 0.94 ± 0.49   | 0.54 ± 0.12  | 0.56 ± 0.33 | 3.7 hours |       |\n\n\nAs shown in the table above, our method achieves a 6× speedup over conventional molecular dynamics while producing W2 and $\\Delta F$ metrics that closely match the ground truth. Visually, the reweighted energy histograms show near-perfect overlap, and the predicted free energy projections align closely with the ground truth. These results support the conclusion that our method generalizes well and is transferable for inference beyond the training data.\n\n**Reweighting Improves Results** \n> \"reporting energy W2, torus W2, etc for both the proposal and the reweighted samples would help support the effectiveness of the EBM\"\n\nThis is a very interesting idea and provides a great opportunity to show case the strength of our EBM model. We report energy W2 and torus W2 for both the proposal and reweighted samples in the table below. In all cases, we see that reweighting significantly improves generated distributions. The EBM (BoltzNCE Vector field) is able to efficiently reweight samples from the GVP vector field model and beats the ECNF baseline on these metrics\n\n\n| File                          | Proposal E-W2 | Reweighted E-W2 | Proposal T-W2 | Reweighted T-W2 |\n|:-----------------------------|-----------------------------:|---------------------------------:|------------------------------:|---------------------------------:|\n| ECNF        | 8.08±0.56                    | 0.37±0.02                        | 1.10±0.01                     | 0.59±0.00                        |\n|GVP Endpoint           | **6.19±1.08**                    | 2.88±0.01                        | 1.12±0.01                     | 0.58±0.01                        |\n|GVP Vector field                 | 7.20±0.13                    | 0.46±0.05                        | **1.09±0.01**                   | 0.60±0.00                        |\n|BoltzNCE Endpoint        | 6.24±0.52                    | 2.78±0.04                        | 1.12±0.01                     | 0.59±0.01                        |\n| BoltzNCE Vector field         | 7.12±0.15                    | **0.27±0.02**                        | 1.12±0.00                     | **0.57±0.00**                        |\n\n**3. Questions**\n> It is not immediately clear to me how the statement made on line 54 relates to the density chasm problem.\n    \nThank you for the question. The density chasm problem arises when the samples drawn from the 2 distributions are easily distinguishable. Such a case results in vanishing gradients[3] under the NCE objective. Therefore to narrow the gap between the distributions that are compared in the contrastive loss we employ annealing through the use of stohcastic interpolant and ensure that the negative samples are in closely proximity to the positive sample. \n\n> Algorithm 1 in appendix B states the EBM uses data samples but the method states the samples at t=0 comes from the generative model. \n\nGood catch! We have fixed it to reflect that we train the EBM on generated samples.\n\n>I am curious to see how does the generative model perform if the score and vector field are parameterized using $\\nabla E_\\theta$ from the EBM?\n    \nThis is an interesting point that is worth exploring further. Our method takes advantage of a good sampler (Flow Matching) and a good likelihood estimator (EBM). We leverage both these properties to get the best of both worlds. For the scope of this rebuttal, however, we mainly focused on experiments related to transferability and leave sampling with $\\nabla E_\\theta$ to future work.\n\n\nIn conclusion, we thank the reviewer for their great suggestions of extra experiments. We believe that the experiments they proposed have truly strengthend our paper and we hope that they see our novel experiments on transferability as even stronger proof of our method. We also show that our EBM is well correlated with likelihoods from the Jacobian trace integral with orders of magnitude speed up. With this, we hope the reviewer considers improving their score on our paper.\n\n[1]Leon Klein and Frank Noé. Transferable boltzmann generators. arXiv preprint arXiv:2406.14426, 2024.\n \n[2]Tan, Charlie B., Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. \"Scalable equilibrium sampling with sequential boltzmann generators.\" arXiv preprint arXiv:2502.18462 (2025).\n\n[3]Liu, B., Rosenfeld, E., Ravikumar, P., & Risteski, A. (2021). Analyzing and improving the optimization landscape of noise-contrastive estimation. arXiv preprint arXiv:2110.11271."}, {"review_id": "CpSYgR4poH", "reviewer": "Reviewer_UPgQ", "summary": "The paper introduces BoltzNCE, a way to train an EBM and continuous normalizing flow simultaneously. One caveat of cnfs is that although they can be trained simulation free tend to have unstable/inefficient likelihood computations, making them infeasible for reweighting. The authors propose to combine a FM loss together with training a family of time-parameterized EBMs that should approximate the distribution of the cnfs. Luckily, this can be done using score matching and Info-NCE, which both do not require solving an ode. The approach is then compared for free energy calculations on alanine dipeptide.", "questions": "1) The INFO NCE loss is not clear: while I am family with the idea of NCE losses in EBMs, i do not understand why that is equivalent to the NLL (is that somehow like 1d integration??). please add some more explanation on that. why is it reasonable to sample negative points around t with a small variance?\n\n2) In the contributions you write \"We show that the learned likelihoods can be used to reweight conformations to match the Boltzmann distribution. To the best of our knowledge, this is the first method to recover the Boltzmann distribution without requiring exact likelihood computations.\" I do not understand this, there have been a lot of papers recovering Boltzmann dists using approximate likelihoods and CNFs?\n\n3) I am not super familiar with molecular simulations, so please explain for me (not necessarily in the paper), why the free energy of a certain angle is such a good proxy for likelihood evals.\n\n4) Does training the EBM on the side deteriorate the sampling (i.e, W_2) performance of the FM method?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Stengths: \n\n1) The approach is to the best of my knowledge new. While many papers have tried training EBMs  on their own, guiding EBMs by a CNF seems smart. \n\n2) The paper is also (up to a few problems that can be misunderstandings on my end) very clearly written, and the approach simple to understand. \n\n3) The motivation is clear and the experiments are encouraging. \n\n4) The ablations on why both the score and NCE loss are needed are very neat. \n\nWeaknesses: \n\n1) I would like to see experimental evaluation on a simple imaging dataset. One can easily train a FM model on MNIST, and then see if one can reuse the EBM from your approach as a regularizer (i.e. prior) in a simple inverse problem such as deblurring or inpainting. This would greatly strengthen the paper, as the usefulness would span to a different subfield and show that the method can also work there. Given that this is a method/experimental paper, I would expect more comparisons. \n\n2) I had a lot of trouble understanding some bits of the paper: more precisely the INFO NCE loss needs more clarity. \n\n3) As a baseline one could also consider distilling the FM model, i.e., train a model which has similar likelihoods/odes but with much fewer steps. This would essentially serve the same purpose as your EBM. A discussion of this (if you do not consider it a strong baseline) or a small experiment would also strengthen the paper. \n\n4) While the paper makes a very good experimental point, that both the NCE and the SM loss are needed, I would also like some (heuristic) explanation why that is. \n\n5) Citations are unfortunately very sparse in the text. While every paper used is cited in the intro/related work, I would also like to see that when introducing the methods so it is very clear where the method is from and where one can read it up. Otherwise, one could confuse some of those as your contributions.", "was_revised": true, "final_justification": "My questions are answered, this paper deserves an acceptance imo.", "rebuttal": "We thank the reviewer for their time and their helpful comments. We are glad that the reviewer found our paper novel and well written. We would also like to mention that in addition to results on alanine dipeptide, we now include new experiments demonstrating transferability For transferability we train on a set of dipeptides and infer on dipeptide systems not present in the training set. With a bit of fine-tuning of the EBM on the unseen dipeptide, we obtain accurate free energy estimates and low energy/torus W-2 scores. Visually our energy histograms have nearly perfect overlap and free energy projections have great agreement with the test set.\n\n**Imaging and other experiments:** \n\n>I would like to see experimental evaluation on a simple imaging dataset. \n\nWe thank the reviewer for their suggestion on solving the inverse problems on images. We definitely think this could be an interesting direction of research and experiment/algorithm design with an EBM deserves further thought. However, for the rebuttal period, we had to focus on other requested experiments on molecular systems (transferability and scalability) as efficiently solving the Boltzmann distribution is the primary purpose of this work. Further details on this can be seen in our reply to reviewer 3 vcq2.\n\n**InfoNCE objective:** \n\n>The INFO NCE loss is not clear...\n\nThank you for your comments and questions, we have improved the writeup for the InfoNCE objective in our manuscript. The InfoNCE objective is an orthogonal objective function that maximizes the conditional likelihood $\\rho_\\theta(t|x_t)$.\n\n$$\\rho_\\theta(t|x) = \\frac{\\rho_\\theta(t,x_t)}{\\rho_\\theta(x_t)} = \\frac{\\rho_\\theta(t,x_t)}{\\int\\rho_\\theta(t',x_t)dt'} = \\frac{\\exp(E_\\theta(t,x_t))}{\\int \\exp(E_\\theta(t',x_t))dt'}$$\n\nwhere $E_\\theta$ is the energy based model. This leads us to minimize the negative log likelihood of the conditional:\n\n$$\\mathcal{L_{\\text{NLL}}}(\\theta) = E_{x_t}\\left[-\\log \\frac{\\exp(E_\\theta(t,x_t))}{\\int \\exp(E_\\theta(t',x_t))dt'}\\right]$$\n\nThe intractable integral in the denominator is approximated by appropriately sampling a set of *negative* time points $t'$ resulting in the InfoNCE loss:\n\n$$\\mathcal{L_{\\text{InfoNCE}}}(\\theta) = E_{x_t}\\left[ -\\log \\frac{\\exp(E_\\theta(t,x_t))} {\\sum_{t' \\in \\{\\tilde{t}\\} \\cup t } {\\exp(E_\\theta(t',x_t))}}\\right]$$\n\n>please add some more explanation on that. why is it reasonable to sample negative points around t with a small variance?\n\nIn practice, on the interpolant process, only time points $t'$ close to $t$ will have a non-negligible likelihood $\\rho_\\theta(t', x_t)$, since $t$ also determines the amount of noise added to the clean sample. Therefore, it is reasonable to sample $t'$ using a distribution with small variance centered around $t$.\n\n> While the paper makes a very good experimental point, that both the NCE and the SM loss are needed, I would also like some (heuristic) explanation why that is.\n\nThank you for your question, we expand on it further below:\n\nThe score matching loss minimizes the Fisher divergence between the model and data distributions, and therefore also implicitly minimizes the KL divergence. However, as demonstrated in our 2D experiments, while the score matching objective effectively identifies the modes of the data distribution, it does not reliably capture their relative weights. The InfoNCE loss, being orthogonal to score matching, serves as a complementary regularizer that helps the model learn more accurate relative weighting between modes.\n\n**Distillation of FM model as a baseline**\n\n> As a baseline one could also consider distilling the FM model, i.e., train a model which has similar likelihoods/odes but with much fewer steps\n\nThank you for the question. Most FM distillation methods we are familiar with fall under the category of consistency models. Due to their algorithmic design, these methods do not support likelihood estimation. The only method we are aware of that potentially reduces the number of integration steps is rectified flows, however, even this approach does not address the challenge of scaling inference with dimensionality. Furthermore, since our FM models are already trained with mini batch OT, rectified flow distillation may not reduce the transport cost by much.\n\n**Citations**\n\n> I would also like to see ['citations'] when introducing the methods\n\nThank you for your comment. We have now added citations throughout the Methods and Results sections, including references related to stochastic interpolants, vector field and score matching objectives, endpoint parameterization, optimal transport, GVP, Graphormer, and more. Additionally, we have included more citations to several seminal works on score matching and flow matching in the main text. \n\n**Free Energy**\n\n>please explain for me (not necessarily in the paper), why the free energy of a certain angle is such a good proxy for likelihood evals.\n\nWe believe the free energy projection of a system is a relevant baseline for the following two reasons:\n\n\n* **It represents a high dimensional integral of probability**: The equation of a free energy projection along a reaction coordinate is given by:\n\n  $$\n  F(r') = -K_BT \\ln \\rho(r')\n  $$\n\n  where $\\rho(r')$ is the probability density of observing the system at position $r'$ along the coordinate,\n\n  $$\n  \\rho(r') = \\frac{1}{Z} \\int_x \\delta(r(x) - r') e^{(-U(x)/K_BT)} \\, dx\n  $$\n  \n  In principle, if we solve the above integral along all points of the reaction coordinate, we solve a ($D-1$) dimensional integral, where $x \\in R^D$. Thus, this serves as a good metric on how well the model matches the ground truth $R^D$ distribtuion.\n  \n  For dipeptides the process of going from the negative $\\varphi$ angle to the positive $\\varphi$ angle is considered a slow process as there are regions of high energy/low probability in between the two. Therefore, this serves as an ideal reaction coordinate to study dipeptide systems.\n  \n* **Domain Relevance:**  Applied studies in the biophysical/ biochemical/ structural biology domain work on elucidating the free energy projection along a *reaction coordinate*. This helps the researchers identify the relative stability of different modes along the coordinate as well as the rate of reaction along the coordinate.\n\n**Other Questions**\n\n> In the contributions you write \"We show that the learned likelihoods can be used to reweight conformations to match the Boltzmann distribution. To the best of our knowledge, this is the first method to recover the Boltzmann distribution without requiring exact likelihood computations.\" I do not understand this, there have been a lot of papers recovering Boltzmann dists using approximate likelihoods and CNFs?\n\nThank you for your question. We do want to emphasize here that, at the time of submission, we were only aware of Boltzmann generator methods that compute likelihoods using the change of variable equation (both discrete and continuous versions). For example, CNFs use the continuous change of variable equation to compute likelihoods. We call these *exact likelihoods* methods. As of now, we are only aware of only one other (concurrent) work [1] that does not use exact likelihoods.  We will revise our claim accordingly and clarify that our method also demonstrates strong performance in transferabile settings.\n\n> Does training the EBM on the side deteriorate the sampling (i.e, W_2) performance of the FM method?\n\nThank you for your question, the EBM does not interfere with the training of the FM model and is only trained to match the likelihoods of samples generated by the FM model. Therefore training it on the side does not affect the performance of the FM model. \n\n\n[1] T. Akhound-Sadegh et al., “Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities,” arXiv.org, 2025. https://arxiv.org/abs/2506.16471 (accessed Jul. 30, 2025)."}]

Author rebuttal:

Meta-review: This paper introduces BoltzNCE, a method that pairs a generative model with an energy-based model to efficiently model Boltzmann distributions for molecular systems. The reviewers collectively recommend accepting the paper which I agree with. The core idea of using an EBM for rapid likelihood estimation is novel and address a significant computational bottleneck. The rebuttal and discussion phase addressed concerns around transferability, strengthening the recommendation for acceptance.
Decision: Accept (poster)