Molecule Design by Latent Prompt Transformer

We sincerely thank you for your thorough review and positive assessment of our paper. We appreciate your recognition of our novel model LPT, the framework for property-conditioned molecule sequence generation, and the strong experimental results across multiple benchmarks. We are pleased that you found our writing clear and well-organized. We would like to address the weaknesses and questions you raised:

W1: Validity rate of generated SMILES strings:

Thank you for this important suggestion. In our molecule design experiments, we used the SELFIES representation (Krenn et al., 2020) instead of SMILES, as briefly mentioned in Section A.1.1. SELFIES is a 100% robust molecular string representation that ensures all generated strings correspond to valid molecules. As a result, we achieved 100% validity in all our molecule design experiments. We acknowledge that we should have made this point more explicit in the main text, and we will clearly state this in the revised paper.

We agree that the validity checking of SMILES string generation is a good metric for model sanity check. Although not reported in the current version of the paper, we did perform this sanity check by evaluating SMILES generation with our model pretrained on ZINC. We achieved a validity rate of 0.999 for 10k generated molecules, which validates our model's capability to generate valid molecules and the model can potentially capture the chemical rules. We will include this result in the revised paper.

W2: Dimension/size of latents and its impact on performance:

Thank you for highlighting this important aspect. We conducted additional ablation studies on the impact of latent variable dimensionality in the challenging single-objective PHGDH experiment with the same number of oracle calls (50k):

We observe a trade-off between performance and computational time. Increasing the number of latent variables improves results but increases optimization time. The performance gain diminishes from 4 to 8 variables, while computational cost continues to rise. Based on these findings, we chose 4 latent vectors for z (each size 256, total dimension 1024) in our main experiments, balancing performance and computational efficiency for practical applications. We shall include this discussion in our revised paper.

Once again, we thank you for your constructive feedback. We are committed to improving our paper based on your suggestions.

Thank you for your detailed and constructive review. We appreciate the time and effort you’ve put into evaluating our work. We will address your comments and questions point by point:

W1: Motivation for framework:

Thank you for your question. We kindly refer you to the Global Response Point 1 for a detailed explanation of our motivation.

W2: Dependence between molecules/properties:

Our approach ensures the dependence between molecules and properties through:

1.Joint probability modeling: We learn $p(x,y,z) = p(z)p(x|z)p(y|z)$, where $z = U_\alpha(z_0)$ and $z_0 \sim \mathcal{N}(0, I_d)$, tying molecules and properties through the shared latent space $z$.

2.MLE learning of joint model (Equation 5, Algorithm 1): a) Sample $z$ from $p(z|x,y)$ using Langevin dynamics: $z_{\tau+1} = z_\tau + s\nabla_z \log p(z_\tau|x,y) + \sqrt{2s}\epsilon_\tau$ The gradient for sampling $z$ is (as in line 137): $\nabla_{z_0} \log p(z_0|x,y) = \nabla_{z_0} \log p_0(z_0) + \nabla_{z_0} \sum_{t=1}^T \log p(x^{(t)}|x^{(<t)},z) + \nabla_{z_0} \log p(y|z)$ This ensures $z$ is sampled from regions consistent with both the observed molecule structure and properties, while the noise term $\epsilon_\tau$ enables exploration of the latent space. b) Update model parameters to maximize likelihood given sampled $z$ above.

W3: Minor issues:

We’ll spell out MCMC on first use and correct the typo in Equation 1.

Q1: Decoupling of generative model and property optimization:

Thank you for this important point. Our original statement was too broad. In our introduction, we primarily focused on the online molecule design. We will revise our statement to reflect this and makes sure to cite the paper you mentioned:

"While some approaches to molecule design have treated generative modeling and property optimization separately, recent work such as [1] and [2] has shown promising integration of these aspects. Our approach builds upon these ideas, extending them in the context of online molecule design."

Q2: Comparsion of VAE and MCMC

In the offline setting (Sec. 3.2), both MCMC-based and VAE-based algorithms are valid for approximating MLE learning. However, VAE-based methods with autoregressive decoding can suffer from posterior collapse without careful design. We conducted additional experiments on single-objective QED optimization with 25K molecule-property pairs in offline settings. We parameterized the encoder as Transformer models with the same layers but with full attention rather than causal attention. Performance degradation for our model was observed:

In the online setting (Sec. 3.4), which is the main focus of our paper, MCMC-based LPT is more compatible with our needs. It can adapt to distribution shifts with a learned prior model and flexibly control exploration and exploitation to improve sample efficiency. MCMC-based optimization should be compared to other online optimization techniques such as Bayesian optimization (BO) with a pre-trained VAE. Our detailed comparisons on the PMO benchmark show that our methods outperform VAE-based methods (Table is shown in Global Response Point 3), even when strong BO techniques are used.

Q3: Computational efficiency of MCMC based methods.

We clarify that with careful design, MCMC is not our computational bottleneck. We’ll include this discussion in the revised paper.

In offline pretraining, we sample $z \sim p(z|x,y) \propto p(z)p(x|z)p(y|z)$, where $z = U_\alpha(z_0)$ with 15-step Langevin dynamics. In online learning, we use a persistent Markov chain that amortizes sampling across optimization steps, reducing steps to 2, thus enhancing real-world applicability.

To address your concern comprehensively, we compare our method with strong VAE-based online optimization methods in PMO benchmark:

Comparison of model size and optimization time: We compare optimization time after pretraining generative models on molecules with baselines in PMO bechmark with multiple-property objective (MPO) experiments. LPT shows comparable computation time to VAE-based Bayesion Optimization(BO) method.

Q4: Sampling $p(z_0|y)$ using Langevin dynamics:

This process is explained in Equation 7 of the main paper. The posterior distribution is given by: $z_0 \sim p_\theta(z_0|y) \propto p_0(z_0)p_\gamma(y|z = U_\alpha(z_0))$. To sample from this distribution using Langevin dynamics, we iterate: $z_0^{\tau+1} = z_0^\tau + s\nabla_{z_0} \log p_\theta(z_0|y) + \sqrt{2s}\epsilon^\tau$ where: $\nabla_{z_0} \log p_\theta(z_0|y) = \nabla_{z_0} \log p_0(z_0) + \nabla_{z_0} \log p_\gamma(y|z = U_\alpha(z_0))$

Here, $p_0$ is a standard Gaussian distribution, so its gradient is straightforward to compute. The second term is also computable. For example, if we assume $p_\gamma(y|z)=\mathcal{N}(s_\gamma(z),\sigma^2)$ as in line 121, the second term becomes $1/\sigma^2(y-s_\gamma(z=U_\alpha(z_0)))\nabla_{z_0}s_\gamma(z=U_\alpha(z_0))$ using the chain rule. This can be implemented by pytorch autograd.

Thank you for your valuable feedback on our submission. If you have any additional concerns, please let us know. We hope our responses have addressed your queries satisfactorily. If so, we would appreciate your consideration in raising your rating of our submission.

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.

Thank you for your thoughtful and comprehensive review of our paper. We greatly appreciate your positive feedback on the originality, quality, clarity, and significance of our work. We're particularly grateful for your recognition of LPT's novelty and significance in de-novo molecule design.

Regarding the weaknesses you identified:

W1. Novelty of conditional generative models:

We appreciate your insightful observation about the novelty of our approach. Indeed, while conditional generative models have been applied in other domains, we believe our work makes a significant contribution by being the first to combine latent variable modeling with autoregressive molecule generation specifically for molecule design. This combination leverages the expressive power of Transformer-based autoregressive modeling for molecule generation and the efficiency of low-dimensional latent space sampling and design. This synergy is particularly advantageous in addressing the complexities of molecular structures and navigating the vast chemical space. Furthermore, it enhances our online learning algorithm, improving both design outcomes and sample efficiency. This integrated approach facilitates more effective exploration and optimization in molecular design.

W2. Comparison with Latent-space Bayesian Optimization (LS-BO):

We appreciate this suggestion and have added additional experiments comparing LPT with prominent BO-based optimization methods, which will be included in the revised version of our paper. We report the multi-property objectives (MPO) experiments under a limited oracle budget of 10K queries on the practical molecular optimization (PMO) benchmark [Gao et al., 2022]. We include the strongest BO-based baseline Gaussian process BO (GP-BO) [Tripp et al., 2021], LS-BO baselines including VAE-BO with both trained with SMILES [Gómez-Bombarelli et al., 2018] and SELFIES strings, and LSTM HC SMILES, which is the best generative molecule design method in PMO and is the same as reported in Table 5 of our main paper. Note that GP-BO optimizes the GP acquisition function with graph genetic algorithm in an inner loop.

As shown, LPT performs comparably to the strong GP-BO baseline and LSTM HC SMILES, while significantly outperforming LS-BO (VAE based) methods across multiple tasks. LPT achieves the highest overall sum score, demonstrating its robust performance across diverse molecular optimization objectives. Since BO-based baselines often require careful parameter tuning, we directly report the numbers from the PMO benchmark that is carefully tuned. In future work, we plan to conduct more thorough comparisons in binding affinity experiments to further validate LPT's performance against these baselines.

References:

[1] W. Gao et al. Sample Efficiency Matters: A Benchmark for Practical Molecular Optimization. NeurIPS, 2022.

[2] A. Tripp et al. A fresh look at de novo molecular design benchmarks. NeurIPS 2021 AI for Science Workshop, 2021.

[3] R. Gómez-Bombarelli et al. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science, 2018.

We believe these additions further solidify our work's contribution to de-novo molecular design. Thank you again for your valuable feedback and recommendation for acceptance. We look forward to the opportunity to present a strengthened version of our work that incorporates your valuable suggestions.

Thank you for your thorough review. We appreciate your recognition of our work's alignment with real-world scenarios, comprehensive experiments, and efficiency improvements. We'll address your comments and questions point by point.

W1. Clarity improvements:

We'll add a diagram illustrating the training and optimization process in the revised version and provide more detailed derivations (See Q1).

W2. MCMC sampling practicality:

Thank you for raising this important point. We hope to clarify that with careful design, MCMC is not our computational bottleneck. We'll include this discussion in the revised paper.

In offline pretraining, we sample $z \sim p(z|x,y) \propto p(z)p(x|z)p(y|z)$, where $z = U_\alpha(z_0)$ with 15-step Langevin dynamics. In online learning or optimization, we use a persistent Markov chain that amortize the sampling across all optimization steps, reducing steps to 2. This speedup focuses on the optimization phase for real-world applicability.

To address your concern comprehensively, we add:

1. Comparison of model size and optimization time: We compare optimization time after pretraining generative models on molecules. LPT shows comparable computation time to Bayesian optimization (BO) methods with much lower model size for LPT. This is for multiple-property objective (MPO) experiments in Practical Molecular Optimization (PMO) from Sec. 5.4.

2. Breakdown of major computational costs for a single batch of 256 samples:

Thanks for the question. We shall consider improving the offline model learning speed in the future work and discuss this explicitly in the limiations.

W3. Minor issues:

• We'll correct the typo "2rd" to "2nd" in Table 1.
• Good point about "Kd" values. We'll clarify these are computational predictions, not wet lab results.
• You're correct, it should be Algorithm 1 in Algorithm 2 step $c$. We'll fix this error.

Q1. Derivation from equation 4 to 5:

We'll provide a detailed derivation in the revision. Please refer to Global Response Point 1.

Q2. Formulating $p(y|z)$ and $\sigma$ choice:

Thank you for this insightful question. You're correct that $p(y|x)$ is deterministic by nature. However, introducing a small error term $\sigma$ in generative models is a common and necessary practice, aligning with various regression models and latent space optimization techniques.

• In probabilistic regression models, including Gaussian Process Regression, $y = f(x) + \epsilon$, where $\epsilon \sim N(0, \sigma^2)$. Here, $\sigma$ represents observation noise or model uncertainty as in Gómez-Bombarelli et al. (2018).
• In Variational Autoencoders (VAE) for continuous variables, the property encoder typically outputs both a mean $\mu$ and a standard deviation $\sigma$ for each dimension, modeling $p(z|y)$ as $N(\mu, \sigma^2)$ such as Jiang et al.(NeurIPS 2020).
• Latent Space Bayesian Optimization methods like LSBO (Tripp et al., NeurIPS 2020) use similar noise terms when fitting its Gaussian Process surrogate model.

In our case, modeling $p(y|z)$ with a small $\sigma$ serves several purposes:

• Account for potential inaccuracies in the oracle function or property prediction.
• Provide smoothness to the optimization landscape, potentially aiding in convergence.
• Allow for a degree of uncertainty in the latent space, which can be beneficial for exploration (in online learning LPT).

You raise an excellent point about the challenges in sampling $p_\theta(z_0|y)$ when $\sigma$ is small. One possible solution is to anneal $\sigma$, by starting from a big value and gradually reducing it.

References:

[1] Gómez-Bombarelli et al. Automatic chemical design using a data-driven continuous representation of molecules. (ACS central science 2018)

[2] Jiang et al. Multi-objective Deep Data Generation with Correlated Property Control. (NeurIPS 2020)

[3] Tripp et al. "Sample-Efficient Optimization in the Latent Space of Deep Generative Models via Weighted Retraining" (NeurIPS 2020)

Q3. Model robustness to Oracle quality:

We've added experiments varying Oracle noise levels to assess robustness on single-objective QED optimization tasks under the budget of 25K. We use noised oracles as $y_{noise}=y_{true}+e$, where $e\sim N(0, \sigma^2)$, varying $\sigma$ as a percentage of the property range. The results show our model is robust to the noised oracles.

Thank you for these insightful comments. They will help improve the paper's clarity and completeness. We hope our responses have addressed your queries satisfactorily. If so, we would appreciate your consideration in raising your rating of our submission.

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