Original: M OL EBM: M OLECULE G ENERATION AND D ESIGN BY L ATENT S PACE E NERGY -B ASED M ODELING
Revised: M OL EBM: M OLECULE G ENERATION AND D ESIGN BY L ATENT S PACE E NERGY -B ASED M ODELING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Generation of molecules with desired chemical and biological properties such as high drug-likeness, high binding affinity to target proteins is critical in drug dis covery
Revised: Generation of molecules with desired chemical and biological properties such as high drug-likeness, high binding affinity to target proteins is critical in drug dis covery

Original: In this paper
Revised: In this paper

Original: , we propose a probabilistic generative model to capture the joint distribution of molecules and their properties
Revised: , we propose a probabilistic generative model to capture the joint distribution of molecules and their properties

Original: Our model assumes an energy based model
Revised: Our model assumes an energy based model

Original: EBM) in the latent space
Revised: EBM) in the latent space

Original: Given the latent vector sampled from the latent space
Revised: Given the latent vector sampled from the latent space

Original: EBM
Revised: EBM

Original: , both molecules and molecular properties are conditionally sam
Revised: , both molecules and molecular properties are conditionally sam

Original: - pled via a molecule generator model and a property regression model respectively.
Revised: - pled via a molecule generator model and a property regression model respectively.

Original: The EBM in a low dimensional latent space allows our model to capture complex chemical rules implicitly but efficiently and effectively.
Revised: The EBM in a low dimensional latent space allows our model to capture complex chemical rules implicitly but efficiently and effectively.

Original: Due to the joint model- ing with chemical properties
Revised: Due to the joint model- ing with chemical properties

Original: , molecule design can be conveniently and naturally achieved by conditional sampling from our learned model given desired proper
Revised: , molecule design can be conveniently and naturally achieved by conditional sampling from our learned model given desired proper

Original: - ties, in both single
Revised: - ties, in both single

Original: -objective and multi-objective optimization settings.
Revised: -objective and multi-objective optimization settings.

Original: The latent space
Revised: The latent space

Original: EBM molecule generator and property regression model are learned jointly by approximate maximum likelihood
Revised: EBM molecule generator and property regression model are learned jointly by approximate maximum likelihood

Original: , while optimization of properties is accom
Revised: , while optimization of properties is accom

Original: - plished by gradual shifting of the model distribution towards the region supported by molecules with high property values.
Revised: - plished by gradual shifting of the model distribution towards the region supported by molecules with high property values.

Original: Our experiments show that our model outperforms state-of-the-art models on various molecule design tasks
Revised: Our experiments show that our model outperforms state-of-the-art models on various molecule design tasks

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In drug discovery, it is of vital importance to find or design molecules with desired pharmacologic or chemical properties such as high drug-likeness and binding affinity to a target protein.
Revised: In drug discovery, it is of vital importance to find or design molecules with desired pharmacologic or chemical properties such as high drug-likeness and binding affinity to a target protein.

Original: It is challenging to directly optimize or search over the drug-like molecule space since it is discrete and enormous, with an estimated size is on the order of 10 33 (Polishchuk et al., 2013).
Revised: It is challenging to directly optimize or search over the drug-like molecule space since it is discrete and enormous, with an estimated size is on the order of 10 33 (Polishchuk et al., 2013).

Original: Recently, a large body of work attempts to tackle this problem.
Revised: Recently, a large body of work attempts to tackle this problem.

Original: The first line of work leverages deep generative models to map the discrete molecule space to a continuous latent space, and optimizes molecular properties in the latent space with methods like Bayesian optimization (G ´ omezBombarelli et al., 2018; Kusner et al., 2017; Jin et al., 2018).
Revised: The first line of work leverages deep generative models to map the discrete molecule space to a continuous latent space, and optimizes molecular properties in the latent space with methods like Bayesian optimization (G ´ omezBombarelli et al., 2018; Kusner et al., 2017; Jin et al., 2018).

Original: The second line of work recruits reinforcement learning algorithms to optimize properties in the molecular graph space directly (You et al., 2018; De Cao & Kipf, 2018;
Revised: The second line of work recruits reinforcement learning algorithms to optimize properties in the molecular graph space directly (You et al., 2018; De Cao & Kipf, 2018;

Original: Zhou et al., 2019; Shi et al., 2020; Luo et al., 2021).
Revised: Zhou et al., 2019; Shi et al., 2020; Luo et al., 2021).

Original: A number of other efforts have been made to optimize molecular properties with genetic algorithms (Nigam et al., 2020), particle-swarm algorithms (Winter et al., 2019)
Revised: A number of other efforts have been made to optimize molecular properties with genetic algorithms (Nigam et al., 2020), particle-swarm algorithms (Winter et al., 2019)

Original: specialized MCMC methods (Xie et al., 2021).
Revised: specialized MCMC methods (Xie et al., 2021).

Original: In this work, we propose a method along the first line mentioned above, by learning a probabilistic latent generative model of molecule distributions and optimizing chemical properties in the latent space.
Revised: In this work, we propose a method along the first line mentioned above, by learning a probabilistic latent generative model of molecule distributions and optimizing chemical properties in the latent space.

Original: Given the central role of latent variables in this approach, we emphasize that it is critical to learn a latent space model that captures the data regularities of the molecules.
Revised: Given the central role of latent variables in this approach, we emphasize that it is critical to learn a latent space model that captures the data regularities of the molecules.

Original: Thus, instead of assuming a simple Gaussian distribution in the latent space as in prior work (G´omez-Bombarelli et al., 2018; Jin et al., 2018), we assume a flexible and expressive energy-based model (EBM) (LeCun et al., 2006; Ngiam et al., 2011; Kim & Bengio, 2016; Xie et al., 2016; Kumar et al., 2019; Nijkamp et al., 2019; Du & Mordatch, 2019; Grathwohl et al., 2019;
Revised: Thus, instead of assuming a simple Gaussian distribution in the latent space as in prior work (G´omez-Bombarelli et al., 2018; Jin et al., 2018), we assume a flexible and expressive energy-based model (EBM) (LeCun et al., 2006; Ngiam et al., 2011; Kim & Bengio, 2016; Xie et al., 2016; Kumar et al., 2019; Nijkamp et al., 2019; Du & Mordatch, 2019; Grathwohl et al., 2019;

Original: Finn et al., 2016) in latent space.
Revised: Finn et al., 2016) in latent space.

Original: This leads to a latent space energy-based model (LSEBM) as studied in Pang et al.
Revised: This leads to a latent space energy-based model (LSEBM) as studied in Pang et al.

Original: ; Nie et al.
Revised: ; Nie et al.

Original: (2021), where LSEBM has been shown to model the distributions of natural images and text well.
Revised: (2021), where LSEBM has been shown to model the distributions of natural images and text well.

Original: For molecule modeling, without any explicit validity constraints in generation, our model generates molecules with high validity with simple SMILES representation (Weininger, 1988).
Revised: For molecule modeling, without any explicit validity constraints in generation, our model generates molecules with high validity with simple SMILES representation (Weininger, 1988).

Original: Given our goal of property optimization, we learn a joint distribution of molecules and their properties.
Revised: Given our goal of property optimization, we learn a joint distribution of molecules and their properties.

Original: Our model consists of 1) an EBM in a low-dimensional continuous latent space, 2) a generator mapping from the latent space to the observed molecule space, and 3) a property regression model mapping from the latent space to the property values (see Figure 1).
Revised: Our model consists of 1) an EBM in a low-dimensional continuous latent space, 2) a generator mapping from the latent space to the observed molecule space, and 3) a property regression model mapping from the latent space to the property values (see Figure 1).

Original: We call our model as MolEBM .
Revised: We call our model as MolEBM .

Original: All three components in our model are learned jointly by approximate maximum likelihood.
Revised: All three components in our model are learned jointly by approximate maximum likelihood.

Original: A learned model generates a molecule with a high property value in two steps: 1) given the property value, sample the latent vector; 2) given the sampled latent vector, generate a molecule (see the topto-bottom path in Figure 1a).
Revised: A learned model generates a molecule with a high property value in two steps: 1) given the property value, sample the latent vector; 2) given the sampled latent vector, generate a molecule (see the topto-bottom path in Figure 1a).

Original: Since the learned model approximates the data distribution, directly sampling from the learned model conditional on a high property value does not work well since a molecule with a high property value is most likely not in the original data distribution.
Revised: Since the learned model approximates the data distribution, directly sampling from the learned model conditional on a high property value does not work well since a molecule with a high property value is most likely not in the original data distribution.

Original: We thus design a method to gradually shift the learned distribution towards the region supported by molecules with high property values, and sample molecules with desirable properties from the shifted distribution.
Revised: We thus design a method to gradually shift the learned distribution towards the region supported by molecules with high property values, and sample molecules with desirable properties from the shifted distribution.

Original: In drug discovery, most often we need to consider multiple properties simultaneously.
Revised: In drug discovery, most often we need to consider multiple properties simultaneously.

Original: Our model can be extended to this setting straightforwardly.
Revised: Our model can be extended to this setting straightforwardly.

Original: With our method, we only need to add a regression model for each property, while the learning and sampling methods remain the same.
Revised: With our method, we only need to add a regression model for each property, while the learning and sampling methods remain the same.

Original: Learning the model involves inferring the latent vector given both the molecule and the property value, and we recruits Langevin dynamics instead of amortized inference network for inference computation.
Revised: Learning the model involves inferring the latent vector given both the molecule and the property value, and we recruits Langevin dynamics instead of amortized inference network for inference computation.

Original: This design makes our approach versatile in dealing with varying number of properties.
Revised: This design makes our approach versatile in dealing with varying number of properties.

Original: We evaluate our method in various settings including single-objective optimization and multiobjective optimization.
Revised: We evaluate our method in various settings including single-objective optimization and multiobjective optimization.

Original: Our method outperforms prior methods by significant margins.
Revised: Our method outperforms prior methods by significant margins.

Original: In summary, our contributions are as follows:
Revised: In summary, our contributions are as follows:

Original: • We propose to learn a latent space energy-based model for the joint distribution of molecules and molecular properties.
Revised: • We propose to learn a latent space energy-based model for the joint distribution of molecules and molecular properties.

Original: • We develop a sampling with gradual distribution shifting method, enabling us to extrapolate the data distribution and sample from the region supported by molecules with high property values.
Revised: • We develop a sampling with gradual distribution shifting method, enabling us to extrapolate the data distribution and sample from the region supported by molecules with high property values.

Original: • Our methods are versatile enough to be extended to optimizing multiple properties together.
Revised: • Our methods are versatile enough to be extended to optimizing multiple properties together.

Original: • Our model achieves state-of-the-art performances on a wide range of molecule optimization
Revised: • Our model achieves state-of-the-art performances on a wide range of molecule optimization

Original: tasks.
Revised: tasks.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Optimization with Generative Models.
Revised: Optimization with Generative Models.

Original: Deep generative models approximate the distribution of molecules with desired biological or non-biological properties.
Revised: Deep generative models approximate the distribution of molecules with desired biological or non-biological properties.

Original: Existing approaches for generating molecules include applying variational autoencoder (VAE) (Kingma & Welling, 2014) and generative adversarial network (GAN) (Goodfellow et al., 2014) etc. to molecule data (G ´ omez-Bombarelli et al., 2018; Jin et al., 2018; De Cao & Kipf, 2018; Honda et al., 2019; Madhawa et al., 2019; Shi et al., 2020; Zang & Wang, 2020; Kotsias et al., 2020; Chen et al., 2021; Fu et al., 2020; Liu et al., 2021; Bagal et al., 2021; Eckmann et al., 2022; Segler et al., 2018).
Revised: Existing approaches for generating molecules include applying variational autoencoder (VAE) (Kingma & Welling, 2014) and generative adversarial network (GAN) (Goodfellow et al., 2014) etc. to molecule data (G ´ omez-Bombarelli et al., 2018; Jin et al., 2018; De Cao & Kipf, 2018; Honda et al., 2019; Madhawa et al., 2019; Shi et al., 2020; Zang & Wang, 2020; Kotsias et al., 2020; Chen et al., 2021; Fu et al., 2020; Liu et al., 2021; Bagal et al., 2021; Eckmann et al., 2022; Segler et al., 2018).

Original: After learning continuous representations for molecules, they are further able to optimize using different methods.
Revised: After learning continuous representations for molecules, they are further able to optimize using different methods.

Original: (Segler et al., 2018) proposes to optimize by simulating design-synthesis-test cycles.
Revised: (Segler et al., 2018) proposes to optimize by simulating design-synthesis-test cycles.

Original: (G´omez-Bombarelli et al., 2018; Jin et al., 2018; Eckmann et al., 2022) propose to learn a surrogate function to predict properties, and then use Bayesian optimization to optimize the latent vectors.
Revised: (G´omez-Bombarelli et al., 2018; Jin et al., 2018; Eckmann et al., 2022) propose to learn a surrogate function to predict properties, and then use Bayesian optimization to optimize the latent vectors.

Original: However, the performance of this latent optimization is not satisfactory due to three major issues.
Revised: However, the performance of this latent optimization is not satisfactory due to three major issues.

Original: First, it is difficult to train an accurate surrogate predictor especially for those novel molecules with high properties along the design trajectories.
Revised: First, it is difficult to train an accurate surrogate predictor especially for those novel molecules with high properties along the design trajectories.

Original: Second, as the learned latent space tries to cover the fixed data space, its ability to explore the targets out of the distribution is limited (Brown et al., 2019; Huang et al., 2021).
Revised: Second, as the learned latent space tries to cover the fixed data space, its ability to explore the targets out of the distribution is limited (Brown et al., 2019; Huang et al., 2021).

Original: Third, those methods are heavily dependent on the quality of learned latent space, which requires non-trivial efforts to design encoders when dealing with multiple properties.
Revised: Third, those methods are heavily dependent on the quality of learned latent space, which requires non-trivial efforts to design encoders when dealing with multiple properties.

Original: To address above issues, (Eckmann et al., 2022) use VAE to learn the latent space and train predictors separately using generated molecules, and then leverage latent inceptionism, which involves the decoder solely, to optimize the latent vector with multiple predictors.
Revised: To address above issues, (Eckmann et al., 2022) use VAE to learn the latent space and train predictors separately using generated molecules, and then leverage latent inceptionism, which involves the decoder solely, to optimize the latent vector with multiple predictors.

Original: In this paper, we propose an encoder-free model in both training and optimization to learn the joint distribution of molecules and properties, and make it possible to obtain several adequate predictors.
Revised: In this paper, we propose an encoder-free model in both training and optimization to learn the joint distribution of molecules and properties, and make it possible to obtain several adequate predictors.

Original: We then design an efficient algorithm to shift the learned distribution iteratively.
Revised: We then design an efficient algorithm to shift the learned distribution iteratively.

Original: x z y p γ ( y | z ) p β ( x | z )
Revised: x z y p γ ( y | z ) p β ( x | z )

Original: Controlled
Revised: Controlled

Original: Generation p α ( z ) (a) Single-Objective Optimization.
Revised: Generation p α ( z ) (a) Single-Objective Optimization.

Original: x z · · · yy m p β ( x | z ) p γm ( y m | z ) p γ 1 ( y 1 | z ) p
Revised: x z · · · yy m p β ( x | z ) p γm ( y m | z ) p γ 1 ( y 1 | z ) p

Original: α ( z ) (b) Multi-Objective Optimization.
Revised: α ( z ) (b) Multi-Objective Optimization.

Original: Optimization with Reinforcement Learning and Evolutionary Algorithms.
Revised: Optimization with Reinforcement Learning and Evolutionary Algorithms.

Original: Reinforcement learning (RL) based methods directly optimize and generate molecules in an explicit data space (You et al., 2018;
Revised: Reinforcement learning (RL) based methods directly optimize and generate molecules in an explicit data space (You et al., 2018;

Original: Zhou et al., 2019; Jin et al., 2020; Gottipati et al., 2020).
Revised: Zhou et al., 2019; Jin et al., 2020; Gottipati et al., 2020).

Original: By formulating the property design as a discrete optimization task, they can modify the molecular substructures guided by an oracle reward function.
Revised: By formulating the property design as a discrete optimization task, they can modify the molecular substructures guided by an oracle reward function.

Original: However, the training of those RL-based methods can be viewed as rejection sampling which is difficult and inefficient due to the random-walk search behavior in the discrete space.
Revised: However, the training of those RL-based methods can be viewed as rejection sampling which is difficult and inefficient due to the random-walk search behavior in the discrete space.

Original: Evolutionary algorithms (EA) also formulate the optimization in a discrete manner (Nigam et al., 2020; Jensen, 2019; Xie et al., 2021;
Revised: Evolutionary algorithms (EA) also formulate the optimization in a discrete manner (Nigam et al., 2020; Jensen, 2019; Xie et al., 2021;

Original: Fu et al., 2021a;b).
Revised: Fu et al., 2021a;b).

Original: By leveraging carefully-crafted combinatorial algorithms, they can search the molecule graph space in a flexible and efficient way.
Revised: By leveraging carefully-crafted combinatorial algorithms, they can search the molecule graph space in a flexible and efficient way.

Original: However, the design of those algorithms is non-trivial and domain specific.
Revised: However, the design of those algorithms is non-trivial and domain specific.

Original:  M ETHODS 3. P ROBLEM S ETUP AND O VERVIEW
Revised:  M ETHODS 3. P ROBLEM S ETUP AND O VERVIEW

Original: We use the SELFIES representation for molecules (Krenn et al., 2020).
Revised: We use the SELFIES representation for molecules (Krenn et al., 2020).

Original: It encodes each molecule as a string of characters and ensures validity of all SELFIES strings.
Revised: It encodes each molecule as a string of characters and ensures validity of all SELFIES strings.

Original: Let x = ( x (1) , ..., x ( t ) , ..., x ( T ) ) be a molecule string encoded in SELFIES,
Revised: Let x = ( x (1) , ..., x ( t ) , ..., x ( T ) ) be a molecule string encoded in SELFIES,

Original: where x ( t ) ∈
Revised: where x ( t ) ∈

Original: V is the t -th character and V is the vocabulary.
Revised: V is the t -th character and V is the vocabulary.

Original: Suppose y ∈ R represents a molecular property of interest.
Revised: Suppose y ∈ R represents a molecular property of interest.

Original: Then the problem we attempt to tackle is to optimize x such that its property y =
Revised: Then the problem we attempt to tackle is to optimize x such that its property y =

Original: y ∗ where y ∗ is some desirable value for y .
Revised: y ∗ where y ∗ is some desirable value for y .

Original: We take a probabilistic approach and treat the optimization problem as a sampling problem, that is,
Revised: We take a probabilistic approach and treat the optimization problem as a sampling problem, that is,

Original: [Equation]
Revised: [Equation]

Original: This is a single-objective optimization problem since only one property is targeted.
Revised: This is a single-objective optimization problem since only one property is targeted.

Original: In real-world drug design settings, we are more likely to need to optimize multiple properties simultaneously, that is, multi-objective optimization .
Revised: In real-world drug design settings, we are more likely to need to optimize multiple properties simultaneously, that is, multi-objective optimization .

Original: Suppose we optimize for { y j ∈ R } mj =1 , then our task is to sample,
Revised: Suppose we optimize for { y j ∈ R } mj =1 , then our task is to sample,

Original: [Equation]
Revised: [Equation]

Original: To address these problems, we propose a solution under a unified probabilistic framework.
Revised: To address these problems, we propose a solution under a unified probabilistic framework.

Original: As a first step, we need to model the data distribution of molecules, p data ( x ) .
Revised: As a first step, we need to model the data distribution of molecules, p data ( x ) .

Original: To this end, we recruit latent space energy-based model (Pang et al., 2020; Nie et al., 2021) for its expressiveness.
Revised: To this end, we recruit latent space energy-based model (Pang et al., 2020; Nie et al., 2021) for its expressiveness.

Original: LSEBM assumes latent vector z ∈ R d in a low dimensional latent space follows an energy-based model, while the observed x is generated by a generator conditional z , that is, p ( x, z )
Revised: LSEBM assumes latent vector z ∈ R d in a low dimensional latent space follows an energy-based model, while the observed x is generated by a generator conditional z , that is, p ( x, z )

Original: =
Revised: =

Original: p ( x | z ) p ( z ) .
Revised: p ( x | z ) p ( z ) .

Original: We introduce LSEBM within the context of molecule data in §3.2.
Revised: We introduce LSEBM within the context of molecule data in §3.2.

Original: For the purpose of property optimization, we propose to model the joint distribution of molecules and molecular properties (§3.3).
Revised: For the purpose of property optimization, we propose to model the joint distribution of molecules and molecular properties (§3.3).

Original: We first consider the single-objective optimization problem (eq. 1).
Revised: We first consider the single-objective optimization problem (eq. 1).

Original: Given the expressiveness of LSEBM, we assume that the latent vector z captures data regularities in x (evaluated in §4.2 and §A.4).
Revised: Given the expressiveness of LSEBM, we assume that the latent vector z captures data regularities in x (evaluated in §4.2 and §A.4).

Original: Thus, we can learn a simple regression model from the lowdimensional latent space, that is, p ( y | z ) , and the joint distribution is p ( x, y, z )
Revised: Thus, we can learn a simple regression model from the lowdimensional latent space, that is, p ( y | z ) , and the joint distribution is p ( x, y, z )

Original: = p ( z ) p ( x | z ) p ( y | z ) .
Revised: = p ( z ) p ( x | z ) p ( y | z ) .

Original: The model is learned with maximum likelihood (see §3.4 and Algorithm 1).
Revised: The model is learned with maximum likelihood (see §3.4 and Algorithm 1).

Original: With the learned model, we can optimize x given y =
Revised: With the learned model, we can optimize x given y =

Original: y ∗ by ancestral sampling z ∗ ∼ p
Revised: y ∗ by ancestral sampling z ∗ ∼ p

Original: z | y =
Revised: z | y =

Original: y ∗ ) and x ∗ ∼ p ( x | z = z ∗ ) .
Revised: y ∗ ) and x ∗ ∼ p ( x | z = z ∗ ) .

Original: However, if y ∗ deviates from the observed data distribution of y , this naive solution involves sampling in an extrapolated regime (or out of distribution regime) where y ∗ is not in the effective support of the learned distribution.
Revised: However, if y ∗ deviates from the observed data distribution of y , this naive solution involves sampling in an extrapolated regime (or out of distribution regime) where y ∗ is not in the effective support of the learned distribution.

Original: We propose a Sampling with Gradual Distribution
Revised: We propose a Sampling with Gradual Distribution

Original: Shifting (SGDS) approach where we 1) sample on the boundary of the effective support of the learned distribution, and 2) gradually shift the learned distribution with these boundary samples to a region where it is supported by high property values (see §3.5 and Algorithm 2).
Revised: Shifting (SGDS) approach where we 1) sample on the boundary of the effective support of the learned distribution, and 2) gradually shift the learned distribution with these boundary samples to a region where it is supported by high property values (see §3.5 and Algorithm 2).

Original: Our model is designed to be versatile such that it admits straightforward extension to multi-objective optimization.
Revised: Our model is designed to be versatile such that it admits straightforward extension to multi-objective optimization.

Original: To optimize x given { y j = y ∗ j } mj =1 , we can simply augment the joint distribution with more regression models, i.e., p ( x, z, y 1 , ..., y m )
Revised: To optimize x given { y j = y ∗ j } mj =1 , we can simply augment the joint distribution with more regression models, i.e., p ( x, z, y 1 , ..., y m )

Original: = p ( z ) p ( x | z ) (cid:81) mj
Revised: = p ( z ) p ( x | z ) (cid:81) mj

Original: =1 p ( y j | z ) .
Revised: =1 p ( y j | z ) .

Original: The sampling procedure follows the same SGDS approach.
Revised: The sampling procedure follows the same SGDS approach.

Original: See §3.6 for more details on multi-objective optimization.
Revised: See §3.6 for more details on multi-objective optimization.

Original: 3. L ATENT S PACE E NERGY -B ASED M ODEL
Revised: 3. L ATENT S PACE E NERGY -B ASED M ODEL

Original: Suppose x = ( x (1) , ..., x ( t ) , ..., x ( T ) ) is a molecule string in SELFIES and z ∈ R d is the latent vector.
Revised: Suppose x = ( x (1) , ..., x ( t ) , ..., x ( T ) ) is a molecule string in SELFIES and z ∈ R d is the latent vector.

Original: Consider the following model,
Revised: Consider the following model,

Original: [Equation]
Revised: [Equation]

Original: where p α ( z ) is a prior model with parameters α , and p β ( x | z ) is a generation model with parameters β .
Revised: where p α ( z ) is a prior model with parameters α , and p β ( x | z ) is a generation model with parameters β .

Original: In VAE (Kingma & Welling, 2014), the prior is simply assumed to be an isotropic Gaussian distribution.
Revised: In VAE (Kingma & Welling, 2014), the prior is simply assumed to be an isotropic Gaussian distribution.

Original: In our model, p α ( z ) is formulated as an energy-based model,
Revised: In our model, p α ( z ) is formulated as an energy-based model,

Original: [Equation]
Revised: [Equation]

Original: where p 

Revised: where p

Original: ( z ) is a reference distribution, assumed to be isotropic Gaussian as in VAE.
Revised: ( z ) is a reference distribution, assumed to be isotropic Gaussian as in VAE.

Original: f α : R d → R is the scalar-valued negative energy function and is parameterized by a small multi-layer perceptron (MLP) with parameters
Revised: f α : R d → R is the scalar-valued negative energy function and is parameterized by a small multi-layer perceptron (MLP) with parameters

Original: α .
Revised: α .

Original: Z ( α ) =
Revised: Z ( α ) =

Original: (cid:82) exp( f α ( z ))
Revised: (cid:82) exp( f α ( z ))

Original: p 0 ( z ) dz = E p 

Revised: p 0 ( z ) dz = E p

Original: [exp( f α ( z ))] is the partition function.
Revised: [exp( f α ( z ))] is the partition function.

Original: The generation model p β ( x | z ) is a conditional autoregressive model,
Revised: The generation model p β ( x | z ) is a conditional autoregressive model,

Original: [Equation]
Revised: [Equation]

Original: which is parameterized by a one-layer LSTM with parameters β .
Revised: which is parameterized by a one-layer LSTM with parameters β .

Original: Note that the latent vector z controls every step of the autoregressive model.
Revised: Note that the latent vector z controls every step of the autoregressive model.

Original: 3. J OINT D ISTRIBUTION OF M OLECULE AND M OLECULAR P ROPERTY
Revised: 3. J OINT D ISTRIBUTION OF M OLECULE AND M OLECULAR P ROPERTY

Original: Given a molecule x , suppose y is the chemical property of interest, such as QED or protein affinity binding.
Revised: Given a molecule x , suppose y is the chemical property of interest, such as QED or protein affinity binding.

Original: The property value can be computed from an input x via open-sourced software RDKit (Landrum et al., 2013) or AutoDock-GPU (Santos-Martins et al., 2021).
Revised: The property value can be computed from an input x via open-sourced software RDKit (Landrum et al., 2013) or AutoDock-GPU (Santos-Martins et al., 2021).

Original: We assume that given z , x and y are conditionally independent.
Revised: We assume that given z , x and y are conditionally independent.

Original: [Equation]
Revised: [Equation]

Original: where p α ( z ) is the EBM prior, p β ( x | z ) is the generation model, and p γ ( y | z ) is the property regression model, and
Revised: where p α ( z ) is the EBM prior, p β ( x | z ) is the generation model, and p γ ( y | z ) is the property regression model, and

Original: θ = ( α, β, γ ) .
Revised: θ = ( α, β, γ ) .

Original: We use the model p θ ( x, y, z ) to approximate the data distribution of ( x, y ) .
Revised: We use the model p θ ( x, y, z ) to approximate the data distribution of ( x, y ) .

Original: See Appendix §A.1 for details.
Revised: See Appendix §A.1 for details.

Original: The property regression model can be written as
Revised: The property regression model can be written as

Original: [Equation]
Revised: [Equation]

Original: (7) where s γ ( z ) is a small MLP, with parameters γ , predicting y based on the latent z .
Revised: (7) where s γ ( z ) is a small MLP, with parameters γ , predicting y based on the latent z .

Original: The variance σis set as a constant or hyperparameter in our work.
Revised: The variance σis set as a constant or hyperparameter in our work.

Original: 3. L EARNING A LGORITHM
Revised: 3. L EARNING A LGORITHM

Original: Suppose we observe training examples { ( x i , y i ) , i = 1 , ..., n } .
Revised: Suppose we observe training examples { ( x i , y i ) , i = 1 , ..., n } .

Original: The log-likelihood function is
Revised: The log-likelihood function is

Original: L ( θ ) = (cid:80) ni
Revised: L ( θ ) = (cid:80) ni

Original: =1 log p θ ( x i , y i ) .
Revised: =1 log p θ ( x i , y i ) .

Original: The learning gradient can be calculated according to
Revised: The learning gradient can be calculated according to

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: For the prior model, ∇ α log p
Revised: For the prior model, ∇ α log p

Original: α
Revised: α

Original: ( z ) =
Revised: ( z ) =

Original: [Equation]
Revised: [Equation]

Original: z
Revised: z

Original: [Equation]
Revised: [Equation]

Original: z )] .
Revised: z )] .

Original: Thus the learning gradient given an example ( x, y ) is
Revised: Thus the learning gradient given an example ( x, y ) is

Original: [Equation]
Revised: [Equation]

Original: α is updated based on the difference between z inferred from empirical observation ( x, y ) , and z sampled from the current prior.
Revised: α is updated based on the difference between z inferred from empirical observation ( x, y ) , and z sampled from the current prior.

Original: For the generation model,
Revised: For the generation model,

Original: [Equation]
Revised: [Equation]

Original: Similarly, for the regression model,
Revised: Similarly, for the regression model,

Original: [Equation]
Revised: [Equation]

Original: Estimating expectations in equations 17, 18, and 19 requires MCMC sampling of the prior model
Revised: Estimating expectations in equations 17, 18, and 19 requires MCMC sampling of the prior model

Original: p α ( z ) and the posterior distribution p θ ( z | x, y ) .
Revised: p α ( z ) and the posterior distribution p θ ( z | x, y ) .

Original: We recruit Langevin dynamics (Neal, 2011).
Revised: We recruit Langevin dynamics (Neal, 2011).

Original: For a target distribution π ( z ) , the dynamics iterates
Revised: For a target distribution π ( z ) , the dynamics iterates

Original: [Equation]
Revised: [Equation]

Original: where k indexes the time step of the Langevin dynamics, s is a small step size, and ϵ k ∼ N (0 , I d ) is the Gaussian white noise.
Revised: where k indexes the time step of the Langevin dynamics, s is a small step size, and ϵ k ∼ N (0 , I d ) is the Gaussian white noise.

Original: π ( z ) can be either p α ( z ) or p θ ( z | x, y ) .
Revised: π ( z ) can be either p α ( z ) or p θ ( z | x, y ) .

Original: In either case, ∇ z log π ( z ) can be efficiently computed by back-propagation.
Revised: In either case, ∇ z log π ( z ) can be efficiently computed by back-propagation.

Original: See Appendix §A.1 for more details.
Revised: See Appendix §A.1 for more details.

Original: [Table]
Revised: [Table]

Original: { s
Revised: { s

Original: , s 1 } output: θ T = α T β γ for t = 0 : T − 1 do
Revised: , s 1 } output: θ T = α T β γ for t = 0 : T − 1 do

Original: γ t
Revised: γ t

Original: γ
Revised: γ

Original: 3. S AMPLING WITH G RADUAL D ISTRIBUTION S HIFTING
Revised: 3. S AMPLING WITH G RADUAL D ISTRIBUTION S HIFTING

Original: To tackle the single-objective optimization problem (eq. 1), one naive approach is to perform ancestral sampling with two steps, given some desirable property value y ∗ ,
Revised: To tackle the single-objective optimization problem (eq. 1), one naive approach is to perform ancestral sampling with two steps, given some desirable property value y ∗ ,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where (1) is an application of Bayes’ rule, with p α ( z ) as the prior and p γ ( y | z ) as the likelihood.
Revised: where (1) is an application of Bayes’ rule, with p α ( z ) as the prior and p γ ( y | z ) as the likelihood.

Original: Our model p θ ( x, y, z ) is learned to capture the data distribution.
Revised: Our model p θ ( x, y, z ) is learned to capture the data distribution.

Original: In real-world settings, y ∗ might not be within the support of the data distribution.
Revised: In real-world settings, y ∗ might not be within the support of the data distribution.

Original: Therefore, sampling following equation 14 does not work well since it involves extrapolating the learned distribution.
Revised: Therefore, sampling following equation 14 does not work well since it involves extrapolating the learned distribution.

Original: We propose a method called sampling with gradual distribution shifting (SGDS) to address this issue.
Revised: We propose a method called sampling with gradual distribution shifting (SGDS) to address this issue.

Original: In particular, we find examples from the training data on the boundary of the distribution support by sorting them according to the values of y and taking top- k samples, { ( x (old) i , y (old) i ) } k
Revised: In particular, we find examples from the training data on the boundary of the distribution support by sorting them according to the values of y and taking top- k samples, { ( x (old) i , y (old) i ) } k

Original: i =1 , with high property values.
Revised: i =1 , with high property values.

Original: We shift the support slightly by adding some small ∆ y to all y ’s, and sample x conditional on shifted y ’s, following equation 14.
Revised: We shift the support slightly by adding some small ∆ y to all y ’s, and sample x conditional on shifted y ’s, following equation 14.

Original: Given the generated x ’s, we compute their groundtruth y ’s with
Revised: Given the generated x ’s, we compute their groundtruth y ’s with

Original: RDKit or AutoDock-GPU.
Revised: RDKit or AutoDock-GPU.

Original: We then shift the learned model by finetuning it with the new data, { ( x (new) i , y (new) i ) } ki =1 , for a few steps (e.g., 10 ).
Revised: We then shift the learned model by finetuning it with the new data, { ( x (new) i , y (new) i ) } ki =1 , for a few steps (e.g., 10 ).

Original: This completes one iteration of distribution shift.
Revised: This completes one iteration of distribution shift.

Original: The molecule with desired high property value is sampled after T shifting iterations.
Revised: The molecule with desired high property value is sampled after T shifting iterations.

Original: We summarize the algorithm in Algorithm 2.
Revised: We summarize the algorithm in Algorithm 2.

Original: For constrained optimization in step we only keep the sampled molecules that satisfy the given constraints.
Revised: For constrained optimization in step we only keep the sampled molecules that satisfy the given constraints.

Original: 3. M ULTI -O BJECTIVE O PTIMIZATION
Revised: 3. M ULTI -O BJECTIVE O PTIMIZATION

Original: We next consider the multi-objective optimization problem.
Revised: We next consider the multi-objective optimization problem.

Original: Suppose we optimize for a set of properties { y j } mj =1 , then we learn a property regression model for each property y j ,
Revised: Suppose we optimize for a set of properties { y j } mj =1 , then we learn a property regression model for each property y j ,

Original: [Equation]
Revised: [Equation]

Original: where each s γ j is a small MLP with parameters γ j .
Revised: where each s γ j is a small MLP with parameters γ j .

Original: Then the joint distribution is,
Revised: Then the joint distribution is,

Original: [Equation]
Revised: [Equation]

Original: m
Revised: m

Original: [Equation]
Revised: [Equation]

Original: Under our framework, the learning algorithm and the sampling algorithm for the single-objective problem can be straightforwardly extended to the multi-objective setting.
Revised: Under our framework, the learning algorithm and the sampling algorithm for the single-objective problem can be straightforwardly extended to the multi-objective setting.

Original: In both settings, the same types of properties are provided in both the initial training stage and the distribution shifting stage.
Revised: In both settings, the same types of properties are provided in both the initial training stage and the distribution shifting stage.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: To demonstrate the effectiveness of our proposed model, we compare our model with previous
Revised: To demonstrate the effectiveness of our proposed model, we compare our model with previous

Original: SOTA methods for unconditional generation (§4.2) and molecule design including single-objective optimization (§4.3) and multi-objective optimization (§4.4).
Revised: SOTA methods for unconditional generation (§4.2) and molecule design including single-objective optimization (§4.3) and multi-objective optimization (§4.4).

Original: In molecule design experiments, we consider both non-biological and biological properties.
Revised: In molecule design experiments, we consider both non-biological and biological properties.

Original: 4. E XPERIMENTAL S ETUP
Revised: 4. E XPERIMENTAL S ETUP

Original: Datasets.
Revised: Datasets.

Original: For the unconditional generation task, we report results on ZINC (Irwin et al., 2012) and
Revised: For the unconditional generation task, we report results on ZINC (Irwin et al., 2012) and

Original: MOSES datasets (Polykovskiy et al., 2020).
Revised: MOSES datasets (Polykovskiy et al., 2020).

Original: ZINC consists of around 250 k molecules, and MOSES comprises around 2 million molecules.
Revised: ZINC consists of around 250 k molecules, and MOSES comprises around 2 million molecules.

Original: For optimization tasks, we conduct experiments with ZINC.
Revised: For optimization tasks, we conduct experiments with ZINC.

Original: Training Details.
Revised: Training Details.

Original: There are three modules in our model, the top-down generation model p β ( x | z ) , the EBM prior p α ( z ) , and the regression model p γ ( y | z ) .
Revised: There are three modules in our model, the top-down generation model p β ( x | z ) , the EBM prior p α ( z ) , and the regression model p γ ( y | z ) .

Original: p β ( x | z ) is parameterized by a single-layer
Revised: p β ( x | z ) is parameterized by a single-layer

Original: LSTM with 1024 hidden units.
Revised: LSTM with 1024 hidden units.

Original: The dimension of latent vector z is 100 .
Revised: The dimension of latent vector z is 100 .

Original: p α ( z ) is a 3 -layer MLP with 100 input dimension and 200 hidden units.
Revised: p α ( z ) is a 3 -layer MLP with 100 input dimension and 200 hidden units.

Original: The property regression model p γ ( y | z ) is a 3layer MLP with 100 input dimension and 100 hidden units.
Revised: The property regression model p γ ( y | z ) is a 3layer MLP with 100 input dimension and 100 hidden units.

Original: In multi-objective optimization settings, we recruit one MLP for each property and these MLPs have the same architecture as mentioned above.
Revised: In multi-objective optimization settings, we recruit one MLP for each property and these MLPs have the same architecture as mentioned above.

Original: It is worth mentioning that compared to most previous models, our model is characterized by its simplicity without adding inference networks, RL-related modules, and graph neural networks.
Revised: It is worth mentioning that compared to most previous models, our model is characterized by its simplicity without adding inference networks, RL-related modules, and graph neural networks.

Original: Adam (Kingma & Ba, 2015) optimizer is used to train our models with learning rates 0 .
Revised: Adam (Kingma & Ba, 2015) optimizer is used to train our models with learning rates 0 .

Original: 0001 for
Revised: 0001 for

Original: EBM and 0 .
Revised: EBM and 0 .

Original: 001 for the rest.
Revised: 001 for the rest.

Original: We train our models for 30 epochs.
Revised: We train our models for 30 epochs.

Original: 10 , 000 (non-biological) and 2 ,
Revised: 10 , 000 (non-biological) and 2 ,

Original: (biological) boundary examples are selected for distribution shifting.
Revised: (biological) boundary examples are selected for distribution shifting.

Original: The numbers of shift iterations are 50 (non-biological single-objective), 20 (non-biological single-objective) and 10 (multiobjective).
Revised: The numbers of shift iterations are 50 (non-biological single-objective), 20 (non-biological single-objective) and 10 (multiobjective).

Original: Within each SGDS iteration, the model is finetuned with only 10 parameter updates.
Revised: Within each SGDS iteration, the model is finetuned with only 10 parameter updates.

Original: All experiments are conducted on Nvidia Titan XP GPU.
Revised: All experiments are conducted on Nvidia Titan XP GPU.

Original: See more details in Appendix §A.6.
Revised: See more details in Appendix §A.6.

Original: 4. U NCONDITIONAL G ENERATION
Revised: 4. U NCONDITIONAL G ENERATION

Original: Three types of encoding systems are used to encode molecules in prior work: SMILES (Weininger, 1988), SELFIES (Krenn et al., 2020), and graph.
Revised: Three types of encoding systems are used to encode molecules in prior work: SMILES (Weininger, 1988), SELFIES (Krenn et al., 2020), and graph.

Original: SMILES and SELFIES linearize a molecule graph into a string of characters.
Revised: SMILES and SELFIES linearize a molecule graph into a string of characters.

Original: Most previous models using SMILES struggle to generate molecules with high validity, which is the percentage of molecules that satisfy the chemical valency rule.
Revised: Most previous models using SMILES struggle to generate molecules with high validity, which is the percentage of molecules that satisfy the chemical valency rule.

Original: Thus graph representations become popular since explicit valency constraints can be imposed.
Revised: Thus graph representations become popular since explicit valency constraints can be imposed.

Original: However, perfect validity in this approach does not imply that the model captures the chemical rules since it is achieved with external constraints.
Revised: However, perfect validity in this approach does not imply that the model captures the chemical rules since it is achieved with external constraints.

Original: Recently, SELFIES is developed where every SELFIES string corresponds to a valid molecule due to the nature of the encoding system.
Revised: Recently, SELFIES is developed where every SELFIES string corresponds to a valid molecule due to the nature of the encoding system.

Original: Thus, validity for generated SMILES strings is a good indicator on how well a learned model captures the basic chemical rules implicitly.
Revised: Thus, validity for generated SMILES strings is a good indicator on how well a learned model captures the basic chemical rules implicitly.

Original: Besides validity, we also compare models on uniqueness (the percentage of unique molecules in all generated samples) and novelty (the percentage of generated molecules that are not in the training set).
Revised: Besides validity, we also compare models on uniqueness (the percentage of unique molecules in all generated samples) and novelty (the percentage of generated molecules that are not in the training set).

Original: Following previous work, we randomly sample 10 k molecules for ZINC and 30 k for MOSES, and compare on the three aforementioned metrics.
Revised: Following previous work, we randomly sample 10 k molecules for ZINC and 30 k for MOSES, and compare on the three aforementioned metrics.

Original: Generations results on ZINC and MOSES are shown in Table 1 and Table 2 respectively.
Revised: Generations results on ZINC and MOSES are shown in Table 1 and Table 2 respectively.

Original: str-smi denotes string-based SMILES representations and str-sfi denotes string-based SELFIES representations.
Revised: str-smi denotes string-based SMILES representations and str-sfi denotes string-based SELFIES representations.

Original: In Table 1, we show generation results for both SMILES and SELFIES.
Revised: In Table 1, we show generation results for both SMILES and SELFIES.

Original: SMILES does not have a validity constraint during generation.
Revised: SMILES does not have a validity constraint during generation.

Original: However, our model, MolEBM, still achieves 95 .
Revised: However, our model, MolEBM, still achieves 95 .

Original: 5% validity, which outperforms other SMILESbased methods and is also comparable to those with valency check.
Revised: 5% validity, which outperforms other SMILESbased methods and is also comparable to those with valency check.

Original: This result demonstrates that our model can capture those valency rules effectively and implicitly (also see Appendix §A.4.).
Revised: This result demonstrates that our model can capture those valency rules effectively and implicitly (also see Appendix §A.4.).

Original: Samples from our model also achieve perfect uniqueness and novelty.
Revised: Samples from our model also achieve perfect uniqueness and novelty.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: 4. S INGLE -O BJECTIVE O PTIMIZATION
Revised: 4. S INGLE -O BJECTIVE O PTIMIZATION

Original: Non-Biological Property Optimization.
Revised: Non-Biological Property Optimization.

Original: For non-biological properties, we are interested in Penalized logP and QED, both of which can be calculated by RDKit (Landrum et al., 2013).
Revised: For non-biological properties, we are interested in Penalized logP and QED, both of which can be calculated by RDKit (Landrum et al., 2013).

Original: Since we know the Penalized logP scores have a positive relationship with the length of the molecules, we maximize Penalized logP either with or without maximum length limit.
Revised: Since we know the Penalized logP scores have a positive relationship with the length of the molecules, we maximize Penalized logP either with or without maximum length limit.

Original: Following (Eckmann et al., 2022)
Revised: Following (Eckmann et al., 2022)

Original: [Table]
Revised: [Table]

Original: the maximum length is set to be the maximum length of molecules in ZINC using SELFIES.
Revised: the maximum length is set to be the maximum length of molecules in ZINC using SELFIES.

Original: From Table
Revised: From Table

Original: [Table]
Revised: [Table]

Original: we can see that with length limit
Revised: we can see that with length limit

Original: [Table]
Revised: [Table]

Original: MolEBM outperforms previous methods by a large margin.
Revised: MolEBM outperforms previous methods by a large margin.

Original: Maximizing penalized logP without length limit leads to the highest Penalized logP, 158 . 0 .
Revised: Maximizing penalized logP without length limit leads to the highest Penalized logP, 158 . 0 .

Original: We also achieve the highest QED with/without length limit.
Revised: We also achieve the highest QED with/without length limit.

Original: These observations demonstrate the effectiveness of our method.
Revised: These observations demonstrate the effectiveness of our method.

Original: We also illustrate our distribution shifting method in Appendix §A.
Revised: We also illustrate our distribution shifting method in Appendix §A.

Original: [Table]
Revised: [Table]

Original: Biological Property Optimization.
Revised: Biological Property Optimization.

Original: ESR1 and ACAA1 are two human proteins.
Revised: ESR1 and ACAA1 are two human proteins.

Original: We aim to design ligands (molecules) that have the maximum binding affinities towards those target proteins.
Revised: We aim to design ligands (molecules) that have the maximum binding affinities towards those target proteins.

Original: ESR is well-studied, which has many existing binders, while ACAA1 does not.
Revised: ESR is well-studied, which has many existing binders, while ACAA1 does not.

Original: However, we did not use any binder-related information in the design process.
Revised: However, we did not use any binder-related information in the design process.

Original: Binding affinity is measured by the estimated dissociation constants K D , which can be computed with AutoDock-GPU (Santos-Martins et al., 2021) given a molecule.
Revised: Binding affinity is measured by the estimated dissociation constants K D , which can be computed with AutoDock-GPU (Santos-Martins et al., 2021) given a molecule.

Original: Large binding affinities corresponds to small K D .
Revised: Large binding affinities corresponds to small K D .

Original: That is, we aim to minimize K D .
Revised: That is, we aim to minimize K D .

Original: Table 4 shows that our model outperforms previous methods on both ESR1 and
Revised: Table 4 shows that our model outperforms previous methods on both ESR1 and

Original: ACAA1 binding affinity maximization tasks.
Revised: ACAA1 binding affinity maximization tasks.

Original: Producing those ligands with high binding affinity plays a vital role in the early stage of drug discovery
Revised: Producing those ligands with high binding affinity plays a vital role in the early stage of drug discovery

Original: [Table]
Revised: [Table]

Original: 4. M ULTI -O BJECTIVE O PTIMIZATION
Revised: 4. M ULTI -O BJECTIVE O PTIMIZATION

Original: We next consider optimizing binding affinity, QED and SAS simultaneously for multi-objective optimization.
Revised: We next consider optimizing binding affinity, QED and SAS simultaneously for multi-objective optimization.

Original: Following Eckmann et al.
Revised: Following Eckmann et al.

Original: , we exclude molecules with abnormal behaviorsin the generation process to make sure the learned distribution shifts towards a desirable region in terms of pharmacologic and synthetic properties.
Revised: , we exclude molecules with abnormal behaviorsin the generation process to make sure the learned distribution shifts towards a desirable region in terms of pharmacologic and synthetic properties.

Original: Table 5 shows our multi-objective binding affinity maximization results comparing to LIMO (Eckmann et al., 2022) and GCPN (You et al., 2018).
Revised: Table 5 shows our multi-objective binding affinity maximization results comparing to LIMO (Eckmann et al., 2022) and GCPN (You et al., 2018).

Original: From the results, we can see that MolEBM is able to find the ligands with desired properties while keeping the pharmacologic structures.
Revised: From the results, we can see that MolEBM is able to find the ligands with desired properties while keeping the pharmacologic structures.

Original: For ESR1, we have two existing binders on the market, Tamoxifen and Raloxifene.
Revised: For ESR1, we have two existing binders on the market, Tamoxifen and Raloxifene.

Original: Our designed ligands have similar QED and SA, with rather low K D .
Revised: Our designed ligands have similar QED and SA, with rather low K D .

Original: Compared to existing methods, MolEBM obtains better results in overall adjustments.
Revised: Compared to existing methods, MolEBM obtains better results in overall adjustments.

Original: For ACAA1, we do not have any existing binders.
Revised: For ACAA1, we do not have any existing binders.

Original: Compared with prior SOTA methods, our optimized ligands outperform those by a large margin.
Revised: Compared with prior SOTA methods, our optimized ligands outperform those by a large margin.

Original: When comparing ACAA1 multi-objective setting with its corresponding single-objective one, we find multiobjective results even outperforms the single-objective one, which may be counter-intuitive since multi-objective optimization is assumed to be a harder task than single-objective optimization.
Revised: When comparing ACAA1 multi-objective setting with its corresponding single-objective one, we find multiobjective results even outperforms the single-objective one, which may be counter-intuitive since multi-objective optimization is assumed to be a harder task than single-objective optimization.

Original: However, our current results indicate that with proper prior knowledge (e.g. multiple objectives are positively aligned), multi-objective settings could be more plausible in de novo design, and those complicated prior knowledge can indeed be captured by our expressive latent space EBM.
Revised: However, our current results indicate that with proper prior knowledge (e.g. multiple objectives are positively aligned), multi-objective settings could be more plausible in de novo design, and those complicated prior knowledge can indeed be captured by our expressive latent space EBM.

Original: While we still need domain expertise to determine the effectiveness of those ligands discovered by MolEBM, we believe our model could shed light on the optimization procedure in the early drug discovery.
Revised: While we still need domain expertise to determine the effectiveness of those ligands discovered by MolEBM, we believe our model could shed light on the optimization procedure in the early drug discovery.

Original: [Table]
Revised: [Table]

Original: Figure 2 shows generated molecules with high binding affinities.
Revised: Figure 2 shows generated molecules with high binding affinities.

Original: See Appendix for more examples.
Revised: See Appendix for more examples.

Original:  C ONCLUSION AND DISCUSSION
Revised:  C ONCLUSION AND DISCUSSION

Original: We propose a deep generative model, MolEBM, which models the joint distribution of molecules and molecular properties.
Revised: We propose a deep generative model, MolEBM, which models the joint distribution of molecules and molecular properties.

Original: It assumes an energy-based prior for a low-dimensional continuous latent space, which effectively captures data regularities of the discrete molecule data.
Revised: It assumes an energy-based prior for a low-dimensional continuous latent space, which effectively captures data regularities of the discrete molecule data.

Original: We then design a distribution shifting method (SGDS) to shift the learned distribution to a region with high property values.
Revised: We then design a distribution shifting method (SGDS) to shift the learned distribution to a region with high property values.

Original: Molecule design can then be achieved by conditional sampling.
Revised: Molecule design can then be achieved by conditional sampling.

Original: Our experiments demonstrate that our method outperforms previous SOTA methods by a significant margin.
Revised: Our experiments demonstrate that our method outperforms previous SOTA methods by a significant margin.

Original: A limiting factor is that the sampling with gradual distribution shifting (SGDS) requires that the properties of interest can be easily computed.
Revised: A limiting factor is that the sampling with gradual distribution shifting (SGDS) requires that the properties of interest can be easily computed.

Original: In our future work, we shall explore semi-supervised learning methods using our model where the properties are available only for a small sample of molecules.
Revised: In our future work, we shall explore semi-supervised learning methods using our model where the properties are available only for a small sample of molecules.

Original: Another limitation is that, for our optimization method, good property scores may not translate to useful molecules for drug design.
Revised: Another limitation is that, for our optimization method, good property scores may not translate to useful molecules for drug design.

Original: This problem may be partially addressed by multi-objective optimization as studied in this paper.
Revised: This problem may be partially addressed by multi-objective optimization as studied in this paper.

Original: In our work, we focus on optimizing given objectives.
Revised: In our work, we focus on optimizing given objectives.

Original: Designing good objectives or metrics is an equally or perhaps even more important problem that deserves careful investigation.
Revised: Designing good objectives or metrics is an equally or perhaps even more important problem that deserves careful investigation.

Original: A A PPENDIX A. D ETAILS ABOUT MODEL AND LEARNING
Revised: A A PPENDIX A. D ETAILS ABOUT MODEL AND LEARNING

Original: Our model is of the form p α ( z ) p
Revised: Our model is of the form p α ( z ) p

Original: [Equation]
Revised: [Equation]

Original: ( x | z ) p γ ( y | z ) .
Revised: ( x | z ) p γ ( y | z ) .

Original: The marginal distribution of ( x, y ) is
Revised: The marginal distribution of ( x, y ) is

Original: [Equation]
Revised: [Equation]

Original: We use p θ ( x, y ) to approximate the data distribution of ( x, y ) .
Revised: We use p θ ( x, y ) to approximate the data distribution of ( x, y ) .

Original: For the data distribution of ( x, y ) , y is a deterministic function of x .
Revised: For the data distribution of ( x, y ) , y is a deterministic function of x .

Original: However, a machine learning method usually cannot learn the deterministic function exactly.
Revised: However, a machine learning method usually cannot learn the deterministic function exactly.

Original: Instead, we can only learn a probabilistic p θ ( y | z ) .
Revised: Instead, we can only learn a probabilistic p θ ( y | z ) .

Original: Our model p θ ( x, y ) seeks to approximate the data distribution p ( x, y ) by maximum likelihood.
Revised: Our model p θ ( x, y ) seeks to approximate the data distribution p ( x, y ) by maximum likelihood.

Original: A learnable and flexible prior model p α ( z ) helps to make the approximate more accurate than a fixed prior model such as that in VAE.
Revised: A learnable and flexible prior model p α ( z ) helps to make the approximate more accurate than a fixed prior model such as that in VAE.

Original: Let the training data be { ( x i , y i ) , i = 1 , ..., n } .
Revised: Let the training data be { ( x i , y i ) , i = 1 , ..., n } .

Original: The log-likelihood function is L ( θ )
Revised: The log-likelihood function is L ( θ )

Original: =
Revised: =

Original: (cid:80) ni
Revised: (cid:80) ni

Original: =1 log p θ ( x i , y i ) .
Revised: =1 log p θ ( x i , y i ) .

Original: The learning gradient is L ′ ( θ )
Revised: The learning gradient is L ′ ( θ )

Original: = (cid:80) ni
Revised: = (cid:80) ni

Original: =1 ∇ θ log p θ ( x i , y i ) .
Revised: =1 ∇ θ log p θ ( x i , y i ) .

Original: In the following, we provide details for calculating ∇ θ log p θ ( x, y ) for a single generic training example ( x, y )
Revised: In the following, we provide details for calculating ∇ θ log p θ ( x, y ) for a single generic training example ( x, y )

Original: (where we drop the subscript i for notation simplicity).
Revised: (where we drop the subscript i for notation simplicity).

Original: [Equation]
Revised: [Equation]

Original: For the prior model,
Revised: For the prior model,

Original: [Equation]
Revised: [Equation]

Original: Thus the learning gradient for α given an example ( x, y ) is
Revised: Thus the learning gradient for α given an example ( x, y ) is

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The above equation has an empirical Bayes nature.
Revised: The above equation has an empirical Bayes nature.

Original: p θ
Revised: p θ

Original: ( z | x, y ) is based on the empirical observation ( x, y )
Revised: ( z | x, y ) is based on the empirical observation ( x, y )

Original: [Equation]
Revised: [Equation]

Original: while p α is the prior model.
Revised: while p α is the prior model.

Original: For the generation model,
Revised: For the generation model,

Original: [Equation]
Revised: [Equation]

Original: Similarly, for the regression model,
Revised: Similarly, for the regression model,

Original: [Equation]
Revised: [Equation]

Original: Estimating expectations in the above equations requires Monte Carlo sampling of the prior model p α ( z ) and the posterior distribution p θ ( z | x, y ) .
Revised: Estimating expectations in the above equations requires Monte Carlo sampling of the prior model p α ( z ) and the posterior distribution p θ ( z | x, y ) .

Original: If we can draw fair samples from the two distributions, and use these Monte Carlo samples to approximate the expectations, the the gradient ascent algorithm based on the Monte Carlo samples is the stochastic gradient ascent algorithm or the stochastic approximation algorithm of Robbins and Monro (Robbins & Monro, 1951), who established the convergence of such an algorithm to a local maximum of the log-likelihood.
Revised: If we can draw fair samples from the two distributions, and use these Monte Carlo samples to approximate the expectations, the the gradient ascent algorithm based on the Monte Carlo samples is the stochastic gradient ascent algorithm or the stochastic approximation algorithm of Robbins and Monro (Robbins & Monro, 1951), who established the convergence of such an algorithm to a local maximum of the log-likelihood.

Original: For MCMC sampling using Langevin dynamics, the finite step or short run Langevin dynamics may cause bias in Monte Carlo sampling.
Revised: For MCMC sampling using Langevin dynamics, the finite step or short run Langevin dynamics may cause bias in Monte Carlo sampling.

Original: The bias was analyzed in Pang et al.
Revised: The bias was analyzed in Pang et al.

Original: The resulting algorithm is an approximate maximum likelihood learning algorithm.
Revised: The resulting algorithm is an approximate maximum likelihood learning algorithm.

Original: A. A BLATION S TUDIES
Revised: A. A BLATION S TUDIES

Original: We conduct ablations on the key components of our method: 1) EBM Prior (EBM vs. standard
Revised: We conduct ablations on the key components of our method: 1) EBM Prior (EBM vs. standard

Original: Gaussian N ( 0 , I ) ), (2) SGDS (shift distributions with SGDS vs. no shift), and (3) joint training (joint training of molecule and molecular property distribution vs. train molecule with LSEBM and learn property regression model in a second step).
Revised: Gaussian N ( 0 , I ) ), (2) SGDS (shift distributions with SGDS vs. no shift), and (3) joint training (joint training of molecule and molecular property distribution vs. train molecule with LSEBM and learn property regression model in a second step).

Original: The ablation results are displayed in Table 6.
Revised: The ablation results are displayed in Table 6.

Original: It is clear that all the proposed components contribute significantly to the good performance of our method.
Revised: It is clear that all the proposed components contribute significantly to the good performance of our method.

Original: [Table]
Revised: [Table]

Original: A. I LLUSTRATION OF S AMPLING WITH G RADUAL D ISTRIBUTION S HIFTING (SGDS)
Revised: A. I LLUSTRATION OF S AMPLING WITH G RADUAL D ISTRIBUTION S HIFTING (SGDS)

Original: Figure 3 shows property score densities of sampled molecules from our MolEBM in the distribution shifting process.
Revised: Figure 3 shows property score densities of sampled molecules from our MolEBM in the distribution shifting process.

Original: We can see the model distribution is gradually shifting towards the region supported by molecules with high property values.
Revised: We can see the model distribution is gradually shifting towards the region supported by molecules with high property values.

Original: A. U NCONDITIONAL G ENERATION
Revised: A. U NCONDITIONAL G ENERATION

Original: For unconditional generation task, we use uniqueness, novelty and validity to compare our generation results with existing methods.
Revised: For unconditional generation task, we use uniqueness, novelty and validity to compare our generation results with existing methods.

Original: Meanwhile, for SELFIES-based ZINC dataset, we split this dataset into train split ( 240 k ) and test split ( 10 k samples).
Revised: Meanwhile, for SELFIES-based ZINC dataset, we split this dataset into train split ( 240 k ) and test split ( 10 k samples).

Original: Here, we randomly sample 10 k molecules from the learned latent EBM and calculate their logP, QED and SA scores using RDKit.
Revised: Here, we randomly sample 10 k molecules from the learned latent EBM and calculate their logP, QED and SA scores using RDKit.

Original: We compare these property densities with molecule property densities in test split.
Revised: We compare these property densities with molecule property densities in test split.

Original: The results are shown in Figure 4.
Revised: The results are shown in Figure 4.

Original: We can see that the marginal distributions (property score densities) of the learned model match those of the data quite well, implying that our model indeed captures the chemical rules implicitly.
Revised: We can see that the marginal distributions (property score densities) of the learned model match those of the data quite well, implying that our model indeed captures the chemical rules implicitly.

Original: A. G ENERATION WITH D IFFERENT L ENGTH OF M ARKOV C HAIN
Revised: A. G ENERATION WITH D IFFERENT L ENGTH OF M ARKOV C HAIN

Original: In our experiments, we use short-run MCMC (i.e. s = 20 ) in Equation 13 for all experiments.
Revised: In our experiments, we use short-run MCMC (i.e. s = 20 ) in Equation 13 for all experiments.

Original: From
Revised: From

Original: Figure 5, we can see with the increasing length of Markov chains, the molecules change accordingly, showing that Markov Chain doesn’t get stuck in the local mode.
Revised: Figure 5, we can see with the increasing length of Markov chains, the molecules change accordingly, showing that Markov chain doesn’t get stuck in the local mode.

Original: A. T RAINING T IME
Revised: [Figure]

Original: Maximum likelihood learning of MolEBM takes 4 hours, to train 30 iterations on a single Nvidia Titan XP GPU with batch size 2048 .
Revised: The joint training of MolEBM takes 4 hours with 30 iterations on a single Nvidia Titan XP GPU with batch size 2048 .

Original: For non-biological single-objective property optimization, it takes around 0 . 5 hours to do 50 distribution shifting (SGDS) iterations.
Revised: For non-biological single-objective property optimization, it takes around 0 .hours to do 50 distribution shifting (SGDS) iterations.

Original: For biological binding affinity maximization, the optimization time is mainly dependent on the number of queries of AutoDockGPU.
Revised: For biological binding affinity maximization, the optimization time is mainly dependent on the number of queries of AutoDock-GPU.

Original: We do 20 and 10 SGDS iterations for the single-objective and multi-objective tasks, respectively, which cost 10 hours and 5 hours.
Revised: We doand 10 SGDS iterations for the single-objective and multi-objective tasks, respectively, which cost 10 hours and 5 hours.

Original: For biological property optimization tasks, we use two Nvidia Titan XP GPUs, one for running our code and another one for running AutoDock-GPU.
Revised: For biological property optimization tasks, we use two Nvidia Titan XP GPUs, one for running our code and another one for running AutoDock-GPU.

Original: We have added a table to compare with previous methods.
Revised: We have added a table to compare with previous methods.

Original: [Table]
Revised: [Table]

Original: Even if we intensively use MCMC based sampling methods, our training speed is affordable comparing to existing methods.
Revised: Even if we intensively use MCMC sampling-based methods, our training speed is affordable comparing to existing methods.

Original: That’s due to our designed latent space EBM is low-dimensional  and we use short-run MCMC (i.e. with fixed iteration steps s = 20 ) in our experiments.
Revised: That’s due to our designed latent space EBM is low-dimensional (i.e. dim ( z ) =100) and we use short-run MCMC (i.e. with fixed iteration steps s = 20 ) in our experiments.

Original: The sampling results with respects to the length of Markov chain is discussed in previous section.
Revised: The sampling results with respect to the length of Markov chain is discussed in previous section.

Original: A. O PTIMIZATION WITH C ONSTRAINTS
Revised: A. O PTIMIZATION WITH C ONSTRAINTS

Original: Inthe unconditional generation task, even with SMILES representation, our model can capture the valency constraints in chemical space.
Revised: In unconditional generation task, even with SMILES representation, our model can capture the valency constraints in chemical space.

Original: This idea can further be extended to constrained optimizations, and those constraints in the chemical space can be directly imposed byonly keeping the sampled molecules that satisfy the constraints during the SGDS molecule sampling step in Algorithm 2.
Revised: This idea can further be extended to constrained optimizations, and those constraints in the chemical space can be directly imposed by keeping the sampled molecules that satisfy the constraints during the SGDS molecule sampling step in Algorithm 2.

Original: We include logP targeting and similarity-constrained optimization experiments in the following.
Revised: We include logP targeting and similarity-constrained optimization experiments in the following.

Original: For example, in logP targeting experiments, during each iteration, we only keep the molecules within the target logP range; in similarity-constrained experiments, weonly use those molecules which satisfy similarity constraints.
Revised: For example, in logP targeting experiments, during each iteration, we only keep the molecules within the target logP range; in similarity-constrained experiments, we use those molecules which satisfy similarity constraints to update the model.

Original: By iteratively updating the model with those selected molecules, we shift the joint distribution towards the target that satisfies the constraints.
Revised: By iteratively updating the model with those selected molecules, we shift the joint distribution towards the region that satisfies the constraints.

Original: A.7. LOG P T ARGETING
Revised: A.7. LOG P T ARGETING

Original: Comparing to previous methods, MolEBM is able to get second best diversity scores with significantly better success rate in both ranges.
Revised: Comparing to previous methods, MolEBM is able to get competitive diversity scores with significantly better success rate in both ranges.

Original: That’s because after SGDS [Table] our model is shifted towards the region that is supported by molecules satisfying the logP constraints.
Revised: That’s because after SGDS, our model is shifted towards the region that is supported by molecules satisfying the logP constraints.

Original: Due to the flexibility of our EBM prior [Table] MolEBM achieves rather high diversity scores while keeping most of the sampled molecules within the logP range.
Revised: Due to the flexibility of our EBM prior, MolEBM achieves rather high diversity scores while keeping most of the sampled molecules within the logP range

Original: [Table]
Revised: [Table]

Original: A.7. S IMILARITY - CONSTRAINED O PTIMIZATION
Revised: A.7. S IMILARITY - CONSTRAINED O PTIMIZATION

Original: Following previous procedures for JT-VAE (Jin et al., 2018)
Revised: Following previous procedures in JT-VAE (Jin et al., 2018)

Original: [Table]
Revised: [Table]

Original: we select 800 molecules with the lowest penalized-logP scores in ZINC250k dataset.
Revised: we select 800 molecules with the lowest penalized-logP scores in ZINC250k dataset.

Original: This experiments aim to generative novel molecules with high penalized-log while similarity to the target molecules.
Revised: This experiment aims to generate novel molecules with high penalized-log while similarity to the target molecules.

Original: We first randomly select one molecule as the target [Table] and then optimize the p-logP. Within each SGDS step, we only keep the molecules that have similarity score greater than minimum value δ.
Revised: We first randomly select one molecule as the target, and then optimize the p-logP. For each SGDS step, we only keep the molecules that have similarity score greater than minimum value δ

Original: [Table]
Revised: [Table]

Original: A. A DDITIONAL E XPERIMENTS ON G UACA M OL B ENCHMARKS
Revised: A. A DDITIONAL E XPERIMENTS ON G UACA M OL B ENCHMARKS

Original: We further evalute our MolEBM on several distribution learning and goal-directed optimization tasks in GuacaMol benchmark (Brown et al., 2019).
Revised: We further evaluate our MolEBM on several distribution learning and goal-directed optimization tasks in GuacaMol benchmark (Brown et al., 2019).

Original: To be specific, for distribution learning [Table] we use the validity, uniqueness and novelty to evaluate our model.
Revised: To be specific, for distribution learning, we use the validity, uniqueness and novelty to evaluate our model.

Original: The results are shown in Table 10.
Revised: The results are shown in Table 10.

Original: Comparing to existing methods, our MolEBM trained using SELFIES representations achieve highest scores among all three tasks.
Revised: Comparing to existing methods, our MolEBM trained using SELFIES representations achieves highest scores among all three tasks.

Original: For goal-directed benchmarks, we select five multiple property optimization tasks (MPO).
Revised: For goal-directed benchmarks, we select five multiple property optimization tasks (MPO).

Original: Due to  time limit, we just run the experiments without carefully fine-tuning the hyper-parameters.
Revised: Due to the time limit, we did not have time to tune our code for the new experiments. We expect to obtain improved results with fine-tuning.

Original: However, our MolEBM is also able to achieve comparable results to the strong baselines as Graph GA.
Revised: However, our MolEBM is still able to achieve comparable results to the strong baseline as Graph GA.

Original: Again, GuacaMol is an extremely useful benchmark since it provides the normalized scoring functions which is a weighed sum of multiple diverse properties of interest.
Revised: Again, GuacaMol is an extremely useful benchmark since it provides the normalized scoring functions which are weighed sum of multiple diverse properties of interest.

Original: The normalized scoring functions enable the fast convergence of MolEBM.
Revised: The normalized scoring functions enable the fast convergence of MolEBM.

Original: With GuacaMol, we can also investigate the design choice between single property regression network to predict the predefined scores and multiple property regression networks to predict original chemical properties individually.
Revised: With GuacaMol, we can also investigate the design choice between single property regression network to predict the pre-defined scores and multiple property regression networks to predict original chemical properties individually.

Original: We leave this  for future studies [Table]
Revised: We leave this question for future studies.

Original: A.8. D ISTRIBUTION -L EARNING B ENCHMARKS
Revised: A.8. D ISTRIBUTION -L EARNING B ENCHMARKS

Original: [Table]
Revised: [Table]

Revised: [Table]

Original: A.8. G OAL - DIRECTED B ENCHMARKS
Revised: A.8. G OAL - DIRECTED B ENCHMARKS

Original: [Table]
Revised: [Table]

Original: A. G ENERATED S AMPLES
Revised: A. G ENERATED S AMPLES

Original: Figure 6 and Figure 7 show generated molecules with high binding affinities towards ESR1 and
Revised: Figure 6 and figure 7 show generated molecules with high binding affinities towards ESR1 and

Original: ACAA1 respectively in single-objective property design.
Revised: ACAA1 respectively in single-objective property design experiments.

Original: Figure 8 and Figure 9 show generated molecules with high binding affinities towards ESR1 and
Revised: Figure 8 and Figure 9 show generated molecules with high binding affinities towards ESR1 and

Original: ACAA1 respectively in multi-objective property design.
Revised: ACAA1 respectively in multi-objective property design.

Original: Comparing to the previous state-of-the-art methods, our MolEBM is able to produce more high quality molecules than top-3 only because after sampling with gradual distribution shifting (SGDS), MolEBM is able to locate at the area supported by molecules with high binding affinities.
Revised: Comparing to the previous state-of-the-art methods, our MolEBM is able to produce more high quality molecules than top-3 molecules because after sampling with gradual distribution shifting (SGDS), MolEBM locates at the area supported by molecules with high binding affinities.

Original: In single-objective design, we find that few generated molecules may be of less practical use due to undesired properties (e.g. the first one in Figure 7 has a large circle).
Revised: In single-objective design, we find that few generated molecules may be of less practical use due to undesired properties (e.g. the first one in Figure 7 has a large circle).

Original: This observation is in accordance with (Eckmann et al., 2022), which is the case where the single-objective optimization is not sufficient.
Revised: This observation is in accordance with (Eckmann et al., 2022), which is the case where the single-objective optimization is not sufficient.

Original: Thus we need multi-objective binding affinities design settings since in contrast to non-biological properties, binding affinities are hard to compute and optimize.
Revised: Thus we need multi-objective binding affinities design settings because in contrast to non-biological properties, binding affinities are hard to compute and optimize.

Original: In multi-objective design settings, we find that those issues mentioned above can be largely solved by optimizing binding affinities, QED and SA at the same time.
Revised: In multi-objective design settings, we find that those issues mentioned above can be partially addressed by optimizing binding affinities, QED and SA at the same time.

Original: Meanwhile, compared to previous generative model based methods, we use Langevin dynamics to infer the posterior distribution p ( z | x, y 1 , . . .
Revised: Meanwhile, compared to previous generative model based methods, we use Langevin dynamics to infer the posterior distribution p ( z | x, y 1 , . . .

Original: , y n ) without bothering to design different encoders which is not trivial.
Revised: , y n ) without bothering to design different encoders when facing different combination of properties.

Original: A. T OP -3 M OLECULES IN P - LOG P AND QED O PTIMIZATION
Revised: A. T OP -3 M OLECULES IN P - LOG P AND QED O PTIMIZATION

Original: A. R EPRODUCIBILITY
Revised: A. R EPRODUCIBILITY

Original: Our code and saved checkpoints can be found here 2 .
Revised: Our code and saved checkpoints can be found here 2 .