Molecule Design by Latent Prompt Transformer

This work explores the challenging problem of molecule design by framing it as a conditional generative modeling task, where target biological properties or desired chemical constraints serve as conditioning variables. We propose the Latent Prompt Transformer (LPT), a novel generative model comprising three components: (1) a latent vector with a learnable prior distribution modeled by a neural transformation of Gaussian white noise; (2) a molecule generation model based on a causal Transformer, which uses the latent vector as a prompt; and (3) a property prediction model that predicts a molecule's target properties and/or constraint values using the latent prompt. LPT can be learned by maximum likelihood estimation on molecule-property pairs. During property optimization, the latent prompt is inferred from target properties and constraints through posterior sampling and then used to guide the autoregressive molecule generation. After initial training on existing molecules and their properties, we adopt an online learning algorithm to progressively shift the model distribution towards regions that support desired target properties. Experiments demonstrate that LPT not only effectively discovers useful molecules across single-objective, multi-objective, and structure-constrained optimization tasks, but also exhibits strong sample efficiency.

Paper

Similar papers

Peer review

Reviewer vrpK7/10 · confidence 3/52024-07-06

Summary

This paper presents LPT, a novel transformer model for conditional molecule sequence design and generation. LPT first generates latent vectors from a learnable prior distribution, then autoregressively generates molecule sequence taking the latent vector as prompt. Comprehensive experiments show that LPT achieves state-of-the-art performance on a variety of molecule, protein and DNA design benchmarks.

Strengths

- Propose a novel transformer model LPT and a novel framework to achieve property conditioned generation of molecule sequence. - Strong and solid experiment results on multiple benchmark datasets. - Good, clear and well-organized writing.

Weaknesses

- As SMILES strings of molecules have grammars so not every SMILES string can be decoded to a real molecule. Authors are encouraged to report validity rate of the generated SMILES strings. - What is the dimension/size of the latents in experiments? How will the dimension/size of the latents impact the performance? Authors are suggested to give a discussion or conduct ablation study experiments about this problem.

Questions

See Weaknesses part.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Reviewer rWJi4/10 · confidence 4/52024-07-12

Summary

The paper introduces an approach for molecule design, by leveraging recent advancements in conditional generative models for language and image generation. It contains three steps: (1) learnable prior distribution, (2) molecule generation model and (3) property prediction model. The experiments are comprehensive and convincing to me.

Strengths

1. The task that the proposed method tries to solve is of great importance in a lot of real-world applications such drug design and protein design. 2. The formulation of multi-objective optimization task is quite novel and well defined. 3. The experiments are comprehensive and convincing to me.

Weaknesses

1. Motivation to use the current framework is unclear to me. Why we use Langevin dynamics to sample latent variables rather than directly using the VAE framework to introduce the prior? 2. I think the dependence between molecules/properties cannot be secured as there is no reconstruction loss on molecules. As you may see in Eq. (5) and the equation under Eq. (6), the information when predicting y from z cannot be back-propagated to x. Then how to align x and y? 3. Minor: (a) line 42: would suggest to write down the full name of MCMC before using its abbreviation for general audience. (b) typo at equation (1): subscript of o^p_m(x).

Questions

1. At line 29, I don't think existing methods decouple the training of the generative model from the property conditioned optimization, such as [1] and [2]. 2. What's the difference between the proposed method and variational autoencoder (VAE)? I see both methods assume the prior on the latents, and there are works that generated molecules from VAE such as [1] and [2]. 3. I noticed that the paper proposes Langevin dynamics to get z, but why not directly use an encoder to encode z from data like VAE does? MCMC-based method seems more computationally heavy to me. 4. At line 154, I'm still confused about how to sample z | y using Langevin dynamics. Especially how can we compute the gradient of the p(z_0|y)? [1] Multi-objective Deep Data Generation with Correlated Property Control. NeurIPS 2022. [2] Property controllable variational autoencoder via invertible mutual dependence. ICLR 2020.

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes. Limitations are well discussed.

Authorsrebuttal2024-08-12

Thank you for your continued feedback. We'd like to address two key points: >1. Posterior collapse in VAE. For the posterior collapse issue in VAEs, the early stages of learning can present a significant mismatch between the encoder and generator models. This mismatch causes the latent variable $z$, inferred by the encoder, to be highly inaccurate. Consequently, the autoregressive decoder model $p(x_t | x_{<t}, z)$ tends to disregard this imprecise latent variable. Instead, it models the input observation primarily using its own parameters and the input $x_{<t}$. This occurs because $p(x_t|x_{<t})$ is often strong enough to generate the molecule on its own, with minimal reliance on $z$. This behavior ultimately leads to posterior collapse. In contrast, Langevin dynamics approaches the problem differently. It consistently samples from the true posterior distribution of the latent variable, bypassing the need for a learned encoder and directly utilizing the autoregressive decoder. This method tends to be more accurate, especially in the initial stages of the learning process. As a result, the sampled $z$ consistently contributes to modeling the input observation through the decoder, mitigating the risk of posterior collapse. > 2. joint probability of $p(x,y,z)$ The joint distribution of $(x, y, z)$ is $p(x, y, z) = p(z) p(x|z) p(y|z).$ The joint distribution of $(x, y)$ is $p(x, y) = \int p(z) p(x|z) p(y|z) dz$, so that their dependency is captured by the sharing of $z$. $z$ plays the role of information bottleneck, i.e., $x$ is predicted from $y$ via $z$, i.e., $$p(x|y) = \int p(x|z) p(z|y) dz.$$ Given $(x, y)$, $z $ can be sampled from $p(z|x, y) \propto p(z) p(y|z) p(x|z)$ as a function of $z$ with $x$ and $y$ fixed. For molecule design, $z$ can be sampled from $p(z|y) \propto p(z) p(y|z)$ as a function of $z$ with $y$ fixed and we generate $x$ conditional on $z$, i.e. $x\sim p(x|z)$ as a form of ancestral sampling using derived $p(x|y)$ above. The sampling from both $p(z|x, y)$ and $p(z|y)$ can be accomplished by Langevin dynamics. We shall make it more explicit in the revised version of our paper. Again, thanks for your insightful comments. Please don't hesitate to let us know if you have any questions.

Authorsrebuttal2024-08-13

Have we addressed your concerns?

Dear Reviewer, Thank you again for your follow-up questions. With only a short time left in the discussion period, we wanted to ensure we've adequately addressed all of your concerns. If you have any additional questions or if any points require further clarification, please don't hesitate to let us know. We're eager to provide any necessary information to fully address your review. We sincerely appreciate your time and valuable insights and hope you will consider supporting and championing our work. Thank you for your thoughtful evaluation. Authors

Reviewer jGf37/10 · confidence 4/52024-07-12

Summary

The authors introduce a new conditional generative model for molecules. The model is called the Latent Prompt Transformer (LPT). A conditional model capable of generating new molecules with desired target properties is very useful in de-novo molecular design since we often want to design new molecules with some set of target properties. The authors initially train the LPT model on existing pairs of known molecules and property values. They then iteratively shift the distribution of the model towards regions with desired target properties. This results in a model capable of generating new molecules that are likely to have the desired target properties. In the experimental results section, the authors apply LPT to optimization tasks where the goal is to find new molecules that have high objective values, where the objective value here is some measurable desirable property of the molecule (i.e. the molecule’s binding affinity to some target of interest). Results show that the LPT model succeeds at these tasks, generating new molecules with higher binding affinity than baseline approaches. Additionally, the authors show that LPT can succeed at multi-objective optimization (generating new molecules with more than one desired target property). In particular, they show that LPT can generate new molecules that achieve high binding affinity, as well as high QED (Quantitative Estimate of Druglikeness) and minimal SA (synthetic accessibility) scores. In this case, results show that LPT is still able to generate molecules with higher binding affinity scores than baseline methods, while also having QED scores above 0.4 and SA scores below 5. Finally, the authors also show that LPT can be successfully applied to a new structure-constrained optimization task.

Strengths

Originality: LPT is novel as it is (as far as I am aware) the first autoregressive conditional generative model designed specifically for the task of conditional molecule generation. Quality: The paper is very well-written. Additionally, the figures and tables are all of good quality - they are both easy to parse and do a nice job of displaying results. Figures 1 and 2 do a very nice job of illustrating the author’s method. I especially like the use of Figure 2 to illustrate the distribution shift towards molecules with higher binding affinity. Clarity: The paper is clear and easy to follow from start to finish. The figures and tables are clear and easy to read. The way the authors set up, trained, and applied the LPT model is clear. Significance: De-novo molecular design is one of the most relevant/significant tasks in computational biology. In particular, finding molecules that bind to targets of interest is at the core of de-novo drug design. The paper shows that the author’s LPT model outperforms baseline approaches for relevant molecular design tasks such as finding new molecules with high binding affinity to targets of interest. This paper is therefore significant and of interest to the community.

Weaknesses

Conditional generative models are not themselves novel as they have been successfully applied to a variety of text and image generation tasks. Much of the methodology in this paper involves taking existing methods in conditional generative modeling and applying them in a new domain: molecule generation. However, as far as I am aware this is the first time that one of these types of autoregressive conditional generative model has been designed for molecules, so this approach is novel from the perspective of methods for de-novo molecular design. I therefore think that this is a minor weakness and that this paper should be accepted. Latent-space Bayesian optimization (LS-BO) approaches are mentioned in related work as an alternative approach, but not directly compared to in experiments. I actually don’t think that this direct comparison is strictly necessary for the paper to be accepted because LS-BO is an orthogonal approach and the authors do compare LPT to a good number of state-of-the-art generative modeling approaches. However, a direct comparison showing that LPT performs as well as state-of-the-art LS-BO methods would strengthen the paper.

Questions

See Weaknesses section above for suggestions/points of discussion.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Reviewer D3n34/10 · confidence 4/52024-07-14

Summary

This work proposes a novel molecular optimization framework Latent Prompt Transformer (LPT),modeling latent distribution, molecule sequences and properties distributions conditioned on the latent distribution. They uses MCMC in MLE training and conditional generation. Additionally, they propose an online learning algorithm to extrapolate to feasible regions supporting the desired property. The framework demonstrates promising results across various molecular design tasks.

Strengths

1. The online learning approach aligns well with real-world scenarios. 2. The experiments are comprehensive, and the introduction of a new task for conditional generation and design trajectory comparison is expressive. 3. The paper discusses and makes efforts to alleviate the bottleneck of computational efficiency. It also improves sample efficiency by reweighting, making the method more feasible.

Weaknesses

- The clarity of the paper could be improved in several aspects: 1. Adding an illustration or diagram of the training and optimization process would make the methodology easier to understand. 2. Some formulas lack derivations and explanations, as raised in Question 1. - I am concerned that using MCMC sampling for training and inference may impact the model's practicality. Could you please compare the training and sampling speed with the baselines and discuss this issue in the limitations section? - Minor Issues: 1. Typo error: "2rd" in Table 1. 2. Using "Kd" values might lead readers to mistakenly think they are wet lab results, while in paper they represent docking scores. 3. In Algorithm 2, 'step (c) Update LPT on synthetic dataset using **Algorithm 2** using MLE.' Should it be Algorithm 1 here?

Questions

1. How is the transition from equation 4 to equation 5 derived? It would be beneficial to provide a detailed derivation for this step, as it is not immediately apparent. 2. When formulating p(y∣z), how do you choose σ\sigmaσ and how does it affect performance? Since properties can be uniquely determined by the molecule, i.e., p(y∣x) is deterministic, this suggests that in the probability model shown in Fig. 1 (Left), p(y∣z) should be very deterministic. However, this creates difficulties for sampling p_\theta(z_0|y) in eq.7, because when \sigma is small, \nabla\log p_\theta(z_0|y) is very small for most z_0, leading to very slow MCMC convergence, especially with poor initialization. 3. How robust is your model to the quality of the Oracle?

Rating

4

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes.

Authorsrebuttal2024-08-13

Have we addressed your concerns?

Dear Reviewer, With only a short time left in the discussion period, we wanted to ensure we've fully addressed your concerns. If you have any additional questions or require further clarification on any points, please don't hesitate to reach out. We greatly appreciate your time and valuable feedback. Best regards, Authors

Authorsrebuttal2024-08-11

Have we addressed your concerns?

Dear Reviewer, As the discussion period nears its end, we wanted to check if we've adequately addressed your concerns. If you have any additional questions or if there are any points that require further clarification, please don't hesitate to let us know. We appreciate your time and valuable feedback. Best, Authors

Authorsrebuttal2024-08-11

Have we addressed your concerns?

Dear Reviewer, As we approach the end of the discussion period, we wanted to check if we've adequately addressed your concerns. If you have any additional questions or if there are any points that require further clarification, please don't hesitate to let us know. We appreciate your time and valuable feedback. Thank you for your consideration. Authors

Reviewer rWJi2024-08-12

response

Thank author for the response which partially addressed my concern. Therefore I have increased my score. However, I'm still confused about several parts of the rebuttal. 1. I still cannot tell how Langevin dynamics can circumvents posterior collapse. I think it's just a method to sample from a distribution by the score function. It's pretty like VAE samples from Gaussian. The reviewer's rebuttal looks quite away from the insights of Langevin dynamics. 2. I think Eq. (4) only indicates that $x$ and $y$ are independent given $z$. From here we can say the alignment between $z$ and $x$, $z$ and $y$, but how to secure the dependence either between $z$ and joint $(x, y)$ or between $x$ and $y$? Simply, since in the end we aim to sample $x~p(x|y)$, how $p(x|y)$ be learned from paper's setting?

Authorsrebuttal2024-08-12

Thank you for reviewing our work.

Dear Reviewer, As we approach the end of the discussion period, we wanted to ensure we've adequately addressed all of your concerns. If you have any additional questions or if any points require further clarification, please don't hesitate to let us know. We're eager to provide any necessary information to fully address your review. We sincerely appreciate your time and valuable insights. We hope you will consider supporting and championing our work. Thank you for your thoughtful evaluation. Best, Authors

Reviewer jGf32024-08-12

Rebuttal Acknowledgement

I would like to thank authors for their response and addressing the points I raised in my review. I am happy to keep my assessment of their work the same.

Program Chairsdecision2024-09-25

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC