CoarsenConf: Equivariant Coarsening with Aggregated Attention for Molecular Conformer Generation

We are glad that you noticed that we achieve state-of-the-art performance in a number of important tasks, including structural precision, property prediction, and docking binding affinity.

We want to clarify that all of the references brought up here have no direct overlap with our method or task, with most being examples of methods for 2D molecule generation, which is a completely different task. We address the comments specifically below.

> Question: Problem Significance of MCG

• Molecular Conformer Generation (MCG) remains a significant task, as many real-world drug development workflows employ lead optimization, which is a molecule-conditioned task by definition [1]. MCG is directly used for drug discovery and is often the first step to protein docking workflows such as AutoDock Vina and DiffDock [2] (both of which do not generate the ligand from scratch). To use any Vina-based software, a conformer must first be generated. It is more realistic to do this from a 2D graph, as generating molecules from scratch still possesses significant validity issues [3]. This would cause immediate failures to any off-the-shelf docking method, such as those which are currently used in industrial drug discovery.
• Furthermore, most methods that generate 3D molecules from scratch are 10-1000x times slower than CoarsenConf (CC), making it infeasible for many domain use cases [4]. The scratch-generation methods pay little attention to the underlying energetics of the generated structures, which is an extreme importance for chemists and domain scientists more generally.
• Improvements to MCG are directly applicable to current polymer, catalyst, and nanoparticle design as it is a conditional task, whereas generating molecules from scratch cannot be as feasibly adapted due to the size and complexity of the domain shift [10].
• It is also important to note that ML methods are only recently barely beating RDKit for conformer generation, which has been around since 2015 [9], so even beyond CC, more work must be done to generate accurate and low-energy conformers. Furthermore, we added additional evaluations that showcase CC can generate better structures for protein docking and energy prediction.

> Question: Novelty

• We explain the difference between the given references and our method in more detail below:

• JT-VAE [5], uses a 2D pre-computed vocabulary, whereas we learn the fragments directly from their torsion angles for a completely different 3D task. Both methods are trying to accomplish different modeling goals and cannot be directly compared in any form of evaluation. The only things these two methods have in common is that they are VAEs and utilize fragments differently, with [5] not being SE(3)-equivariant.

• Zhang et al. [6] was designed for the structure-based drug discovery benchmarks. It relies on a fragment vocabulary (which we do not), and as a result, [6] cannot generalize to any molecule and is limited to those covered by the fixed vocabulary. CC can operate on any drug-like molecule, no matter the fragment breakdown, and can be extended to any coarsening criteria (see Appendix A).

  • Furthermore, CC learns a latent representation of the CG molecule (structure, and atomic features) that can be used in both autoregressive and non-autoregressive generation techniques.
  • [6] requires an iterative approach, where a motif binding site is selected and filled in by choosing one of the available fragments from its vocab. CoarsenConf also learns FG and CG positions, as well as distances and angles, leaving nothing fixed.
  • We note we cannot compare with the SBDD benchmarks as presented in [5] directly, as they generate a molecule conditioned on the protein, whereas we target the off-the-shelf docking workflow where the ligand is known, and then must create the 3D structure.
  • We are also the first to employ fragment-based generation in MCG.

• RetMol [7] is a 2D task with embedding cross attention, whereas our Aggregated Attention specifically uses conditional 3D queries to predict the desired 3D shape and coordinates. Furthermore, our attention implementation operates over learned CG coordinates, not invariant features.

  • Unlike to prior CG work, we utilize the entire learned representation, as discussed in Section 2. We note that information fusing is a central point of any cross-attention mechanism, and we use attention in a completely different way.
  • We have additional novelty in how we leverage attention to enable variable-length coarse graining without a vocabulary, and to learn the FG positions from the entire latent space (which prior CG methods are unable to do). [7] is another example of another entirely different task, where the only overlap is the molecular domain and the use of attention. Our attention scheme and usage have never been used before, as the queries are constructed to yield the shape necessary for equivariant 3D construction.

The problem significance is well-illustrated by the authors. For novelty, although CoarsenConf is the first work to study coarse-graining in molecular conformation generation, and it is exciting to hear that the proposed method is not limited to a fixed vocabulary, the main body of coarse-graining workflow largely follows previous works, and the key idea behind the aggregated attention - utilizing a reference library to enhance the quality of generation - is not quite new. But overall, I am convinced that the authors have made some unique contributions, so I updated my rating to 5.

We are glad to see that the reviewer found our improvements for MCG novel. We respond to all comments below, and discuss a number of ablation studies that we conducted, as well as the differences between CoarsenConf and Torsional Diffusion.

> Question: Finding the source of improvement in the architecture

This is a very interesting question as the architecture is complex, and we have done our best to run key ablations. We agree that testing this hypothesis is quite challenging, but important, as removing the coarsening/pooling layers of our architecture would strip the model of our main inductive biases. An example of a conditional VAE for MCG with no pooling/torsion angle information, nor equivariant updates, is found in CVGAE [1], which does not perform well (but was a crucial building block for all MCG methods). We conducted several ablations where we removed key components of our model to see if it was necessary for performance:

• 1. We softened the definition of what a torsion angle is, as done in Torsional Diffusion. Specifically, we allowed for non-physical rotations like rotations around double and triple bonds. Using this coarsening strategy, CoarsenConf had poor conformer reconstruction, and training would not converge. Our method is sensitive to the coarsening strategy being physically plausible, which is crucial to this specific task of low-energy conformer generation. It is possible that a different coarse-grained strategy may prove useful for different modeling objectives.
• 2. We also removed the distance-based auxiliary loss. This resulted in the interatomic distance error steadily growing during training resulting in infeasible sampled structures. Since unlike TD, CoarsenConf can update angles, coordinates, and distances at will, it is very important to correctly optimize all modalities. This is because if one modality is left uncontrolled, it tends to turn into a main source of error.
• 3. Along with experimenting with both autoregressive and non-autoregressive architectures for 3D molecule tasks, our main focus was to introduce more evaluation methods that prior MCG methods did not evaluate on, which mimic real-world docking workflows. Due to the SBDD docking evaluations taking over 2 weeks per method, we did not do an extensive hyperparameter sweep, but provided an initial hyperparameter ablation based on RMSD in the appendix.
• 4. To get closer to the question of what part of the architecture makes the most impact, we also ran a study on the amount of equivariant feature mixing in our decoder architecture (Eqn 10b). We found that increasing or decreasing the magnitude of mixing had little impact on the loss performance. From this, we believe that while the ability to break the problem down into fragments is beneficial, especially for downstream docking, there exist many equivalent versions of the architecture that differ with the weightings of the equivariant updates.

We are glad that you found the idea straightforward and novel, and pointed out that we had very good performance on multiple downstream applications. We answer all comments below, and have also made a number of updates to the paper to reflect our responses (including new algorithms to describe the encoder architecture, training, and inference), which are in red.

> Question: Variable-length Coarse-to-fine

• Thank you for bringing this paper to our attention! As it was published only 2 months (July ‘23 at ICML) before the ICLR deadline, we had not seen it at the time of the deadline. We have updated our paper to include this paper in Appendix A, and to clarify that we are the first to employ variable-length coarsening without a fragment vocabulary (as done in Qiang et al.) This is an important difference, as our method has more flexibility and can work on any given molecule, not only those that can be parsed by the provided vocab. When [1] and CoarsenConf are limited to the same atom types, if a new fragment not in the vocab were to be encountered at test time, [1] would have to be retrained from scratch whereas CoarsenConf can handle any molecular substructure.

> Question: Encoder Architecture clarity

• We have updated our paper to add new algorithms and clarify the description of the encoder architecture in Appendix D, including what the inputs and outputs are of the three modules, and getting outputs based on inputs. The encoder is set up similarly to a bi-LSTM, but in our case, we have three equivariant message-passing modules that pass information from the fine module to the pooling module, to the final coarse-grain module in each layer. We bring up the analogy to the LSTM, as the product of one module is used as the input for the next.
• At a high level, our encoder takes in (X, h) and (X^, h), the ground truth and RDkit approximated conformer (equivariant 3D positions/features, invariant atomic features), and produces Z and Z^, the CG latent equivariant representation for the ground truth and RDKit inputs.
• Each module is an EGNN update, which takes in equivariant features and invariant features, and produces a message-passing update for both, given by equations 6, 7, 8, and 9 for each of the provided modules. These equations detail the method of achieving outputs from inputs, but not the exact methodology for how the information is passed from one module to another, which we go into further detail next.
• In more detail, as part of our preprocessing for each molecule, we initialize a fine-grain (FG), pooling, and coarse-grain (CG) graph scaffolds (correct shape zero init) with a coarsening strategy that we know a priori (in our case, by torsion angles). Thus, we know the n, n+N, and N nodes for each graph, along with their chemical bond and distance-based auxiliary edges during preprocessing. Then, we can initialize graphs to conduct the hierarchical message passing. We refer to these as the FG, pooling, and CG graphs in the appendix, which provide the edges, and nodes that are used in their respective equations. In each module, the respective graph features are updated and then used to pass information from one module to another. This is shown in Algorithm 1 in Appendix D, which demonstrates the forward pass of a single batch.

> Question: Algorithm for method (Training and Inference)

• We have added training and inference algorithms in Appendix D, which are labeled Algorithm 2 and Algorithm 3. To summarize, we encode the data into our coarse-grain latent representation, sample from the latent space, use Aggregated Attention to decode from CG to FG, and run our FG decoder to produce the desired structures.

We appreciated that the reviewer found our work to be novel, and noted that our work hits all three areas of efficiency, flexibility, and quality of the generated conformers. We answer all of your comments below, add more baselines, and have also made updates to the paper, which are in red.

> Question: Performance on GEOM-DRUGS

• We note that CoarsenConf (CC) is very competitive with Torsional Diffusion (TD), and outperforms all other methods (as seen in Table 1). We are better than Torsional Diffusion on Median AMR, and are very close on the other GEOM-DRUGS metrics (for example, 52.0 to 52.1 on mean coverage). Importantly, as seen in Figure 4, CoarsenConf has the lowest RMSD distribution of any model. This means that in general, CC samples structures with lower error compared to all other models, including TD.
• Additionally, we include multiple experiments that test the quality of our generated DRUGS structures that go beyond RMSD. Note that RMSD is not frequently used by domain scientists, so we also did comprehensive testing of the quality of our generated conformers on downstream, real-world applications:
  • 1. property prediction, where CoarsenConf generates conformers with the most accurate energy: we reduce the $E_{min}$ error of TD by 50% (Table 2).
  • 2. Two downstream real-world applications, where our generated DRUGS structures are tested on how they would operate in practice for protein docking. CC outperforms all prior methods by a large margin, with improvements of up to 53% compared to the next best method (as you noted, in Tables 3 and 4, Figure 5).

> Question: Performance on baseline different from the literature.

• The numbers in our table are from generating molecules by running the different codebase models, following their public instructions. When available, we also used their published pre-trained weights. These numbers do have a slight discrepancy from the original papers. (We note that in running these baselines, some metrics, such as TD performance on AMR, are a little better than what was previously published in the paper).

> Question: Recall

• We publish recall values for all evaluated methods for 7 different sampling budgets in Appendix J, Figure 7. This issue of recall scores being very different based on the sampling budget was discussed in "Advantages and limitations of RMSD-based metrics” on page 7 (section 4.1).
• Unlike Precision, Recall is extremely dependent on the chosen sampling budget, with performances degrading by over 50% when restricted to more feasible sampling budgets (i.e., less computationally expensive). As seen in Figure 7, CC is comparable, and better in Coverage, on Recall compared to all other models (including TD) when the sampling budget is lower (and more realistic for how sampling would be done in the real world).

> Question: Only a few baselines are provided

• We initially provided the same baselines as in Torsional Diffusion. We went and added more model baselines from prior papers (including GraphDG, CGCF, ConfVAE, ConfGF, and DMCG), that were evaluated on a different random subset of GEOM molecules for RMSD evaluation. The paper has been updated to reflect this. CoarsenConf outperforms all of these additional baselines.
• One comment is that there is a train/test data mismatch issue, and it has been encountered before (more details can be found at: https://openreview.net/forum?id=w6fj2r62r_H&noteId=eTl4eNd2IDw). Other methods, like DMCG (that do not compare to TD), use a different but similar test set, and also only publish their DRUGS results for a coverage threshold of 1.25, compared to our 0.75 angstroms (a higher coverage threshold gives better results). Given this, we left blank coverage results, which are a function of AMR. AMR values can be roughly compared as both test sets are just different equal-sized subsets (see the updated Tables 5 and 6). As noted in the appendix, we tried very hard, but were not able to run the GeoDiff code for downstream applications because of dependency conflicts with torch-scatter and torch-geometric. However, the TD authors shared their DRUGS-generated conformers for GeoDiff evaluation.