Transferable Boltzmann Generators

Authors: leon  klein
frank  noé
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:52:36

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.

Main paper: 
 introduction: 
 Generative models have demonstrated remarkable success in the physical sciences, including protein structure prediction [1, 2, 3] , generation of de novo molecules [4, 5, 6, 7] , and efficiently generating samples from the Boltzmann distribution [8, 9, 10] . In this work, we will focus on the latter for molecular systems, which represents a promising avenue for addressing the sampling problem. The sampling problem refers to the long-standing challenge in statistical physics to generate samples from equilibrium Boltzmann distributions µ(x) ∝ exp (-U (x)/k B T ), where U (x) is the potential energy of the system, k B 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 to the necessity of performing small update steps, in the order of femtoseconds, for stability. This is especially challenging in the 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, 11] . One such method is the Boltzmann Generators (BG) [8] . In this work, we refer to BGs as a model that allows for the approximate sampling of the Boltzmann distribution of interest and the subsequent reweighting to the unbiased target distribution. If the model is only capable of generating approximate samples, which may stem from a subset of the Boltzmann distribution, we refer to them as Boltzmann Emulators 1 . Boltzmann Generators transform a typically simple prior distribution into an approximation of the target Boltzmann distribution through a normalizing flow [12, 13, 14, 15] . Once generated, samples can be reweighted to align with the unbiased target distribution. The effectiveness of this reweighting hinges on how closely the generated distribution approximates the target. As a result, it is possible to obtain uncorrelated and unbiased samples from the target Boltzmann distribution, potentially achieving significant speed-ups compared to classical MD simulations.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) [16, 17] , as opposed to coupling flows [18] . Recently, flow matching [19, 20, 21, 22] has emerged as an alternative training method for CNFs. This approach is simulation-free, enabling more efficient training of CNFs.Thus far, Boltzmann Generators have been limited by the requirement to train them on the specific system of interest. This training process demands a significant amount of time, making it difficult to achieve substantial speed-ups over classical MD simulations. Consequently, there is a 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 a framework for transferable Boltzmann Generators based on 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 transferability 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 [11] , which employs large time steps instead of generating independent samples as our method does. 
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 to generate samples from Boltzmann distributions of molecules [25, 26, 27, 28, 29, 30] , as well as lattice systems [26, 31, 32, 33] . Most BGs for molecular systems require system-specific featurizations such as internal coordinates [8, 34, 26, 27, 35, 36, 30] . Only recently, BGs for small molecular systems in Cartesian coordinates were introduced [23, 24] , using CNFs and coupling flows, respectively. Equivariant normalizing flows [37, 38, 39, 28, 23, 40] 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 [41, 36, 42] , which makes the learning task simpler. However, all previous Boltzmann Generators are not transferable. Arguably, the work of [7] represents an exception, as they are able to generate samples from unseen conditional (Boltzmann) distributions in 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 [11] . Instead of generating independent samples, they learn large time steps and combine these with Metropolis-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 [43] , who propose a three stage transferable CNF model to learn peptide ensembles. [44] use flow matching to learn distributions of proteins, while [45] utilize diffusion models. [46] build a transferable Boltzmann Emulator for small molecules. Others aim to additionally also capture the correct dynamics of the molecular systems, such as [47] , who use a diffusion model to predict transition probabilities. Scaling to larger systems often requires coarse graining [48, 49, 44] , e.g. describing amino acids by a single bead rather than the individual atoms. However, this approach precludes the possibility of reweighting to the Boltzmann distribution.A 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, 6, 50, 51, 52, 7] . 
boltzmann generators and normalizing flows: 
 Here, we describe Boltzmann Generators and normalizing flows, which are a central part of our proposed transferable Boltzmann Generator framework. We follow the notation of [23] .Boltzmann Generators (BGs) [8] combine an exact likelihood deep generative model and a reweighting algorithm to reweight the generated distribution to the target Boltzmann distribution. The exact likelihood deep generative model is trained to generate samples from a distribution p(x) that is close to the target Boltzmann distribution µ(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 ∼ p(x) with the exact likelihood model and then computing corresponding importance weights w(x) = µ(x)/p(x) for each sample. These allow to reweight generated samples to the target Boltzmann distribution µ(x). It is possible to estimate observables of interest (asymptotically unbiased) using the weights w(x) with importance sampling viaEQUATIONFurthermore, these reweighting weights can be employed to assess the efficiency of trained BGs by computing the effective sample size (ESS) with Kish's equation [53] . In this work, we will compute the relative ESS, rather than the absolute one, and refer to it as ESS.Normalizing flows [15, 54] are a type of deep generative model used to learn complex probability densities µ(x) by transforming a prior distribution q(x) through an invertible transformation f θ : R n → R n , resulting in the push-forward distribution p(x). Continuous Normalizing Flows (CNFs) [16, 17] are a specific kind of normalizing flow. For CNFs, the invertible transformation f t θ (x) is defined by the ordinary differential equationEQUATIONwhere v θ (t, x) : R n × [0, 1] → R n is a time-dependent vector field. The solution to this initial value problem provides the transformation equationEQUATIONwith f 1 θ (x) = pt (x). The corresponding change in log density from the prior to the push-forward distribution is described by the continuous change of variable equation log p(x) = log q(x) -EQUATIONEquivariant flows The energy of molecular systems is typically invariant under permutations of interchangeable particles and global rotations and translations. Consequently, it is advantageous for the push-forward distribution of a Boltzmann Generator to possess the same symmetries as the target system. In [37, 38] the authors demonstrate that the push-forward distribution p(x) of a permutation and rotation equivariant normalizing flow with a permutation and rotation invariant prior distribution, is again rotation and permutation invariant. Furthermore, [37] present a method to construct such equivariant CNFs by using an equivariant vector field v θ .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 v θ (t, x) throughEQUATIONThere are many possible parametrizations for the conditional vector field u t (x|z) and the conditional probability path p t (x|z). One of the most simple, but powerful possible parametrization isEQUATIONEQUATIONwhich we use in this work to train our models. For a more detailed description refer to [19, 22, 23, 50] . 
transferable boltzmann generators: 
 This section presents our proposed framework for transferable Boltzmann Generators (TBGs).Our proposed transferable Boltzmann Generator is based on a CNF. The corresponding vector field v θ (t, x) is parametrized by an O(D)and S(N )-equivariant graph neural network (EGNN) [39, 55, 43] , as commonly used in prior work, e.g. [23, 39] . Although, less expressive than other equivariant networks such as [56, 57, 58, 59] , it is faster to evaluate, which is important for CNFs as there can be hundreds of vector field calls during inference.The vector field v θ (t, x) consists of L consecutive layers. The position of the i-th particle x i is updated according to the following equations:EQUATIONEQUATIONEQUATIONEQUATIONwhere ϕ α represents different neural networks, d ij is the Euclidean distance between particle i and j, t is the time, a i is an embedding for the particle type, b i for the amino acid, and c i 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 O(D)and S(N )-invariant, as demonstrated in [60] .The embedding of each atom is constructed from three parts. The first part is the atom type a i , 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 is divided into 20 classes. The third part is the position of the amino acid in the peptide sequence. This embedding is similar to the embedding used in [43] for the rotamer embeddings.The amino acid and positional embeddings are only used for the transferable experiments. For more details see Appendix B.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 a special case of our architecture, wherein b i and c i are omitted and a i encodes the atom type or a backbone atom. Hence, there are 13 classes for a i . We refer to this architecture as TBG + backbone. Furthermore, we refer to the specific architecture employed in [23] as BG + backbone for the alanine dipeptide experiments. Note that using only a i 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 as simply TBG.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.Sampling with a transferable Boltzmann Generator, especially on unseen peptides, poses multiple challenges: (i) Some generated samples may not correspond to 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.4. 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 a reference bond graph. If the two graphs are isomorphic, we can conclude that the configuration is correct. For more details, see Appendix B.1. For (iii), we employ the code of [11] 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 configuration and chirality are considered valid samples from the Boltzmann distribution of interest. Table 1 : 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 [23] . -127.06 ± 0.12 6.03 ± 1.34%Model NLL (↓) ESS (↑) 
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 Appendix B.2. The training trajectory was generated by sampling with respect to a classical force field, and subsequently, 10 5 random samples were relaxed with respect to the semi-empirical GFN2-xTB force-field [61] for 100fs each. The objective is to train a Boltzmann Generator capable of sampling from the equilibrium Boltzmann distribution defined by the semi-empirical GFN2-xTB force-field efficiently and to recover the free energy surface along the slowest transition, i.e. the φ angle. Following the methodology outlined in [23] , the training data is biased towards the less probable (positive) φ 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 1 ). This is achieved with nearly the same amount of parameters and maintaining comparable training and inference times (see Appendix B.3). Furthermore, the correct free energy difference is recovered, as demonstrated in Appendix A.1.In [23] the authors utilized a 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 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 1 and Figure 1 . 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 in Figure 1 , both models recover the free energy landscape correctly. In our second experiment, we evaluate our model on dipeptides and show transferability. The dataset was introduced in [11] . The training set consists of 200 dipeptides, which were simulated each with a 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 2 and Figure 4a ), energies (Figure 2d ), the ratio of correct configurations (Table 2 ), 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 A.5).The results for the well-performing TBG + full model are presented for two exemplary peptides from the test set in Figure 2 and Figure 3 . 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 an exemplary Boltzmann Emulator, as it is capable of capturing all metastable states of the target Boltzmann distribution (Figure 2b ,c and Figure 3b, c ). However, it is furthermore also a capable Boltzmann Generator, as it allows for efficient reweighting (Figure 2d ,e,f and Figure 3d ,e,f) with good ESSs (Table 2 ). To identify different metastable states, we employ time-lagged independent component analysis (TICA) [62] , a 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 as [23] indicate that it can be advantageous to bias the training data towards states that are less probable, such as positive φ states, to recover free energy landscapes. Therefore, we bias 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 a 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 4bc ). 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 4a ). Nevertheless, these findings show that transferable Boltzmann Generators can be effectively trained using very small datasets, even when individual trajectories lack metastable states.In this section, we compare our TBG model with the Timewarp model [11] on the dipeptide dataset. Unlike our approach, which generates independent samples, the Timewarp model predicts large time steps, which can be cobined with Metropolis-Hastings acceptance steps to ensure asymptotically unbiased sampling. Additionally, Timewarp employs a coupling flow rather than a continuous flow. Similar to our TBG + full model, the Timewarp model also treats most particles as 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 Ramachandran 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 in Figure 5 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. 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 Metropolis-Hastings acceptance (exploration mode). Additional results and details can be found in Appendix A.7. 
discussion: 
 For the first time, we demonstrated the feasibility of training transferable Boltzmann Generators. We introduced a general framework for training and evaluating transferable Boltzmann Generators based on continuous normalizing flows. Furthermore, we developed a 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 on 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 to 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 e.g. shown in [11] , where they are able to train their transferable model on 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 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 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 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 [11] . 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, 50] 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 is to use a Harmonic prior distribution [63, 64] , 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 [65] , indicating that chemically informed priors do not enhance performance substantially compared to simpler uninformed priors for flow matching in molecular systems. Instead, the network architecture and inductive bias play a more crucial role.Despite conducting a series of ablation studies, we did not pursue the impact of a 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 a semi-empirical force field and training on this modified version, similar to 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 semi-empirical 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 as those proposed by [56, 57, 58, 59, 66, 67, 64] , 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. 
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 a known convergence criterion, making it impossible to confirm that all potential configurations have been identified, even with an infinite number of samples. This could lead to false claims about the results, potentially affecting subsequent applications.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] TBG + full (ours) Free energy difference / k B T 4.10 ± 0.26 4.10 ± 0.08 4.09 ± 0.05 test set. The dipeptides are randomly selected, but it 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 are in good agreement with the first ones as presented in Table 4 and Figure 10 .We report individual results for the different architectures in Figure 10c, d .To illustrate the performance of the TBG + full model, we present additional examples of dipeptides from the test set in Figure 11a -f and Figure 12a-f . Furthermore, we also again show results for training on the biased dataset in the same figures (Figure 11g ,h,i and Figure 12g ,h,i). As observed previously, the TBG + full (biased) model recovers the free energy landscape better than the TBG + full model, especially for the φ projections. We present additional examples of dipeptides from the test set for the TBG + full model in Figure 13 , Figure 14 , Figure 15 , Figure 16 , and Figure 17 .Furthermore, we present results for two example peptides from the test set for the TBG + full (smaller) model, which is trained on trajectories that are tenfold smaller than those used for the TBG + full model. These results are shown in Figure 18 and Figure 19 .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 all models (see Appendix B.2). However, since sampling without the Jacobian trace is less expensive and we do not require as many samples as for estimating the ESS, we also employ the TBG + full (smaller) model as a 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 set of test peptides, and nearly always finds all metastable states within less than one hour of wall clock time. This is a notable improvement over MD simulations, which often take longer to explore due to the iterative nature of MD. Some examples are shown in Figure 20 . This experiment shares similarities with the exploration mode of [11] , where they employ their model without the acceptance step and therefore also explore a 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, we do not apply reweighting, whereas in the Timewarp model, samples are rejected only if consecutive samples show an significant energy increase to prevent divergence [11] .Although reweighting is 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 of 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 to those from the TBG + full model. Consequently, the distribution produced by the TBG + full model more accurately approximates the target Boltzmann distribution.For all Timewarp experiments, we used a proposal batch size of 100, as recommended in [11] for A-100 GPUs. We primarily use the following code libraries: PyTorch (BSD-3) [68] , bgflow (MIT license) [8, 37] , torchdyn (Apache License 2.0) [69] , and NetworkX (BSD-3) [70] for validating graph isomorphisms. Additionally, we use the code from [39] (MIT license) for EGNNs, as well as the code from [11] (MIT license) and [23] (MIT license) for datasets and related evaluation methods.Our code is available here: https://osf.io/n8vz3/?view_only= 1052300a21bd43c08f700016728aa96e.The investigated benchmark systems were created in prior studies [23, 11] .The alanine dipeptide datasets were created in [23] (CC BY 4.0), we refer to them for detailed simulation details. The classical trajectory was created at T = 300K with the classical Amber ff99SBildn force-field. The subsequent relaxation was performed with the semi-empirical GFN2-xTB force-field [61] .Dipeptides (2AA dataset) The original dipeptide dataset as introduced in [11] (MIT License) is available here: https://huggingface.co/datasets/microsoft/timewarp. As this includes a lot of intermediate saved states and quantities, like energies, we create a smaller version with is available here: https://osf.io/n8vz3/?view_only=1052300a21bd43c08f700016728aa96e. For a comprehensive overview of the simulation details, refer to [11] . All dipeptides were simulated with a classical amber-14 force-field at T = 310K. The simulation of the training peptides were run for 50ns, while the test set peptides were run for 1µs.Choice of test set peptides Inference is a costly process with CNFs (see Appendix B.7), and extensive sampling is 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 at least once. Only the best performing model (TBG + full) is evaluated on the whole test set of 100 dipeptides.We report the model hyperparameters for the different model architectures as describes in Section 4.1 in Table 5 . As in [23] all neural networks ϕ α have one hidden layer with n hidden neurons and SiLU activation functions. The input size of the embedding n embedding depends on the model architecture. We report training hyperparameters for the different model architectures in Table 6 . It 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 [71] . For the dipeptide training, each batch consists of three samples for each peptide.As introduced in [23] , it can be beneficial to bias the training data in such a way that unlikely states are more prominent. For alalnine dipeptide and many dipeptides, the positive φ states at φ = 1 are often the unlikely ones and transition between the positive and negative φ states are slow. For the alanine dipeptide dataset, the biasing methodology proposed in [23] is employed. Similarly, we bias the dipeptides based on the von Mises distribution f vM . The weights ω are computed along the φ dihedral angle as where r is computed based on the ratio of positive and negative φ states, such that both have nearly the same weight after the biasing.EQUATIONThe atom type embedding a i is a one-hot vector of 54 classes. The classes are mostly defined by the atom type 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. 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 as two CH3 groups bound to the same carbon atom. The relative effective sample sizes (ESS) are computed with Kish's equation [53] as in prior work.For the alanine dipeptide experiments we use 2 × 10 5 samples for the forward ESS and 1 × 10 4 for the negative log likelihood computation. A total of × 10 4 samples were used for each dipeptide in the forward ESS, while 4.5 × 10 3 samples were employed for the negative log likelihood computation. 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 a significant amount of time was required for hyperparameter tuning. It is estimated that at least ten times the compute time reported in Appendix B.3 was necessary to identify suitable hyperparameters. Furthermore, inference with CNFs is expensive, especially if one requires the reweighting weights. Generating 3 × 10 4 samples with the large transferable models for the dipeptides requires approximately four days, whereas generating 2 × 10 5 samples for the alanine dipeptide experiments takes less than one day. However, generating samples without corresponding weights significantly accelerates the sampling process. In the case of the dipeptides, the generation of 2 × 10 5 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. • 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. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there 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 a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our models do not pose such a risk, as they are for molecular data.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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.1 and Appendix B.2. Our assets will be available under the MIT / CC BY 4.0 license.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: No crowdsourcing nor human subjects Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of 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. 
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?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. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.To the best of our knowledge Bowen Jing introduced the name first. 38th Conference on Neural Information Processing Systems (NeurIPS2024).ExperimentsIn this section, we compare our model with similar previous work on equivariant Boltzmann Generators[23] for alanine dipeptide. Moreover, we show the transferability of our model on dipeptides, where we compare our 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. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
- 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.
- 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.
- Through comprehensive analyses, the authors showed that proposed model is accurate, data efficient, and has potential to sample unseen metastable states.
strength_weakness: strength_weakness: 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.
questions: 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:
    - TBG+full: atom 54 encoding, unique encoding for most atoms.
    - TBG/BG+backbone: unique encoding for backbone atoms + atom element types for side chain atoms.
    - TBG: 5 atom element types
    
    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?
    

1. 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? 

1. 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?

1. Minor typo: eq 11 in the summation i → j for removing the center of mass
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

***

In 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.

***

[1] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.

[2] Osama Abdin and Philip M Kim. Pepflow: direct conformational sampling from peptide energy landscapes through hypernetwork-conditioned diffusion. bioRxiv, pages 2023–06, 2023.

[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.
questions: questions: Please see Weaknesses.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. 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.
4. 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.
questions: questions: Did the authors compare the proposed method against the TimeWarp method?
limitations: 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.
paper_summary: paper_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.
The authors conclude that transferable Boltzmann Generators can be trained on some cases, when trained using normalizing flows.
strength_weakness: strength_weakness: It is not clear how different their approach is from ref[22]. Maybe the authors could comment more on that.
questions: questions: Can the authors comment on experiments on other types of data ?
limitations: limitations: The limitations are discussed adequately in the article.

Comments: comment: We are glad that we could address all the reviewer's questions.
We are happy that we could address the reviewer's concerns and would like to thank them again for suggesting the additional experiments for Boltzmann Emulators.
Thank you for the authors' thoughtful response, especially the newly conducted comparison with Timewarp. Most of my concerns have been resolved, so I will be raising the review score.
I thank the authors for the clarification on the points raised - my questions have been resolved and happy to maintain the original rating.
We are pleased that the reviewer is satisfied with our response.
We would like to thank the reviewer once again for suggesting the Timewarp experiments. We will include the results, along with the suggested discussion of the embedding, in the final version of the paper.
I thank the authors for their precise answer which clarify the point that I have raised. I'm changing "contribution" from 2->3.
The provided comparisons to the TimeWarp method clearly improve the paper and underline the merit of TBGs over existing approaches.

In particular, the reasoning behind the different embedding requirements for Boltzmann Generators vs methods that have access to an initial structure (such as TimeWarp), is noteworthy and should be included in a more thorough discussion of the embeddings in the paper.