Summary
This paper presents a novel gradient flow-based method to traverse the latent space of molecular generation models, known as ChemFlow. The authors instantiate their framework with a number of different flows inspired by dynamical systems. They also investigate the use of supervised and unsupervised guidance for the flow methods with the goal of optimising molecular properties of interest. The authors perform many experiments focused on molecular optimisation in order to evaluate their method. Specifically, they investigate unconstrained optimisation, similarity-constrained optimisation and multi-property optimisation.
Strengths
- The authors present a novel formulation for exploring chemical latent spaces and optimising molecular properties which is an important and relevant task within the pharmaceutical and molecular design domains. They present a few different versions of their method, including an implementation which can use a surrogate property prediction model to guide the flow to optimised chemical space, as well as an unsupervised implementation which aims to maximise structural changes to the molecule.
- The authors perform an extensive evaluation on molecule optimisation related tasks, comparing a number of different instantiations of their framework. They also benchmark against a previously introduced model for latent space traversal and a random traversal strategy.
- ChemFlow methods show very promising results in comparison to baselines, especially when only looking at top performing molecules or applying similarity-constrained optimisation.
- The authors also present a useful analysis of the latent changes under the random traversal strategy and an explanation for why random traversal can work reasonably well.
- The paper is mostly very well written and, the evaluations in particular, are very clearly presented and easy to follow.
Weaknesses
- I find the methodology section on its own quite unclear since it's not clear how to actually use the objective functions that are outlined. Appendix sections D4 and D5 are helpful but ideally it would be possible to follow the main text on its own. Particularly, I think the text would benefit from showing the full loss function in the methodology and including a short outline of the training and sampling procedures and referring to the appendix.
- The baselines for some tasks are a bit weak, particularly for the unconstrained molecular optimisation. For this task techniques like evolutionary algorithms and reinforcement learning fine-tuning have been proposed and widely used before. It would be very useful to see a comparison of ChemFlow with methods like these, as well as an evaluation of the training and sampling time for each.
Questions
- 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?
- For the unsupervised cases did you experiment with different values of k? Which values of k were used? It seems to me that k might need to be very large in general in order to find a flow which matches with an arbitrary property.
Other suggestions: I assume equation 8 should have $\phi^k$ instead of $\phi$? Additionally, the ordering of t and z in $\phi(.,.)$ or $\phi^k(.,.)$ is inconsistent - compare figure 1, the caption for figure 1 and equation 11 with all the other places were $\phi$ is used.
Limitations
- 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.
- 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, which are much more commonly used in practice for molecular generation.
- As far as I can tell 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.