Transferable Boltzmann Generators

We thank the reviewer for their review and questions.

It is not clear how different their approach is from [22]. Maybe the authors could comment more on that.

We agree with the reviewer that this aspect of our work is crucial and should be more prominently highlighted.

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

Moreover, [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).

Thus, 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.

We will incorporate these points into the final version of the paper.

Can the authors comment on experiments on other types of data ?

Our 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

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

Additionally, 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.

[22] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In NeurIPS 2023.

[70] Berend Zwartsenberg et al. Conditional permutation invariant flows. arXiv preprint arXiv:2206.09021, 2022.

[71] Marin Biloš and Stephan Günnemann. Equivariant normalizing flows for point processes and sets, 2021.

We thank the reviewer for their detailed review, insightful questions, and valuable suggestions. We will now address their comments and questions individually.

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.

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

Real world impact of Boltzmann Generators and scalability.

We 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. It remains to be explored to what extent Boltzmann Generators can be scaled to larger systems. 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 [10]. One potential real world application would be the simulation of small ligands with expensive force fields.

We will expand our discussion of these points in the limitations section of the final version.

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.

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

Comparison to Timewarp [10] for different computational budgets

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

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

Timewarp already uses a similar embedding scheme as our proposed methods, as it takes the current position and types of atoms as an input. Therefore, Timewarp already has all the information about the bond graph. We 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

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

[22] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In NeurIPS 2023.

We thank the reviewer for their review and questions. We now address their comments individually.

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.

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

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

Moreover, [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).

Thus, 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.

We will incorporate these points into the final version of the paper.

Comparison with Boltzmann Emulators

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

Additionally, 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.

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

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.

Lack of experiments on the transferability in multiple systems and scalability

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

Can the authors clarify the differences between the three “architectures” of TBG-full, TBG-backbone, and TBG?

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

Architecture that respects chirality

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

Classical vs semi-empirical force field performance

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

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

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

Typo in eq. 11

Thanks! That is a good catch.

[22] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In NeurIPS 2023.