Summary
This paper proposes a text-conditional generative model of molecules in 3D space by pairing a diffusion model with a text encoder used to guide the denoising process. The authors then perform a series of experiments to validate the proposed approach, showing how it can generate molecules conditioned on one or several properties, as well as handle more free-form textual descriptions.
Strengths
(S1): This work explores an important topic of molecule generation. While 2D-based generative models have long been adopted in the pharma industry, models operating in 3D directly are a newer frontier, which many practitioners are excited about, and so developing such models is worthwhile.
(S2): The high-level design seems sensible, and it makes use of relatively modern DL components. The text conditioning idea is interesting from an ML point of view (even if I'm not sure about its practicality – see (W2)).
(S3): The aspects of the method that the authors choose to talk about in depth are explained quite clearly (however, many other aspects seem to be not talked about – see (W1)).
Weaknesses
(W1): Many aspects of this work are not clear to me.
- (a) Many existing diffusion-based models for generating molecules (or point clouds more generally) have a caveat around number of atoms (points), which has to be fixed beforehand. Is it also the case here? When sampling from the model, do you sample the number of atoms separately? Is that conditioned on the text?
- (b) How does $\Gamma$ work? As I understand, it is a model mapping from text to a molecular conformation, i.e. the output is a set. How is this set decoded? Are there any issues in pretraining this component related to the fact that the atoms do not have a predefined ordering?
- (c) What are the different training stages involved here? I understand $\Gamma$ is trained beforehand and frozen; is the unconditional EDM also pretrained (without involving text) and frozen? Are the parameters of the linear transformation $\theta$ then trained separately, in a third stage? I can't fully match this with Figure 2 which suggests two stages, not three.
- (d) Why is the form of conditioning (linear combination of reference geometry and the current latent) so constrained? It seems the text encoder has very limited ability to "send information" to the diffusion process, as it cannot "send" arbitrary activations and rather something that is constrained to be an actual geometry. Is this design informed by having very little paired data (i.e. samples having both molecular conformation and text description at the same time)?
(W2): While the text conditioning is theoretically nice and flexible, I fear it may not be fully practical in a real drug design setting. In those scenarios, the generation would typically be also conditioned on a property that is specific to a given drug discovery target; for those properties there would normally be very few samples available [1], and they wouldn't be known upfront, so it's not clear how such properties would fit in the text-conditioned model (one may have to retrain the model for each project?). On the other hand, this issue is not specific to this work, and it is perhaps why simple 2D-based models (utilizing autoencoders or genetic algorithms) were so far more widely adopted in pharma than 3D-based ones, as the former are easier to condition on arbitrary properties not known beforehand. That being said, including 3D in the generative process is generally an important direction, and those classes of models are less mature than the 2D-based ones, so a smaller practicality is perhaps fine at this stage. At the very least, I would hope for some more discussion of this limitation in the "Limitations" section i.e. of the fact that the properties the model can condition on have to be known beforehand when training the text encoder, and this assumption may not fully reflect the real-world usecase.
=== Other comments ===
(O1): Table 2 presents some results, bolding the best numbers, however, the differences between most of these results are tiny, and they don't look statistically significant. So, first, some significance test would be good here, and second, could the authors discuss why the differences are so small? It seems the novelty/stability percentages are influenced much more strongly by the property being conditioned on than the generative model itself, which I found somewhat surprising and counter-intuitive.
(O2): Given that the related works discussion mentions [2], it would also make sense to refer to more modern extensions of that work [3,4].
=== Nitpicks ===
Below I list nitpicks (e.g. typos, grammar errors), which did not have a significant impact on my review score, but it would be good to fix those to improve the paper further.
- "By “unconditional”, some studies (…) craft atom coordinates and types without external constraints. By “conditional”, (…)" - the use of the word "by" seems off to me
- "searching suitable molecules in the drug design" -> I would drop "the" in this context
- "$\mathcal{N}(O,I_n)$" - should the $O$ be $0$?
- "rests on the assumption that, with the model distribution $p(G) = p(x, h)$ remains invariant" - I would either say "if (…) remains invariant" or "with (…) remaining invariant"
- "Experiment" (title of Section 4) - maybe "Experiments", as there are several
- missing space before the citation of Igashov et al
- "novelty is enhanced instead in comparison to the baseline" - the use of the word "instead" seems slightly off to me
- "Mvaluation Metrics" (title of Appendix B.1)
=== References ===
[1] "FS-Mol: A Few-Shot Learning Dataset of Molecules"
[2] "Junction Tree Variational Autoencoder for Molecular Graph Generation"
[3] "Learning to Extend Molecular Scaffolds with Structural Motifs"
[4] "Hierarchical Generation of Molecular Graphs using Structural Motifs"