Any-Property-Conditional Molecule Generation with Self-Criticism using Spanning Trees

Thank you for your review. We address your questions and comments below.

Q1) STGG originally used a causal Transformer. We could use non-causal Transformer but for this we would need to switch to diffusion or use BERT-like masking. But in this case, we generate one atom and vertex after another in a random order, so we must follow causality in order to autoregressively generate the molecule. A limitation of diffusion is that we would not be able to use the powerful masking from STGG which masks future invalid tokens before the softmax. STGG masking is very powerful, so we pushed into this direction.

Q2) We were not very familiar with GANs for molecule generation, so we couldn’t say for sure. But in practice, discrete data (text, molecules, etc.) tends to work better with autoregressive models than with GANs. Hence why GANs are very rarely used for chatbots anymore.

Q3) We did initially consider Mamba, RKWV, and similar recurrent/SSSM models to improve inference speed. However, after consulting with chemists, we learned that it is extremely hard to synthesize molecules of very large sizes (>120 atoms for example), so the context length is limited (e.g., the largest molecule had 511 tokens on ChromophoreDB compared to the >=4096 context-length we normally see in LLMs). As long as context length is smaller than let say 1024-2048, FlashAttention is fast enough that there is no inference speed benefit for using these recurrent/SSSMs models.

Q4) Each time we sample a training molecule, we choose a random number t of properties to mask uniformly between 0 and T, then we randomize the order of the properties and mask the first t properties (thus, a random subset is chosen).The code looks like this: batch_choices = torch.arange(n_properties).unsqueeze(0).repeat(batch_size,1) < torch.randint(0, n_properties+1, (batch_size,1)) # random choose how many properties to keep (equal-prob of each amount of properties) batch_choices = shufflerow(batch_choices, 1) # shuffle [b, n_properties] to randomize which properties are selected

We addressed your comments: 1) We added a shortened table comparing all 3 best methods for unconditional generalization to be more quantitative in how we compare them; we still left the full table in the Appendix. This is more clear and easy for the reader to follow. 2) We changed the header to FCD to clarify that Distance stands for FCD.

Question:

We give a brief explanation of the path that led to this direction in the Background section. To be more specific: We chose STGG after doing an extensive literature review of molecules generation papers. We found that many papers recently focused on diffusion models respecting equivariance. However, the older SMILES-based models were performing just as well, which was surprising to us given all this time spent by researchers on newer diffusion methods. Out of the SMILES autoregressive methods, we found a slight variant: STGG, which is effectively a slightly modified SMILES vocabulary with the powerful masking of invalid tokens which improves performance. We found STGG to perform the best or among the best among many molecule generation datasets (see “Jang, Yunhui, Seul Lee, and Sungsoo Ahn. "A Simple and Scalable Representation for Graph Generation." The Twelfth International Conference on Learning Representations. 2023.” and “Ahn, Sungsoo, et al. "Spanning tree-based graph generation for molecules." International Conference on Learning Representations. 2021. “). This was the clear winner in our point of view. So we seeked to extend it, solve some of its remaining limitations, and improve it so that we could push it to a practical level. We also found STGG+ can perform very well in real-world applications (on proprietary datasets).

Thank you for your review. We address your questions and comments below.

Weaknesses

While our approach is intuitive, it also introduces several novel improvements. We propose many improvements, such as guidance in autoregressive models (extremely rarely used in the autoregressive literature), random guidance (we are unaware of prior works using this technique), self-property-predictor (we are not aware of molecule generation models using this, but there are likely similar ideas in the LLM world).

Questions

1. We updated Figure 2 to make self-criticism more clear (per your suggestion). Explanation: During training we predict the properties of the molecule. At generation we generate K molecules conditional on the desired properties. Then we mask the properties, effectively removing this information, and the model predict the properties. Out of the K molecules, we keep the best one, i.e., the molecule whose predicted properties are closest to the desired properties.

2. To clarify, we propose many novelties (as mentioned above) and we use them in order to improve property conditioning performance. Table 1 (especially the full version in Table 9, Appendix A.9) shows that our model generally has the best Validity, coverage, Frechet distance and property accuracy as well. Can you elaborate on why you believe that in Table 2, STGG+ outperforms by design?

3. They are inherently different approaches to solve a similar problem. We agree that Table 3 comparison is thus somewhat of an apples-to-oranges comparison. Even if the settings are different, we performed this comparison to see if generative models could perform well, as they have the advantage of being able to generate molecules for different property combinations without needing to retrain with a different reward. We believe that having this comparison, even if imperfect, is better than not.

4. We have ablations in Table 8 of Appendix A.8.

5. In theory it would be possible to train GFlowNet models on the 6 different OOD properties for Table 2. However, this would require a property predictor (increasing the complexity of this comparison), because if GFlowNet were to use RDKit directly to measure the true property, then it is an unfair advantage in terms of the effective dataset size. We only made the RL/GFlowNet comparison in Table 3 in order to have one small comparison with these methods. It would be more suitable to have a separate paper focusing on comparing RL/GFlowNet to generative models as this is not our main goal.

Paper suggestions

a) We initially thought that Figure 3 (supplementary) was a bit too big and could confuse the reader, which is why we had put it in the Appendix. But since you suggested putting it in the main text, we gladly changed the main Figure to Figure 3 (supplementary).

b) Per your suggestion, we made Figure 2 much simpler and more clear. This significantly improves readability.

Thank you for your review. We address your questions and comments below.

Weaknesses

1. The method still performs very well without the self-property-predictor (see k=1 in the Tables). In general, we found the self-property-predictor to perform well on most OOD properties, but in some cases, less so. As mentioned in lines 440-443 and 534-535, we found performance good except for OOD high logP conditioning values.

2. Just to be clear, we evaluated coverage, diversity, and similarity to training molecules in Table 1. We evaluated the % of valid, novel, and unique molecules (efficiency) in Table 2 and 4 (novelty means that the generated molecules are not found in the training dataset). We also look at Tanimoto diversity in Table 3. In all of these Tables, we see that our model performs well both in diversity metrics and property at the same time (efficiency close to 1 on Table 2 and Table 3, and highest coverage with high diversity in Table 1). It is possible to balance diversity and property fidelity trade-offs by changing the guidance and temperature (lower temperature (e.g., 0.7) leads to higher diversity, and lower guidance (e.g., 0.8) pushes less towards the property conditioning).

Questions:

1. In all experiments we compare STGG and STGG+ showing improvements. We also have ablations in Table 8 (Appendix A.8) starting from STGG and incrementally adding features to make it STGG+. The generation time is generally unimportant as it is sufficiently fast; there is no significant speed difference in our experience.

2. The compute cost is directly proportional to k, so k=5 makes the generation 5 times slower. But ultimately generation time is not a big factor since we are dealing with objects of at most 511 tokens (the largest molecule in Chromophore Db). We are very far from the LLM world which deals with > 4096 tokens and large models (we use only 3 layers!).

3. Can you point us to the table's column heads that are unclear? Another reviewer asked for clarifications to Table 1 header, so we will update these ones.

4. To clarify, Table 1 is the MAE (instead of the MinMAE) so it compares all generated molecules properties to the conditioning property. Also the Distance in Table 1 is the Frechet Distance which measures a distributional distance between training and generated molecules. For Table 3 we agree that it would be preferable to also show the MAE over top-10 or top-100 molecules in addition to the MinMAE. However, we reached out to the author of the VAE baselines and they didn’t have access to the code anymore. Replicating is non-trivial, so we focused on the top-1 which is the only metric used in the paper with the VAE baselines.

Thank you for your review. We address your questions and comments below.

Weaknesses:

1. These choices are applicable to molecules in general, we generally do not change most hyperparameters or modeling decisions between datasets. Making such choices is common to many generative models across domains, such as the need to set a maximum size or number of tokens. As such, these decisions do not affect the generality of the method to molecules. Also note you can ignore the self-criticism through the property predictor by simply setting k=1, which we evaluate throughout the paper.

2. Please refer to Table 8 of Appendix A.8 for ablation analysis, which already addresses most of your suggestions. We made this more prominent in the updated version.

3. We agree that it would be preferable to show the MAE with percentiles. However, we reached out to the author of the VAE baselines and they didn’t have access to their code anymore. Replicating is non-trivial, so for comparison, we focused on the MinMAE which is the only metric used in the paper with the VAE baselines. In most cases, it's impossible for the model to memorize extreme molecules, see the % of molecules below and above mean + 4SD below:

For Zinc:

MolWt, number of observations above u+4sd = 0

MolWt, number of observations below u-4sd = 0

MolLogP, number of observations above u+4sd = 0.0004% (1 case)

MolLogP, number of observations below u-4sd = 0.068%

QED, number of observations above u+4sd = 0

QED, number of observations below u-4sd = 0.019%

For Chromophore:

ExactMolWt, number of observations above u+4sd = 0.67%

ExactMolWt, number of observations below u-4sd = 0

MolLogP, number of observations above u+4sd = 0.78%

MolLogP, number of observations below u-4sd = 0

QED, number of observations above u+4sd = 0

QED, number of observations below u-4sd = 0

Thank you for your review. We address your questions and comments below.

The STGG+ representation of molecules is within the capabilities of SMILES representation

We want to reiterate that the core contributions of our approach are property-conditioning, improved architecture, improved spanning-tree,

1. We want to clarify that the “if-else conditions” is the main contribution of STGG over SMILES. It's through this masking with if-else conditions that STGG significantly improves the quality of generated molecules over SMILES. The representation is not necessarily better, but the STGG masking prevents invalid choices during inference rather than needing to discard invalid molecules after generation. The authors of STGG showed that it performs better than SMILES. We further extend and enhance STGG given its good performance. Your concerns are thus targeted specifically at STGG, which is the foundation we use. To reiterate, we made 5 sets of contributions (property conditioning, improved architecture, many improvements to STGG, auxiliary prediction loss for improved generalization and allowing us to do self-criticism, and classifier-free guidance with random sampling).

2. We agree that canonical SMILES could use the STGG masking, the STGG creator simply modified the SMILES slightly to make the masking a bit easier to define. Again, this is a concern targeted specifically at STGG, not our work which just leverages and extends STGG. We use STGG because its performance was shown to be better than SMILES (See “Ahn, Sungsoo, et al. "Spanning tree-based graph generation for molecules." International Conference on Learning Representations. 2021.“).

As a side note: The main novelty of STGG is the masking of tokens using a SMILES-like vocabulary for improved molecule validity. This is why we extend it to make it even better.

This work lacks clarity on the spanning-tree representation such as explicit/implicit hydrogen atoms and canonical representation.

1. Following STGG, the vocabulary uses explicit H’s, but after converting STGG to SMILES, we pass the SMILES to RDKit which can decide to ignore the explicit H and change the number of H. This is why the plots may be different since they are RDKit plots. We added this information in Appendix A.3.

2. There are many possible orderings in which to traverse the graph. Contrary to STGG, we do not canonicalize molecules, we use a random ordering of the molecules (a different random ordering is sampled for each molecule during training). Doing so improves generalization (see Table 8 in Appendix A.8). It is only mentioned on line 85 in ‘contributions’ that we use random ordering; we will mention it in the text because it should be made more clear.

3. We convert our STGG molecules to canonical SMILES after generation, thus novelty and uniqueness is correct. Thank you for the great suggestion; we added an example of the transition from SMILES to STGG to Canonical SMILES in Appendix A.3.

4. Training with only implicit H’s is expected to work, but we haven’t done this ourselves. We added these details to Appendix A.3.

The reported property MAE of the conditional generation needs more explanation.

1. We use RDKIT to evaluate the properties, except for Property Acc. in Table 1 which is based on a property-predictor as mentioned in the footnote of the table. We had also forgotten to mention that Table 3 uses MXMNet to evaluate the HOMO-LUMO Gap. We now mention all of this information more explicitly in the paper. We also revised the headers of Table 1 to make it more clear what the metrics are (as suggested by a reviewer).

• MinMAE means that for each generated molecule, we get the MAE between the generated molecule properties and the real properties (mean over all properties), then we report the minimum MAE (for the molecule closest to the true property). We prefer the MinMAE over the MeanMAE since our goal is to find new molecules with desired OOD properties. It is not important to our needs that the average molecule is close to the property, what matters most for material discovery is to find even just one such molecule that has the correct difficult-to-obtain set of properties. Finding the needle in the haystack is the ultimate goal.
• On another note, we would like to point out that the VAE baselines have no open-source code; we reach out to the authors and they responded that they don’t have access to the code anymore, so replication would be extremely difficult, which is why we ended up only using the same metric that they used (MinMAE).
• Note that Table 1, contrary to other tables, compares the pairwise generated molecules properties to real properties and average over all generated molecules so this specific table does not use the MinMAE and instead uses the MeanMAE.

[1/2]

I am still not convinced of the idea that "STGG improves over SMILES through the masking to prevent invalid choices". The "STGG+" seems more of "SMILES+". In other words, in my opinion, similar improvements can be achieved with SMILES.

• We agree that STGG and SMILES are very similar, and that our improvements could have been based on a SMILES representation instead of STGG. We chose to improve STGG instead of SMILES because it is the state-of-the-art on many molecule generation datasets [1,2]. This does not change the fact that our proposed model shows improvements on a wide variety of tasks.

The chosen metrics such as molWt, logP, and QED are not exactly challenging - they are not challenging enough to learn or predict. HOMO-LUMO gap is arguably the most interesting property in the paper but not many results are shown.

See our response to Reviewer GMmJ:

• We fully agree that the focus on simple properties by much of the existing literature in general is an issue for obtaining models that are directly useful in practice. However in Table 1, each of the 3 datasets have one experimental property that is non-trivial: 1) HIV: HIV virus replication inhibition, 2) BBBP: blood-brain barrier permeability, and 3) BACE: human β-secretase 1 inhibition. For other datasets, we follow standard protocol with a standard set of properties. From a machine learning perspective, the problems tackled by the literature are not solved considering the massive gap in performance between all methods compared to STGG+ and Graph DiT (for example, see Table 2 and 3).

Again, we fully agree that the focus on basic properties is not useful from a chemists perspective, but this is a limitation of the typical datasets and benchmarks used by much of the machine learning community in general. In this work, we chose these datasets because they allow for better comparability to existing work.

MinMAE is a weak metric - MeanMAE reflects the true capability of a model - Sure, the STGG+ decreased the error to a few magnitudes lower, but this seems more of a float-point accuracy improvement

• Table 2 experiments with the 3 datasets evaluate the Property Accuracy on HIV, BBBP, BACE (the meaningful properties mentioned above) and the MAE (i.e., MeanMAE) on the synthetic accessibility (SAS) property. We show strong performance improvements over recent state-of-the-art methods (far beyond float-point accuracies). As mentioned, we do not use the MeanMAE in Table 3 because we cannot replicate the baselines by Kwon et al. (2023). The authors of this paper themselves told us that they don’t have access to the code anymore, thus replication is not possible so we used their previous metrics which are Efficiency and MinMAE.
• Following your suggestion, in Appendix A.12, we added the average MAE over the top100 molecules for Table 3 and 5. Table 5 use 100 molecules, so this is the MeanMAE. Since Kwon et al. (2023) did not release code, we can only report on STGG and STGG+ models.

the authors picked some values that can be exactly achieved by certain molecules and STGG+ found one.

We want to clarify that we used +-4 standard-deviation, as mentioned in line 433. These are the exact same values as used by Kwon et al. (2023), we did not pick the values manually.

Sure, the STGG+ decreased the error to a few magnitudes lower, but this seems more of a float-point accuracy improvement

• Since a few reviewers were wondering about the difficulty of the OOD task, we added a row in both Tables 3 and 5 with the training data closest sample where we see that except for high QED, all +-4 SD OOD properties are far away from the closest sample in the training data: our model generate molecules much closer to the desired OOD property than the closest training data molecules.

Table 3 (STGG+ is much better except for QED-high where we perform slightly worse):

From 5.7e+1 to 1.0e−3 and 7.3e+1 to 6.1e−3 are more than floating point accuracy improvements.

Table 5 (STGG+ is much better except for QED-high where we perform equally):

From 1.40 to 0.47 and 9.62 to 0.35 are more than floating point accuracy improvements.

[1] Jang, Yunhui, Seul Lee, and Sungsoo Ahn. "A Simple and Scalable Representation for Graph Generation." The Twelfth International Conference on Learning Representations. 2023.

[2] Ahn, Sungsoo, et al. "Spanning tree-based graph generation for molecules." International Conference on Learning Representations. 2021.

Thank you for your review. We address your questions and comments below.

Weaknesses:

1. The main goal of our work is to show that our model is better able to generate molecules of the specified properties to condition on. One such property could be synthesizability. We agree that synthesizability is an important problem in molecule generation, but good metrics of synthesizability are difficult to obtain. Therefore, we focus on improving the property conditioning of properties in general, and with a suitable dataset that contains synthesizability as one of the properties, we should be able to improve generation of synthesizable molecules as well.

We will add comparisons of synthesizability (determined by AiZynthFinder) between training molecules (since synthesizability depends on the specific choice of reactants and starting molecules chosen, it's likely that many training molecules are labeled as not synthesizable, especially for ChromophoreDB), generated molecules in-distribution and out-of-distribution in time for the camera-ready-version.

2. We accounted for GraphGA, LSTM-HC in Table 1 (we had to trim the table as it was huge, the full table with many more baseline methods is in Appendix A.9) and they are not particularly good on the 3 datasets of that table. We used numbers from “Liu, Gang, et al. Graph Diffusion Transformers for Multi-Conditional Molecular Generation. arXiv preprint arXiv:2401.13858 (2024).”, so implementing those baselines into our code for our other tasks would take significantly more work. We believe that the much worse performance on 3 datasets is enough of a filter to justify not pursuing these methods further in the other datasets.

3. We fully agree that the focus on simple properties by much of the existing literature in general is an issue for obtaining models that are directly useful in practice. However in Table 1, each of the 3 datasets have one experimental property that is non-trivial: 1) HIV: HIV virus replication inhibition, 2) BBBP: blood-brain barrier permeability, and 3) BACE: human β-secretase 1 inhibition (mentioned in line 361-363). For other datasets, we follow standard protocol with a standard set of properties. From a machine learning perspective, the problems tackled by the literature are not solved considering the massive gap in performance between all methods compared to STGG+ and Graph DiT (for example, see Table 1 and 3).

4. Thank you for catching these issues. We added more details, cross-references, and fixed incorrect statements to the items that you mentioned.

Questions:

1. Contrary to STGG, we do not canonicalize molecules, we use a random ordering of the molecules (a different random ordering is sampled for each molecule during training).. Doing so improves generalization (see Table 8 in Appendix A.8). It was only mentioned on line 85 in ‘contributions’ that we used random ordering; we added a mention in the text because it should have been made more clear.

2. The non-STGG results are from “Jain, Moksh, et al. "Multi-objective gflownets." International conference on machine learning. PMLR, 2023”; they use 1M molecules, they do not show performance over time. The search space is parametrized differently, but all methods have the same output space.

Thank you for your review. We address your questions and comments below.

To clarify, we already provide property prediction performance of the self-property-predictor in Table 6 of Appendix A.6. For generation, we always compare with and without self-criticism in all experiment tables (k=1 vs k>1).

To answer your question,we found that the self-property-predictor was not always optimal at OOD and can make mistakes when choosing the best-out-of-k=5; we found this to happen for high logP conditioning values (we mention this in lines 440-443 and 534-535).