Transferable Boltzmann Generators
Authors: Leon Klein
Frank Noe
Venue: NeurIPS 2024 poster (2024)
Keywords: Boltzmann Generators
Normalizing Flows
Sampling Problem
Flow Matching
Molecular Dynamics
Primary area: machine_learning_for_physical_sciences
TL;DR: We introduce transferable Boltzmann Generators that allow efficient sampling on unseen small peptide systems

Abstract: The generation of equilibrium samples of molecular systems has been a long-standing problem in statistical physics. Boltzmann Generators are a generative machine learning method that addresses this issue by learning a transformation via a normalizing flow from a simple prior distribution to the target Boltzmann distribution of interest. Recently, flow matching has been employed to train Boltzmann Generators for small molecular systems in Cartesian coordinates. We extend this work and propose a first framework for Boltzmann Generators that are transferable across chemical space, such that they predict zero-shot Boltzmann distributions for test molecules without being retraining for these systems. These transferable Boltzmann Generators allow approximate sampling from the target distribution of unseen systems, as well as efficient reweighting to the target Boltzmann distribution. The transferability of the proposed framework is evaluated on dipeptides, where we show that it generalizes efficiently to unseen systems.
Furthermore, we demonstrate that our proposed architecture enhances the efficiency of Boltzmann Generators trained on single molecular systems.

Transferable Boltzmann Generators
Leon Klein Freie Universität Berlin leon. #lein@fu-berlth de
Frank Noé
Microsoft Research A14Science
Freie Universität Berlin
Rice University Saanknoo@@iccoooofffoo com
Abstract
The generation of equilibrium samples of molecular systems has been long- standing problem in statistical physics. Boltzmann Generators are generative machine learning method that addresses this issue by learning a transformation via a normalizing flow from a simple prior distribution to the target Boltzmann distri- bution of interest. Recently, flow matching has been employed to train Boltzmann Generators for small molecular systems Cartesian coordinates. We extend this work and propose a first framework for Boltzmann Generators that are transferable across chemical space, such that they predict zero-shot Boltzmann distributions for test molecules without being retrained for these systems. These transferable Boltzmann Generators allow approximate sampling from the target distribution of unseen systems, well efficient reweighting the target Boltzmann distri- bution. The transferability of the proposed framework is evaluated on dipeptides, where show that it generalizes efficiently unseen systems. Furthermore, we demonstrate that our proposed architecture enhances the efficiency of Boltzmann Generators trained on single molecular systems.
1 Introduction
Generative models have demonstrated remarkable success in the physical sciences, including protein structure prediction 3 generation of de novo molecules 4 45 6 7] and efficiently generating samples from the Boltzmann distribution 8 g [0]. In this work, we will focus on the latter for molecular systems, which represents promising avenue for addressing the sampling problem. The sampling problem refers to the lonn-standing challenge in statistical physics to generate samples from equilibrium Boltzmann distributions u(x) ) с exp (-p(-((x)) (-U(x) (x)/kBT)) where U ((x) is the potential energy of the system, kB the Boltzmann constant, and T the temperature. Traditionally. samples are generated with sequential sampling algorithms such as Markov Chain Monte Carlo and Molecular Dynamics (MD) simulations. However, these algorithms require a significant amount of time to generate uncorrelated samples from the target distribution. This is due the necessity of performing small update steps, in the order of femtoseconds, for stability. This is especially challenging in presence of well-separated metastable states, where transitions are unlikely due to high energy barriers. In recent years, numerous machine learning methods have emerged to address this challenge 8 One such method is the Boltzmann Generators (BG) 8 In this work, we refer to BGs model that allows for the approximate sampling of the Boltzmann distribution of interest and the subsequent reweighting to the unbiased target distribution. the model is only capable of generating approximate samples, which may from subset of the Boltzmann distribution, we refer to them as Boltzmann Emulator Boltzmann Generators transform a typically simple prior distribution into approximation of the target Boltzmann distribution through normalizing flow 12 13 14 15). Once generated, samples can be reweighted to align with the unbiased target
'To the best of our knowledge Bowen Jing introduced the name first.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

The effectiveness of this reweighting hinges on how closely the generated distribution approximates the target. As result, is possible obtain uncorrelated and unbiased samples from the target Boltzmann distribution, potentially achieving significant speed-ups compared to classical MD simulations.
distribution.
There are numerous ways to construct a Boltzmann Generator due to the variety of realizations of normalizing flows available. In this work, we concentrate on continuous normalizing flows (CNFs)
[6 17, as opposed to coupling flows [8). Recently, flow matching [99 20 22 has emerged as
an alternative training method for CNFs. This approach is simulation freel, enabling more efficient training of CNFs.
Thus far, Boltzmann Generators have been limited by the requirement to train them the specific system of interest. This training process demands a significant amount of time, making it difficult to achieve substantial speed-ups over classical MID simulations. Consequently, there is strong desire for a transferable Boltzmann Generator that can be trained on one set of molecules and effectively generalize to another, enabling efficient generation of Boltzmann samples at inference time without the need for retraining. Recently, reliable Boltzmann Generators in Cartesian coordinates for molecules have been introduced [23] 24 paving the way for transferable Boltzmann Generators. This advancement is particularly significant because these models do not rely on molecule-specific internal coordinate representations, which have traditionally made the construction of transferable models challenging.
In this work, we present framework for transferable Boltzmann Generators based CNFs, enabling effective sample generation from previously unseen Boltzmann distributions. Transferable Boltzmann Generators are particularly advantageous as they do not require retraining for similar systems and can be trained on shorter trajectories that may not fully capture all metastable states.
We make the following main contributions:
1. To the best of our knowledge, we introduce the first transferable Boltzmann Generator. We demonstrate its transferabiliity on dipeptides, successfully generating unbiased samples from the Boltzmann distributions of unseen dipeptides.
2. We outline a general framework for training and sampling with transferable Boltzmann Generators based on continuous normalizing flows, which also includes the post-processing of generated samples.
3. We conduct several ablation studies to investigate the effects of different architectures, training set sizes, and biasing of the training data. The results reveal that even small training sets can suffice to train transferable Boltzmann Generators. Additionally, we compare our model with Timewarp (1]) which employs large time steps instead of generating independent samples as our method does.
2
Related work
The initial work on Boltzmann Generators 8 has led to a great deal of subsequent research. The most common application of BGs is generate samples from Boltzmann distributions of molecules
23
26 27
28
29 30, as well as lattice systems 26 B1 32 33 Most BGs and related methods
for molecular systems require system-specifie featurizations such as internal coordinates [8] 34 26
27
35
36
30,
37
38). Only recently, BGs for small molecular systems in Cartesian coordinates
were introduced 23
24 using CNFS and coupling flows, respectively. Equivariant normalizing
played a pivotal role in the success of Boltzmann Generators in Cartesian coordinates. not only for molecular systems. The majority of BGs employ a Gaussian prior distribution, but it is also possible to start from prior distributions close to the target distribution which makes the learning task simpler. However, all previous Boltzmann Generators are not transferable. Arguably, the work of 7 45 represents an exception, as they are able to generate samples from unseen conditional (Boltzmann) distributions torsion space. However, the distribution is conditioned on a single local structure for each molecule, namely fixed bonds and angles. Consequently, in contrast to our work, they are unable to generate samples from the full Boltzmann distribution in Euclidean space. The first transferable deep generative model that was able to generate asymptotically unbiased samples from the Boltzmann distribution is [T.. Instead of generating independent samples, they learn large time steps and combine these with Metropolis-
flows 39
28
23
42
43
36
44

---

Hastings acceptance steps, to ensure asymptotically unbiased samples. However, in contrast to our work, they do not generate uncorrelated samples.
Boltzmann Emulators are analogous to Boltzmann Generators, yet they are not designed to generate unbiased equilibrium samples from the target Boltzmann distribution. Instead, they are intended to generate approximate samples that do not undergo reweighting, Furthermore, the generation of all metastable states may not be a necessary requirement, depending on the system. Boltzmann Emulators do not need to be based on flow models, as they do not aim to do reweighing to the target distribution. They are often similar to Boltzmann Generators and use normalizing flows or diffusion models for the architecture, but due to removing the constraint of sampling the unbiased Boltzmann distribution, they can target significantly larger systems or are transferable. One example is 44). who propose three stage transferable CNF model to learn peptide ensembles. 47 use
flow matching to learn distributions of proteins, while [48] utilize diffusion models. 49 build
transferable Boltzmann Emulator for small molecules. Others aim to additionally also capture the correct dynamics of the molecular systems, such as [50]. who use a diffusion model to predict
transition probabilities. Scaling to larger systems often requires coarse graining [5] 52 47 e.g.
describing amino acids by single bead rather than the individual atoms. However, this approach precludes the possibilityy of reweighting to the Boltzmann distribution.
distinct, though related, learning objective is to generate novel molecular conformations. However, approximations from the Boltzmann distribution are not necessary; it is sufficient to generate a few (or even a single) conformation per molecule. The utilized architectures are once again analogous, as flow and diffusion models are employed 5 53
55 7
3 Boltzmann Generators and Normalizing Flows
Here, we describe Boltzmann Generators and normalizing flows, which are central part of our proposed transferable Boltzmann Generator framework. We follow the notation of (23).
3.1 Boltzmann Generators
Boltzmann Generators (BGs) 8 combine exact likelihood deep generative model and reweight- ing algorithm to reweight the generated distribution to the target Boltzmann distribution. The exact likelihood deep generative model is trained to generate samples from distribution p(x:) that is close to the target Boltzmann distribution p(x). A common choice for the exact likelihood model are normalizing flows.
The Boltzmann Generator can be used to generate unbiased samples by first sampling x ~ ((x) with the exact likelihood model and then computing corresponding importance weights w(x) = (lz) ((x)) for each sample. These allow to reweight generated samples to the target Boltzmann distribution u(x). It is possible to estimate observables of interest (asymptotically unbiased) using the weights wxx) with importance sampling via
(0) Ea- w())
E E [w(x)]
Furthermore, these reweighting weights can be employed assess the efficiency of trained BGs by computing the effective sample size (ESS) with Kish's equation (56). In this work, we will compute the relative ESS, rather than the absolute one, and refer as ESS.
3.2 Continuous Normalizing Flows (CNFs)
Normalizing flows 15 57] are type of deep generative model used learn complex probability densities p(x) by transforming a prior distribution ((x) through an invertible transformation fo RT > "R, resulting in the push-forward distribution p(z))
Continuous Normalizing Flows (CNFs) 16, 17
are a specific kind of normalizing flow. For CNFs, the invertible transformation (x) is defined by the ordinary differential equation
(x) up (t. f((x)) [z) Zo.
(2)
dt
3

---

where (t(t, IR" x [0, R" is time denendendd vector field. The solution to this initial value problem provides the transformation equation
at fz(z) - zo + f.atro (e.56
(3)
with (z) = (x)).. The corresponding change in log density from the prior to the push-forward distribution is described by the continuous change of variable equation
((g)()) = log g(x) - -f. di - ue (L))).
The energy of molecular systems is typically invariant under permutations of interchangeable particles and global rotations and translations. Consequently, it advantageous for the push-forward distribution of Boltzmann Generator possess the same symmetries the target system. (39 40 authors demonstrate that the push-forward distribution p(x:) of permutation
Equivariant flows
and rotation equivariant normalizing flow with a permutation and rotation invariant prior distribution, again rotation and permutation invariant. Furthermore, [39] present a method to construct such equivariant CNFs by using equivariant vector field vg.
3.3 Flow matching
Flow matching 19
20 21 22 enables efficient, simulation free training of CNFs. The conditional
flow matching training objective allows for the direct training of the vector field (t, x)) through
LCPM (0) 1,=-00.1),--xx(cll
[|uo 2> ue (x)x))]1
(5)
There are many possible parametrizations for the conditional vector field uy (a)2) and the conditional probability path pa (x]2). One of the most simple, but powerful possible parametrization is
z = (Iro, 21)
and p(=) ((20))((11))
(6)
u (x|z)) = 20 and p=(212) = M((ztt + (1-t) (0,))))),0)))
(7)
which we use in this work to train our models. For a more detailed description refer to 19 22
Transferable Boltzmann Generators
This section presents proposed framework for transferable Boltzmann Generators (TBGs).
4.1
Architecture
Our proposed transferable Boltzmann Generator is based on a CNE, The corresponding vector field (t 2) is parametrized by O(D)- and S(N) -qquivvaaaa graph neural network (EGNN) [1 (8 (46) as commonly used in prior work, c.g. 23 41). Although, less expressive than other equivariant networks such as [59] 60 61 62). it is faster to evaluate, which is important for CNFs as there can hundreds of vector field calls during inference.
The vector field vo a)) consists of L consecutive layers. The position of the i-th particle is updated according to the following equations:
mlj = de (h/s h/. hy.d)))) h = (t, aq, b,, ci,
(8)
a 2. (I) 24))
dij + oe (e(mly))
ip1
h/41 = Φn mH),
m Σ
Dm (mi
(10)
0o (t, x = 2. x/
Σ
(nt - =2))
(11)

---

where da represents different neural networks, dy is the Euclidean distance between particle and j. è is the time, o an embedding for the particle type, b; for the amino acid, and e or the amino acid position in the peptide. In the final step, the geometric center is subtracted to ensure that the center of positions is conserved. When combined with a symmetric mean-free prior distribution, the push-forward distribution of the CNF will be OD)- and ((N))invarian) as demonstrated in [63].
The embedding of each atom is constructed from three parts. The first part the atom type a; which is a one-hot vector of 54 classes. The classes are defined based on the atom types in the peptide topology. Therefore, only a few atoms are indistinguishable such as hydrogen atoms that are bound to the same carbon or nitrogen atom. The second part is the amino acid to which the atom belongs, which divided into 20 classes. The third part the position of the amino acid in the peptide sequence. This embedding is similar to the embedding used in [46] for the rotamer embeddings. The amino acid and positional embeddings are only used for the transferable experiments. For more details see Appendix [.5] In this study, we refer to this transferable Boltzmann Generator architecture as TBG + full, and we use this name even when we apply it to a non-transferable setting.
The proposed architecture in 23 uses distinct encodings for all backbone atoms and the atom types for all other atoms. This represents special case of our architecture, wherein b; and c; are omitted and ag encodes the atom type or backbone atom. Hence, there are 13 classes for ag We refer to this architecture as TBG backbone. Furthermore, we refer the specific architecture employed in [23 as BG + backbone for the alanine dipeptide experiments. Note that using only a, causes problems for transferability, see Appendix A.4 for more details.
Moreover, we employ a model that utilizes the atom type as the sole encoding (there are only five distinct atom types). This model is referred to simply TBG.
4.2 Training transferable Boltzmann Generators
All transferable Boltzmann Generators utilize flow matching for training. Given the variation in peptides across batches, the flow matching loss for each peptide is normalized by the number of atoms it contains. This training procedure is applied to all different architectures. For further details, refer to Appendix B
4.3 Inference with transferable Boltzmann Generators
Sampling with a transferable Boltzmann Generator, especially on unseen peptides, poses multiple challenges: (i) Some generated samples may not correspond the molecule of interest, but rather to a molecule that contains the same atoms but has a different bonding graph. Some of these configurations might even be valid molecules. For some examples see Appendix A.44 However, as we are in this work interested in sampling from the equilibrium Boltzmann distribution for a given molecular bonding graph, rather than sampling distinct molecules, we would like to avoid these cases. Nevertheless, this effect can be largely mitigated by our proposed TBG + full architecture. (ii) When working with classical force fields, the correct ordering with respect to the topology is crucial for evaluating energies. This is not a concern for semi-empirical force fields, as they respect the permutation symmetry of particles of the same type. As we typically use a Gaussian prior distribution, it is common that the generated samples are not arranged according to the topology. Consequently, in order to evaluate the energy, it is necessary to reorder the generated samples according to the topology. (iii) It is possible that the chirality of generated samples differs from that of the peptide of interest.
We resolve (i) and (ii) by generating a bond graph for the generated samples, based on empirical bond distances and atom types. This graph is then compared with reference bond graph. If the two graphs are isomorphic, we can conclude that the configuration is correct. For more details, see Appendix B.I For (iii), we employ the code of [I] to check all chiral centers. If all chirality centers of a peptide are flipped, this can be resolved by mirroring. Otherwise, these samples are assigned high energies, as they are not from the target Boltzmann distribution of interest. It is important to note that only generated samples with the correct confi anfiguration and chirality are considered valid samples from the Boltzmann distribution of interest.
5

---

*)
Blased MD
Boltzmann Generator
b)
Classical energy distribution
Free energy projection
80
No
baskore
energy koT
RD becktine snviieteed
Energy
Figure 1: Results for the alanine dipeptide system simulated with classical force field (a) Ramachan- dran plots for the biased MD distribution (left) and for samples generate with the TBG + full model (right). (b) Energies of samples generated with different methods. (c) Free energy projection along the slowest transition (f angle), computed with different methods.
Table Comparison of Boltzmann Generators with different architectures for the single molecular system alanine dipeptide. Errors are computed over five runs. The results for the Boltzmann Generator and backbone encoding (BG + backbone) for the semi-empirical force field are taken from 223..
Model
NLL (4)
ESS (1)
Alanine dipeptide semi-cmpirical force field
BG + backbone
23
107.56 0.09
0.50 + 0.13%
TBG + full (ours)
124.71 + 0.08
1.03 + 0.17%
Alanine dipeptide classical force field
BG + backbone
[23
-109.02 0.01
1.56 + 0.30%
TBG + full (ours)
-127.06 = 0.12
6.03 + 1.34%
Experiments
In this section, compare model with similar previous work equivariant Boltzmann Genera- tors (23) for alanine dipeptide. Moreover, we show the transferability of our model on dipeptides, where compare model with the transferable Timewarp model (11). More experimental details,
such as dataset details, the specifics of the employed models, and the utilized computing infrastructure can be found in Appendix B
5.1 Alanine dipeptide
In our first experiment, we investigate the single molecule alanine dipeptide in implicit solvent, described in Cartesian coordinates. The dataset was introduced in [23], for more details see Ap- pendix B.2 The training trajectory was generated by sampling with respect to a classical force field, and subsequently. 10a random samples were relaxed with respect to the semi-empirical GFN2-xTB force-field (64) for 100Is each. The objective is to train Boltzmann Generator capable of sampling from the equilibrium Boltzmann distribution defined by the semi-empirical GFN2-xTT force-field efficiently and to recover the free energy surface along the slowest transition, i.e. the ip angle. Following the methodology outlined in (23), the training data is biased towards the less probable (positive) ip state. It is evident that any trained model on this set will be biased in comparison to the true Boltzmann distribution defined by the semi-empirical energy. However, the reweighting technique allows for the debiasing of the samples. The model is trained in the same way as described in [23). Overall, the likelihoods and ESS values observed for the TGB full model are superior to those reported in [23 (Table [p. This is achieved with nearly the same amount of parameters and maintaining comparable training and inference times (see Appendix [.3). Furthermore, the correct free energy difference is recovered, as demonstrated in Appendix A.
In 23 the authors utilized semi-empirical potential to avoid the required ordering of the atoms to the topology for classical force fields. As the prior distribution of the Boltzmann Generator is usually a multivariate standard Gaussian distribution, generated samples will almost certainly not have the correct ordering. As we have introduced an efficient way to reorder samples in Section 4.3 we can now also evaluate alanine dipeptide for a classical force field. Therefore, we retrain the model in 23 on the classical MD trajectory and compare with our TBG + full architecture. We bias the training

---

MD
TBG full
MD
TB0 M
d)
energy distribution
energy projection
energy projection
MD
Free energy Als
223
TDC weggitd
1 free
103
Energy
kg
T1C0
Figure 2: Results for the KS dipeptide (a) Sample generated with the TBG + full model (b) Ra- machandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full model (right). (d) Energies of samples generated with different methods and architectures. (e) Free energy projection along the angle. (f) Free energy projection along the slowest transition (TICO).
a)
b)
TBG
MD
TBC Mil
TICO
d).
Classical distribution
Free projection
Free projection TICA
ND
Free eneroyy Kg)
X05
TBC Mi
tetibern
mogtitd
Free 2000o
TB0 nt nenpptte
300
-23 -26
Energy
-22 -26
Figure 3: Results for the GN dipeptide (a) Sample generated with the TBG full model (b) Ra- machandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (d) Energies of samples generated with different methods and architectures. (e) Free energy projection along the & angle. Free energy projection along the slowest transition (TICO).
data as before towards the unlikely state. As expected, the likelihood and ESS for the classical force field are much better than for the semi-empirical one, as the training data stems from the target distribution. Our proposed architecture again performs significantly better, as shown in Table and F The majority of generated samples with the TBG full model and the BG + backbone sample nearly exclusively correct configurations, i.e. configurations with the correct bond graph, namely nearly 100% and about 98%, respectively. As presented Figure] both models recover the free energy landscape correctly.
5.2 Dipeptides (2AA)
In our second experiment, we evaluate our model on dipeptides and show transferability. The dataset was introduced in []] The training set consists of 200 dipeptides, which were simulated each with

---

Dipertidas- effective samaling 8izes
b)
Free energy projection 9
Free anergy projection
30-
TBS bucksane
Free energy kk7
seed
10C tallaaa TDC
TB2 ful
2
TBc fal is seed remeggted
fall laleed rawagrsed
T92 full bmeleo
emergy
10
13 074
152 145 18
Popiide
Figure 4: (a) Effective samples sizes (ESS) for the first test peptides for different transferable architectures and training sets. (b) Free energy projection along the yp angle for the TBG + full model trained on the biased dataset for the KS dipeptide. The weighted free energy projection demonstrates a superior fit compared to the TBG full model (see Figure 2 (c) Free energy projection along the op angle for the TBG full model trained on the biased dataset for the GN dipeptide. The weighted free energy projection demonstrates superior fit compared to the TBG full model (see Figure BD).
classical force field for 50 ns and, therefore, may not have reached convergence. Nevertheless, as previously demonstrated, it is not necessary to train on unbiased data in order to obtain unbiased samples with a Boltzmann Generator.
We compare the three different transferable architectures described in Section 4.1 and use the same training procedure for all of them. Similar to the alanine dipeptide experiments, we obtain significantly better results for the TBG full model in terms of ESS (Table and Figure A.) energies (Figure 2d), the ratio of correct configurations (Table and likelihoods of test set samples (Appendix(A.5) In particular, the extremely low number of correct configurations for numerous test peptides for the TBG and TBG + backbone models renders them unsuitable as Boltzmann Generators for this setting (Table 2 and Appendix A.5)) Furthermore, the TBG full model always finds all metastable states for unseen test peptides (see also Appendix
The results for the well-performing TBG + full model are presented for two example peptides from the test set in Figure 2 and Figure They were chosen as all architectures sample relevant amounts of valid configurations. Detailed results for other evaluated test peptides are shown in Appendix A.5
The TBG full model is exemplary Boltzmann Emulator, as it is capable of capturing all metastable states of the target Boltzmann distribution (Figurepp.c and Figure33.c). However, it is furthermore also a capable Boltzmann Generator, allows for efficient reweighting (Figure f and Figure fe.,f) with good ESSs (Table 2 To identify different metastable states, we employ time-lagged independent component analysis (TICA) (65) dimensionality reduction technique that separates metastable states. We show this analysis in addition to the Ramachandran plots for the dihedral angles.
Moreover, we investigate the influence of the training set in two ablation studies.
Training on a biased training set Our alanine dipeptide results as well (3 indicate that it can be advantageous to bias the training data towards states that are less probable, such as positive yp states, to recover free energy landscapes. Therefore, we hias the training data by weighting positive states for each training peptide, such that they have nearly equal weight to the negative states (see also Appendix [B.4). We show that TBG + full model trained on this dataset (TBG + full (biased)) produces even more accurate free energy landscapes for both the Ramachandran and TICA projections (Figure (bc). Notably, the unweighted projection shows a clear bias, as expected. However, as the training data is now biased, the effective sample size (ESS) is generally lower (Table 2 and Figure 41).
Training on a smaller training set Additionally, we examine the impact of smaller training sets on the generalization results. To this end, we train the TBG + full model on two smaller datasets with shorter simulation times: (i) 5ns for each training simulation and (ii) only 500ps of each training simulation. Consequently, the training trajectories are 10 times and 100 times smaller than before. As we utilize only the initial portion of each trajectory. a greater number of metastable states are missed during the brief simulations, as illustrated in Appendix A.. Training on the tenfold smaller trainings set, we refer to the model as TBG + full (smaller), shows slightly worse results compared to training on the whole trainings set (Table 2 and Appendix (A.5). The even smaller trainings set leads to inferior results, with several metastable states being missed presented inAppendix A.3
8

---

Table 2: Effective samples size and correct configuration rate for unseen dipeptides across different transferable Boltzmann Generator (TBG) architectures. The values are computed for test dipeptides. See Appendix A.5 for more results.
Model
ESS (f)
Correct configurations (1) Mean
Mean
Range
Range
TBG
0.48 = 0.59%
(0.0%, 1.47%)
13 = 18%
(1%,48%%
TBG + backbone
0.58 = 1.04%
(0.0%, 3.24%
17 21%
(1%, 52%)
TBG + full
15.29 9.227%
(4.08% 31.98%)
98 - 2%
(94%, 100%)
TBG + full (smaller)
6.13 = 3.13%
(1.93%, 11.16%)
96 = 3%
(88%, 100%)
TBG + full (biased)
10.24 + 7.14%
(3.21%, 22.66%)
98 + 2%
(93%, 100%)
a)
Dipeptides 30000 energy uations
b)
Dipeptides sampling budget
6)
Dipeotides sampling budget
mean 338
Temeworr -mw.--t-2250000221 TBG + 2r6anm0.10020.220
iinewara -maared 1/985 213 TBG + fu1 mma.-m:8822000 2.009
Tec Al maar-01100s..
0.05
10.0
12.5
0.
30.5
0.4
0.2
6.0
AT GNLW HT ET
KS AT GNEWNY MF ACGPKO
GNUMINY RD RLETAC
Figure 5: Timewarp MCMC 23 sampling experiments. Wasserstein distance between the generated Ramachandran plot and the MD Ramachandran plot for different computational budgets for dipeptides. Lower is better. (a) After 30000 energy evaluations (b) After 12h wall-clock-time. (c)
After 24h wall-clock-time.
Nevertheless, these findings show that transferable Boltzmann Generators can be effectively trained using very small datasets, even when individual trajectories lack metastable states.
Comparison with Timewarp
In this section, we compare TBG model with the Timewarp model ||| the dipeptide dataset. Unlike our approach, which generates independent samples, the Timewarp model predicts large time steps, which can be cobined with Metropolis-Hoe Hasttigss acceptance steps ensure asymptotically unbiased sampling. Additionally, Timewarp employs coupling flow rather than a continuous flow. Similar our TBG full model, the Timewarp model also treats most particles distinguishable. It does so by conditioning the learned probability on the current state, implicitly capturing the topology information.
We compare the Wasserstein distance between the generated Ramachandran plot and the MD Ra- machandran plot for different computational budgets for dipeptides: (a) 30,000 energy evaluations, (b) 12 hours, and (c) 24 hours of wall-clock simulation time on an A-100 GPU. As shown Figure 3 the TBG full model outperformed the Timewarp model across all budgets, particularly in terms of energy evaluations. Energy evaluations are especially critical, as they become major computational bottleneck with more complex force fields. Moreover, the unweighted samples of the TBG + full model have lower energies, and are therefore closer to the actual target Boltzmann distribution than samples generated with Timewarp without Metmopllis- acceptance (exploration mode). Additional results and details can be found Appendix A.7
6 Discussion
For the first time, we demonstrated the feasibility of training transferable Boltzmann Generators. We introduced general framework for training and evaluating transferable Boltzmann Generators based on continuous normalizing flows. Furthermore, we developed transferable architecture based on equivariant graph neural networks and demonstrated the importance of including topology information in the architecture to enable efficient generalization to unseen, but similar systems. The transferable Boltzmann Generator was evaluated on dipeptides, where significant effective sample sizes were demonstrated unseen test peptides and accurate sampling of physical properties, such

---

as the free energy difference between metastable states, was achieved. Moreover, we have shown in ablation studies that transferable Boltzmann Generators can be extremely data efficient, with even small training trajectories being sufficient. Future research will determine whether and how transferable Boltzmann Generators can be scaled larger systems.
Limitations / Future work
Scaling transferable Boltzmann Generators to larger systems remains for future work. Notably, this usually requires large amounts of computational resources, as c.g. shown in [[T], where they are able to train their transferable model tetrapeptides, but use more than 100 times more parameters than us. Scaling to larger systems often involves coarsegraining, which typically results in the loss of an explieit energy function. In such cases, the transferable Boltzmann Generator would effectively become a transferable Boltzmann Emulator. Depending on the specific application, samples from distribution close to the target Boltzmann distribution may be sufficient. Our results show that unweighted samples from the TBG full model already closely resemble the target Boltzmann distribution.
Additionally, small systems can still be highly relevant, especially when paired with more expensive force fields, such semi-empirical or first-principles guantum-mechanical force fields. In these scenarios energy evaluations become a primary bottleneck. Boltzmann Generators are particularly advantageous in this context, they require several orders of magnitude fewer energy evaluations compared to MD simulations or other iterative methods, such as Timewarp [1]. One potential real world application would therefore be the simulation of small molecules with expensive force fields.
Instead of flow matching, one could use optimal transport flow matching 22 or equivariant optimal transport flow matching [23 53 for training. but as indicated in (23 the gains for molecular systems, especially in the presence of many distinguishable particles, are small.
Throughout our work, we utilize a standard Gaussian prior distribution. However, as recently introduced, an alternative to use Harmonic prior distribution 66 67). where atoms that are close in the bond graph are sampled in the vicinity of each other. Notably, We experimented with this Harmonic prior but found no significant improvements for our transferable model. This aligns with findings by [68], indicating that chemically informed priors do not enhance performance substantially compared to simpler uninformed priors for flow matching molecular systems. Instead, the network architecture and inductive bias play more crucial role. However, such priors might become more important for larger systems.
Despite conducting a series of ablation studies, we did not pursue the impact of training set comprising a smaller number of peptides. Instead, we opted for investigating shorter trajectories. Another potential direction could be relaxing the 2AA dataset using semi-empirich force field and training on this modified version, similar the alanine dipeptide experiment. However, this approach incurs additional computational costs, as the entire 2AA dataset would need to be relaxed with respect to the aamieemiricall force field.
We used the EGNN architecture for the vector field due to its fast evaluation capabilities. Future research could explore alternative architectures for the vector field, such those proposed by [59] 60 61 62 69 77, 67]. to determine if they improve performance and enable scaling to larger systems. We hope that our framework will facilitate the scaling of transferable Boltzmann Generators to larger systems in future research.
8 Broader Impact
This work represents foundational research with no immediate societal impact. However, if our method is scalable to larger, more relevant systems, it could facilitate the acceleration of drug and material discovery by replacing or enhancing MD simulations, which often play a crucial part in the process. A potential risk is that this method then might be used to identify new diseases or develop biological weapons, Another risk is the lack of known convergence criterion, making impossible to confirm that all potential configurations have been identified, even with infinite number of samples. This could lead to false claims about the results, potentially affecting subsequent applications.
10

---

Acknowledgements
We gratefully acknowledge support by the Deutsche Forschungseennensscaatt (SFB1114, Projects No. CO3, No. A04, and No. B08), the European Research Council (ERC CoG 772230 "ScaleCell"), the Berlin Mathematics center MATH+ (AA1-10), and the German Ministry for Education and Research (BIFOLD- Berlin Institute for the Foundations of Learning and Data). We gratefully acknowledge the computing time made available to them on the high-performance computer "Lise" at the NHR Center NHR @ ZIB. This center is jointly supported by the Federal Ministry of Education and Research and the state governments participating in the NHR (www -hrr-verein de/ unsere- partner). We thank Andreas Krämer, Michele Invermizzi, Jonas Köhler, Atharvar Kelkar, Yaoyi Chen, Max Schebek, Iryna Zaporozhets, and Hannes Stärk for insightful discussions. Moreover, we thank Andrew Foong and Yaoyi Chen for providing the datasets from the Timewarp project [11].

---

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, pages 1-3, 2024.
[2] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589, 2021.
[3] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos Santos Costa, Maryam protein sequences at the scale of evolution enable accurate structure prediction. BioRxiv, 2022:500902, 2022.
Tom Sereu, Sal Candido, et al. Language models of
[4] Niklas Gebauer, Michael Gastegger, and Kristof Schüitt. Symmetry-adapee generation of 3d point sets for the targeted discovery of molecules. Advances in neural information processing systems, 32, 2019.
[5]
Emiel Hoogeboom, Víctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8867-8887. PMLR, 17-23 Jul 2022.
[6] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[7] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Tor- sional diffusion for molecular conformer generation. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 22253.. Curran Associates, Inc., 2022.
[8] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators sampling equilibrium states many-body systems with deep learning. Science, 365:eaaw 147, 2019.
[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 urXiv:2404.16555 2024.
[10] Grant M Rotskoff. Sampling thermodynamic ensembles of molecular systems with generative neural networks: Will integrating physics-basee models close the generalization gap? Current Opinion in Solid State and Materials Science, 30: 10:101158: 2024.
[11] Leon Klein, Andrew Y. K. Foong, Tor Erlend Fjelde, Bruno Kacper Mlodozeniee, Mare Brockschmidt, Sebastian Nowozin, Frank Noe, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[12] Esteban G Tabak, Eric Vanden- et al. Density estimation by dual ascent of the log- likelihood. Communications in Mathematical Sciences, (1):21172333 2010.
[13] Esteban G Tabak and Cristina V Turner. A family of nonparametric density estimation algo- rithms. Communications on Pure and Applied Mathematics, 6602):145-164 2013.
[14] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. CoRR, 2019.
[15] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference machine learning, pages 1530-1538. PMLR, 2015.
[16] Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances neural information processing systems, pages 6571-6583, 2018.
12

---

[17] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable re- versible generative models with free-form continuous dynamics. In International Conference on Learning Representations, 2019.
[18] Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. CoRR, 2014.
[19] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. The Eleventh International Conference on Learning Representations, 2023.
[20] Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
[21] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
[22] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Dector-Brooks Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint wxXi::2302.00482. 2023.
[23] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[24] Laurence Blnng Midgley, Vincent Stimper, Javier Antoran, Emile Mathicu, Bernhard Schölkopf, and José Miguel Hernández-Lobato. SE(3) equivariant augmented coupling flows. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
[25] Peter Wirnsherger, Andrew J Ballard, George Papamakarios, Stuart Abercrombie, 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):144112, 2020.
[26] Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and Boltzmann generators. Phys. Rev. Res., 4:L042005, Oct 2022.
[27] Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. In M. Ranzato,
A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 2796-2809. Curran Associates, Inc., 2021.
[28] Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023.
[29] Xinqiang Ding and Bin Zhang. Deepbar: A fast and exact method for binding free energy computation. Journal of Physical Chemistry Letters, 12:2509-2518 3 2021.
[30] Joseph C. Kim, David Bloore, Karan Kapoor, Jun Feng. Ming-Hong Hao, and Mengdi Wang. Scalable normalizing flows enable boltzmann generators for macromolecules. In International Conference Learning Representations (ICLR), 2024.
[31] Rasool Ahmad and Wei Cai. Free energy calculation of crystalline solids using normaliz- ing flows. Modelling and Simulation in Materials Science and Engineering, 30(6):065007, September 2022.
[32] Kim A Nicoli, Christopher J Anders, Lena Funcke, Tobias Hartung, Karl Jansen, Pan Kessel, Shinichi Nakajima, and Paolo Stomati.. Estimation of thermodynamie observables in lattice field theories with deep generative models. Physical review letters, 126(3):032001, 2021.
[33] Maximilian Schebek, Michele Invernizzi, Frank Noé, and Jutta Rogal. Efficient mapping of phase diagrams with conditional boltzmann generators. Machine Learning: Science and Technology, 2024.
13

---

[34] Hao Wu, Jonas Köhler, and Frank Noe. Stochastic normalizing flows. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 9933-5944. Curran Associates, Inc., 2020.
[35] Xingiang Ding and Bin Zhang. Computing absolute free energy with deep generative models. Biophysical Journal, 1203))195a, 2021.
[36] Andrea Rizzi, Paolo Carloni, and Michele Parrinello. Multimap targeted free energy estimation, 2023.
[37] Samuel Tamagnone, Alessandro Laio, and Marylou Gabrié. Coarse-grained molecular dynamics with normalizing flows. Journal of Chemical Theory and Computation, 20(18):7796-7805, 2024.
[38] Christoph Schönle, Marylou Gabrié, Tony Lelièvre, and Gabriel Stoltz. Sampling metastable systems using collective variables and jarzynski-cron -rookks paths. arXiv preprint orXiv:2405.18160 2024.
[39] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pages 5361-5370. PMLR, 2020.
[40] Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant Hamiltonian flows. arXiv preprint arXiv: (tt.:::999.73330 13739, 2019.
[41] Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 4181 4191-4192. Curran Associates, Inc., 2021.
[42] Jonas Köhler, Michele Invernizzi, Pim de Haan, and Frank Noé. Rigid body flows for sampling molecular crystal structures. In International Conference on Machine Learning. ICML 2023, volume 202 of Proceedings of Machine Learning Research, pages 17301-17326. PMLR, 2023.
[43] Andrea Rizzi, Paolo Carloni, and Michele Parrinello. Targeted free energy perturbation revisited: Accurate free energies from mapped reference potentials. Journal of Physical Chemistry Letters, 12:9449-9454, 2021.
[44] Michele Invernizzi, Andreas Krämer, Cecilia Clementi, and Frank Noé. Skipping the replica exchange ladder with normalizing flows. The Journal of Physical Chemistry Letters, 13: 13:11 1 11649, 2022.
[45] Kirk Swanson, Jake Lawrence Williams, and Eric Jonas. Von mises mixture distributions for molecular conformation generation. In International Conference on Machine Learning, pages PMLR, 2023.
[46] Osama Abdin and Philip M Kim. Pepflow: direct conformational sampling from peptide energy landscapes through bypernetwork-condiioon diffusion. bioRxiv, pages 2023-06, 2023.
[47] Bowen Jing. Bonnic Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402. 048455 2024.
[48] Shuxin Zheng, Jiyan He, 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, pages 1-10, 2024.
[49] 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.
[50] Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
14

---

[51] Nicholas E Charron, Felix Musil, Andrea Guljas, Yaoyi Chen, Klara Bonneau, Aldo S Pasos- Trejo, Jacopo Venturin, Daria Gusew, Iryna Zaporozhets, Andreas Krämer, al. Navigating protein landscapes with machine-learned transferable coarse-grained model. arXiv preprint arXiv:2310. /8278, 2023.
[52] Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining of molecular dynamics without forces. Journal of Chemical Theory and Computation, 19(3)992-9552, 2023.
[53] 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. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[54] 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 arXiv:2403. 14088, 2024.
[55] Felix Draxler, Peter Sorrenson, Lea Zimmermann, Armand Rousselot, and Ullrich Köthe. Free-form flows: Make any architecture normalizing flow. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li, editors, Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pages 2197-2205. PMLR, 02-04 May 2024.
[56] H Wiegand. Kish, I.: Survey sampling. john wiley & sons, inc., new york, london 1965, ix+ 643 31 abb., 56 tab., preis 83 s., 1968.
[57] George Papamakarios, Eric T Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. J. Mach. Learn. Res., 22657):1-64, 2021.
[58] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[59] Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations, 2021.
[60] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pages 9377-9388. PMLR, 2021.
[61] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2023.
[62] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790 680000668022222 2021.
[63] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning, pages 5361-5370. PMLR, 2020.
[64] Christoph Bannwarth, Sebastian Ehlert, and Stefan Grimme. Gfn2-xttb an accurate and broadly parametrized self-consistent electrostatics and density-dependent dispersion contributions. Journal of Chemical Theory and Computation, 15(3):6652-161) 1671, 2019, PMID: 30741547.
quantum chemical method with multipole
[65] Guillermo Pérez-Hernández, Fabian Paul, Toni Giorgino, Gianni De Fabritiis, and Frank Noé. Identification of slow molecular order parameters for Markov model construction. The Journal of chemical physics, 139( 1,
[66] Bowen Jing, Ezra Erives, Peter Pao-Huang, Gabriele Corso, Bonnie Berger, and Tommi Jaakkola. Figenfold: Generative protein structure prediction with diffusion models. arXiv preprint 2xX1v:2304.02198, 2023.
15

---

[67] 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:23 2023.
[68] Dina A Sharon, Yining Huang, Motolani Oyewole, and Sammy Mustafa. How to go with the flow: an analysis of flow matching molecular docking performance with priors of varying information content. In ICLR 2024 Workshop Generative and Experimental Perspectives for Biomolecular Design, 2024.
[69] Weitao Du, Yuanqi Du, Limei Wang, Dieqiao Feng, Guifeng Wang, Shuiwang Ji, Carla Gomes, and Zhi-Ming Ma. A new perspective on building efficient and expressive 3d equivariant graph neural networks. arXiv preprini arXiv:2304.0 200:2334.0447,, 2023.
[70] Han Yang, Chenxi Hu, Yichi Zhou, Xixian Liu, Yu Shi, Jielan Li, Guanzhi Li, Zekun Chen, Shuizhou Chen, Claudio Zeni, et al. Mattersim: A deep learning atomistic model across elements, temperatures and pressures. arXiv preprint ArXiv:2405.04990 2024.
[71] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, nigh-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-8033 Curran Associates, Inc., 2019.
[72] Michael Poli, Stefano Massaroli, Atsushi Yamashita, Hajime Asama, Jinkyoo Park, and Stefano Ermon. Torchdyn: Implicit models and neural numerical methods in pytorch. In Neural Information Processing Systems, Workshop on Physical Reasoning and Inductive Biases for the Real World, volume 2, 2021.
[73] Aric Hagberg, Pieter Swart, and Daniel S Chult. Exploring network structure, dynamics, and function using networkx. Technical report, Los Alamos National Lab,(LANL), Los Alamos, NM (United States), 2008
[74] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6900 2014.
16

---

Table 3: Dimensionless free energy differences for the slowest transition of alanine dipeptide estimated with various methods. Umbrella sampling yields a converged reference solution, Errors are calculated over five runs. Values for BG + backbone and Umbrella sampling are taken from 23
Umbrella sampling
BG + backbone 23 4.10 + 0.08
TBG + full (ours)
Free energy difference / KyT
4.10 0.26
4.09 0.05
b)
Training data
10x smaller
100x sma
Training data
10x smaller
100x smaller
Figure 6: Example Ramachandran plots for different trajectory lengths for the training data. It can be observed that as the trajectory length decreases, the number of metastable states that are missed increases, thereby making the learning task more challenging. (a) AY dipeptide (b) IH dipeptide.
Appendix
A Additional results and experiments
A.1
Semi-empirical force field for alanine dipeptide
We report the free energy differences for the slowest transitions of alanine dipeptide for a semi- empirical force field in Table 3, See Section 5 for more details.
A.2 Dipeptide training data
When training on smaller training sets, i.e. with shorter trajectories, additional metastable states will not be visited during the short simulation times. We show this for two example training peptides in Figure 6 Nevertheless, the TBG + full (smaller) model trained on 10 times shorter trajectories, is nearly as good as the model trained on the full trajectories, see Section 5) However, for the 100 times smaller trajectories, the TBG full model performs significantly worse, see Appendix A.3
A.3 Smaller dataset
We investigate the effect 100 times smaller trainings trajectories, i.e. simulation time of only 500ps. As shown in Appendix A.2] these trajectories miss many metastable states. This can be also observed for the so trained models, which we refer as TBG + full (smaller500), as they do not capture especially unlikely metastable states well presented in Figure trained on larger trajectories find all metastable states and allow for efficient reweighting, as discussed in Section 5 and Appendix A.5
aan and Figure & In contrast, models
A.4 Sampled dipeptide configurations
For some amino acid combinations, both the TBG and TBG backbone models sample only a small number of correct configurations. Although the generated configurations potentially valid molecular configurations, they are not the one of the target dipeptide as shown in Figure 9 This often due the encoding, which e.g. cannot distinguish between different orderings of amino acids a peptide. Only the various TBG + full architectures samples nearly exclusively correct configurations.
A.5 Additional results for dipeptides
Inference a costly process, and extensive sampling is necessary to obtain reliable estimates for the expected sample size (ESS). Therefore, we only evaluate the transferable models subset of the
17

---

b)
MD
TBG + full (smaller500)
MD
TBC (smaller500)
d)
energy distribution
energy projection
energy projection
Mo
"ao lemallasion
Froe enens Av
vg
temallrsion
Free energy 1ky
MD.
200
300
Energy
kg
T1C0
Figure 7: Results for the KS dipeptide for TBG + full model trained on 100 times smaller training trajectories. As can be seen in Figure the results for the TBG full model trained on the whole trajectories are much better. (a) KS dipeptide (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the model (right). (d) Energies samples generated with the model. (e) Free energy projection along the yp angle. (f) Free energy projection along the slowest transition (TICO).
a)
b)
MD
TBG + full (smaller500)
MD
TBG full (mmaller50))
d)
Classical energy distribution
energy projection
energy projection
kD
VD
MD
tomalle500
tomile-s0m
Free ene ko]
lamalletoo:
Frod 2009
304
300
Energy
/kgT
TIC0
Figure 8: Results for the GN dipeptide for TBG full model trained on 100 times smaller training trajectories. As can be seen in Figure the results for the TBG full model trained on the whole trajectories are much better. (a) GN dipeptide (b) Ramachandran plot for weighted MD distribution (left) and for samples generate with the model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the model (right). (d) Energies samples generated with the model. (c) Free energy projection along the yp angle. (f) Free energy projection along the slowest transition (TICO).
set. The dipeptides are randomly selected, but is ensured that all amino acids are represented at least once. However, we evaluate the best-performing model, namely TBG + full, for all 100 test peptides. The results for the additional test peptides good agreement with the first ones as presented in Table 4 and Figure 10
We report individual results for the different architectures in Figure 10
To illustrate performance from the test set in Figure hT and Figure 12a-f. Furthermore, we also again show results for training on the biased dataset the same figures (Figure TTg.h, and Figure T2g.h,i, As observed previously, the TBG full (biased) model recovers the free energy landscape better than the TBG +
TBG full model, we present additional examples of dipeptides
18

---

a)
b)
Figure 9: Sampled molecules with the TBG and TBG + backbone models, which do not have the correct topology. (a) NY dipeptide reference (b) Generated molecule with NY atoms by the TBG model. (c) IM dipeptide reference (d) Generated molecule with IM atoms by the TBG model. (e) TD dipeptide reference (f) Generated molecule with TD atoms by the TBG backbone model.
-))
b)
d)
8
Ettective Sameling Sizee
Cornect Contiguratire
Dioeptides correct configurations
Dipeptides
5
Rate
Conood cotiigrrraaatt atten rates
80
E15
to
T50 teddone
Tike baccheee TDC
150 eal 24see)
Tik sond
78
181- Tall
IBG Tull
121
377
10
H1
15
A1
120
57
10
Figure 10: (a) Effective samples sizes for all 100 test dipeptides for the TBG full model. (b) Percentage of correct configurations for all 100 test dipeptides sampled with the TBG + full model. (c,d) Performance comparison for different transferable architectures and training sets on 8 dipeptides from the test (c) Correct configuration rate (d) Likelihood per particle.
full model, especially for the projections. We present additional examples of dipeptides from the test set for the TBG full model Figure Figure 14
Figure
15
Figure T6 and Figure 17
Furthermore, we present results for two example peptides from the test set for the TBG full (smaller) model, which is trained trajectories that tenfold smaller than those used for the TBG + full model. These results are shown in Figure 18 and Figure 19
A.6 Transferable Boltzmann Generators Boltzmann Emulators
Given the high cost of sampling with CNFs, which necessitates integrating the Jacobian trace along the positions, we did not evaluate all available test peptides for models (see Appendix [B.2). However, since sampling without the Jacobian trace is less expensive and do not require many samples as for estimating the ESS, we also employ the TBG + full (smaller) model as Boltzmann Emulator to ascertain whether we have identified all metastable states, despite the fact that it was only trained on the 10 times smaller training set. The Boltzmann Emulator is evaluated on a diverse of test peptides, and nearly always finds all metastable states within less than hour of wall clock time. This is a notable improvement over MD simulations, which often take longer to explore due the iterative nature of MD. Some examples are shown Figure 20 This experiment shares similarities with the exploration mode of [I]] where they employ their model without the acceptance
Table Effective samples size and correct configuration rate for unseen dipeptides for the TBG + full architecture for different numbers of test peptides.
Model
ESS (t)
Correct configurations (t) Mean
Mean
Range
Range
TBG full (8 test peptides)
15.29 9.27%
(4.0%% 31 93%
98 + 2%
(94%. 100%)
TBG full (100 test peptides)
13.11 + 8.59%
(0.82%, 36.766.778 78%)
98 - 2%
(92%. 100%)
19

---

MD
TBG full
MD
TBG f
d)
energy
energy projection
energy projection
No
Free energy Aot
HD
Free energy AAa
TDO Aal reoegtte
Energy ke
TICO
g)
Classical energy distribution
h)
Free energy projection
Free energy projection TICA
3:0
Free anargy Aat
Free energy kal
Tad All Doordi manigna
Thd Doordi govedi manignted
100
160 (AMad)
Energy
Figure 11: Results for the IM dipeptide (a) Sample generated with the TBG full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (d) Energies samples generated with the TBG + full model. (e) Free energy projection along the angle. (f) Free energy projection along the slowest transition (TICO). (g) Energies samples generated with the TBG full (biased) model. (h) Free energy projection along the &p angle for the TBG + full (biased) model. (i) Free energy projection along the slowest transition (TICO) for the TBG + full (biased) model.
step and therefore also explore potentially biased distribution rather than the unbiased Boltzmann distribution.
A.7 Comparison with the Timewarp model 11
Here, we provide additional results comparing our TBG + full model to the Timewarp model 11. as discussed in Section 5.3 In this second scenario, the objective is to explore all meta-stable states as given in the Ramachandran plot, aligning with the goals of a Boltzmann Emulator, where only approximate sampling from the Boltzmann distribution is required. For the TBG full model, do not apply reweighting, whereas in the Timewarp model, samples are rejected only consecutive samples show significant energy increase prevent divergence [1]
Although reweighting not strictly needed, we evaluate the energy of all samples generated by the TBG full model, filtering out any rare, high-energy samples. We then compare both models based on the mean number energy evaluations and the mean wall-clock time needed to explorae all states. As shown in Figure 21] the Timewarp model generally finds all states more quickly but requires more energy evaluations. Additionally, as illustrated in Figure 22 the Timewarp model's generated samples tend to exhibit higher energies compared those from the TBG full model. Consequently. the distribution produced by the TBG + full model more accurately approximates the target Boltzmann distribution.
For Timewarp experiments, we used a proposal batch size of 100, as recommended in for A- 100 GPUs.
20

---

MD
TBG full
MD
TBG fill
d)
energy distribution
energy projection
energy projection
MD
Free energyy Mot
101
Free S
179
TDG Md reneigtsod
Energy kg Classical energy distribution
TICO
h)
Free energy projection
i)
Free projection TICA
Free enera Doy
M reeehere
M reeehiee
Bessel
Free energp
200
33
-290
Energy
-20
-45 -20
Figure 12: Results for the NY dipeptide (a) Sample generated with the TBG full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (d) Energies of samples generated with the TBG full model. (c) Free energy projection along the ip angle. (f) Free energy projection along the slowest transition (TICO). (g) Energies of samples generated with the TBG + full (biased) model. (h) Free energy projection along the yp angle for the TBG + full (biased) model. (i) Free energy projection along the slowest transition (TICO) for the TBG + full (biased) model.
B Technical details
B.1
Code libraries
We primarily use the following code libraries: PyTorch (BSD-3) [71, bgflow (MIT license) 39
torchdyı (Apache License 2.0) [72], and NetworkX (BSD-3) (73 for validating graph isomorphisms. Additionally, use the code from 4T] (MIT license) for EGNNs, as well the code from (MIT license) and 23 (MIT license) for datasets and related evaluation methods.
Our
code
is
available
here:
https://osf htt.://3/...//////22////vv/v wtti://///vvieoooool
10520000221144400887001171788882
B.2 Benchmark systems
The investigated benchmark systems were created in prior studies A
Alanine dipeptide The alanine dipeptide datasets created in 23
(CC BY 4.0), refer
to them for detailed simulation details. The classical trajectory was created at T = 300K with the classical Amber l/99BBiln force-field. The subsequent relaxation was performed with the semi-empirical GFN2-x7B force-field 64
Dipeptides (2AA dataset) The original dipeptide dataset as introduced in [I] (MIT License) is As this includes lot of intermediate saved states and quantities, like energies, we create a smaller version with is available here: https:
available here: https://hugginffaet
:/op::://osf. //osf. io//i//////vvvve //vvew 1tta://///ieeeooool1102222200000 wwwp::/00020002e2244000000000000
For comprehensive overview of the simulation details, refer to [I]] All dipeptides were simulated
21

---

MD
TBG full
MD
TBG f
d),
MD energy distribution
Free enengy 147 0
energy projection
energy projection
309
XM
vg
Free energy INs
20 Energy
kg
11C0
Figure 13: Results for the ET dipeptide (a) Sample generated with the TBG full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right), (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full model (right). (d) Energies of generated samples (e) Free energy projection along the to angle. (f) Free energy projection along the slowest transition (TICO).
b)
MD
TBG full
MD
TBC M
d)
Classical energy distribution
energy projection
f).
energy projection
Free energy 1k!
PD
Free energy KyT
MC
100
320
Energy
KgT
TICO
Figure 14: Results for the RV dipeptide (a) Sample generated with TBG full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full model (right). (d) Energies of generated samples (c) Free energy projection along the o angle. (f) Free energy projection along the slowest transition (TICO).
with classical amber-14 force- force-field T 310K. The simulation of the training peptides were run for 50ns, while the test set peptides were run for 1ps.
Choice of test set peptides Inference is costly process with CNFs (see Appendix B.7 ). and extensive sampling necessary to obtain reliable estimates for the relative effective sample size (ESS). Therefore, we evaluate most of the transferable models on a subset of the test set. The dipeptides are randomly selected, but it is ensured that all amino acids are represented least once. Only the best performing model (TBG full) is evaluated on the whole test set of 100 dipeptides.
B.3 Hyperparameters
We report the model hyperparameters for the different model architectures describes in Section 4. 5 As 23) neural networks d., have one hidden layer with Thiiden neurons and SiZ.U activation functions. The input size of the embedding Rembedding depends on the model architecture.
22

---

b)
MD
TBG full
MD
TBG f
d)
energy distribution
energy projection
energy projection
104
MD
vp
Free energyerrgyy AkT
MD.
Free energy
33
300
Energy
kg7
T1C0
Figure 15: Results for the AC dipeptide (a) Sample generated with the TBG full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full model (right). (d) Energies of generated samples (e) Free energy projection along the to angle. (f) Free energy projection along the slowest transition (TICO).
b)
MD
TBG full
MD
TBG M
d)
Classical energy
energy projection
energy projection
MD
Free emergy AgT
2.0
Free emerge
300
TDO 1a mosigttd
Energy keT
TICO
Figure 16: Results for the NF dipeptide (a) Sample generated with the TBG full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full model (right). (d) Energies of generated samples (c) Free energy projection along the &o angle. (f) Free energy projection along the slowest transition (TICO).
report training hyperparameters for the different model architectures in Table 6 should be noted that all TBG models are trained in an identical manner if the training set is identical. We use the ADAM optimizer for all experiments (74). For the dipeptide training, each batch consists of three samples for each peptide.
B.4 Biasing target samples
As introduced in 23 can be beneficial to bias the training data in such way that unlikely states are more prominent. For alalnine dipeptide and many dipeptides, the positive states at = are often the unlikely ones and transition between the positive and negative &p states slow. For the alanine dipeptide dataset, the biasing methodology proposed in [23] is employed. Similarly, we bias the dipeptides based the von Mises distribution fvM. The weights are computed along the ip dihedral angle as
w(p)) -.v (f/1 1, 10) +
(12)
23

---

b)
MD
TBG full
MD
TBG M
d)
energy distribution
enetty brojection
energy projection
MD
Free energy kaT
Mo
105
Free energy
300
TEG ronoghted
300
Energy
kg
T1C0
Figure 17: Results for the GP dipeptide (a) Sample generated with the TBG + full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full model (right). (d) Energies of generated samples (e) Free energy projection along the angle. (f) Free energy projection along the slowest transition (TICO).
a)
b)
TBG full (smaller)
(smaller)
TICO
TICO
d)
Classical energy distribution
Free energy projection
Free energy projection TICA
304
Tel fill rovalle)
Te0 rovally anggitd
Ta0 Al (crake)
tevading menigood
Free energy kol
Thi -ull (cnater
(crol ent masiigtteed
300
Free energy
300
Energy
TICO
Figure 18: Results for the LW dipeptide for the TBG + full (smaller) model, which is trained on tenfold smaller trajectories than the TBG + full model. (a) Sample generated with the TBG + full (smaller) model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full (smaller) model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full (smaller) model (right). (d) Energies of generated samples. (e) Free energy projection along the y angle. (f) Free energy projection along the slowest. transition (TICO).
where r is computed based on the ratio of positive and negative yp states, such that both have nearly the same weight after the biasing.
B.5 Encoding of atom types
The atom type embedding a; one-hot vector of 54 classes. The classes are mostly defined by the atom type in the peptide topology. Therefore, only few atoms are indistinguishable such as hydrogen atoms that are bound to the same carbon or nitrogen atom. Moreover, we also treat oxygen atoms bound to the same carbon atom as indistinguishable, unless they are in the carboxyl group. Notably, we never treat particle groups as indistinguishable, such two CH3 groups bound the same carbon atom.
24

---

a)
MD
TBG Tull (smaller)
MD
TBC full (smaller)
d)
Classical distribution
Free projection
Free projection TICA
Free 1 19
Free enorgy kkT
180 (imiloe)
180 tom nowe glace
220
133
Energy
TICO
Figure 19: Results for the TD dipeptide for the TBG full (smaller) model, which is trained on tenfold smaller trajectories than the TBG + full model. (a) Sample generated with the TBG + full (smaller) model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full (smaller) model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG full (smaller) model (right). (d) Energies of generated samples (e) Free energy projection along the angle. (f) Free energy projection along the slowest transition (TICO).
Table 5: Model hyperparameters
Model
L
Thidden
Nembedding
Num. of parameters
alanine dipeptide
BG backbone
5
64
8
147599
TBG + full
5
64
15
149147
Dipeptides (2AA)
TBG
128
5
1044239
TBG backbone
128
13
1046295
TBG full
9
128
76
1062486
B.6 Effective samples sizes
The relative effective sample sizes (ESS) are computed with Kish's equation [56] as in prior work. For the alanine dipeptide experiments we use 2 x 10 samples for the forward ESS and 1 x 10 for the negative log likelihood computation. total 3 x 10 samples were used for each dipeptide the forward ESS, while 4.5 103 samples were employed for the negative log likelihood computation.
Table 6: Training hyperparameters
Mdoel
Batch size
Learning rate
Epochs
Training time
Alanine dipeptide
BG + backbone
256
5e-4/5e-5 5e-4/5e-5
500/500
3.5h
TBG full
256
500/500
3.5h
Dipeptides (2AA)
TBG
600
5e-4/5e-5/5e-6
7/7/7
4d
TBG backbone
600
5e-4/5e-5/5e-6
7/7/7
4d
TBG + full
600
5e-4/5e-5/5e-6
7/7/7
4d
TBG full (smaller)
600
5e-4/5e-5/5e-6
30/30/30
2.5d
25

---

MD 3h
TBG + full (sma 1h b)
HD
KD. 10
TBG full (ame le))
11.0
Go
T1C0
TICO
1100
MD
MD 1h
Tea a fall (sma 1h
d)
ND
HD-18
Mi
T1C0
T1CO
ND-
(smaller)
Figure 20: Comparison classical MD runs for hour (MD 1h) and the sampling with the TBG + full (smaller) model without weight computation for I hour (TBG full (smaller) 1h). The TICA plots of different peptides from the test set are shown. is important to note that the TICA projection is always computed with respect to the long MD trajectory (MD), All peptides stem from the test set. (a) CS dipeptide (b) EK dipeptide (c) KI dipeptide (d) LW dipeptide (e) RL dipeptide (f) TF dipeptide.
a) Boltzmann energy evaluations
b)
Boltzmann Emulator exploration time c)
Boltzmann Emulator energy differences
hum evehuttors
60k
TEC + masnsstaLi31
Tireworp moaniisnessit
353
-I- Thoaws mear. 56-
TBs lat monne
Eneray ditennnce An?
Timear---mn 14a--------
150 ful mearwd 22n1 lp?
250
200
150
100
16
50
AT
NY TO KU BL KQ destide
Pepride
Pepride
Figure 21: Boltzmann Emulator / Timewarp exploration experiments. Errors over 5 runs. Lower is better. (a) Number of energy evaluations until all states are found. (b) Wall-clock time until all
states found. (c) Energy difference between the mean MD energy and the mean sample energy generated by the two methods.
B.7 Computing resources
All training and inference was performed on single NVIDIA A100 GPUs with 80GB of RAM.
The training time for the models is reported in Appendix B.3 although it should be noted that significant amount of time was required for hyperparameter tuning. is estimated that least times the compute time reported in Appendix [.3] was necessary to identify suitable hyperparam- eters. Furthermore, inference with CNFs is expensive, especially if one requires the reweighting weights. Generating 3 x 10t samples with the large transferable models for the dipeptides requires approximately four days, whereas generating takes less than one day. However, generating samples without corresponding weights significantly accelerates the sampling process. the case of the dipeptides, the generation of 2 x 10 samples can be completed in less than one day. However, it should be noted that sampling can be done fully in parallel, as Boltzmann Generators generate independent samples.
10 samples for the alanine dipeptide experiments
26

---

a)
Classical energy distribution KS
b)
Classical energy distribution AT
3.05
MD
0.06
MD
Timewarp TBG full
3.04
0.05
TBG full
0.04
0.03
0.03
220
Energy kkT
Figure 22: Comparison of energies generated samples with long MD simulation, Timewarp exploration, and TBG full without reweighting. The energy distribution of the TBG full model matches the Boltzmann distribution generated with MD better. (a) For the dipeptide KS. (b) For the dipeptide AT.
27

---

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

---

Justification:
We only utilize theorems of prior work, which we reference accordingly.
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 referenced in 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 are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We explain how we performed our experiments in Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not 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 novel architecture, describing the architecture fully might suffice, or if the contribution is a 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 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) If the contribution is primarily new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is a new model (e.g., large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path 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?
29

---

reproducible, they Answer: [Yes]
Justification: Code, data, model checkpoints as well as detailed instructions are available here: https://osf
as stated in Appendix [BI| Nevertheless, high level training, evaluation and implementation details are described Appendix B and Section
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https://nnis..c)
souuddes/CoeeSubissssioooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
otw.a//sseeeeeeSSbbieaaamemooooo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We discuss the dataset details Appendix and refer to related work for datasets introduced in prior work.
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 a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide error bars for the alanine dipeptide experiments. In contrast, for the much more expensive transferable experiments, we utilize our computational resources to sample multitude of different peptides from the test set, rather than training and sampling distinct instances of the same architecture for the same peptide. Consequently, we obtain error bounds by averaging results over different test dipeptides rather than runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
30

---

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.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог 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 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: We discuss the required computational resources for this work for the training and inference in Appendix B.3 and Appendix B.7
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs 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 NeurIPS Code of Ethics httppp/://errrpp...pppppbtttttaa
Answer: [Yes]
Justification: The research conducted conforms with the NeurlPS Code of Ethics. All authors have read 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 to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impact of our work in Section g
Guidelines:
The answer NA means that there is no societal impact of the work performed.
31

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess 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 tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional 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 to 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 models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our models do not pose such risk, they are for molecular data.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
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: We give credit to used assets in this work in Appendix B.I and Appendix B.2 Our assets will available under the MIT CC BY 4.0 license.
Guidelines:
The answer NA means that the paper does not 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 used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
32

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We document our models in Appendix B
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing nor human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: No crowdsoureing nor human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
33

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
34


Reviews:
[{"review_id": "DRRltcwlPS", "reviewer": "Reviewer_C3Gj", "summary": "- This work builds upon Equivariant flow matching (Klein, 2023) and proposes a transferable Boltzmann Generator that sample Boltzmann distribution for molecules outside the training set. It is a common scenario that simulation data are often scarce for the system of interest and model transferability is highly desired.\n- The key contribution of this work is expanding an existing system-dependent framework to the transferable setting, and pinpointing the importance of proper topology encoding in previous models.\n- Together with the advantages of normalizing flows trained with score matching, the authors showed the model have superb iid sampling capability of Boltzmann distributions on the model system of alanine dipeptide (single protein) and strong transferability between different dipeptides.\n- Through comprehensive analyses, the authors showed that proposed model is accurate, data efficient, and has potential to sample unseen metastable states.", "questions": "1. Can the authors clarify the differences between three “architectures” of TBG-full, TBG-backbone, and TBG? This question remains throughout the reading and I current understand that the main difference is the atom encoding (i.e., topology informations) while all other model specifications are remain the same. Specifically:\n    - TBG+full: atom 54 encoding, unique encoding for most atoms.\n    - TBG/BG+backbone: unique encoding for backbone atoms + atom element types for side chain atoms.\n    - TBG: 5 atom element types\n    \n    If so, can the authors elaborate on the reason of the performance difference (albeit small) between BG-Backbone and TBG-full in Alanine dipeptide experiment?\n    \n\n1. As the isomorphism problem (chemical bonding and permutation symmetry) can be largely mitigated by unique atom encoding, the chirality remains a problem. While the authors uses a post-processing to filter valid molecules, one might wonder if the treatment scales as the system becomes larger and includes more chiral centers. Is there any other means to avoid chirality issues? \n\n1. Line 234 - 235: the authors state that the performance for the classical force field is much better than the semi-empirical one because the training data stems from the target distribution. While data was simulated using a classical force field AND subsequently relaxed with respect to the semi-empirical force field; and the objectives if to model the Boltzmann distribution defined by the semi-empirical force field (line 216 - 219). Can the authors help to better understand which is the “target distribution” and which distribution does the training data follow?\n\n1. Minor typo: eq 11 in the summation i → j for removing the center of mass", "limitations": "The authors extensively discussed the limitations and further justifications of current work, including 1) scaling to larger systems; 2) different flow matching forms; 3) different priors; 4) training peptides diversity; and 5) alternative architectures with enhanced performance.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- While being an followup work on Equivariant flow matching (Klein, 2023), this work pioneered in providing the first proof-of-the-concept work that Boltzmann Generators learned through flow matching has potential transferability to unseen systems.\n- During this exploration, they note the importance of encoding topology information through atom encoding to avoid indistinguishable atoms and transferability to unseen molecules.\n- Additionally, they conducted comprehensive analysis and ablation studies to answer key questions on models ability of 1) recovering the exact Boltzmann distribution and free energies along the reaction coordinates, 2) sample unknown metastable states, and 3) the influence of training on limited data. These information are valuable for communities to better understand model’s capability.\n- Overall, it is a pioneering, valuable, and well-written work that improving the solution towards an important question in AI for science.", "weaknesses": "The main weakness is the lack of experiments on the transferability in multiple systems. The authors only demonstrated the transferability using a small system of dipeptides, which has limited complexities. Despite being a proof-of-the-concept work, more efforts on improving model scalability and verifying on larger transferable systems would have improved the contribution and significance of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and questions. We now address their comments individually. To keep the rebuttal within the character limit, we often only cite parts of each question.\n\n> Lack of experiments on the transferability in multiple systems and scalability\n\nWe acknowledge the reviewer's point that testing our methods on larger systems or systems of varying sizes would be valuable. However, due to the significant computational resources required, we did not conduct these experiments. We hope that future research will explore these areas and that our proposed TBG framework will prove useful in such investigations.\n\n> Can the authors clarify the differences between the three “architectures” of TBG-full, TBG-backbone, and TBG?\n\nThe reviewer describes the differences mostly correctly. The TBG-full architecture includes not only encoding for the atom type but also for the amino acid to which the atom belongs. This additional encoding allows the TBG-full model to perform better on alanine dipeptide and other dipeptides, as it can better differentiate between atom pairs and is, therefore, more expressive. The EGNN generates updates based on pairwise inputs. If input pairs share the same embeddings as other pairs, the update contributions follow the same function. This uniform treatment may not be ideal, as atoms in different parts of the molecules can behave differently, yet the model must treat them the same if their embeddings are identical.\n\n> Architecture that respects chirality\n\nWe are unsure if this is feasible, as we start from Gaussian noise rather than from a configuration with the correct chirality. Therefore, a chirality-conserving architecture would not be beneficial. However, it is important to note that even the current TBG+full model can distinguish between different chiralities because the distances between atoms vary based on the chirality. The model is only equivariant with respect to global reflections, which we do not consider as representing different chiralities, since they can be resolved by mirroring.\n\n> Classical vs semi-empirical force field performance\n\nWe conducted two experiments with alanine dipeptide. In the first experiment, the target is the Boltzmann distribution defined by the classical force field. Here, the training data originates from an MD trajectory run with the classical force field, so the training distribution matches the target distribution.\n\nIn the second experiment, the target is the Boltzmann distribution specified by the semi-empirical force field. However, the training data was not generated from a long MD simulation with the semi-empirical force field. Instead, it was obtained by relaxing samples from the classical MD simulation for a few steps with respect to the semi-empirical force field. As a result, these training samples do not stem from the Boltzmann distribution defined by the semi-empirical force field. This discrepancy is evident from the free energy differences between the positive and negative phi states, as shown in Table 2 of [22], where the free energy difference in the training data significantly differs from the value obtained through Umbrella sampling with the semi-empirical force field.\n\nWe acknowledge that this may not have been clearly communicated in the current version of the paper and will provide a more detailed explanation in the final version.\n\n> Typo in eq. 11\n\nThanks! That is a good catch.\n\n[22] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In NeurIPS 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PbSWoIiBEo", "reviewer": "Reviewer_AbVU", "summary": "The authors proposes Transferable Boltzmann Generators (TBGs), which are transferrable for approximating the target distribution of unseen molecular datasets. TBGs are based on the graph-based continuous normalizing flow, and trained by using the simple flow matching. The authors experimentally demonstrate that the proposed model outperforms its competitors.", "questions": "Please see Weaknesses.", "limitations": "The authors adequately address the limitation, including the lack of evaluation on large-scale systems, and future directions, e.g., the use of other training objectives and prior distributions, for this work.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The motivation of the paper is very clear and relevant to the field of the AI + Science. The experimental results seem to be convincing, supported by thorough ablation studies (though it seems there is a lack of comparison with some previous works, please see Weaknesses).", "weaknesses": "I think the authors should put a more effort on explaining why the proposed TBGs are more transferrable than the previous models, for examples, coordinates-based Boltzmann generators [1]. I found that the proposed architecture is heavily based on [1], and the clear difference between [1] and the proposed one is the use of the auxiliary inputs $b_i$ and $c_i$ (and possibly the modification of $a_i$, though I am not sure what is “the topology for a classical force field”, which seems to be the added information compared to the previous $a_i$ of [22], i.e., “simply the atom types of the backbone encoding”). In the current version of this paper, the description about the proposed framework is only approximately one-page, and it is not unclear why the model can be transferred to the untrained dataset. To hedge this, the authors might explain the clear difference between BGs and TBGs, with emphasizing the usefulness of the incorporated inputs and modifications, with particular consideration for readers who are not familiar with physical chemistry.\n\n***\n\nIn Section 2, the authors introduce the Boltzmann Emulators, which are transferable due to removing the constraint of weighted sampling [2, 3]. However, the authors do not compare the proposed TBGs with these methods for the transferability experiment. I understand the Boltzmann Emulators lack the unbiased estimations, but I believe the experimental evidence should be provided for the completeness of the authors’ claim.\n\n***\n\n[1] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.\n\n[2] Osama Abdin and Philip M Kim. Pepflow: direct conformational sampling from peptide energy landscapes through hypernetwork-conditioned diffusion. bioRxiv, pages 2023–06, 2023.\n\n[3] 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and questions. We now address their comments individually.\n\n> I think the authors should put a more effort on explaining why the proposed TBGs are more transferrable than the previous models, for examples, coordinates-based Boltzmann generators (...) The authors should explain the clear difference between BGs and TBGs, with emphasizing the usefulness of the incorporated inputs and modifications, with particular consideration for readers who are not familiar with physical chemistry.\n\nWe agree with the reviewer that this aspect of our work is crucial and should be more prominently highlighted and presented in a way that is accessible to readers less familiar with physical chemistry.\n\nAs noted by the reviewer, our work builds upon the research presented in [1]. A major difference is in how we encode molecular structures in our model. In [1], the encoding methods involve either using solely the atom type (TBG) or adding different encodings for each atom in the peptide backbone (TBG + backbone).  The backbone of a peptide consists of a repeating sequence of atoms, which form the sequence. Each amino acids contributes the same atoms to the backbone. Therefore, there are not many different encodings for the atoms, and many atoms share the same encoding for both models.\nIn contrast, our proposed model assigns embeddings based on the positions of an atom in the peptide as well as the corresponding amino acid it belongs to. This results in different embeddings for nearly all atoms. This embedding makes the model more expressive, as the EGNN generates updates based on pair wise inputs. If input pairs share the same embeddings with other pairs, the update contributions follow the same function.  This uniform treatment may not be ideal, as atoms of the same atom type in different regions of the molecule can behave differently; however, the model treats them the same if their embeddings are identical.\n\nMoreover, [1] does not explore transferability to unseen systems. Our experiments demonstrate that their proposed architecture fails to generate significant effective sample sizes (ESS) for unseen dipeptides, as detailed in Section 5.2. In contrast, our TBG + full model achieves significant ESS across the entire test set (see Section 5.2 and Figure R4cd in the PDF of the global rebuttal).\n\nThus, we provide the first demonstration of a transferable Boltzmann Generator, which is highly relevant to the field of AI and Science, as noted by the reviewer. Additionally, we propose a framework for transferable Boltzmann Generators using continuous normalizing flows,  that can be used with different vector fields, e.g. different equivariant models, and also includes post-processing of generated samples. We hope this framework will facilitate future research in this area.\n\nWe will incorporate these points into the final version of the paper.\n\n> Comparison with Boltzmann Emulators\n\nWe thank the reviewer for the suggestion and agree that this comparison is valuable. We have conducted additional experiments using the TBG + full model as a Boltzmann Emulator and compared it to Timewarp [10], which is also suitable as a transferable Boltzmann Emulator for small peptide systems. We selected Timewarp over [2] and [3] because [2] is designed for larger systems, and [3] is not currently applicable to molecules with rings, which are common in many of the dipeptides studied in our work.\n\nAdditionally, we compared the TBG + full model used as a Boltzmann Generator with the Timewarp model combined with Metropolis-Hastings. Our findings demonstrate that the TBG + full model performs better in both scenarios. Please refer to the global rebuttal for detailed results.\n\n[10] Leon Klein et al. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In NeurIPS 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AS4NP3g3P3", "reviewer": "Reviewer_84N3", "summary": "The paper tackles the challenging and high-impact problem of sampling from the high-dimensional Boltzmann distribution of molecular systems. The proposed method allows to train a Boltzmann generator that is applicable to systems it has not been trained on as demonstrated for a dataset of different Dipeptides. The proposed method leverages flow matching, equivariant neural networks and reweighting to sample from the target Boltzmann distribution in an unbiased manner and without auto-correlation between generated samples.", "questions": "Did the authors compare the proposed method against the TimeWarp method?", "limitations": "The authors discuss the limitations of the proposed approach, including the limited size of molecular systems studied, as well as directions for further research, such as using more advanced flow matching approaches, informative prior distributions and more different model architectures.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The proposed method represents a significant improvement over the state-of-the-art Boltzmann generators given that prior approaches are only applicable to systems they have been trained on, limiting their applicability in practice.\n2. The experiments are appropriate to support the claim of transferability. The ablation studies based on limited and biased training data are relevant to judge practical applicability of the proposed method.\n3. The paper clearly highlights the difference between Boltzmann generators and emulators and the importance of the former to predict downstream quantities in a physically rigorous manner.", "weaknesses": "1. The main methodological innovation consists of including topology information in the embedding. Otherwise, standard methods such as flow matching and reweighting are used. This seems like an incremental improvement.\n2. In order to achieve real-world impact, a Boltzmann generator needs to be applicable to large molecular systems, but the dataset of Dipeptides in gas phase is not representative of real-world applications. Given that the proposed method depends strongly on reweighting generated samples to the true Boltzmann distribution and reweighting efficiency decreases exponentially with the system size, scaling the approach to large system of practical relevance appears challenging.\n3. The employed EGNN architecture is not considered state-of-the-art anymore and results with more recent models for the vector field would be interesting.\n4. The authors did not compare their method to the recent TimeWarp method, which tackles the same problem (even using the same dataset), but leverages an existing initial structure to sample in an unbiased, but autocorrelated manner using Hamiltonian Monte Carlo. To judge the merits of the proposed method, it seems important to compare these two methods (using the same topological embedding scheme) in terms of sampling efficiency for a given computational budget and transferability.  In addition, the TimeWarp paper also benchmarks their method against Tetrapeptides. It would be interesting to see which of the two methods generalizes better to larger systems and how this impacts sampling efficiency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review, insightful questions, and valuable suggestions. We will now address their comments and questions individually.\n\n> The main methodological innovation consists of including topology information in the embedding. Otherwise, standard methods such as flow matching and reweighting are used. This seems like an incremental improvement.\n\nWe agree with the reviewer that our proposed method builds upon prior work on Boltzmann Generators and flow matching, particularly the work of [22]. However, we demonstrate for the first time that a transferable Boltzmann Generator capable of achieving relevant effective sample sizes is feasible. This capability was not possible with previous architectures, as evidenced by our experiments. Therefore, we do not consider the improvement to be incremental. Additionally, we introduce a framework for transferable Boltzmann Generators using continuous normalizing flows. This framework is adaptable to various architectures, including different equivariant models, and incorporates post-processing of generated samples. We hope this framework will facilitate further research in this area.\n\n> Real world impact of Boltzmann Generators and scalability. \n\nWe agree with the reviewer that scaling Boltzmann Generators to larger systems remains a significant challenge. However, we see several potential pathways forward. Scaling to larger systems often involves coarse-graining, which typically results in the loss of an explicit energy function. In such cases, the transferable Boltzmann Generator would effectively become a transferable Boltzmann Emulator. Depending on the specific application, samples from a distribution close to the target Boltzmann distribution may be sufficient. Our results in the general rebuttal show that unweighted samples from the TBG + full model closely resemble the target Boltzmann distribution.\nIt remains to be explored to what extent Boltzmann Generators can be scaled to larger systems.\nAdditionally, small systems can still be highly relevant, especially when paired with more expensive force fields, such as semi-empirical or first-principles quantum-mechanical force fields. In these scenarios, energy evaluations become a primary bottleneck. Boltzmann Generators are particularly advantageous in this context, as they require several orders of magnitude fewer energy evaluations compared to MD simulations or other iterative methods, such as Timewarp [10]. One potential real world application would be the simulation of small ligands with expensive force fields. \n\nWe will expand our discussion of these points in the limitations section of the final version.\n\n> The employed EGNN architecture is not considered state-of-the-art anymore and results with more recent models for the vector field would be interesting.\n\nWe agree that there are more expressive EGNN architectures available. However, the one used in our work is particularly efficient to evaluate, which is crucial given the hundreds of vector field evaluations required for inference. Since our results are already quite promising—e.g., a median effective sample size (ESS) of over 10% on the test set (see Figure R4c) and nearly all generated samples have correct configurations (see Figure R4d)—we have focused on this efficient architecture. Experiments with different vector field architectures are left for future research, but we have provided a framework to facilitate such investigations.\n\n> Comparison to Timewarp [10] for different computational budgets\n\nWe thank the reviewer for this suggestion and agree that it is an important method to compare against. We present the results in the global rebuttal, where we demonstrate that our TBG + full model outperforms Timewarp, particularly in terms of energy evaluations, which are crucial as discussed above.\n\nWe use the Wasserstein distance between generated distributions as our evaluation metric rather than effective sample size (ESS). The reported ESS for Timewarp is based on autocorrelations of a specific collected variable, which contrasts with Kish's ESS used for Boltzmann Generators. This discrepancy makes direct comparison of ESS values challenging.\n\nTimewarp already uses a similar embedding scheme as our proposed methods, as it takes the current position and types of atoms as an input. \nTherefore, Timewarp already has all the information about the bond graph.\nWe concur with the reviewer that experiments with tetrapeptides would be valuable. However, due to the significant computational resources required for such larger experiments, we did not conduct them. We hope future research will explore these systems\n\n[10] Leon Klein et al. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In NeurIPS 2023.\n\n[22] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In NeurIPS 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CPOaxwwtPe", "reviewer": "Reviewer_GrZN", "summary": "The article \"Transferable Boltzmann Generators\" deals with the elaboration of a generative model based on Boltzmann Generators in order to estimate Boltzmann Distributions of molecules on which it has not been trained. The training procedure is based on continuous-time normalizing flow, where the trained model is then \"transferred\" to unseen molecules. Numerical experiments on dipeptides are provided to show how their method works.\nThe authors conclude that transferable Boltzmann Generators can be trained on some cases, when trained using normalizing flows.", "questions": "Can the authors comment on experiments on other types of data ?", "limitations": "The limitations are discussed adequately in the article.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The method introduced in this article is strongly evaluated on peptides. They illustrate with details how the samples generated by their methods (TBG+full or TBG+full reweighted) can work well (the latter model being better) and reproduce accurately the Ramachandran plots and the free energy projections.\n\nThe work provide a full set of details of the numerical implementation (number of parameters, batch size, learning rate, etc) in the appendices and many different experiments. In addition, their code is available online.", "weaknesses": "It is not clear how different their approach is from ref[22]. Maybe the authors could comment more on that.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and questions.\n\n> It is not clear how different their approach is from [22]. Maybe the authors could comment more on that.\n\nWe agree with the reviewer that this aspect of our work is crucial and should be more prominently highlighted.\n\nAs noted by the reviewer, our work builds upon the research presented in [22]. A major difference is in how we encode molecular structures in our model. In [22], the encoding methods involve either using solely the atom type (TBG) or adding different encodings for each atom in the peptide backbone (TBG + backbone). Therefore, there are not many different encodings for the atoms, and many atoms share the same encoding for both models.\nIn contrast, our proposed model assigns embeddings based on the positions of an atom in the peptide as well as the corresponding amino acid it belongs to. This results in different embeddings for nearly all atoms. This embedding makes the model more expressive, as the EGNN generates updates based on pair wise inputs. If input pairs share the same embeddings with other pairs, the update contributions follow the same function. This uniform treatment may not be ideal, as atoms of the same atom type in different regions of the molecule can behave differently; however, the model treats them the same if their embeddings are identical.\n\nMoreover, [22] does not explore transferability to unseen systems. Our experiments demonstrate that their proposed architecture fails to generate significant effective sample sizes (ESS) for unseen dipeptides, as detailed in Section 5.2. In contrast, our TBG + full model achieves significant ESS across the entire test set (see Section 5.2 and Figure R4cd in the PDF of the global rebuttal).\n\nThus, we provide the first demonstration of a transferable Boltzmann Generator, which is highly relevant to the field of AI and Science. Additionally, we propose a framework for transferable Boltzmann Generators using continuous normalizing flows,  that can be used with different vector fields, e.g. different equivariant models, and also includes post-processing of generated samples. We hope this framework will facilitate future research in this area.\n\nWe will incorporate these points into the final version of the paper.\n\n> Can the authors comment on experiments on other types of data ?\n\nOur method is specifically tailored for molecular systems, particularly peptides and proteins. However, it could be extended to other small molecules or peptides simulated with more expensive force fields, such as semi-empirical ones are even first principle quantum mechanic force fields.  In these scenarios, energy evaluations become a primary bottleneck. Boltzmann Generators are particularly advantageous in this context, as they require several orders of magnitude fewer energy evaluations compared to MD simulations or other iterative methods\n\nThere is potential to apply our approach to other types of data that exhibit symmetries, such as intersection traffic data [70] and point clouds [71]. Transferable models could be beneficial in these contexts as well. However, the choice of embeddings would need to be adapted based on the specific system of interest. It is challenging to predict how well our method would perform in these scenarios or identify potential additional obstacles.\n\nAdditionally, while we did not apply our model to a different dataset, we evaluated it for a different task—namely, as a Boltzmann Emulator. The results of this evaluation can be found in the global rebuttal.\n\n[22] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In NeurIPS 2023.\n\n[70] Berend Zwartsenberg et al. Conditional permutation invariant flows. arXiv preprint arXiv:2206.09021, 2022.\n\n[71] Marin Biloš and Stephan Günnemann. Equivariant normalizing flows for point processes and sets, 2021.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We appreciate the reviewers for their time and insightful feedback on our paper.

In response to the reviewers' request, we have included a comparison of our Transferable Boltzmann Generator (TBG) with the Timewarp model [10]. Unlike our approach, which generates independent samples, the Timewarp model learns to predict large time steps, which are combined with Metropolis-Hastings acceptance steps to ensure asymptotically unbiased samples. For this comparison, we used the publicly available weights for the Timewarp model, which was trained on the same dipeptide dataset we employed in our experiments. Notably, the Timewarp model required nearly three weeks of training on four A-100 GPUs, representing a training budget approximately 30 times greater than that of the TBG + full model.

We conducted a comparison of the two methods using 16 peptides from the test set. The corresponding figures can be found in the attached PDF.

1. In the first scenario, the goal is to generate samples from the target Boltzmann distribution for dipeptides unseen during training.     Timewarp accomplishes this by iteratively proposing samples and accepting or rejecting them using the Metropolis-Hastings algorithm, resulting in an asymptotically unbiased Markov chain. This approach aligns with the typical objective of a Boltzmann Generator.
We generated samples using both methods under fixed computational budgets and evaluated the results by comparing the Wasserstein distance between the generated Ramachandran plot and one generated from a long MD simulation. The computational budgets were: (a) 30,000 energy evaluations, (b) 12 hours, and (c) 24 hours of wall-clock simulation time on an A-100 GPU. As shown in Figure R1, the TBG + full model outperformed the Timewarp model across all budgets, particularly in terms of energy evaluations.
Energy evaluations are especially critical, as they become a major computational bottleneck with more complex force fields, such as semi-empirical or even quantum mechanical force fields. Furthermore, we present a comparison of the decay of the Wasserstein distance in Figure R2a ad R2b.
Notably, Timewarp occasionally fails to explore all states within the 24-hour computational budget, as shown in Figure R2d, whereas the TBG + full model consistently identifies all states within the same budget (see Figures in the main paper and appendix). We still compute the Wasserstein distance in these cases.

2. In the second scenario, the objective is to explore the most unlikely state, indicated by an orange circle in Figure R2d, which is the most unlikely state for most dipeptides. This goal aligns with that of a Boltzmann Emulator, as only approximate sampling from the Boltzmann distribution is required. For the TBG + full model, we do not apply reweighting, while for the Timewarp model, generated samples are only rejected if the energy increases by more than 300 $k_B T$ in consecutive samples, which is essential to prevent divergence [10]. Although not strictly necessary, we evaluated the energy of all samples generated with the TBG + full model and filtered out high-energy samples.
We compared the mean number of energy evaluations and the mean wall-clock time required to discover the state. Timewarp generally finds the state more quickly but requires more energy evaluations, as shown in Figure R3a and R3b. Additionally, the energies of samples generated with the Timewarp model tend to be higher than those generated with the TBG + full model, as illustrated in Figure R3c and R4a and R4b. Therefore, the distribution generated by the TBG + full model more closely approximates the target Boltzmann distribution.

For all Timewarp experiments, we used a proposal batch size of 100, as recommended in [10] for A-100 GPUs.

We will include these experiments in the final version of the paper. 

[10] Leon Klein et al. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In NeurIPS 2023.

Meta-review: The article presents a generative method for equilibrium samples of molecular systems that allows approximate sampling from the target distribution of unseen systems. 

Reviewers indicate the article deals with a challenging and high impact problem and expands an existing system-dependent framework to the transferable setting. They find that it provides an important proof of concept, presents comprehensive analysis and ablation studies, and provides a full set of details on the numerical implementation. They praised it as a pioneering, valuable, and well-written work towards an important question in AI for science, and a significant improvement over the state of the art. Some weaknesses were listed including on the methodology, comparisons, and scalability. The rebuttal offered a comparison with Timewarp that was suggested by reviewers and resolved most concerns in initial reviews, prompting raised review scores. 

In view of the positive reception of the article, I am recommending accept. Authors are asked to implement the promised additions in the final version of the article, particularly discussion of embedding and Timewarp experiments.
Decision: Accept (poster)