Molecule Design by Latent Prompt Transformer

Authors: deqian  kong
yuhao  huang
jianwen  xie
edouardo  honig
ming  xu
shuanghong  xue
pei  lin
sanping  zhou
sheng  zhong
nanning  zheng
ying nian wu
biomap  research
akool  research
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:39:27

Abstract: 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.

Main paper: 
 introduction: 
 Molecule design plays a pivotal role in drug discovery, focusing on identifying or creating molecules with desired pharmacological or chemical properties. However, the vast and complex space of potential drug-like molecules presents significant challenges for systematically designing efficient machine learning models to navigate this space. Latent space optimization (LSO) (Maus et al., 2022; Tripp et al., 2020; Gómez-Bombarelli et al., 2018) , a two-stage procedure, has emerged as a promising approach to address this challenge. The first stage involves training a deep generative model, typically a deep variational auto-encoder (VAE) (Kingma and Welling, 2014) , to map low-dimensional continuous vectors to the data manifold in input space, creating a simplified and continuous analog of the original optimization problem. In the second stage, the objective function is optimized over this learned latent space using a surrogate model. While LSO has demonstrated potential in tackling high-dimensional and structured input spaces, existing LSO-based methods separate the training of the generative model from the property-conditioned optimization. This decoupling can lead to suboptimal performance, as the learned latent space may not be ideally suited to the specific optimization task. Furthermore, LSO depends on an additional inference model during training, increasing the number of model parameters and adding complexity to the model design process.To address the above limitations, we propose a novel generative model, the Latent Prompt Transformer (LPT), which unifies molecule generation and optimization within a single framework. The LPT can be directly trained on observed molecule-property pairs via maximum likelihood estimation (MLE) in an end-to-end fashion. We take advantage of the immediately-available Markov chain Monte Carlo (MCMC) inference engine derived from the posterior distribution, eliminating the need for an auxiliary network for variational inference. The LPT consists of three components: (1) a learnable prior model of the latent vector based on a neural transformation of Gaussian noise, (2) a molecule generation model that generates molecule sequences using a causal transformer with the latent vector serving as the prompt, and (3) a property predictor model that estimates the target property values given the latent vector. The latent prompts serve as shared representations for both molecules and properties. By incorporating the MLE of the LPT, we employ an online learning algorithm of the LPT for property optimization, formulated as a conditional generative modeling task. This algorithm progressively shifts the model distribution towards regions associated with desired properties.To enhance the practical application of designed molecules and enable a comparative analysis of our generated designs versus those made by human experts, we introduce a new task, which is the conditional generation of molecules that bind to the NAD binding site of Phosphoglycerate dehydrogenase (PHGDH). In addition to single-objective optimization for this specific protein, we also introduce structure-constrained optimization to analyze the design pathways of learning-based models in comparison with human experts. This additional experiment aims to provide valuable insights for improving learning-based model design. We manually process the data to ensure compatibility with docking software, which is essential for facilitating precise score computation. Compared to other widely used evaluation metrics, the use of PHGDH offers greater practical significance and enables more accurate calculations in real-world drug discovery scenarios.Our work makes the following key contributions: (1) We propose the LPT, a novel generative framework for jointly modeling molecule sequences and their target properties. This framework leverages a learnable informative prior distribution on a latent space, conceptualized as an intrinsic design representation. (2) We propose an MCMC-based MLE to train the LPT without needing an auxiliary network for variational inference. (3) We develop a novel online learning algorithm for LPT, allowing the model to gradually extrapolate to feasible regions associated with desired properties, improving computational and sample efficiency. (4) We present a new task for the conditional generation of molecules that bind to the NAD binding site of PHGDH, and introduce structure-constrained optimization to compare the design pathways of learning-based models with those of human experts.(5) Our model achieves state-of-the-art performance across a variety of molecule-based optimization tasks, including single-objective design, multi-objective design, and biological sequence design.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: See Sec. 3 and Appendix.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
background: 
 Let x = (x (1) , ..., x (t) , ..., x (T ) ) ∈ X be a sequence representation of molecules such as SMILES (Weininger, 1988) or SELFIES (Krenn et al., 2020; Cheng et al., 2023) , where x (t) ∈ V is the t-th element in the vocabulary V, and X is the space of molecules. In molecule design, the objective is to generate targeted molecules that optimize several properties of interest, represented by y = {y i = o p i (x) ∈ R} for i = 1, . . . , m, while satisfying specific constraints, c = {c j = o c j (x) ∈ Z 2 } for j = 1, . . . , n. Here, o p (x) and o c (x) denote oracle functions determining property values and constraint satisfaction, respectively, and these values can be obtained either by querying existing software or by conducting wet lab experiments.The multi-objective multi-constraint optimization (MMO) problem can be formulated as follows:EQUATIONwhere o c (x) = 1 denotes constraint satisfaction and o c (x) = 0 denotes constraint violation. From a probabilistic modeling perspective, the MMO problem can be viewed as a conditional sampling problem, x * ∼ p(x|y = y * , c = 1), where y * represents the desired values of the target properties, and c = 1 indicates that all constraints are satisfied.In generative molecule design, we assume access to a dataset of molecule sequences and their associated properties for offline pretraining of the generative model. During the design phase, the model queries existing software (oracle functions) to obtain single or multiple metrics for the generated molecules, with these oracle functions providing ground-truth property values to be optimized. It is important to acknowledge that developing proper software for accurate evaluation of molecule properties is of equal or even greater importance than designing molecules based on those software. However, it is beyond the scope of this paper to address the development of such software.Here, we treat the property values obtained from oracle functions as ground-truths for optimization.3 Latent Prompt Transformer1 /n log p θ (x i , y i ).3. Setting the total number of shifting iterations as 1.4. Replacing the Unet prior with a Gaussian N (0, I d ).As shown in Tab. 9, each component was essential, and removing or altering any of them might lead to performance degradation, underscoring their importance in achieving the model's high efficacy. In addition, we investigate the effect of the exploitation scheme 1 /σ 2 in Eq. ( 8). Experiments are conducted on the practical molecular optimization (PMO) benchmark, focusing on optimizing the multi-property objective (MPO) for amlodipine under a 10K-query oracle budget. As shown in Tab. 10, increasing 1 /σ 2 enhances LPT's performance, indicating a stronger bias toward exploiting the sampled posterior z and leading to improved optimization efficiency.Furthermore, we study the effect of oracle budget size on performance, conducting experiments on the single-objective ESR1 binding affinity optimization task. As shown in Tab. 11, LPT demonstrates robust performance even under a limited oracle query budget of 10k, outperforming most existing methods. As the budget increases, LPT continues to exhibit significant performance enhancements, further highlighting its efficiency and scalability. Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: See Sec. 6.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
model: 
 Suppose x = (x (1) , ..., x (t) , ..., x (T ) ) is a molecule sequence, (e.g. a string-based representation like SMILES (Weininger, 1988) or SELFIES (Krenn et al., 2020; Cheng et al., 2023) ), where x (t) ∈ V is the t-th element of sequence in the vocabulary V. The latent vector is z ∈ R d and y ∈ R is the target property value, or y ∈ {0, 1} indicates constraint satisfaction. The term property refers to both real-valued properties and binary-valued constraints in the following sections. The joint distribution of a molecule and its property is defined as p(x, y).z z 0 x y z = Uα(z0) pα(z) pβ(x|z) pγ(y|z) Cross-attention Causal Transformer ×N z x (0) , . . . , x (T -1)x (1) , . . . , x (T ) The latent variable z, conceptualized as a design representation, decouples the molecule generation from property prediction. Specifically, given z, we assume x and y are conditionally independent, making z the information bottleneck. With this assumption, our Latent Prompt Transformer (LPT) is defined as,EQUATIONwhere θ = (α, β, γ). p α (z) is a prior model with parameters α. Here, z serves as the latent prompt for the generation model p β (x|z) parameterized by a causal Transformer with parameters β. p γ (y|z) is the predictor model with parameters γ. As shown in Fig. 1 , LPT defines the generation process as,EQUATIONFor the prior model, p α (z) is formulated as a learnable neural transformation from an uninformative distribution, such as an isotropic Gaussian, z = U α (z 0 ), where z 0 ∼ N (0, I d ). U α (•) is parameterized by an expressive neural network such as a Unet (Ronneberger et al., 2015) with parameters α. In this way, p α (z) can be viewed as an implicit generator model. 0) , ..., x (t-1) , z), which is realized by a causal Transformer with parameters β. The latent vector z, serving as the prompt, controls every step of the autoregressive molecule generation. We incorporate these latent prompts z into the p β (x|z) through cross-attention, as shown in Fig. 1 .p β (x|z) is a conditional autoregressive model, p β (x|z) = T t=1 p β (x (t) |x (The real-valued property predictor is a non-linear regression model p γ (y|z) = N (s γ (z), σ 2 ), where s γ (z) is a small multi-layer perceptron (MLP) predicting y based on the latent prompt z. The variance σ 2 is treated as a hyper-parameter that balances the exploitation-exploration trade-off. For binary-valued constraints, p γ (y|z) = s γ (z) y (1 -s γ (z)) 1-y . For multi-objective tasks, we can use heuristic-based combinations of multiple objectives to form a special single-objective function.D = {(x i , y i ), i = 1, ..., n}. The log-likelihood function is L(θ) = 1 /n n i=1 log p θ (x i , y i ).Since z = U α (z 0 ), we can write the model asEQUATIONwhere p 0 (z 0 ) ∼ N (0, I d ). The learning gradient can be calculated as follows,EQUATIONFor the prior model, the learning gradient isδ α (x, y) = E p θ (z0|x,y) [∇ α log p β (x|U α (z 0 )) + ∇ α log p γ (y|U α (z 0 )))].The learning gradient for the molecule generation model isδ β (x, y) = E p θ (z0|x,y) [∇ β log p β (x|U α (z 0 ))].The learning gradient for the predictor model isδ γ (x, y) = E p θ (z0|x,y) [∇ γ log p γ (y|U α (z 0 ))].Estimating these expectations requires MCMC sampling of the posterior distribution p θ (z 0 |x, y). We use Langevin dynamics (Neal, 2011) . For a target distribution π(z), the dynamics iterates as follows,EQUATIONwhere τ indexes the time step, s is step size, and ϵ τ ∼ N (0, I d ) is the Gaussian white noise.Here, π(z) is instantiated by the posterior distribution p θ (z 0 |x, y). With z = U α (z 0 ), we have p(z 0 |x, y) ∝ p 0 (z 0 )p β (x|z)p γ (y|z). The gradient is Pretrain LPT In practical applications involving multiple molecular generation tasks, each characterized by a different target property y, each model p θ (x, y) may require separate training. To enhance efficiency, we adopt a pretaining strategy focusing solely on molecule sequences. In this case, we aim to maximize∇ z0 log p θ (z 0 |x, y) = ∇ z0 log p 0 (z 0 ) prior +∇ z0 T t=1 log p β (x (t) |x (<t) , z)n i=1 log p θ (x i ) = n i=1 log p θ (x i , z i )dz i . The learning gradient is ∇ θ log p θ (x) = E p θ (z0|x) [∇ β log p θ (x|z = U α (z 0 ))].After pretraining LPT, we finetune the model with target properties using Eq. ( 5) for a small number of epochs. This two-stage approach is adaptable for semi-supervised scenarios where property values are limited.After MLE learning from a collected (offline) dataset, LPT is capable of generating molecules with desired properties. Since the latent prompt is designed as the information bottleneck, we can generate the molecule given property value as p(x|y) = p(z 0 |y)p(x|z 0 )dz 0 . We first infer the latent prompt via posterior sampling using Bayes' rule,EQUATIONThis posterior sampling is performed using Langevin dynamics similar to the training process. Specifically, we replace the target distribution in Eq. ( 6) with p θ (z 0 |y) and run MCMC for a fixed number of steps, i.e.,z τ +1 0 = z τ 0 + s∇ z0 log p(z 0 |y) + √ 2sϵ τ ,where∇ z0 log p(z 0 |y) = ∇ z0 (log p(z 0 ) + log p(y|z 0 )) = -z 0 + 1 σ 2 (y -s γ (z))∇ z0 s γ (z). (8)This process of sampling p(z 0 |y) iteratively refines the latent prompts z, increasing their likelihood given the desired property or constraints. This optimization of molecules is achieved in the latent space by gradient-based sampling as a form of test-time computation. Once we generate the latent prompt z, the molecule generation model uses this latent prompt to sample the next element x (t) ∼ p β (x (t) |x (<t) , z = U α (z 0 )) until termination.MLE learning of the joint model is equivalent to minimizing the D KL (p data (x, y)∥p θ (x, y)). Note that Eq. ( 7) is reliable when condition y is supported by p data (y). However, in real-world molecule design, desired property values often lie far from those in the collected offline dataset. For such cases, LPT should be applied in an online learning setting, as discussed in the next section.When the desired property value y is not supported within the learned distribution p θ (z, x, y), we propose an online learning approach to gradually shift the model distribution towards regions supporting desired properties. We assume access to the oracle functions o(x) via software such as RDKit (Landrum et al.) This process is repeated until the model converges or budget limits are reached. This online learning method maintains a latent space generative model (LPT) and a synthetic dataset and gradually shifts both towards regions of desired targets.In step (a), to account for small extrapolation, we use property-conditioned generation as mentioned in Sec. 3.3 by setting the desired property as y = y * + δ y , where y * lies on the boundary of the learned LPT in the previous shifting iteration (e.g. y * can be the maximum value) and δ y is a small increment representing extrapolation. The hyper-parameter δ y quantifies the shifting speed. Specifically, since the latent prompts decouple the molecule generation and property prediction, we first sample the latent prompts z 0 ∼ p θ (z 0 |y = y * + δ y ) before the molecule generation and then use x ∼ p β (x|z = U α (z 0 )) to generate the novel molecules. To sample molecules satisfying binary constraints, we use z 0 ∼ p θ (z 0 |y = 1). This latent space sampling scheme relieves the expensive autoregressive sampling in the data space, and allows efficient gradient-based sampling in low-dimensional latent space, such as Langevin dynamics. In step (b), after generating the novel molecules in step (a), we use the oracle function o(x) to relabel them, obtaining a synthetic dataset of molecule-property pairs with ground-truth values. This synthetic dataset serves as a replay buffer, storing generated molecules along the optimization trajectory, and is used for MLE training of the LPT in step (c). The illustration of the online learning of LPT is depicted in Fig. 2 .Compared to population-based methods such as genetic algorithms (Nigam et al., 2020) and particleswarm algorithms (Winter et al., 2019) , our method maintains both a synthetic dataset (which can be considered a small population) and a generative model to fit the dataset, enabling the improvement of generated molecules from the model. The model itself can be seen as an infinite population, as it can generate an unlimited number of new samples.In molecule design, computational efficiency and sample efficiency are crucial, with varying priorities depending on the objectives. Computational efficiency measures the ability to learn quickly with limited resources, while sample efficiency assesses the capacity to learn a good model with minimal oracle function interactions. Generative models are applied to molecular design for two purposes: obtaining starting points and optimizing them to meet specific objectives. For objectives accessible via existing software, computational efficiency is prioritized, aiming to train the model and query software as quickly as possible until convergence. For properties that are time-consuming and expensive to simulate, such as wet-lab experiments, sample efficiency is more important, minimizing the number of oracle queries for practical usage. In real-world scenarios, a combination of both may be necessary, depending on the specific requirements. We propose the following techniques for LPT to address efficiency issues accordingly.Computational efficiency The computational overhead for online learning of LPT primarily arises from the iterative MCMC sampling procedure. In the posterior sampling stage of p θ (z 0 |y) and p θ (z 0 |x, y) in step (a) and (c), rather than using a Gaussian noise-initialized MCMC chain with a fixed number of MCMC steps (N = 15) for each learning iteration, we employ the Persistent Markov Chain (PMC) method (Tieleman, 2008; Xie et al., 2016; Han et al., 2017; Xie et al., 2019) , which amortizes sampling across shifting iterations. Specifically, for the starting point z τ =0 0 in Eq. ( 6) of each MCMC chain, its initialization is drawn from a Gaussian distribution at the first learning iteration, then subsequently from the previous iteration's sampling output. With this approach, the number of MCMC updates can be reduced to N = 2 steps per sampling stage, achieving an approximate 5× speedup in posterior sampling.We aim to improve sample efficiency in both steps (a) and (c) in Sec. 3.4. In step (a), property-conditioned generation, we refine the exploration strategy during test-time computation. This approach demonstrates a clear advantage of posterior inference over direct use of VAE or GAN. By guiding the learned LPT towards exploitation during generation, we encourage the latent prompts z to converge on the modes of the posterior distribution in Langevin dynamics. This concept is analogous to the intuition behind classifier guidance in conditional diffusion models (Dhariwal and Nichol, 2021; Ho and Salimans, 2022) . In Eq. ( 8), we can adjust σ 2 to balance the trade-off between exploitation and exploration: ∇ z0 log p(z 0 |y) = ∇ z0 (log p(z 0 ) + log p(y|z 0 )) = -z 0 + 1 σ 2 (y -s γ (z))∇ z0 s γ (z). When 1 /σ 2 = 1, the sampled latent prompts z represent the density of the posterior distribution, resulting in an efficient exploration scheme. As 1 /σ 2 increases, the sampled posterior z concentrates around the modes of the posterior distribution, indicating increased confidence and a stronger bias towards exploitation. In step (c), we leverage the synthetic dataset to train the LPT on the most informative samples in terms of property values. By training the LPT on a distribution that assigns higher probability mass to high-value points and lower mass to low-value points, the training objective encourages a larger fraction of the feasible region's volume to model high-value points while simultaneously using other data points to learn useful representations and avoid overfitting. We modify the standard objective n i=1 1 /n log p θ (x i , y i ) by assigning an importance weight w i to each molecule-property pair (x i , y i ) in the synthetic dataset, resulting in a biased objective function, n i=1 w i log p θ (x i , y i ), where i w i = 1. In our experiments, we set w i = 1/N if y i is in the top-N property scores, and w i = 0 otherwise. This approach is inspired by prioritized experience replay (Schaul et al., 2015) in online reinforcement learning and weighted retraining (Tripp et al., 2020) in black-box optimization, both of which prioritize learning from the most informative samples to improve sample efficiency.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 7. Experiment Statistical Significance Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: See Sec. 5 and App. A.2. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: See App. A.7. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to 
related work: 
 Latent space optimization Latent space optimization has been widely applied in generative models and high-dimensional data manipulation (Jin et al., 2018; Kusner et al., 2017; Kajino, 2019; Dai et al., 2018) . This method involves representing data in a lower-dimensional space while retaining its essential characteristics. Latent space optimization improves the fidelity and quality of generated data and facilitates more efficient and effective exploration of the latent space. This leads to better generalization and robustness in various applications, such as image synthesis (Karras et al., 2019; Razavi et al., 2019; Song et al., 2020; Dhariwal and Nichol, 2021; Rombach et al., 2022) , data compression (Ballé et al., 2016; Mentzer et al., 2018) , molecular optimization (Kong et al., 2023; Jain et al., 2023; Zhu et al., 2024) , and automatic machine learning (Liu et al., 2018; Zhang et al., 2019) . When the search space is significantly simplified in the latent space, familiar Bayesian optimization (BO) tools can be readily applied (Maus et al., 2022; Tripp et al., 2020) . Unlike BO-based methods, our LPT is based on the explicit probabilistic form of p(z|y), which allows us to perform optimization as conditional generation using Bayes' rule.Generative molecule design This research follows two main approaches. The first uses latent space generative models to translate discrete molecule graphs into continuous latent vectors, enabling optimization of molecular properties within the latent space (Gómez-Bombarelli et al., 2018; Kusner et al., 2017; Jin et al., 2018; Maziarz et al., 2021; Eckmann et al., 2022; Kong et al., 2023) . The second directly employs combinatorial optimization methods, such as reinforcement learning, to fine-tune molecular attributes within the graph data space (You et al., 2018; De Cao and Kipf, 2018; Zhou et al., 2019; Shi et al., 2020; Luo et al., 2021; Du et al., 2022) . Alternative data space methodologies, like genetic algorithms (Nigam et al., 2020) , particle-swarm strategies (Winter et al., 2019) , Monte Carlo tree search (Yang et al., 2020) , and scaffolding trees (Fu et al., 2021) , have also gained traction. 
experiments: 
 We demonstrate the effectiveness of our approach across a wide range of optimization tasks. In the context of molecule design, this includes binding affinity maximization, constrained optimization, and multi-objective optimization (see Sec. 5.2). Additionally, we optimize protein sequences for high fluorescence and DNA sequences for enhanced binding affinity (see Sec. 5.3). Finally, we validate the sample efficiency of LPT by performing optimization with a limited number of oracle function queries (see Sec. 5.4). Additional experiments, including robustness to noisy oracles, online learning from scratch, ablation studies and a detailed discussion of related works and baselines, are provided in Apps. A.2 and A.4.Molecule Sequence Design For molecule design tasks, we use SELFIES representations of the ZINC (Irwin et al., 2012) dataset, which comprises 250K drug-like molecules as our offline dataset D. We utilize RDKit (Landrum et al.) to compute several key metrics, including penalized logP, drug-likeness (QED), and the synthetic accessibility score (SA). Additionally, we use AutoDock-GPU (Santos-Martins et al., 2021) to derive docking scores E, which serve as proxies for estimating the binding affinity of compounds to three protein targets: the human estrogen receptor (ESR1), human peroxisomal acetyl-CoA acyltransferase 1 (ACAA1), and Phosphoglycerate dehydrogenase (PHGDH). The binding affinity is expressed as the dissociation constant K D (nM), which is approximated by the formula K D ≈ e -E/c , where E is the docking score and c is a constant. A lower K D value indicates stronger binding affinity. ESR1 is a well-characterized protein with numerous known binders, making it a suitable reference point for evaluating molecules generated by our model. In contrast, ACAA1 has no known binders, providing an opportunity to test the model's capability for de novo design (Eckmann et al., 2022) . Additionally, we propose to design molecules that bind to PHGDH, an enzyme pivotal in the early stages of L-serine synthesis. Recently, PHGDH has gained attention as a potential therapeutic target in cancer treatment due to to its involvement in various human cancers (Zhao et al., 2020) . The crystal structure of PHGDH (PDB: 2G76) is well-established, and there exists a comprehensive case study on the development of PHGDH inhibitors, showcasing a structure-based progression from simpler to more complex molecules targeting its NAD binding site (Spillier and Frédérick, 2021) , as illustrated in Fig. 3 . Furthermore, we observe that the K D values estimated from wet lab experiments align closely with the trends predicted by AutoDock-GPU. This congruence makes PHGDH-NAD an excellent case study for testing LPT on single-objective, structure-constrained, and multi-objective optimization tasks using AutoDock-GPU. More details can be found in App. A.5.We further apply our method to biological sequence design via two tasks in Design-Bench (Trabucco et al., 2022) : TF Bind 8 and GFP. To be specific, the TF Bind 8 task focuses on identifying DNA sequences that are 8 bases long, aiming for maximum binding affinity. This task contains a training set of 32,898 samples and includes an exact oracle function. The GFP task involves generating protein sequences of 237 amino acids that exhibit high fluorescence. Training Setup The prior model, p α (z), of LPT is a one-dimensional UNet (Ronneberger et al., 2015) where z contains 4 tokens, each of size 256. The sequence generation model, p β (x|z), is implemented as a 3-layer causal Transformer, while a 3-layer MLP serves as the predictor model, p γ (y|z). As described in Sec. 3.2, we pre-train LPT on molecules for 30 epochs and then fine-tune it with target properties for an additional 10 epochs, following the procedure outlined in Alg. 1 in App. A.3. We perform up to 25 iterations of online learning, generating 2,500 samples per iteration, which totals a maximum of 62.5K oracle function queries. We use the AdamW optimizer (Loshchilov and Hutter, 2019; Kingma and Ba, 2014) with a weight decay of 0.1. Training was conducted on an NVIDIA A6000 GPU, requiring 20 hours for pre-training, 10 hours for fine-tuning, and 12 hours for online learning. Additional details can be found in App. A.3.Single-Objective Optimization For the single-objective binding affinity optimization task, we aim to design ligands with optimal binding affinities to ESR1, ACAA1, and PHGDH as de novo design, without any constraints. LPT does not use any prior knowledge of existing binders, and exclusively uses the crystal structures of the aforementioned proteins. Structure-constrained Optimization The structure-constrained optimization task mimics lead optimization in drug discovery, aiming to decorate a fixed core substructure to optimize activity and pharmacological properties. Our model's factorization, p(z, x, y) = p(z)p(x|z)p(y|z), enables the decoupling of molecule generation and property prediction, simplifying conditional generation. Given a substructure x = (x (1) , . . . , x (k) ), we aim to sample from p θ (x, y|x). This is accomplished by sampling z ∼ p(z|y) and then x ∼ p(x|x, z), which only requires rearranging x's sequence to start from the desired atom. In Fig. 3 , Compound 2 (C2) is designed by humans as an extension of Compound 1 (C1), and Compound 3 (C3) is similarly an extension of C2. In Fig. 4 , we show that:(1) given C1 or C2, our model is able to design compounds similar to the human-designed C2 or C3, and (2) our model can identify molecules that outperform those designed by humans. Meanwhile, we confirm that C2, C3, and LPT-generated molecules are novel compared to the ZINC training set, with an average Tanimoto similarity less than 0.5. More generated molecules are shown in App. A.6.2.Our model explores high-affinity PHGDH inhibitors by adding functional groups to an indole backbone, a common scaffold for such inhibitors (Spillier and Frédérick, 2021) . The model identifies aromatic or heteroaromatic groups, such as benzene or pyridine, at the second position as frequently occurring and exhibiting higher binding scores compared to the indole backbone itself. This finding aligns with reported data: a published molecule with a benzothiophene backbone similar to C1 exhibits a binding affinity of 470 µM, while C2, which includes an aromatic group at the second position, shows a significantly improved binding affinity of 1.6 µM (Spillier and Frédérick, 2021; Fuller et al., 2016) . Furthermore, the model introduces a second functional group at the sixth position of the indole in C2, generating inhibitors closely resembling the structure of C3. The observed trend of increasing binding affinities (1.6 µM for C2 and 0.18 µM for C3) aligns with the literature values, providing validation for our proposed method in identifying potential high-affinity inhibitors.Multi-Objective Optimization For multi-objective optimization tasks, we aim to simultaneously maximize binding affinity and QED, while minimizing SA. These objectives are balanced as a weighted combination, with constraints of QED > 0.4 and SA < 5.5. We evaluate our method on three protein targets: ESR1, ACAA1, and PHGDH, comparing the results against two baseline methods, LIMO (Eckmann et al., 2022) and SGDS (Kong et al., 2023) . As shown in Tab. 3, our method, LPT, achieves QED and SA scores comparable to those of SGDS while significantly improving binding affinity across all three protein targets, demonstrating its superior modeling capability. Examples of the generated molecules are provided in App. A.6.1. Our model excels in biological sequence design, a single-objective optimization application, as demonstrated by two benchmarks in Design-Bench (Trabucco et al., 2022) : TF Bind 8 and GFP.Tab. 4 shows that LPT significantly outperforms other methods in these tasks. In the TF Bind 8 task, our approach surpasses the strong competitor GFlowNet-AL (Jain et al., 2022) , while maintaining comparable diversity. For the GFP task, we achieve superior performance with reasonable diversity. We validate LPT's sample efficiency on the Practical Molecular Optimization (PMO) benchmark (Gao et al., 2022) , where multi-property objectives (MPO) are optimized within a limited oracle budget of 10K queries. Tab. 5 shows that our method, LPT, surpasses previous approaches, such as MARS (Xie et al., 2021) , GFlowNet (Jain et al., 2022) and SMILES/SELFIES-VAE (Gómez-Bombarelli et al., 2018; Maus et al., 2022) with Bayesian Optimization (BO). Also, LPT achieves comparable performance to LSTM HC (Brown et al., 2019) , the best generative molecule design method in PMO, and demonstrates performance on par with GP-BO (Tripp et al., 2021) , the best BO-based method in PMO, under the limited budge of oracle function queries. We acknowledge that there remains a performance gap between generative model-based optimization and methods like genetic algorithms when working with a small budget. This is primarily due to the data-intensive requirements of training generative models. To ensure a fair assessment, our comparison focuses on representative generative molecule design methods within PMO. It's worth noting that generative models offer distinct advantages when maintaining relatively large budgets, as the learned model itself can be viewed as infinite populations for further exploration.  
limitation and conclusion: 
 In this work, we presented LPT, a novel generative model for molecule design that achieves strong performance through its offline and online learning algorithms. Contemporary work has extended the similar model to offline reinforcement learning (Kong et al., 2024) . While the model demonstrates significant potential, there are opportunities to better understand how LPT handles the inherent trade-offs in multi-objective optimization scenarios, particularly in characterizing the Pareto front nature of optimal solutions. Future work could also explore alternative architectures to extend LPT's applicability beyond sequence-based optimization problems in science and engineering. 
 
 Appendix: We provide the derivation of Eq. ( 5) in Sec. 3.2.∇ θ log p θ (x, y) = ∇ θ p θ (x, y) p θ (x, y) = 1 p θ (x, y) ∇ θ p θ (x, y, z = U α (z 0 ))dz 0 = p θ (x, y, z = U α (z 0 )) p θ (x, y) ∇ θ log p θ (x, y, z = U α (z 0 ))dz 0 = p θ (z 0 |x, y)∇ θ log p θ (x, y, z = U α (z 0 ))dz 0 = E p θ (z0|x,y) [∇ θ log p θ (x, y, z = U α (z 0 ))] = E p θ (z0|x,y) [∇ θ log p β (x|U α (z 0 )) + ∇ θ log p γ (y|U α (z 0 )) + ∇ θ log p 0 (z 0 )] = E p θ (z0|x,y) [∇ θ log p β (x|U α (z 0 )) + ∇ θ log p γ (y|U α (z 0 ))] .The experiments focus on optimizing the Penalized logP (P-logP) and QED properties, both of which can be calculated using RDKit. Since P-logP scores are positively correlated with the length of a molecule, we maximize P-logP while limiting molecule length to the maximum length of molecules in ZINC using the SELFIES (Krenn et al., 2020; Cheng et al., 2023 ) representation, following Eckmann et al. (2022) . We compare our model with several baseline methods, including JT-VAE (Jin et al., 2018) , MolDQN (Zhou et al., 2019) , LIMO (Eckmann et al., 2022) , GCPN (You et al., 2018) , GraphDF (You et al., 2018) , MARS (Xie et al., 2021) , SGDS (Kong et al., 2023) . Tab. 6 presents the results, demonstrating that LPT outperforms other methods and achieves the highest QED score among methods that limits molecule length. To evaluate the robustness of our model in single-objective QED optimization tasks under the oracle query budget of 25K, we conduct experiments with varying levels of oracle noise. We define noised oracles as y noise = y true + e, where e ∼ N (0, σ 2 ) and σ varies as a percentage of the property range. The minimal degradation in performance seen in Tab. 7 demonstrates that our model is resilient to the noised oracles.We investigate LPT's performance without offline pre-training data, relying solely on online learning with a budget limit of 300K (comparable to the size of the ZINC dataset plus our previous online We conduct ablation studies to investigate the contributions of key components in LPT using a challenging PHGDH single-objective optimization task. The variations in our experiments included:1. Using samples from z ∼ p α (z) instead of p θ (z|y) to generate the proposals P t .As illustrated in Fig. 1 , the prior model in LPT is parameterized by a 1D Unet, with z sequence length of 4, where each of them is size of 256. The molecule generation model, p β (x|z), employs a 3-layer causal Transformer with an embedding size of 256 and maximum input token length of 73.The predictor model is a 3-layer MLP, which takes z as input and outputs predicted property values or classification results. The total number of parameters for LPT is 4.33M.LPT is trained in a two-step process. Initially, it undergoes pre-training solely on molecules for 30 epochs, using cross-entropy loss with a learning rate that varies between 7.5e-4 and 7.5e-5, following a cosine scheduling approach. Subsequently, LPT is fine-tuned for 10 epochs on both molecules and their properties, as outlined in Alg. 1, with the learning rate adjusted between 3e-4 and 7.5e-5. For multi-objective optimization, i.e., simultaneously maximizing binding affinity, QED, and minimizing SA, the predictors for binding affinity and QED/SA are selected as a regression model and a classifier, which are supervised by mean squared error (MSE) and binary cross-entropy (BCE) loss functions, respectively.Algorithm 1 MLE learning of Latent Prompt Transformer (LPT)Input: Number of learning iterations T , initial parameters θ 0 = (α 0 , β 0 , γ 0 ), observed samples D = {x i , y i } n i=1 , posterior sampling step size s, the number of MCMC steps N , and the learning rate η 0 , η 1 , η 2 . Output: θ T for t = 1 to T do 1.Posterior sampling: For each (x i , y i ), sample z 0 ∼ p θt (z 0 |x i , y i ) using Eq. ( 6), where the target distribution π is p θt (z 0 |x i , y i ) with N steps and step size s.2.Learn prior model p α (z), generation model p β (x|z) and predictor model p γ (y|z):α t+1 = α t + η 0 1 n i δ α (x i , y i ); β t+1 = β t + η 1 1 n i δ β (x i , y i ); γ t+1 = γ t + η 2 1 n i δ γ (x i , y i ) as in Sec. 3.2. end forFor LPT's online learning, as detailed in Alg. 2, we establish a maximum number of shifting iterations to be 25, generating 2,500 new samples in each iteration. This results in a total of up to 62.5k oracle function queries. Throughout the training processes, we utilize the AdamW optimizer (Loshchilov and Hutter, 2019) with a weight decay of 0.1. The pre-training, fine-tuning, and online learning phases of LPT require approximately 20, 10, and 12 hours, respectively, on a single NVIDIA A6000 GPU.Algorithm 2 Online learning of LPT Input: Number of proposals m, fixed increment δ y , number of PMC steps N , initial parameters θ t = (α t , β t , γ t ), initial dataset D 0 = {x 0 i , y 0 i } n i=1 , oracle function o(x), maximum number of shifting iterations T . Output: θ T , D T while t < T do step (a): Sample molecules and properties from the learned model {z t i , x t i , y t i } m i=1 , where z t ∼ p θ (z|y = y t-1 + δ y ); x t ∼ p β (x|z = z t ). step (b): Relabel the proposal property values by oracle function query: y t i ← o(x t i ), and update the dataset by D t = {z t i , x t i , y t i } m i=1 ∪ D t-1 . step (c): Update LPT using maximum likelihood on synthetic dataset D t by following Alg. 1 end whileOur model is based on Kong et al. (2023) . The differences are as follows. (1) While Kong et al. (2023) used an LSTM model for molecule generation, we adopt a more expressive causal Transformer model for generation, with the latent vector serving as latent prompt. (2) While Kong et al. (2023) used a latent space energy-based model for the prior distribution of the latent vector, we assume that the latent z is generated by a UNet transformation of a Gaussian white noise vector. This approach allows us to eliminate the need for Langevin dynamics in prior sampling during training, thus simplifying the learning algorithm. (3) Our experimental results are significantly stronger, surpassing those of Kong et al. (2023) and achieving new state-of-the-art performance.For molecule generation, JT-VAE (Jin et al., 2018) utilizes a variational autoencoder (VAE). GCPN (You et al., 2018) and GraphDF (Luo et al., 2021) employ the deep graph model to discover novel molecules. A reinforcement learning framework that fuses chemical domain knowledge with double Q-learning and randomized value functions for molecule optimization was presented by MolDQN (Zhou et al., 2019) . MARS (Xie et al., 2021) develops a Markov molecular Sampling framework targeting multi-objective drug discovery, while LIMO (Eckmann et al., 2022) uses a VAE-generated latent space along with neural networks for property prediction, enabling efficient gradient-based optimization of molecular properties.In contrast to existing latent space generative models (Gómez-Bombarelli et al., 2018; Kusner et al., 2017; Jin et al., 2018; Eckmann et al., 2022) , our approach incorporates a learnable prior model, enabling our model to effectively catch up with the evolving dataset in the optimization process.The landscape of biological sequence design is shaped by a diverse range of computational strategies. DyNAPPO (Angermueller et al., 2019) harnesses active learning with reinforcement learning to facilitate iterative sequence generation. GFlowNet-AL (Jain et al., 2022) capitalizes on GFlowNets for generative active learning within sequence contexts. Model-based optimization techniques like COMs (Trabucco et al., 2021) and AmortizedBO (Swersky et al., 2020) integrate Bayesian Optimization with reinforcement learning, enhancing search efficiency. BO-QEI (Wilson et al., 2017) , a variant of Bayesian optimization, refines the search process using quantile Expected Improvement. Deep generative models, e.g., CBAs (Fannjiang and Listgarten, 2020) and MINs (Brookes et al., 2019) , leverage deep learning for complex pattern discovery. Meanwhile, CMA-ES (Hansen, 2006) excels in high-dimensional optimization, adapting search strategies over generations.A.5 Background of PHGDH and its NAD binding site Phosphoglycerate dehydrogenase (PHGDH) is an enzyme that plays a crucial role in the early stages of L-serine synthesis. Recently, PHGDH has been identified as an attractive therapeutic target in cancer therapy due to its involvement in various human cancers, including breast cancer, melanoma, lung cancer, pancreatic cancer, and kidney cancer. Several studies have been devoted to the exploration of small molecule inhibitors targeting PHGDH including CBR-5884 (IC50 = 33±12 (µM)), NCT-503 (IC50 = 2.5±0.6 (µM)), AZ PHGDH inhibitor (IC50 = 180 (nM)), RAZE PHGDH inhibitor (IC50 = 0.01 ∼ 1.5 (µM)), PKUMDL-WQ-2201 (IC50 = 35.7 (µM)), BI-4924 (IC50 = 2 (nM)), and others. Additionally, the crystal structure of PHGDH (PDB: 2G76) has been elucidated. This makes PHGDH an excellent protein model for conducting docking calculations, aiding in the precise localization of binding sites in order to optimize our algorithms.PHGDH utilizes nicotinamide adenine dinucleotide (oxidized form, NAD+; reduced form, NADH) as a co-factor for enzymatic activity, producing NADH during the synthesis of 3phosphohydroxypyruvate (3PHP) from 3-phosphoglycerate (3PG) (Samanta and Semenza, 2016) .The co-crystallization of PHGDH with NAD has been documented (PDB: 5N6C). The NAD+ pocket is surrounded by hydrophobic residues of P176, Y174, L151, L193, L216, T213, T207 and L210 (Mullarky et al., 2019) . Specifically, nicotinamide moiety exhibited interactions with the protein backbone (specifically A285 and C233) as well as the side chain of D259. Additionally, the hydroxyl groups of the sugar moieties were also involved in hydrogen bonds with both the protein backbone (T206) and the side chain of D174. The phosphate linker demonstrated interactions with the main chain of R154 and I155, along with the side chain of R154 (Fig. 5 ) (Unterlass et al., 2017) . We present visualizations of a subset of generated molecules from both single-objective and multiobjective optimization for ESR1, ACAA1, and PHGDH. Additionally, several conditionally generated compounds, C2 and C3, which have similar characteristics to human-designed ones, are also presented in Fig. 12 and Fig. 13 .A.6.1 Multi-Objective and Single-Objective Optimization We introduce a novel generative model for jointly modeling molecule sequences and their target properties, which potentially leads to more capable and efficient molecule design algorithm. One potential negative impact could come from the misuse of the generative modeling in designing molecules or biological sequences for harmful purposes. Therefore, developing suitable safeguards and regulations will be crucial to mitigate potential negative impacts. 


Decision: The reviewers have appreciated the novelty of the approach, although some of it stemming from the fact that the approach was applied to molecules for the first time; they also liked the fact that the experimental work was comprehensive. The reviewers had concerns about the computational weight of the MCMC sampling, but those were alleviated already by the proposed approach and additional comparisons to SOTAs were performed during the rebuttal period to show practical efficiency of the method. The authors contributed a number of additional results including comparisons to additional methods during the rebuttal. Overall the paper is a solid contribution and of interest to the NeurIPS community.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See Weaknesses section above for suggestions/points of discussion.
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: Yes. Limitations are well discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: See Weaknesses part.
limitations: limitations: Yes.

Comments: title: Have we addressed your concerns?
Thank you!
Have we addressed your concerns?
Rebuttal Acknowledgement
Thank you for reviewing our work.
response
Have we addressed your concerns?
Have we addressed your concerns?
Thank you!
Follow-up Response
comment: 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
Thank you for your positive feedback and thorough review. We're glad all your concerns have been addressed. We really appreciate your continued support.

Best,

Authors
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
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.
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
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 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?
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
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
Thank you for your positive feedback and thorough review. We’re glad all your concerns have been addressed. Have a nice day!
I appreciate authors' efforts in rebuttal. All my concerns have been addressed so I will keep my rating.