Transferable Boltzmann Generators

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.

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

Questions

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

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations are discussed adequately in the article.

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.

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. 2. 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. 3. 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. 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

Did the authors compare the proposed method against the TimeWarp method?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

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.

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.

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

Please see Weaknesses.

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

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.

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.

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. - During this exploration, they note the importance of encoding topology information through atom encoding to avoid indistinguishable atoms and transferability to unseen molecules. - 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. - 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.

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

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.