Paper: Scalable Normalizing Flows Enable Boltzmann Generators for Macromolecules (ICLR-2024)

Reviewer comment (Evaluation):
One of the main goals of Boltzmann generators is to generate samples from the equilibrium Boltzmann distribution. As the output distribution will usually differ from the target distribution, Boltzmann generators employ reweighting to the target distribution (Noe et al. 2019). This allows to generate unbiased samples from the target distribution with the Boltzmann generator. However, efficiently doing reweighting requires that the output distribution is close to the target distribution. This can be measured using Kish's effective sample size (ESS), using the reweighting weights. Hence, the ESS should be reported in the paper. Merely presenting the mean energy values of samples does not provide a comprehensive assessment of the Boltzmann generator's ability to efficiently generate samples from the equilibrium distribution. In addition, also reweighted energies distribution could be shown.

Author response:
Thank you for your suggestion. While efficient Boltzmann reweighting is feasible for a small system like alanine dipeptide, the current work makes several modeling assumptions to scale to larger molecules. We model a distribution on a space with reduced dimensionality, which is not exactly the Boltzmann distribution. To be precise, we model $p(\mathbf{\tau}, \mathbf{\theta}_{bb} | L = \bar{L})$, where $\tau$ are torsion angles, $\theta\_{bb}$ denotes backbone bond angles, and $L$ and $\bar{L}$ denote other internal coordinates and their mean marginals, respectively. In addition, while the MD simulations for the training data were trained with explicit water, we use an implicit water model (for efficiency) when training with the energy function. This results in a wider range of energy values for our training data and generated samples for protein G and villin HP35 (the alanine dipeptide data is open source and is run with an implicit solvent). Due to precision issues, it is difficult to meaningfully compare importance weights.

    Thus, we report the effective sample size (ESS) (Table S.3) and display the reweighted energy distribution according to the energies computed for the training data distribution of protein G (Fig. S.2). As we can see from the table, for larger systems such as protein G and HP35, only the model that utilizes our novel architecture and multi-stage training strategy are capable of capturing a meaningful subset of the data distribution, as measured by ESS. This is primarily due to the fact that the other models generate samples with atomic clashes that dramatically increase their associated energies.

Reviewer comment (Experiments):
Boltzmann generators generally struggle to achieve a speed-up over traditional MD simulations, as they require data from these MD simulations to be trained. The comparison presented in Figure 4 may not be entirely fair, given that generating Boltzmann generator samples necessitates the entire training dataset, which probably took considerably more time to generate than the test data.

Author response:
Indeed, training the Boltzmann generator requires simulation data. However, we note in Sec 4.3 and Fig 4 that our model is capable of generating novel conformations that are not found in the simulation data. In fact, the representative structure highlighted in Fig 4b involves a largely bent hairpin loop that is not found in multiple-seeded MD simulations despite its low energy. In addition, the test data simulation took just under a week to generate while our model generated one million samples in 20 minutes.

Reviewer comment (Novelty):
The authors should consider discussing this limitation in their work, along with citing relevant work on transferable models for MD acceleration, such as references [1, 2, 3], and the recent advancements in Boltzmann generators for molecules in Cartesian coordinates, e.g. [4, 5]. These could be directly trained in a transferable manner, unlike Boltzmann generators operating in internal coordinates.

Author response:
Thank you for your suggestion. We have added some limitations of our model in the Discussion section. We have also added sections discussing transferable models for MD acceleration and recent advances in Cartesian coordinates for BGs in the supplements (Sec A V and A VI).

Reviewer comment (Presentation):
While I think the background section is helpful, sections 2.2 and 2.3 may be helpful to put into context when discussed in the method. While I realize they are described in detail in previous literature, it seems somewhat disjointed when all presented together.

Author response:
Thank you for the suggestion! If space permits, we will move these sections to be more inline with the discussion of the method in the camera-ready version.

Reviewer comment (Novelty):
Generally, I need some justification as to why ML models need to be fit to simulation data. Can’t I just run the simulation? What does this give us that a simulation doesn’t?

Author response:
The primary issue with MD simulations are that they often get stuck in local minima and are unable to sample past large energy barriers. MD simulations are heavily seed-dependent for larger systems, and are computationally expensive. In fact, the total time for training our model and generating one million samples for analysis is less than the time it took to run an independent MD simulation (which we used as our test set in Fig 4). In Sec 4.3, we demonstrate that our model is capable of identifying novel conformations of protein G that are not sampled in all of our simulation data. In particular Fig 4b displays a representative example of a low-energy generated sample not found in the simulation data; this new structure is characterized by a large bent conformation in the hair-pin loop which links beta-strand 3 and 4 of protein G.

Reviewer comment (Experiments):
“For protein G, we use a von Mises base distribution for dihedral coordinates; we noticed that using a von Mises base distribution improved training for the protein G system as compared to a uniform or truncated normal distribution.” Why isn’t this consistent? While I like the Von Mises Distribution, it is odd to me that it isn’t consistent across proteins. As a reader, I want a method I can apply to any protein of interest, or heuristics on which probability distribution to use.

Author response:
We observed an almost negligible difference between the results when using a von MIses distribution, a uniform distribution, or a truncated Gaussian distribution for the base distribution in the ADP and HP35 systems. We hypothesize that this is due to the simplicity of the two systems. In particular, HP35 is characterized primarily by alpha helices, for which backbone dihedrals are relatively fixed. We chose to report the results with the uniform base distribution for these two systems as this is consistent with what has been done in previous works. We hypothesize that for larger, more complex systems, the von Mises base distribution will lead to generally better results. Thus, we suggest that the reader stick with the von Mises base distribution in the general case.

Reviewer comment (Experiments):
How was the training and test set defined for Protein G? I’d prefer for any notion of novel conformations to be independently confirmed by the simulation. For Figure 4b and 4c, would they never be observed during the simulation?

Author response:
The training for protein G is generated by MD as described in Sec 4.1. An independent simulation was run for the test dataset. We confirmed that the novel conformations as described in Sec 4.3 and Fig 4(b,c) are low-energy structures according to the AMBER force field used for all simulations. While we do not know for sure if these conformations will *never* be observed in simulation, we note that the large bend in the hair-pin loop which links beta-strand 3 and 4 in the new conformation involves crossing a large energy barrier from the closest structure in the simulation dataset.

Reviewer comment (Evaluation):
“An overlay of high-resolution, lowest-energy all-atom structures of protein G generated by the BG model. This demonstrates that our model is capable of sampling low-energy conformations at atomic resolution.” - I don’t think this conclusion is validated by the Figure in 4d. The most you can say is that it kind of looks like the same protein?

Author response:
Thank you for your observation. We make this statement because the top ten closest generated structures by RMSD to the ground state structure are within 1.5 Angstroms.

Reviewer comment (Writing):
How do “metastable” conformations differ from stable conformations?

Author response:
Metastable conformations are functionally important  local energy-minimum structures that are different from global energy-minimum conformation (native structure). Both global and local energy-minimum conformations are stable but their equilibrium probabilities are different. The probability of the various conformations change under different experimental conditions.

Reviewer comment (Writing):
For equation 3, for the upper triangle of a distance matrix, over the sum, one can potentially use the (i < j) notation.

Author response:
Thank you for the suggestion! We’ve added this to Definition 3.1.

Reviewer comment (Writing):
“In addition, for larger systems, maximum likelihood training often results in high-energy generated samples.” What does this mean in practice?

Author response:
When we train a flow model with internal coordinates with just maximum likelihood training, as is often done in systems as small as alanine dipeptide, the model will often generate outputs that are characterized by clashes (Fig 3d).

Reviewer comment (Experiments):
Any reason in particular for 58 rotational quadratic spline coupling layers? Or any of the other subsequent number of coupling layers in section 4.1?

Author response:
These are generally arbitrary and chosen based on what fits our GPUs. For our architecture, we utilize 48 GA-RQS backbone coupling layers simply because this is 4 times the number of RQS layers used in [1], and 10 GA-RQS full latent size layers as any more would cause us to run out of memory when training with a batch size of 256. To make the baseline NSF model comparable, we also use 58 RQS layers with a comparable number of parameters for the transformations.

Reviewer comment (Writing):
Table 1: “Results for ∆D and u(·) that are within tolerable range are bold-faced,” What is “tolerable range”?

Author response:
Thank you for pointing this out. This is actually a typo. We meant “best results”.

Reviewer comment (Writing):
Can you define RMSF? I only have a rough understanding of it based of RMSD.

Author response:
The RMSF is a measure of the displacement of a particular atom relative to the reference structure, averaged over the number of atoms. Higher RMSF values indicate greater flexibility during the MD simulation. For atom $i$, where $\mathbf{x}_i$ is the coordinates of atom $i$ and $\langle \mathbf{x}_i \rangle$ is the ensemble average position of $i$, the RMSF is given by $$\text{RMSF}_i = \sqrt{\langle (\mathbf{x}_i - \langle \mathbf{x}_i \rangle )^2 \rangle }.$$

Reviewer comment (Presentation):
Figure 4a - What is the star? I don’t see this defined in the paper.

Author response:
In the caption of Fig 4, for the description of (b), we write “both structures are depicted as stars with their respective structural colors in (a)”. However, there is a typo in (b) regarding the colors. The cyan structure is the novel conformation while the magenta conformation is the closest structure in the training data by RMSD. We have corrected this in the manuscript. We have also updated the zoomed in image that demonstrates the large change in the hair-pin loop between beta strands 3 and 4 to make the conformational difference clearer.

Reviewer comment (Experiments):
Multi-stage training strategy introduces complexity that means significant more effort will be required to tune the algorithm. The third stage of training actually contains two stages. Although the rationale behind the individual losses is clear, the authors do not provide much rationale for each stage of the multi-stage training.

Author response:
1. Thank you for your insight and suggestions. Due to the complexity of conformation sampling, we feel that multi-stage training is necessary; the additional complexity is a cost for scaling up to these larger protein systems. What we observed is that NLL training is crucial in order to stably train with the other loss terms. In the literature [1,2], training of Boltzmann generators is done in two stages: 1) minimizing NLL, i.e., maximum likelihood training, and 2) training with a combination of NLL and reverse KL divergence, i.e., training by energy. In Sec 3.5, we highlight several reasons for this: 

    - Reverse KL divergence training is expensive and can lead to unstable training without a good initialization. 
    - We introduce the 2-Wasserstein training stage as a way to smooth the transition from pure maximum likelihood training to training by energy. 
    - The 2-Wasserstein loss is dropped at the end; this final stage can be seen as an energy “relaxation” stage. 

    We provide a summary of the hyperparameters in the supplements (Table S.1). Indeed, we could try to unify everything into a single training stage with a loss coefficient schedule, e.g., annealing. However, our previous attempts to move in this direction were not as successful as training in this multi-stage manner.

Reviewer comment (Novelty):
The architecture is not novel and I do not think this is a significant contribution - internal coordinates have been commonly used, and fixing bond angles of side chains has also been done in literature.

Author response:
2. Indeed, internal coordinates have been used in the past, as we have mentioned in Sec 3.2. However, they have not been utilized for Boltzmann generators in any system beyond alanine dipeptide. The key contributions of our architecture are: 1) GAUs, 2) rotary positional embeddings and the T5 relative positional bias for learning long range interactions, 3) split backbone and sidechain channels for efficient training. These three contributions combined allowed our model to successfully learn long range interactions for internal coordinates. We have added a table in our supplements (Table S.2) that gives results for ablating the training strategies for the baseline NSF architecture. We see that while our training strategy dramatically improves the performance of the NSF architecture, larger systems like villin HP35 and protein G are still not modeled well with the NSF architecture (see Table 1).

Reviewer comment (Experiments):
What is the rationale/intuition behind each stage of the multistage training? For example, how come the Wasserstein loss is dropped for the final training period? Can this be unified into a single training stage (potentially with annealing of some loss coefficients)?

Author response:
1. Thank you for your insight and suggestions. Due to the complexity of conformation sampling, we feel that multi-stage training is necessary; the additional complexity is a cost for scaling up to these larger protein systems. What we observed is that NLL training is crucial in order to stably train with the other loss terms. In the literature [1,2], training of Boltzmann generators is done in two stages: 1) minimizing NLL, i.e., maximum likelihood training, and 2) training with a combination of NLL and reverse KL divergence, i.e., training by energy. In Sec 3.5, we highlight several reasons for this: 

    - Reverse KL divergence training is expensive and can lead to unstable training without a good initialization. 
    - We introduce the 2-Wasserstein training stage as a way to smooth the transition from pure maximum likelihood training to training by energy. 
    - The 2-Wasserstein loss is dropped at the end; this final stage can be seen as an energy “relaxation” stage. 

    We provide a summary of the hyperparameters in the supplements (Table S.1). Indeed, we could try to unify everything into a single training stage with a loss coefficient schedule, e.g., annealing. However, our previous attempts to move in this direction were not as successful as training in this multi-stage manner.

Reviewer comment (Reproducibility):
The protein-G dataset introduced in the paper seems like it would be useful - will this dataset be released?

Author response:
3. We will indeed release the protein G dataset upon acceptance. Thank you for the suggestion.

Reviewer comment (Novelty):
I didn't fully understand the exact difference between the proposed "new architecture" and NSF. Is GAU the difference, or something else? Could you illustrate the difference more explicitly. Thanks.

Author response:
The new architecture refers to several changes: 1) GAUs, 2) rotary positional embeddings and the T5 relative positional bias for learning long range interactions, 3) split backbone and sidechain channels for efficient training. These three contributions combined allowed our model to successfully learn long range interactions for internal coordinates. In addition, we make use of reduced internal coordinates for representing our systems, which are notoriously more difficult to handle in the BG literature. We demonstrate that our model (architecture + training strategies) is capable of handling the difficulties of working with internal coordinates. The NSF baseline model which we used for comparison is the same as used in previous literature [1,2].

Reviewer comment (Experiments):
The experiment results are very good, but it is better to understand what is the key factor for the improvement by more ablation studies, such as NSF+NLL+KL+W2, Cartesian coordinate+your architecture+NLL+KL+W2, etc.

Author response:
We thank you for recognizing our experiment results. We noticed that our model architecture achieved much better results with NLL training than the baseline NSF model. We have added a full set of ablation results with the NSF architecture in the supplements (Appendix Sec. E, Table S.1). As we can see, while our different training strategies improve upon the baseline model with NLL training, the improvements are not as drastic as for our architecture. As for Cartesian coordinates, we purposefully utilize a reduced set of internal coordinates in order to reduce the dimension of the problem (Sec 3.2).

Reviewer comment (Evaluation):
In Table 1, I notice that in Protein G and HP35, some energy are as large as 10^6 to 10^10. What are the reasons for these extremely high energies? If they are caused by some naive reasons, can we quick-fix them to make the comparison more fair?

Author response:
These large energies are typically caused by van der waal’s interactions, i.e., atomic clashes (Sec 4.2). As we are modeling our inputs with internal coordinates, it is expected for training to result in clashes more than if we were to model with Cartesian coordinates (Sec 3.2). However, a key contribution in this work is the introduction of a training strategy and architecture that is capable of handling such issues that are typical when working with internal coordinates. Before submitting this work, we tried several methods to reduce the clashing in a fair manner such as by incorporating a violation loss as in AlphaFold2 [3]. However, these methods proved ineffective compared to the contributions in this paper for our internal coordinate inputs.

Reviewer comment (Evaluation):
In Table 1, NLL loss of Protein G is negative, why?

Author response:
For protein G, we use a von Mises base distribution for rotatable coordinates as this led to better empirical results (Sec 4.1). The continuous probability densities for computing the log likelihoods in the base distribution result in positive log likelihood values (negative NLL). In comparison, we use a uniform base distribution for ADP and HP35.