Scalable Normalizing Flows Enable Boltzmann Generators for Macromolecules

1. … the ESS should be reported in the paper … In addition, also reweighted energies distribution could be shown…

   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.

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

   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.

3. The authors should consider discussing this limitation in their work, along with citing relevant work on transferable models for MD acceleration

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

4. The code is currently not available

The code will be made available upon acceptance.

5. However, from the notation in Equation (6) it seems as they measure the distance of the backbone bond and torsion angles

In the line after equation (8), we specify that equations (7) and (8) are mean and covariance of the vectorized backbone atom distance matrices.

6. Could the authors comment on how the energy distribution looks like when all high energy samples are discarded?

This is a good point and we thank you for your suggestion. In the caption of Table 1, we note that statistics for the energy are reported after filtering out samples with energy higher than the median. We also took your suggestion in point (1), and in Fig S.2, we present a re-weighted energy distribution for the samples generated for protein G and villin HP35.

7. Can the authors elaborate on the time requirements for equilibrium Molecular Dynamics (MD) simulations in comparison to the time required for sampling with a Boltzmann generator?

The MD simulations for the training data took approximately 1 week to complete. On a single A100 GPU, it takes approximately 20 minutes to sample one million structures. The test set independent simulation took just under 1 week to complete.

1. 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].
2. 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).
3. 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.
4. 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.

[1] Kohler et al. Flow-matching - efficient coarse-graining molecular dynamics without forces. arXiv:2203.11167, 2022. [2] Midgley et al. Flow annealed importance sampling bootstrap. ICLR, 2023. [3] Jumper et al. Highly accurate protein structure prediction with AlphaFold. Nature, 2021.

1. … sections 2.2 and 2.3 may be helpful to put into context when discussed in the method …

   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.

2. Can’t I just run the simulation? What does this give us that a simulation doesn’t?

   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.

3. While I like the Von Mises Distribution, it is odd to me that it isn’t consistent across proteins

   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.

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

   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.

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

   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.

6. How do “metastable” conformations differ from stable conformations?

   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.

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

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

8. “In addition, for larger systems, maximum likelihood training often results in high-energy generated samples.” What does this mean in practice?

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

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

   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.

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.

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

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

[1] Noe et al. Boltzmann generators: sampling equilibrium states of many-body systems with deep learning. Science, 2019.

[2] Midgley et al. Flow annealed importance sampling bootstrap. ICLR, 2023.