Paper: CoarsenConf: Equivariant Coarsening with Aggregated Attention for Molecular Conformer Generation (ICLR-2024)

Reviewer comment (Novelty):
Problem Significance: The authors may need to demonstrate the problem of conformation generation remains significant, in the context of the rapid development of 3D molecule generation from scratch.

Author response:
**> 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.

Reviewer comment (Novelty):
Novelty: There has been a line of work studying coarse-grained molecule generation in the community [1, 2]. The authors may need to further discuss the novelty of their methods in comparison to these existing methods.

Author response:
**> 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.

Reviewer comment (Novelty):
Novelty Again: There has been another work proposing its information fusion attention that is similar to the aggregated attention strategy [3].

Author response:
- 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.

**> Aggregated attention uniqueness**

We emphasize that our aggregated attention mechanism accomplishes something that is entirely unique:
- 1) Our aggregated attention mechanism uses queries in the raw 3D coordinate space to learn molecular structures from the learned latent representation. 
- 2) **Unlike CoarsenConf, all prior work limits all attention operations to the embedded feature space.**
- 3) **CoarsenConf does not use a reference library or fragment vocabulary to enhance the quality of generation**, as it operates directly on the input molecule. This is an important distinction, **as our method has more flexibility and can work on any given molecule**, not only those that can be parsed by the provided vocab. 
 - 4) We set up conformer generation as a translation problem to go from RDKit approximation to low-energy structures (Sec. 3). As a result, we construct attention queries from the raw source language to the target language embedding. This is equivalent to having raw source text being fed into an attention-head of a large language model for machine translation. **This is something that has never been done**, and **cannot be done in prior methods**, as the text domain requires an embedding, or some form of text-to-numerical or byte encoding to enable matrix multiplication. This is important, as learned embeddings come at a cost of loss of information, and our method avoids this issue entirely. **Our Aggregated Attention mechanism leverages the inductive bias of our 3D translation problem to circumvent the need for encoding, and as a result, uniquely operates on raw molecule structures.**

Overall, CoarsenConf is the first attention-based method to not require an embedding for the attention queries, as it mixes molecular structure space with a learned latent representation. We are also the first coarse-graining method to utilize attention to enable variable-length CG, compared to prior methods which force a fixed-length encoding.

Reviewer comment (Novelty):
Authors say that they are the first model to employ variable-length coarse-graining. As far as I know, it has already been used in the molecular field (e.g., Qiang B, Song Y, Xu M, et al. Coarse-to-fine: a hierarchical diffusion model for molecule generation in 3D, ICML2023).

Author response:
- 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.

Reviewer comment (Writing):
The model architecture needs further explanations. For example, the description of encoder architecture in your appendix is unclear. What are the inputs and outputs of the three modules? How to get outputs based on inputs? Please reorganize this section.

Author response:
- 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.**

Reviewer comment (Reproducibility):
The presentation needs further improvement. This paper does not provide any algorithm for the proposed method, making me very confused about a lot of training and inference details.

Author response:
- 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.

Reviewer comment (Evaluation):
As shown in Table 1, Table 5 and Table 6, the performance of this work seems to be suboptimal, especially for Recall.

Author response:
- We note in Tables 1 and 5 that CoarsenConf achieves the best QM9 precision, including with a number of additional baselines added (Table 5 in Appendix).
- For DRUGS, 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.
- 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). 
- 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).

Reviewer comment (Experiments):
Why is there a lack of comparison with Geodiff in many experiments? They have already released their code.

Author response:
- As we noted in Appendix J, we tried very hard but were not able to run the GeoDiff code because of dependency conflicts with torch-scatter and torch-geometric. However, the TD authors shared their GeoDiff test set conformers for us to run our evaluation. Furthermore, unlike Torsional Diffusion which has built-in functionality to generate conformers from a custom set of SMILES, the public GeoDiff code only has the ability to re-run the RMSD evaluations (building out that functionality for GeoDiff would require new dataloaders, preprocessing, and evaluation scripts). For these reasons, we could not evaluate GeoDiff on the protein docking evaluations which take 2+ weeks per method.
- We initially provided the same baselines as in Torsional Diffusion, and have now 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.
- One comment is that there is a train/test data mismatch issue, and it has been encountered before (more details can be found at [2]). 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, and AMR values can be roughly compared as both test sets are just different equal-sized subsets (see the updated Table 6). 
- These additions further demonstrate that CoarsenConf outperforms many prior methods, while using less data and compute, and offers a large improvement on real-world downstream tasks.

Reviewer comment (Evaluation):
Weak performance on GEOM-DRUGS compared to Torsional Diffusion. In addition, the performances of baselines are different from the literature. Also, Recall should be reported as well. Please explain what causes the discrepancy.

Author response:
**> 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).

Reviewer comment (Experiments):
Only few baselines are provided. If more baselines are provided, then it will be better to evaluate the effectiveness of the proposed method compared to recent techniques.

Author response:
**> 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.

Reviewer comment (Experiments):
QM9 and ZINC250 have been used for learning molecular distributions. Also, several representations have been used for generation such as string (SMILES, SELFIES) and (2D) graphs. Is it possible to compare MCG with other graphs or string based methods? Often papers provide other groups of approaches in tables as references with dim fonts.

Author response:
**> Smiles Representation**
- It is not possible to compare MCG with string or SMILES-based approaches (that do not convert it to a 2D graph), as to our knowledge, none exist. The task of MCG is inherently graph-based, as for each node we need a 3D position, and these positions must respect the 2D connectivity (and 3D energetics). This is not found in a SMILES string, due to its issue of enumeration and canonicalization [1]. Many language model approaches rely on SMILES enumeration or other augmentation techniques, which would violate the equivariance of the MCG task, but it is an interesting future research direction.
- We do think that language model representations may be able to be used along with equivariant graph models in a similar way to EquiFold [4], but this has not been done for molecules, and it is an open research area.
- We note there is a method, DMCG [2] (which we provide as a baseline; CoarsenConf has better performance), that uses AlphaFold’s structure module to learn conformers from a pure 2D input, but uses over 100x the model parameters. It is also inefficient due to its loss function enumeration, which is done because of the removal of explicit equivariance in parts of the model architecture.

Reviewer comment (Experiments):
Coarse-to-fine is a popular strategy in many applications. The conditional generation idea can be generalized in other directions. 2D to 3D generation is quite popular in the computer vision domain. Have you ever considered other coarse representations? and how robust/sensitive is the proposed pipeline to the quality of coarse-grain generation?

Author response:
**> Other coarse-grained representations**

- This is a good point that there are multiple coarse-graining strategies, and indeed, the best way to coarse-grain molecules is an active research area in chemistry and physics. We coarsened along contiguous rotatable fragments separated by single bonds, i.e. torsion angles. However, we did an ablation and tested Torsional Diffusion’s torsion angles, which allows 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.

**> How robust/sensitive is the proposed pipeline to the quality of coarse-grain generation**
- The quality of the latent representation was crucial to model performance. If there were any imbalance between the KL divergence in the learned prior and posterior, the model would not converge. As a result, we implemented self-targeting KL tuning. Unlike beta annealing where the KL weight gets larger, we set a target trade-off between the KL loss and reconstruction loss, and the model will adjust its weighting to achieve this goal [3]. By directly controlling the flow of information in our latent space, we created a robust latent representation that was relatively unphased by changes to hyperparameters in the underlying architecture.