Navigating Chemical Space with Latent Flows

C1: Typo: L747 "we first verify if the learned variational poster also follows a Gaussian distribution and we find that it does learn so", poster -> posterior.

A: Thanks for pointing it out. It was a typo and we will fix it in the revised manuscript.

C2: In Table 2 and 3, why did ChemFlow perform best under mild similarity constraint, while suboptimal when the similarity threshold gets higher than 0.4? I think the similarity constraint is important, since usually in practical drug design people would expect to develop new therapeutics based on some drugs whose effects are already known, and the newly designed drug molecules are preferred to be similar so as to keep the effect.

A: Thank you for the question. We did not explicitly trained our methods for similarity constraint optimization. In the future work, we will explicitly encode more constraints to encourage our method to maintain similarity while optimizing the molecular properties. In addition, even though the absolute improvement of our method is not optimal for similarity threshold gets higher than 0.4, it has a optimal success rate compared to baseline methods.

C3: For Figure 4, why is the predictor deviating so much from ground truth? LogP, QED and SA don't seem very hard to learn as far as I know.

A: Thanks for the question. In Figure 4 (original paper), we set up an out-of-distribution scenario where the ground truth/prediction was the ZINC250k dataset (our VAE model was trained on a mixed of MOSES, ChEMBL and ZINC250k dataset where ZINC250k dataset was only a small fraction). However, the training set for the surrogate model was 10k randomly sampled molecular structures from the VAE model. We also only sampled 10k data points from ZINC250k which might not be representative of ZINC250k either. We have uploaded a new Figure 4 (Figure 2 in the uploaded PDF in the general response) showing the comparison on the full ZINC250k dataset (still out-of-distribution but slightly better).

We have further reported the training/test performance for the surrogate model, see Table 1 and 2 (training MAE and test MAE). To further study the out-of-distribution scenario we set up, we include Figure 3 in the uploaded PDF to show our training/test set for the surrogate model and ZINC250k dataset.

Table 1 Training Error (In-Distribution)

Table 2 Test Error (Out-of-Distribution)

C4: Please consider elaborating on why "the learned variational poster[ior] also follows a Gaussian distribution" and "a strong correlation between almost all molecular properties and their latent norms" would contribute to the observed result that "a random latent vector taking a random direction will change the molecular property smoothly and monotonically".

A: Thanks for the question. As VAE enforces the variational posterior to be a Gaussian distribution and we know that the geometry of high-dimensional Gaussian distribution is spheraical and the mass concentrates on the shell (if it is zero-centered, the norm of the sampled data is around $\sqrt{d}$ where d is the dimension of the data). In $R^d$, any random direction would eventually take it to the outer shell of the Gaussian ball with larger norms and further increase or decreases the property value.

Thus the observation of the correlation between property values and latent norms is an emergent geometry of the learned latent space. Under this geometry, it is reasonable that any random direction could lead to monotonic and smooth change of the property.

C5: This paper does not include a discussion on efficiency as compared with ChemSpace. Similar to the table 2 in chemspace

A: Thanks for suggesting discussing the efficiency. Below is a table that summarize the efficiency of all baselines and our methods. It is benchmarked by training the model described in the unconstrained optimization task. The inference time is the time for optimizing 100,000 molecules for 1 step using a batch size of 10,000.

Even though our method has a slower training time compared to ChemSpace because learning the flows requires learning a neural network instead of a linear model (e.g. linearSVM), they have similar inference time. This fast inference time ensures that our method is also capable of conducting high-throughput molecule optimization and screening for drug discovery.

C1: ChemFlow employs multiple approaches to learning different latent flows. However, the experiment's results show that different methods have different specialties, and the paper does not discuss the connection between flow learning and downstream tasks.

A: Thanks for the question. This is indeed a good point. For molecular optimization task, we might have better properties of one flow vs another (e.g. Langevin dynamics vs gradient flow for optimization). However, for the latent traversal or unsupervised setting, we often do not know which flow is better. Nevertheless, we can use this as a prior to perform traversal while simultaneously structuring the latent space. As indicated in Table 5 of the Appendix, we demonstrate that the different flows work as various inductive biases to improve the performance.

C2: The paper mentioned the out-of-distribution generation problem in Appendix D.7 and Sec. 4.2. ChemFlow has encountered such a problem in an unsupervised manner. This could hinder the utility of the learned latent space for scenarios with distribution shifts.

A: We appreciate reviewer points out that out-of-distribution generalization hinder the utility of the learned latent space. The out-of-distribution generalization is not a particular problem for our proposed method, but rather a limitation of the underlying generative models and the surrogate model used. We thus leave this for future study.

C3: Is it possible to visualize the vector field using a tool like t-SNE or UMAP to provide further insight into the entanglement of the molecular properties?

A: Thanks for suggesting visualizing the vector field. We provided the t-SNE visualization of the traversal trajectory for each different property using both supervised and unsupervised wave flow in Figure 1 of the uploaded PDF. The plot shows that almost all trajectories grow towards a unique direction in the t-SNE plot. This implies the disentanglement of learned directions and, thus, molecular properties. In addition, the figures display sinusoidal wave-shape trajectories, indicating the flow is following the wave-like dynamics.

In the unsupervised t-SNE plot, the trajectories of some properties overlap, such as plogP and sa. It is because some properties correlate with the same disentangled direction, so their traversal follows the same direction, thus the same trajectories.

C1: The description of the experimental section lacks detail, such as which software was used to measure the docking scores?

A: Thanks for pointing it out. We will thoroughly revise the experimental section in the revised manuscript to make sure all details are fully explained. Specifically, we used AutoDock [1] to calculate the docking score. We also used RDKit [2] and TDC [3] to calculate other molecular properties and structure similarity.

[1] Morris, G.M. et al. (2009) ‘AutoDock4 and AutoDockTools4: Automated docking with selective receptor flexibility’, Journal of computational chemistry, 30(16), pp. 2785–2791.

[2] RDKit: Open-source cheminformatics. RDKit. Available at: https://www.rdkit.org.

[3] Huang, K. et al. (2021) ‘Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development’, arXiv [cs.LG]. Available at: http://arxiv.org/abs/2102.09548.

C2: Why were only two specific targets evaluated in the docking score experiment?

A: Thanks for the comment. To respect previous literature, we follow the set-up as in [4] that target the binding sites of two human proteins ESR1 and ACAA1. Human estrogen receptor (ESR1) is chosen because it is a well-characterized protein with known disease relevance. Human peroxisomal acetyl-CoA acyl transferase 1 (ACAA1) is chosen to demonstrate the model’s de novo drug design ability as ACAA1 has no known binders.

In addition, computing docking scores requires significant computational resources, as calculating the docking score for 10,000 molecules generated by a single method takes approximately 20 hours on one GPU.

[4] Eckmann, P. et al. (2022) ‘LIMO: Latent Inceptionism for Targeted Molecule Generation’, Proceedings of machine learning research, 162, pp. 5777–5792.

C3: I would like to get some intuitive understanding : What problems arise if training a molecular property predictor in latent space and updating it directly based on its gradient? How is this issue typically addressed in the proposed method?

A: Thanks for the comments. Training a property predictor and updating the latent vector directly based on its gradient is the exact method used in the previous work LIMO [4] as gradient-based optimization can be viewed as discretization of a gradient flow. In our work, we generalize this method as different types of flows. Gradient-based optimization may suffer from several challenges, such as stuck in local minima and poor convergence, especially in a high-dimensional space with noisy gradient guidance.

In our proposed method, these issues can be improved using techniques like Langevin dynamics, which introduces diffusion noise into the gradient updates. Intuitively, this approach helps the model escape local minima by injecting stochasticity into the optimization process, thus promoting better exploration of the latent space.

C4: Table 1 displays properties of several indicators, but showcasing only the top 3 among numerous sampled molecules may lack sufficient persuasiveness. It would be better to include a broader distribution, such as the mean and so on.

Thank you for pointing it out. We have included a new table with mean, median, and standard deviation in the general response.

C1: The baselines for some tasks are a bit weak, particularly for the unconstrained molecular optimisation. It would be very useful to see a comparison of ChemFlow with methods like EA and RL fine-tuning.

A: Thanks for the suggestion. As our method focuses on the latent space of deep generative models, we compared mostly with methods for molecular optimization and traversal in a similar setup. However, we agree with the reviewer that the paper would benefit from additional baselines. We add an evolutionary algorithm-based (EA) approach to optimize molecules in the latent space. The pseudocode is provided as Algorithm 1 in the uploaded PDF in the general response. For a fair comparison, all methods in Table 1 in the uploaded PDF have the same number of oracle calls. The results in Table 1 show that our methods outperform all EA approaches.

It is possible to use reinforcement learning to guide the search in the latent space of molecular generative models but the main reason to use the latent space of generative model is to avoid the discrete nature of molecular structures and instead conduct optimization over a continuous space. We believe it will be nontrivial to propose a new reinforcement learning algorithm and thus leave it as a future study.

C2: For the unsupervised guidance, when you match flows with properties, what the correlations are computed between? Does this require you to have an existing dataset of molecule-property pairs?

A: We compute “the correlation between the property and a natural sequence (from 1 to time step t) along the optimization trajectory” (line 177). Specifically, we compute the following measurement:

$Spearman([P(m_1), P(m_2), …, P(m_t)], [1,2,...,t])$

where $P(\cdot)$ is the function to measure the real chemical property of a given molecule $m_t$ at time $t$.

Even though we do not require access to a dataset of molecule-property pairs, we do rely on minimal supervision to match the direction to the list of properties it may control. In reality, we do this by scoring the molecules using oracle functions.

C3: For the unsupervised cases did you experiment with different values of k? Which values of k were used?

We observe as long as the value of k is larger than the actual number of properties, the results will not be impacted much. We agree having an approximation of the number of properties is a prior knowledge of determining the hyperparameters, but a good practice is always to start with a relatively large k and select the active ones after training.

C4: While the authors evaluate their approach on optimising two properties simultaneously it's unclear how well this would work with a larger number, since their approach relies on simply adding the guidance terms together.

A: Thanks for the constructive advice. We get motivated by the disentanglement literature so we mainly focus on optimizing individual properties. Assuming each property is an energy $\phi^k$ so the stationary distribution is a Boltzmann distribution $p^k(x) \sim exp(\phi^k)$, then the summation of the guidance terms (i.e. energies) corresponds to sampling from the product distribution $\pi \sim exp(\sum \phi^k) = \prod exp(\phi^k)$. Even though this implicitly assumes independence among different properties, it is commonly used (also known as product of expert [1]) in machine learning, e.g. energy-based models [2]. We will leave how to leverage the correlation between different objectives as future work.

[1] Hinton, G.E., 2002. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8), pp.1771-1800.

[2] Du, Y. and Mordatch, I., 2019. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32.

C5: In this study the authors used a VAE with a fixed input size. It's unclear how easy it would be to apply a similar approach to a model with an arbitrary number of elements in an input or latent sequence, such as chemical language models or many diffusion-based models.

A: Thanks for the question. The fixed-length VAE model supports any input with a length less than its maximum limit. Indeed it can be less efficient for varied-length input such as text, however, it has still been widely used in languages and other sequence data. Moreover, despite that we validate the proposed method on a specific problem, molecular design, and select a fixed-length VAE model as the generative model, the method is not limited to fixed-length input.

As long as the model architecture has a well-defined latent space, e.g. [3] uses the bottleneck of U-Net as the latent space for diffusion models and [4] similarly in the attention head of large language models, it is possible to adopt our method in other networks to discover meaningful properties.

[3] Kwon, M., Jeong, J. and Uh, Y., Diffusion Models Already Have A Semantic Latent Space. In The Eleventh International Conference on Learning Representations.

[4] Li, K., Patel, O., Viégas, F., Pfister, H. and Wattenberg, M., 2024. Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36.

C6: The method as it stands requires a latent space optimisation to be done for every sample. This could lead to much longer sampling times than other methods such as RL fine-tuning which allows samples to be generated as normal but from an optimised model.

We focus on pre-trained generative models, where multiple objectives can be composed or removed at any time, while we admit this necessitates the extra sampling time as the model itself does not directly sample from the desired distribution, RL-based fine-tuning would limit the flexibility of the pre-trained model once it has been tuned for a specific task. The extra inference time is also minimal as it only requires a surrogate model evaluation (which is a relatively simple MLP model).