Title: Molecule Design by Latent Prompt Transformer
Venue: NeurIPS 2024
Accepted: yes
Keywords: latent space generative modeling; Latent Prompt Transformer; molecule design
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.

Molecule Design by Latent Prompt Transformer

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.

1 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 (V AE) (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.

2 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:

where 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 Transformer

3.1 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

α

(

z

0

)

p

α

(

z

)

p

β

(

x

|

z

)

p

γ

(

y

|

z

)

Cross-attention

×

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,

where θ = ( α, β, γ ) . 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,

For 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.

The molecule generation model p β ( x | z ) is a conditional autoregressive model, p β ( x | z ) = ∏ T t =1 p β ( x ( t ) | x (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.

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.

3.2 Maximum Likelihood Learning of LPT

Suppose we observe training examples from the dataset with molecule sequence and property pairs 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 as

where p 0 ( z 0 ) ∼ N (0 , I d ) . The learning gradient can be calculated as follows,

For the prior model, the learning gradient is

The learning gradient for the molecule generation model is

The learning gradient for the predictor model is

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,

where τ 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

We initialize z τ =0 0 ∼ N (0 , I d ) , and employ N steps of Langevin dynamics (e.g. N = 15 ) for approximate sampling from the posterior distribution, rendering our learning algorithm as an approximate MLE. See Pang et al. (2020); Nijkamp et al. (2020); Xie et al. (2023) for a theoretical understanding of the learning algorithm based on the finite-step MCMC.

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 ∑ 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 θ ( z 0 | 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.

3.3 Property Conditioned Generation

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,

This 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.,

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.

3.4 Optimization via Online Learning with LPT

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.), AutoDock-GPU (Santos-Martins et al., 2021).

As discussed in Sec. 3.3, generated molecules are reliable when the desired property value is supported by the learned model. To leverage this, we propose an iterative distribution shifting method, as shown in Alg. 2, as a form of online learning for LPT. In each iteration of distribution shifting, we

(a) Sample molecules from the learned LPT, using incremental extrapolation to generate a synthetic dataset with improved desired properties.

(b) Use hindsight relabeling for the sampled molecules' properties based on oracle functions.

(c) Apply MLE learning of LPT based on this synthetic dataset of molecule-property pairs as described in Sec. 3.2.

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.

3.5 Efficiency Analysis

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

9

10

11

12

13

14

0

5

10

15

20

25

Iteration

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 MCMCupdates can be reduced to N = 2 steps per sampling stage, achieving an approximate 5 × speedup in posterior sampling.

Sample efficiency Weaim 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 V AE 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: ∇ z 0 log p ( z 0 | y ) = ∇ z 0 (log p ( z 0 ) + log p ( y | z 0 )) = -z 0 + 1 σ 2 ( y -s γ ( z )) ∇ z 0 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 topN 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.

4 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.

5 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.

5.1 Overview

Molecule Sequence Design For molecule design tasks, we use SELFIES representations of the ZINC (Irwin et al., 2012) dataset, which comprises 250 K 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 AutoDockGPU (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.

Protein and DNA Sequence Design 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.

Autodock Kd = 36.1 μM

Autodock Kd = 2.5 μM

Autodock Kd = 0.97 μM

C1

C2

For this task, we use a subset of 5,000 samples as the training set by following the methodology outlined in Trabucco et al. (2022). Due to the unavailability of an exact oracle function for the GFP task, we follow the same oracle function preparation as in Design-Bench, and train a Transformer regression model on the full dataset with a total of 56,086 samples as the oracle function. We evaluate the design performance and diversity of the generated samples.

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 . 5 Koracle 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.

5.2 Binding Affinity Maximization

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. The predictor model p γ ( y | z ) for this task is a regression model that estimates docking scores. We compare our model with several baseline methods, which are introduced in App. A.4. As shown in Tabs. 1 and 2, our model significantly surpasses other methods across all three binding affinity maximization tasks in terms of K D , often achieving substantial improvement. Lower K D values indicate better performance. Furthermore, in Tab. 2, we report the average performance of the top 50 and top 100 molecules to demonstrate that our model can effectively generate a diverse pool of candidate molecules with the desired properties. Visualizations of generated molecules are provided in App. A.6.1.

Table 2: Single-objective binding affinity maximization results for PHGDH, reporting the 1st, 2nd, and 3rd performance, along with average performance of top 50 and top 100 molecules for each model. Performance is measured by K D (10 -2 nM) . The highest scores are in bold.

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.

Kd = 6.68 μM

Kd = 1.95 μM

Kd = 1.17 μM

Kd = 0.33 μM

Kd = 0.01 μM

Kd = 0.03 μM

Kd = 5.67 nM

Kd = 7.56 nM

C1

C2

C2

1

2

3

4

5

6

7

(a)

(b)

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.

5.3 Biological Sequence Design

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.

5.4 Sample Efficiency

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.

6 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.

References

Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International conference on learning representations (ICLR) , 2019.

Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimized image compression. arXiv preprint arXiv:1611.01704 , 2016.

David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In International conference on machine learning (ICML) , pages 773-782, 2019.

Nathan Brown, Marco Fiscato, Marwin HS Segler, and Alain C Vaucher. Guacamol: benchmarking models for de novo molecular design. Journal of Chemical Information and Modeling , 59(3): 1096-1108, 2019.

Austin H Cheng, Andy Cai, Santiago Miret, Gustavo Malkomes, Mariano Phielipp, and Alán AspuruGuzik. Group selfies: a robust fragment-based molecular string representation. Digital Discovery , 2023.

Hanjun Dai, Yingtao Tian, Bo Dai, Steven Skiena, and Le Song. Syntax-directed variational autoencoder for structured data. In International Conference on Learning Representations (ICLR) , 2018.

Nicola De Cao and Thomas Kipf. Molgan: An implicit generative model for small molecular graphs. arXiv preprint arXiv:1805.11973 , 2018.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS) , volume 34, pages 8780-8794, 2021.

Yuanqi Du, Tianfan Fu, Jimeng Sun, and Shengchao Liu. Molgensurvey: A systematic survey in machine learning models for molecule design. arXiv preprint arXiv:2203.14500 , 2022.

Peter Eckmann, Kunyang Sun, Bo Zhao, Mudong Feng, Michael K Gilson, and Rose Yu. Limo: Latent inceptionism for targeted molecule generation. In International Conference on Machine Learning (ICML) , 2022.

Clara Fannjiang and Jennifer Listgarten. Autofocused oracles for model-based design. In Advances in Neural Information Processing Systems (NeurIPS) , volume 33, pages 12945-12956, 2020.

Tianfan Fu, Wenhao Gao, Cao Xiao, Jacob Yasonik, Connor W Coley, and Jimeng Sun. Differentiable scaffolding tree for molecular optimization. arXiv preprint arXiv:2109.10469 , 2021.

Nathan Fuller, Loredana Spadola, Scott Cowen, Joe Patel, Heike Schönherr, Qing Cao, Andrew McKenzie, Fredrik Edfeldt, Al Rabow, and Robert Goodnow. An improved model for fragmentbased lead generation at astrazeneca. Drug Discovery Today , 21(8):1272-1283, 2016.

Wenhao Gao, Tianfan Fu, Jimeng Sun, and Connor Coley. Sample efficiency matters: a benchmark for practical molecular optimization. In Advances in Neural Information Processing Systems (NeurIPS) , volume 35, pages 21342-21357, 2022.

Rafael Gómez-Bombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS Central Science , 4(2):268-276, 2018.

Tian Han, Yang Lu, Song-Chun Zhu, and Ying Nian Wu. Alternating back-propagation for generator network. In AAAI Conference on Artificial Intelligence (AAAI) , pages 1976-1984, 2017.

Nikolaus Hansen. The cma evolution strategy: a comparing review. Towards a new evolutionary computation: Advances in the estimation of distribution algorithms , pages 75-102, 2006.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

John J Irwin, Teague Sterling, Michael M Mysinger, Erin S Bolstad, and Ryan G Coleman. Zinc: a free tool to discover chemistry for biology. Journal of Chemical Information and Modeling , 52(7): 1757-1768, 2012.

Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning (ICML) , pages 9786-9801. PMLR, 2022.

Moksh Jain, Sharath Chandra Raparthy, Alex Hernández-Garcıa, Jarrid Rector-Brooks, Yoshua Bengio, Santiago Miret, and Emmanuel Bengio. Multi-objective gflownets. In International Conference on Machine Learning (ICML) , pages 14631-14653, 2023.

Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International Conference on Machine Learning (ICML) , pages 2323-2332, 2018.

Hiroshi Kajino. Molecular hypergraph grammar with its application to molecular optimization. In International Conference on Machine Learning (ICML) , pages 3183-3191, 2019.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 4401-4410, 2019.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR) , 2014.

Deqian Kong, Bo Pang, Tian Han, and Ying Nian Wu. Molecule design by latent space energy-based modeling and gradual distribution shifting. In Conference on Uncertainty in Artificial Intelligence (UAI) , volume 216, pages 1109-1120, 2023.

Deqian Kong, Dehong Xu, Minglu Zhao, Bo Pang, Jianwen Xie, Andrew Lizarraga, Yuhao Huang, Sirui Xie, and Ying Nian Wu. Latent plan transformer for trajectory abstraction: Planning as latent space inference. In Advances in Neural Information Processing Systems (NeurIPS) , 2024.

Mario Krenn, Florian Häse, AkshatKumar Nigam, Pascal Friederich, and Alan Aspuru-Guzik. Selfreferencing embedded strings (selfies): A 100% robust molecular string representation. Machine Learning: Science and Technology , 1(4):045024, 2020.

Matt J Kusner, Brooks Paige, and José Miguel Hernández-Lobato. Grammar variational autoencoder. In International Conference on Machine Learning (ICML) , pages 1945-1954, 2017.

Greg Landrum et al. RDKit: Open-source cheminformatics. URL https://www.rdkit.org .

Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv:1806.09055 , 2018.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR) , 2019.

Youzhi Luo, Keqiang Yan, and Shuiwang Ji. Graphdf: A discrete flow model for molecular graph generation. In International Conference on Machine Learning (ICML) , pages 7192-7203, 2021.

Natalie Maus, Haydn Jones, Juston Moore, Matt J Kusner, John Bradshaw, and Jacob Gardner. Local latent space bayesian optimization over structured inputs. In Advances in Neural Information Processing Systems (NeurIPS) , volume 35, pages 34505-34518, 2022.

Krzysztof Maziarz, Henry Jackson-Flux, Pashmina Cameron, Finton Sirockin, Nadine Schneider, Nikolaus Stiefl, Marwin Segler, and Marc Brockschmidt. Learning to extend molecular scaffolds with structural motifs. arXiv preprint arXiv:2103.03864 , 2021.

Fabian Mentzer, Eirikur Agustsson, Michael Tschannen, Radu Timofte, and Luc Van Gool. Conditional probability models for deep image compression. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 4394-4402, 2018.

Edouard Mullarky, Jiayi Xu, Anita D Robin, David J Huggins, Andy Jennings, Naoyoshi Noguchi, Andrea Olland, Damodharan Lakshminarasimhan, Michael Miller, Daisuke Tomita, et al. Inhibition of 3-phosphoglycerate dehydrogenase (phgdh) by indole amides abrogates de novo serine synthesis in cancer cells. Bioorganic & Medicinal Chemistry Letters , 29(17):2503-2510, 2019.

Radford M Neal. MCMC using hamiltonian dynamics. Handbook of Markov Chain Monte Carlo , 2, 2011.

AkshatKumar Nigam, Pascal Friederich, Mario Krenn, and Alán Aspuru-Guzik. Augmenting genetic algorithms with deep neural networks for exploring the chemical space. In International Conference on Learning Representations (ICLR) , 2020.

Erik Nijkamp, Bo Pang, Tian Han, Linqi Zhou, Song-Chun Zhu, and Ying Nian Wu. Learning multi-layer latent variable model via variational optimization of short run mcmc for approximate inference. In European Conference on Computer Vision (ECCV) , pages 361-378, 2020.

Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. In Advances in Neural Information Processing Systems (NeurIPS) , volume 32, 2019.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 10684-10695, 2022.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical Image Computing and ComputerAssisted Intervention (MICCAI) , pages 234-241, 2015.

Debangshu Samanta and Gregg L Semenza. Serine synthesis helps hypoxic cancer stem cells regulate redox. Cancer Research , 76(22):6458-6462, 2016.

Diogo Santos-Martins, Leonardo Solis-Vasquez, Andreas F Tillack, Michel F Sanner, Andreas Koch, and Stefano Forli. Accelerating AutoDock4 with GPUs and gradient-based local search. Journal of Chemical Theory and Computation , 17(2):1060-1073, 2021.

Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprint arXiv:1511.05952 , 2015.

Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. Graphaf: a flow-based autoregressive model for molecular graph generation. arXiv preprint arXiv:2001.09382 , 2020.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

Quentin Spillier and Raphaël Frédérick. Phosphoglycerate dehydrogenase (phgdh) inhibitors: A comprehensive review 2015-2020. Expert Opinion on Therapeutic Patents , 31(7):597-608, 2021.

Kevin Swersky, Yulia Rubanova, David Dohan, and Kevin Murphy. Amortized bayesian optimization over discrete spaces. In Conference on Uncertainty in Artificial Intelligence (UAI) , pages 769-778, 2020.

Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In International Conference on Machine Learning (ICML) , pages 1064-1071, 2008.

Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Conservative objective models for effective offline model-based optimization. In International Conference on Machine Learning (ICML) , pages 10358-10368, 2021.

Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench: Benchmarks for data-driven offline model-based optimization. In International Conference on Machine Learning (ICML) , pages 21658-21676, 2022.

Austin Tripp, Erik Daxberger, and José Miguel Hernández-Lobato. Sample-efficient optimization in the latent space of deep generative models via weighted retraining. In Advances in Neural Information Processing Systems (NeurIPS) , volume 33, pages 11259-11272, 2020.

Austin Tripp, Gregor NC Simm, and José Miguel Hernández-Lobato. A fresh look at de novo molecular design benchmarks. In NeurIPS 2021 AI for Science Workshop , 2021.

Judith E Unterlass, Robert J Wood, Arnaud Baslé, Julie Tucker, Céline Cano, Martin ME Noble, and Nicola J Curtin. Structural insights into the enzymatic activity and potential substrate promiscuity of human 3-phosphoglycerate dehydrogenase (phgdh). Oncotarget , 8(61):104478, 2017.

David Weininger. Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of Chemical Information and Computer Sciences , 28(1):31-36, 1988.

James T Wilson, Riccardo Moriconi, Frank Hutter, and Marc Peter Deisenroth. The reparameterization trick for acquisition functions. arXiv preprint arXiv:1712.00424 , 2017.

Robin Winter, Floriane Montanari, Andreas Steffen, Hans Briem, Frank Noé, and Djork-Arné Clevert. Efficient multi-objective molecular optimization in a continuous latent space. Chemical Science , 10(34):8016-8024, 2019.

Jianwen Xie, Yang Lu, Song-Chun Zhu, and Ying Nian Wu. A theory of generative convnet. In International Conference on Machine Learning (ICML) , pages 2635-2644, 2016.

Jianwen Xie, Ruiqi Gao, Zilong Zheng, Song-Chun Zhu, and Ying Nian Wu. Learning dynamic generator model by alternating back-propagation through time. In AAAI Conference on Artificial Intelligence, (AAAI) , pages 5498-5507, 2019.

Jianwen Xie, Yaxuan Zhu, Yifei Xu, Dingcheng Li, and Ping Li. A tale of two latent flows: Learning latent space normalizing flow with short-run langevin flow for approximate inference. In AAAI Conference on Artificial Intelligence (AAAI) , 2023.

Yutong Xie, Chence Shi, Hao Zhou, Yuwei Yang, Weinan Zhang, Yong Yu, and Lei Li. Mars: Markov molecular sampling for multi-objective drug discovery. In International Conference on Learning Representations (ICLR) , 2021.

Xiufeng Yang, Tanuj Kr Aasawat, and Kazuki Yoshizoe. Practical massively parallel monte-carlo tree search applied to molecular design. arXiv preprint arXiv:2006.10504 , 2020.

Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec. Graph convolutional policy network for goal-directed molecular graph generation. In Advances in Neural Information Processing Systems (NeurIPS) , pages 6410-6421, 2018.

Muhan Zhang, Shali Jiang, Zhicheng Cui, Roman Garnett, and Yixin Chen. D-vae: A variational autoencoder for directed acyclic graphs. In Advances in Neural Information Processing Systems (NeurIPS) , pages 1588-1600, 2019.

Xiaoya Zhao, Jianfei Fu, Jinlin Du, and Wenxia Xu. The role of d-3-phosphoglycerate dehydrogenase in cancer. International Journal of Biological Sciences , 16(9):1495, 2020.

Zhenpeng Zhou, Steven Kearnes, Li Li, Richard N Zare, and Patrick Riley. Optimization of molecules via deep reinforcement learning. Scientific Reports , 9(1):1-10, 2019.

Yiheng Zhu, Jialu Wu, Chaowen Hu, Jiahuan Yan, Tingjun Hou, Jian Wu, et al. Sample-efficient multiobjective molecular optimization with gflownets. In Advances in Neural Information Processing Systems (NeurIPS) , volume 36, 2024.

A Appendix

A.1 Model Learning

We provide the derivation of Eq. (5) in Sec. 3.2.

A.2 Additional Experiments

A.2.1 Model Sanity Check: Penalized logP and QED Maximization

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.

A.2.2 Robustness to Noisy Oracles

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.

A.2.3 Online Learning from Scratch

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

learning budget). Results in Tab. 8 show that LPT can effectively discover high-binding molecules for ESR1 and ACAA1, even without pre-training. However, performance on PHGDH remained suboptimal compared to the version with pre-training, indicating that this target may require additional oracle queries due to its inherent complexity. These findings highlight the potential of pure online learning approaches for future exploration.

A.2.4 Ablation Studies

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:

Using samples from z ∼ p α ( z ) instead of p θ ( z | y ) to generate the proposals P t .

Removing weighted retraining and applying the standard objective ∑ n i =1 1 / n log p θ ( x i , y i ) .

Setting the total number of shifting iterations as 1 .

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.

A.3 Model Architecture and Training Details

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 . 33 M.

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 . 5 e -4 and 7 . 5 e -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 3 e -4 and 7 . 5 e -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

t

+1

+1

γ

t

+1

end for

1

n

1

n

1

n

∑

i

∑

i

∑

i

For 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 . 5 k 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.

δ

δ

α

β

δ

(

(

γ

x

x

(

i

i

x

i

, y

, y

i

i

, y

i

)

)

;

;

)

as in Sec. 3.2.

=

=

=

α

β

t

t

γ

t

+

+

+

η

η

0

1

η

2

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 } i m =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 } i m =1 ∪ D t -1 . step (c): Update LPT using maximum likelihood on synthetic dataset D t by following Alg. 1 end while

A.4 Baselines

Our 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).

T206

C233

D259

A285

D1740

I155S

R154

A.6 Visualization of Generated Molecules

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

Figs. 6 to 8 visualize some molecules produced during the multi-objective optimization for ESR1, ACAA1, and PHGDH, respectively. Meanwhile, Figs. 9 to 11 depict some examples of molecules produced during the single-objective optimization for ESR1, ACAA1, and PHGDH, respectively.

0.038 3.46 0.582

0.047 3.24 0.460

0.048 5.14 0.392

0.049 5.09 0.839

0.049 5.01 0.543

0.062 4.83 0.526

0.063 5.04 0.558

0.065 4.87 0.434

0.065 5.46 0.585

0.065 3.52 0.443

0.066 4.22 0.280

0.066 3.69 0.598

0.069 5.18 0.553

0.073 4.93 0.692

0.077 4.77 0.515

0.078 3.43 0.597

0.078 5.50 0.556

0.078 3.26 0.469

0.079 4.95 0.531

0.081 3.34 0.507

0.181 4.86 0.503

0.214 4.18 0.612

0.214 4.80 0.408

0.246 4.46 0.451

0.258 5.02 0.457

0.276 4.46 0.553

0.281 5.03 0.580

0.301 4.26 0.566

0.306 4.64 0.590

0.3273.84 0.455

0.333 4.02 0.449

0.344 4.95 0.613

0.344 4.33 0.404

0.362 4.74 0.470

0.368 5.37 0.508

0.368 5.32 0.626

0.375 4.60 0.519

0.387 4.71 0.500

0.394 4.63 0.425

0.408 4.45 0.479

0.021 3.11 0.503

0.030 3.22 0.433

0.036 4.95 0.414

0.047 4.91 0.470

0.064 4.98 0.495

0.079 5.23 0.619

0.113 4.27 0.468

0.113 4.70 0.505

0.119 4.72 0.490

0.156 4.58 0.421

0.164 4.79 0.408

0.172 3.57 0.466

0.181 4.22 0.423

0.181 3.91 0.458

0.184 4.75 0.486

0.191 3.78 0.416

0.201 4.08 0.362

0.211 4.78 0.500

0.215 4.73 0.543

0.218 4.07 0.386

89

0.0043

0.0053

0.0091

0.0137

0.0139

0.0152

0.0168

0.0192

0.0195

0.0260

个

0.0319

0.0324

0.0324

0.0330

0.0347

a

0.0353

0.0377

0.0390

0.0411

0.0447

0.0371

0.0411

0.0451

0.0529

0.0566

0.0566

0.0585

0.0605

0.0704

0.0716

0.0729

0.0767

0.0793

0.0820

0.0834

0.0863

0.0863

0.0907

0.0923

0.0939

0.0007

0.0016

0.0019

0.0024

0.0026

0.0027

0.0041

0.0043

0.0051

0.0052

0.0055

0.0055

0.0056

0.0065

0.0066

0.0071

0.0073

0.0076

0.0080

0.0081

A.6.2 Structure-constrained Optimization

6.797 2.52 0.693

1.675 2.71 0.626

1.917 2.65 0.685

1.136 3.58 0.796

1.415 4.11 0.913

1.792 3.90 0.873

1.592 3.10 0.873

2.157 3.10 0.873

2.347 2.91 0.507

2.157 3.10 0.873

0.732 3.85 0.403

2.428 3.03 0.259

0.367 3.81 0.171

2.554 3.37 0.216

0.166 3.60 0.219

0.187 3.67 0.265

0.315 3.11 0.248

1.278 3.18 0.282

1.463 3.14 0.296

1.215 3.40 0.275

A.7 Broader Impact

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.

NeurIPS Paper Checklist

1. Claims

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.

2. Limitations

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: See Sec. 3.2.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: See App. A.3.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Code is released in our project page.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: See App. A.3 and Sec. 5.1.

Guidelines:

The answer NA means that the paper does not include experiments.

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 error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: See App. A.3.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: None.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

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

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [No]

Justification: None.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: None.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [No]

Justification: None.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [No]

Justification: None.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [No]

Justification: None.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n2. The formulation of multi-objective optimization task is quite novel and well defined.\n3. The experiments are comprehensive and convincing to me.\n\nWEAKNESSES:\n1. 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?\n2. 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?\n3. Minor: \n(a) line 42: would suggest to write down the full name of MCMC before using its abbreviation for general audience.\n(b) typo at equation (1): subscript of o^p_m(x).\n\nQUESTIONS:\n1. 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].\n2. 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].\n3. 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.\n4. 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)?\n\n[1] Multi-objective Deep Data Generation with Correlated Property Control. NeurIPS 2022.\n[2] Property controllable variational autoencoder via invertible mutual dependence. ICLR 2020.\n\nLIMITATIONS:\nYes. Limitations are well discussed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOriginality: 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. \nQuality: 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. \nClarity: 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. \nSignificance: 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.\n\nWEAKNESSES:\nConditional 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.\n\nLatent-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.\n\nQUESTIONS:\nSee Weaknesses section above for suggestions/points of discussion.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. The online learning approach aligns well with real-world scenarios.\n2. The experiments are comprehensive, and the introduction of a new task for conditional generation and design trajectory comparison is expressive.\n3. 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.\n\nWEAKNESSES:\n- The clarity of the paper could be improved in several aspects:\n  1. Adding an illustration or diagram of the training and optimization process would make the methodology easier to understand.\n  2. Some formulas lack derivations and explanations, as raised in Question 1.\n- 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?\n- Minor Issues:\n1. Typo error: \"2rd\" in Table 1.\n2. Using \"Kd\" values might lead readers to mistakenly think they are wet lab results, while in paper they represent docking scores.\n3. In Algorithm 2, 'step (c) Update LPT on synthetic dataset using **Algorithm 2** using MLE.' Should it be Algorithm 1 here?\n\nQUESTIONS:\n1. 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.\n2. 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.\n3. How robust is your model to the quality of the Oracle?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}