MolEBM: Molecule Generation and Design by Latent Space Energy-Based Modeling
Authors: Deqian Kong
Bo Pang
Tian Han
Ying Nian Wu
Venue: Submitted to ICLR 2023 (2023)
Keywords: molecule design
energy-based model

Abstract: Generation of molecules with desired chemical and biological properties such as high drug-likeness, high binding affinity to target proteins, is critical in drug discovery. In this paper, we propose a probabilistic generative model to capture the joint distribution of molecules and their properties. Our model assumes an energy-based model (EBM) in the latent space. Given the latent vector sampled from the latent space EBM, both molecules and molecular properties are conditionally sampled via a top-down molecule generator model and a property regression model respectively. The EBM in a low dimensional latent space allows our model to capture complex chemical rules implicitly but efficiently and effectively. Due to the joint modeling with chemical properties, molecule design can be conveniently and naturally achieved by conditional sampling from our learned model given desired properties, in both single-objective and multi-objective optimization settings. The latent space EBM, top-down molecule generator, and property regression model are learned jointly by maximum likelihood, while optimization of properties is accomplished by gradual shifting of the model distribution towards the region supported by molecules with high property values. Our experiments show that our model outperforms state-of-the-art models on various molecule design tasks. 

Under review as a conference paper at ICLR 2023
MOLEBM: MOLECULE GENERATION AND DESIGN BY LATENT SPACE ENERGY-BASED MODELING
Anonymous authors Paper under double-blind review
ABSTRACT
Generation of molecules with desired chemical and biological properties such as high drug-likenese high binding affinity target proteins, is critical in drug dis- covery. In this paper, propose probabilistic generative model to capture the joint distribution of molecules and their properties. Our model assumes an energy- based model (EBM) in the latent space. Given the latent vector sampled from the latent space EBM, both molecules and molecular properties are conditionally sam- pled via molecule generator model and property regression model respectively. The EBM in low dimensional latent space allows our model to capture complex chemical rules implicitly but efficiently and effectively. Due the joint model- ing with chemical properties, molecule design can be conveniently and naturally achieved by conditional sampling from our learned model given desired proper- ties, in both single-obeective and multi-objective optimization settings. The latent space EBM, molecule generator, and property regression model are learned jointly by approximate maximum likelihood, while optimization of properties is accom- plished by gradual shifting of the model distribution towards the region supported by molecules with high property values. Our experiments show that model outperforms stte-offthe-an models on various molecule design tasks.
INTRODUCTION
In drug discovery, it of vital importance find or design molecules with desired pharmacologic or chemical properties such as high drug-likeness and binding affinity to target protein. It is challenging 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 1033 (Polishchuk et al. 2013).
Recently, large body of work attempts to tackle this problem. The first line of work leverages deep generative models map the discrete molecule space to a continuous latent space, and op- timizes molecular properties the latent space with methods like Bayesian optimization (Gómez- Bombarelli et al. 2018 Kusner et al., 2017 Jin et al. 2018). 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 Zhou al. 2019: Shi et al. 2020 Luo et al.) 2021). number of other efforts have been made to optimize molecular properties with genetic algorithms (Nigam) al. 2020), particle-swarm algorithms (Winter et al. 2019) specialized MCMC methods (Xie al. 2021
In this work, we propose method along the first line mentioned above, by learning probabilistic latent generative model of molecule distributions and optimizing chemical properties in the latent space. Given the central role of latent variables in this approach, we emphasize that it is critical to learn latent space model that captures the data regularities of the molecules. Thus, instead of as- suming simple Gaussian distribution in the latent space as prior work (Gómez- Boannzz-Bombarelll baaegy-based model (EBM) (LeCun 1. Kim & Bengio 2016 Xie et al. 2016 Kumar et al. 2019 Nijkamp Du & Mordatch] 2019 Grathwohl et al. 2019 Finn et al. (0166) in latent space. This to latent space mergy-based model (LSEBM) as studied in Pang et al. 2020); Nie et al.
al,,
2018
Jin et al.) 2018). we assume a flexible and expressive
ct al.
2006
Ngiam et al. 2011
et al.
2019
leads
2021), where LSEBM has been shown to model the distributions of natural images and text well. For molecule modeling, without any explicit validity constraints in generation, our model generates molecules with high validity with simple SMILES representation (Weininger) 1988).

---

Under review as a conference paper at ICLR 2023
Given our goal of property optimization, we learn a joint distribution of molecules and their proper- ties. Our model consists of I) an EBM in a cow-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 D.. We call our model as MolEBM. All three components in our model are learned jointly by approximate maximum likelihood. A learned model generates a molecule with a high property value in two steps: 1) given the property value, sample the latent vector; given the sampled latent vector, generate molecule (see the top- to-bottom path in Figure Tap. 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 molecule with high property value is most likely not in the original data distribution. We thus de- sign 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 distribu- tion.
In drug discovery, most often we need to consider multiple properties simultaneously. Our model can be extended to this setting straightforwardly. With our method, we only need to add a regression model for each property, while the learning and sampling methods remain the same. 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. This design makes our approach versatile in dealing with varying number of properties.
We evaluate method in various settings including single-objective optimization and multi- objective optimization. Our method outperforms prior methods by significant margins.
In summary, our contributions are as follows:
We propose to learn latent space energy-based model for the joint distribution of molecules and molecular properties.
We develop a sampling with gradual distribution shifting method, enabling to extrapolate the data distribution and sample from the region supported by molecules with high property values.
Our methods are versatile enough be extended to optimizing multiple properties together.
Our model achieves state-of-the- art performances on a wide range of molecule optimization tasks.
RELATED WORK
Optimization with Generative Models. Deep generative models approximate the distribution of molecules with desired biological or non-biological properties. Existing approaches for generating
molecules include applying variational autoencoder (VAE) (Kingma & Welling] (0014) and genera-
tive adversarial network (GAN) (Goodfellow et al. 2014) etc. to molecule data (Gómez-Bombarell
al. 2018 Jin et al.] 2018 De Cao & Kipf 2018 Honda et al., 2019 Madhawa et al.) 2019) Shi
et al.
2020 Zang & Wang 3020,KKotsia et al. |020 Chen et al. 202 1: Fu et 2020 L.i et al.
2021
| Eckmann et al. 2022] Segler et al. 2018). After learning continuous rep- resentations for molecules, they are further able to optimize using different methods. (Segler al.
Bagal et al.] |021|222]]]1
2018)
cycles. (Gómez-Bombare et al,, 2018 Jin et al. 2018 Eckmann et al. 2022) propose to learn a surrogate function predict prop- erties, and then use Bayesian optimization to optimize the latent vectors. However, the performance of this latent optimization is satisfactory due to three major issues. First, it is difficult train an accurate surrogate predictor especially for those novel molecules with high properties along the design trajectories. Second, as the learned latent space tries to cover the fixed data space, its abil- ity to explore the targets out of the distribution is limited (Brown et al.] 2019; Huang et al.] (02TT) 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. To address above is- sues, (Eckmann et al., (022) 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. In this paper, we propose an encoder-free model in both training and optimization to leam the joint distribution of molecules and properties, and make it possible to obtain several adcquate predictors. We then design an efficient algorithm to shift the learned distribution iteratively.
proposes to optimize by simulating
2

---

Under review as a conference paper at ICLR 2023
y
ym
Control lad
Py (y]=)
Generation
Piy (3n (=)
Prom (Dmm [8]
(2)
Po(=)
(=(2)=)
(a) Single-Objectivy Optimization.
(b) Multi- Optimization.
Figure 1: An illustration of MolEBM. represents molecule, 2 the latent vector, y a molecular property of interest, {yj 1 indicates m properties.
Optimization with Reinforcement Learning and Evolutionary Algorithms. Reinforcement learning (RL) based methods directly optimize and generate molecules in an explicit data space YYou et al.) 2018 Zhou al.) 2019 Jin al.) 2020 Gottipati al., 2020). By formulating the property design as a discrete optimization task, they can modify the molecular substructures guided by an or- acle reward function. However, the training of those RL-based methods can be viewed as rejection sampling which is difficult and inefficient due to the random-wall -aak search behavior in the discrete space. Evolutionary algorithms (EA) also formulate the optimization in discrete manner (Nigam) et al., 2020)) Jensen. 2019 Xie et al. 2021 Fu et al. [b) By leveraging carefully-cratted combinatorial algorithms, they can search molecule graph space flexible and efficient way. However, the design of those algorithms is non-trivial and domain specific.
METHODS
3.1 PROBLEM SETUP AND OVERVIEW
We use the SELFIES representation for molecules (Krenn et al. 2020). It encodes each molecule as a string of characters and ensures validity of all SELFIES strings. Let = (ax1), ,..." ,itt,..., x(7) be a molecule string encoded in SELFIES, where (0 E V the 1-th character and V is the vocabulary. Suppose y E R represents a molecular property of interest. Then the problem attempt tackle is to optimize æ such that its property = y² where y is some desirable value for We take a probabilistic approach and treat the optimization problem sampling problem, that is,
p(ay =-"))) y²")
This is a single-objective optimization problem since only one property is targeted. In real-world drug design settings, we are more likely to need to optimize multiple properties simultaneously. that is, multi- -ojjetivbeeetiveeeee optimization. Suppose optimize for {yz { then our task is to sample,
x p(aly1 = vi, Um = vmm)
To address these problems, we propose a solution under unified probabilistic framework. As first step, we need to model the data distribution of molecules, (data (z). To this end, we recruit latent space energy-based model (Pang al. 2020 Nie al. (20TT) for its expressiveness. LSEBM assumes latent vector E I in low dimensional latent space follows an energy-based model, while the observed x is generated by a generator conditional z, that is, p(x, =) = p(x()))(z). We introduce LSEBM within the context of molecule data in $.2
For the purpose of property optimization, we propose to model the joint distribution of molecules and molecular properties ($3.3) We first consider the singlc-objective optimization problem (eq. IT] Given the expressiveness of LSEBM, we assume that the latent vector x captures data regularities learn simple regression model from the low-
in æ (evaluated in $4.2 and $A.4). Thus, dimensional latent space, that is, p(y|2), and the joint distribution is p(æ, y, 2) = p(=)p(x(=))p()) 2)) The model is learned with maximum likelihood (see $3.4 and Algorithm [
With the learned model, we can optimize x given y = y by ancestral sampling =* ~ p(zly = y*) and x² p[a|= z"). However, if y² deviates from the observed data distribution of y, this naive
3

---

Under review as a conference paper at ICLR 2023
solution involves sampling in extrapolated regime (or out of distribution regime) where y" is not in the effective support the learned distribution. We propose à Sampling with Gradual Distribution Shifting (SGDS) approach where learned distribution, and 2) gradually shift the learned distribution with these boundary samples to region where it is supported by high property values (see $3.3 and Algorithm
1) sample the boundary of the effective support of the
Our model is designed to be versatile such that admits straightforward extension to multi-objective optimization. To optimize given {v, y j= can simply augment the joint distribution with more regression models, i.e., p(r, 2, y1,...,, ym p( (x+() Π5= p(y)(z). The sampling proce- dure follows the same SGDS approach. See 3.6 for more details on multi-objective optimization.
3.2 LATENT SPACE ENERGY- MODEL.
(T) is a molecule string in SELFIES and 2 E Rd is the latent vector. Consider the following model.
Suppose = (z(1),
~~ Pa()),
2 x~ p<(22z)))
(3)
where Pa (z) a prior model with parameters a, and ps(x])) is generation model with parameters 3. In VAE (Kingma & Welling 2014). the prior is simply assumed to be an isotropic Gaussian distribution. In our model, (2(2) formulated as an energy-based model,
(=(2) Z(a)
(4)
where (o(2) is a reference distribution, assumed to be isotropic Gaussian as in VAE. fa Rd + IR the scalar-valued negative energy function and is parameterized by small multi-layer percep- tron (MLP) with parameters a. Z(a) = exp(fa(2))Po())d) Epo (exp(fa(z))] is the partition function.
The generation model p=(x])) is conditional autoregressive model,
Pa (x = II (z " (((-1),2)
(5)
t=1
which parameterized by one-layer LSTM with parameters Note that the latent vector controls every step of the autoregressive model.
3.3 JOINT DISTRIBUTION OF MOLECULE AND MOLECULAR PROPERTY
Given molecule x, suppose y the chemical property of interest, such as QED or protein affinity binding. The property value can be computed from an input via open-sourced software RD- Kit (Landrum et al.) 2013) or AutoDock-GPU (Santos-Martins et al. 2021). We assume that given 2, and y conditionally independent.
po (2r, y, =) = Pa (ul)))
(6)
where Pa (2) is the EBM prior, pa(x 2)) is the generation model, and p, (y|z)) is the property regres- sion model, and = (a, 7). We use the model pa (z, y, 2) to approximate the data distribution of (Ir, y). See Appendix $A.1 for details.
The property regression model can be written as
(vl2)
2 exp
1 20
(y
s,(2))²
where (2) is small MLP, with parameters y, predicting y based on the latent The variance o² is set as a constant or hyperparameter in our work.
3.4 LEARNING ALGORITIM
Suppose observe training examples ((1,,1), 1,...,n) The log- likllieeihoodd function is ((0) = Σm log (o,,, y:). The learning gradient can be calculated according to
Vo Volog Po (z, y) = Epp(= [o po
(8)
(pa(=l=.)) [e Po (2) + logps(<(2) log p. (y 2))]
(9)
4

---

Under review as a conference paper at ICLR 2023
For the prior model, V. log Pa()) = Vafa(z) (a Thus the learning gradient given an example (r,y) v) is (z, v) V. (z, v/ = E (= (V [z) (a)
Vafa())].
(10)
is updated based on the difference between & inferred from empirical observation (æ, y), and sampled from the current prior. For the generation model,
(1 v) = V log (2, u) (= IV (x]2)
(11)
Similarly, for the regression model,
(z, y) = V log lo y) (x V log (og1)
(12)
Estimating expectations in equations 17 18 and 19 requires MCMC sampling of the prior model
(z) and the posterior distribution (= 2). We recruit Langevin dynamics (Neal 2011). For target distribution T (2)), the dynamics iterates
-k+1 24 + vV, log (=*) + v 2SER,
(13)
where k indexes the time step of the Langevin dynamics, 8 is small step size, and cx ~ No0 Ia) the Gaussian white noise. (z) can be either Pa(2) pe(2 22222.0 y). In either case. V. log π (2) can
be efficiently computed by back- pack-propagation. See Appendix A.1 for more details.
Algorithm 1: Learning MolEBM
input Learning iterations T,, learning rates for the prior, generation, and regression model (1o, 1/1, 7/.
initial parameters f = (oa, S0 "oo), observed examples {(x,, 34 y1))= batch size m, number of prior and posterior sampling steps (Ko, K₁ }. and prior and posterior sampling step sizes
{8u, }1
output: Of (aT, Br,
for T T- do
Mini-batch:
Sample observed examples ((r,, y1)
2. Prior sampling: For each i, sample 2, Pure (z) using equation
13
where the target distribution
<(2)= = Po. (z), and s so, K Ko.
3. Posterior sampling: For each (2., 3.). sample 2,- ~ pa, (2/ar, Ma) using equation
13
where the
target distribution (=)) pa, (2|ar, Mt). and 81, K K1.
4. Update prior model: 0x41 = ai + 17o m Σ Wor (n v.Ja, (z7)
Update generation model: Be+1 = Be + m Σm
V logga, (x)l2))-
Update regression model: 7++1 7t + 172
V (u:l=
3.5 SAMPLING WITH GRADUAL DISTRIBUTION SHIFTING
To tackle the single-objective optimization problem (eq. D.. one naive approach is to perform ances- tral sampling with two steps, given some desirable property value yy"
y²1=) (1) (2 = y²) (2(2)) -y= (2) x po(2l2 = =²),
(14)
where (1) is application of Bayes' rule, with (=(2) as the prior and pr(1(2) the likelihood.
Our model po y, 2) is learned to capture the data distribution. In real-world settings, y" might not be within the support of the data distribution. Therefore, sampling following equation 14 does not work well since it involves extrapolating the learned distribution. We propose method called sampling with gradual distribution shifting (SGDS) address this issue. In particular, examples from the training data on the boundary of the distribution support by sorting them accord-
find
ing to the values of y and taking top-k samples, {(z,)old) v. (old))) D with high property values. We shift the support slightly by adding some small A, to all y's, and sample x conditional on shifted y's, following equation [4] Given the generated x compute their groundtruth y' with RDKit or AutoDock-GPP. We then shift the learned model by finetuning with the new data, [(=/nee) (emeew.yewww (new)) )}k }/= for few steps (e.g., 10). This completes iteration of distribution shift. The molecule with desired high property value is sampled after T shifting iterations. We summarize the algorithm in Algorithm 2
5

---

Under review as a conference paper at ICLR 2023
Algorithm 2: Sampling with Gradual Distribution Shifting (SGDS).
input Shift iterations T; initial parameters 8, (0. initial examples { v from the data distribution boundary, shift magnitude Ay. AutoDock-GPP LearningAlgoritiee
RDKit
Algorithm
output: {(x, 3/"") ))
for
0: T
Property shift: For cach vi. 5 = vi Ay.
2. Latent sampling: For each ū! sample 2 ~ po, (=1u)
Molecule sampling: For each 2;-1, sample a/+1 po, [x]]a] (x)=)t1).
4. Property computation: For each x compute vf PropettyoompuceEngioeee
5. Distribution shift: 0 = LearningAlgoritmm(( (/
For constrained optimization, in step 3, we only keep the sampled molecules that satisfy the given constraints.
3.6 MULTI-OBIECTIVE OPTIMIZATION
We next consider the multi-objective optimization problem. Suppose we optimize for a set of prop- erties {y, } then we learn property regression model for each property y/-
1 (y) - 8ny (z))² 20²
Pr (uglz) = V
exp
(15)
2
where each 8y, is small MLP with parameters 7/j. Then the joint distribution is,
Po(x, 2, Vi, Vm) Pa (x(2)ps(<12)
II J=1
Pr, (vilz).
(16)
Under our framework, the learning algorithm and the sampling algorithm for the single-objective problem can be straightforwardly extended to the multi-objective setting. In both settings, the same types of properties are provided in both the initial training stage and the distribution shifting stage.
EXPERIMENTS
To demonstrate the effectiveness our proposed model, we compare our model with previous SOTA methods for unconditional generation ($4.2) and molecule design including single-objective optimization ($4.3) and multi-objective optimization ($4.4). In molecule design experiments, we consider both non-biological and biological properties.
4.1 EXPERIMENTAL SETUP
Datasets. For the unconditional generation task, we report results on ZINC (Irwin et al. 2012) and MOSES datasets (Polykovskiy et al. 2020). ZINC consists around 250k molecules, and MOSES comprises around 2 million molecules. For optimization tasks, we conduct experiments with ZINC.
Training Details. There are three modules in our model, the top-down generation model p&s ( z)) the EBM prior (z), and the regression model p, (8(2). P=(x(2) is parameterized by single-layer L.STM with 1024 hidden units. The dimension of latent vector z is 100. Pa is a 3-layer MLP with 100 input dimension and 200 hidden units. The property regression model (ylv)) =)] is 3- layer MLP with 100 input dimension and 100 hidden units. In multi-objective optimization settings, recruit one MLP for each property and these MLPs have the same architecture mentioned above. It worth mentioning that compared to most previous models, our model characterized by its simplicity without adding inference networks, RL-related modules, and graph neural networks. Adam (Kingma & Bal 2013) optimizer is used to train our models with learning rates 0.0001 for EBM and 0.001 for the We train our models for 30 epochs. 10, 000 (non-biological) and 2, 000

---

Under review as a conference paper at ICLR 2023
(biological) boundary examples are selected for distribution shifting. The numbers of shift itera- tions are 50 (non-biological single-objective, 20 mon-biolooical single-objective) and 10 (multi- objective). Within each SGDS iteration, the model is finetuned with only 10 parameter updates. All experiments are conducted on Nvidia Titan XP GPU. See more details in Appendix AA6
4.2 UNCONDITIONAL GENERATION
Three types of encoding systems are used to encode molecules in prior work: SMILES (Weininger [988]. SELFIES (Krenn al. 2020), and graph. SMILES and SELFIES linearize a molecule graph into string of characters. Most previous models using SMILES struggle to generate molecules with high validity, which is the percentage of molecules that satisfy the chemical valency rule. Thus graph representations become popular since explicit valency constraints can be imposed. However, perfect validity in this approach does not imply that the model captures the chemical rules since it is achieved with external constraints. Recently, SELFIES is developed where every SELFIES string corresponds to a valid molecule due to the nature of the encoding system. Thus, validity for generated SMILES strings is good indicator how well a learned model captures the basic chemical rules implicitly. 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).
Following previous work, we randomly sample 10k molecules for ZINC and 30k for MOSES, and compare on the three aforementioned metrics. Generations results on ZINC and MOSES are shown in Table T and Table 2 respectively. str-smi denotes string- SMILES representations and str-sf fi denotes string-based SELFIES representations. In Table we show generation results for both SMILES and SELFIES. SMILES does not have a validity constraint during generation. However, our model, MolEBM, still achieves 95.5% validity, which outperforms other SMILES- based methods and also comparable to those with valency check. This result demonstrates that model capture those valency rules effectively and implicitly (also see Appendix A.4 Samples from our model also achieve perfect uniqueness and novelty.
Model
Representation
Validity
Novelty
Uniqueness
JT-VAE (Tin et al.
2018
Graph
1,000"
1.000
1.000
GCPN You et al. 2018
Graph
1.000'
1.000
1.000
Graph (aahaaaa et al. 2019
Graph
0.426
1.000
0,948
GraphAF Shu et al. (2020)
Graph
1.000"
1.000
0.991
GraphDF Luo et al. 2021
Graph
1.000"
1.000
1.000
Chem
et al.
2018
str-ami
0.170
0.980
0.310
GrammarVAl (Kusner al. 2017)
str-ami
0.310
1.000
0.108
SDVAE (Dai al. 2018)
str-ami
0.435
MolEBM
str-smi 0.955
1.000
1.000
MolEBM
str- ffi
1.000
1.000
1.000
Table 1: Unconditional generation on ZINC. denotes valency check.
Model
Representation
Validity
Novelty
Uniqueness
JT-VAF (in
2018)
Graph
1.000*
0.914
1.000
GraphAF Shi al. 2020
Graph
1.000*
1.000
0.991
GraphDF (Luo et al. 2021
Graph
1.000*
1.000
1.000
LIMO (Eckmann et al. 2022
str-sfi
1.000
1.000
0.976
MolEBM
sts-sfi
1.000
1.000
1.000
Table Unconditional generation on MOSES. denotes valency check. Results obtained from Polykovskiy
et al. 2020 Eckmann et al. 2022)
4.3
OINGLE-OBIECTIVE OPTIMIZATION
Non- Mon-Biological Property Optimization. For biologial properties, interested in Penal- ized logP and QED, both of which can be calculated by RDKit (Landrum et al. 2013). Since we know the Penalized logP scores have a positive relationship with the length of the molecules,

---

Under review as a conference paper at ICLR 2023
ESR1
0.543
1.711, 0.42, 3.85
2.280, 0.56, 2.57
ACAA1
6.714
6.944
8.503
5.671, 0.60, 4.58
6.602, 0.56, 4.07
Figure 2: Generated molccules with high binding affinitics. Left: Top-3 single-objective optimized molecules. The number denotes Kp in nmol/L. Right: Top-2 multi-objectie optimized molecules. Numbers denote KD in nmol/L. QED and SA respectively.
maximize Penalized logP either with or without maximum length limit. Following (Eckmann et al, 2022). the maximum length is set be the maximum length of molecules in ZINC using SELFIES. From T we can see that with length limit, MolEBM outperforms previous methods by large margin. Maximizing penalized logP without length limit leads to the highest Penalized logP, 158.0. We also achieve the highest QED with/without length limit. These observations demonstrate the effectiveness of our method. We also illustrate our distribution shifting method in Appendix A.3
Penalized logP
(t)
QED (t)
Method
LL
1st
2nd
3rd
1st
2rd
3rd
JT-VAE
X
5.30
4,93
4.49
0.925
0.911
0.910
GCPN
V
7.98
7.85
7.80
0.948
0.947
0.946
MolDQN
V
11.8
11.8
11.8
0.948
0.943
0.943
MARS
X
45.0
44.3
43.8
0.948
0,948
0.948
GraphDF
13.7
13.2
0.948
0,948
0.948
LIMO
10.5
9.69
9.60
0.947
0.946
0.945
MolEBM
26.4
25.1
24.4
0.948
0.948
0.948
MolEBM
X
158.0
157.8
157.5
0.948
0.948
0.948
Table 3: Non-biological single-objective optimization. Report top-3 highest scores found by each model. LL (Length Limit) denotes whether the model has the limit of maximum length. Baseline results obtained from (Bekmann al. 2022 You et al. 2018 Luo et al. 2021 Xie et al. 2021).
Biological Property Optimization. ESRI and ACAAI are two human proteins. We aim to design ligands (molecules) that have the maximum binding affinities towards those target proteins. ESR well-studied, which has many existing binders, while ACAAI does not. However, did not use any binder-related information in the design process. Binding affinity is measured by the estimated dissociation constants KD. which be computed with AutoDock-GPU (Santos-Martins et al., 2021) given a molecule. Large binding affinities corresponds to small KD- That is, we aim to minimize KD. Table 4 shows that our model outperforms previous methods on both ESRI and ACAAI binding affinity maximization tasks. Producing those ligands with high binding affinity plays a vital role in the early stage of drug discovery.
ESRI KD
(L)
ACAAI KD (L)
Method
1st
2rd
3rd
1st
2rd
3rd
GCPN
75
83
84
MolDQN
373
588
1062
240
337
608
MARS
25
47
51
370
520
590
GraphDF
17
64
69
163
203
236
LIMO
0.72
0.89
1.4
37
37
MolEBM
0.52
0.54
0.54
6.71
6.94
8,50
Table 4: Biological single-objective optimization. Report top-3 lowest KD (in nanomoles/liter) found by each
model. Baseline results obtained from [Eckmann et al 2022
8

---

Under review as a conference paper at ICLR 2023
4.4 MULTI-OBIECTIVE OPTIMIZATION
We next consider optimizing binding affinity, QED and SAS simultaneously for multi-objective optimization. Following Eckmann al. 2022). we exclude molecules with abnormal behaviors in the generation process to make sure the learned distribution shifts towards a desirable region in terms of pharmacologic and synthetic properties.
Table 5 shows our multi-objective binding affinity maximization results comparing to LIMO (Eck- mann al. 2022) and GCPN (You et al. 2018). From the results, we can see that MolEBM is able to find the ligands with desired properties while keeping the pharmacologic structures. For E we have two existing binders on the market, Tamoxifen and Raloxifene. Our designed ligands have similar QED and SA, with rather low KD. Compared to existing methods, MolEBM obtains bet- ter results in overall adjustments. For ACAAI, we do not have any existing binders. Compared with prior SOTA methods, our optimized ligands outperform those by a large margin. When com- paring ACAAI multi-objective setting with its corresponding single-objective onc, we find multi- objective 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. How- ever, our current results indicate that with proper prior knowledge (e.g. multiple objectives 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. 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.
ESRI
ACAAI QED (t)
Ligand
KD (L)
QED (t)
SA (L)
KD (L)
SA (L)
GCPN
810
0.43
4.2
8500
0.69
4.2
GCPN nd
2.7 x 10
0.80
3.7
8500
0.54
4.3
LIMO 1
4.6
0.43
4.8
28
0.57
5.5
LIMO 2dd
2.8
0.64
4.9
31
0.44
4.9
Tamoxifen
87
0.45
2.0
Raloxifene
7.9 x 10"
0.32
2.4
MolEBM
1.71
0.42
3.85
5.67
0.60
4.58
MolEBM 2nd
2.28
0.56
2.56
6.60
0.56
4.07
Table 5: Muli Muji-Objective Binding Affinity Maximization for both ESRI and ACAAL Report Top-2 average scores related to KD (in nmol/I.), QED and SA. Baseline results obtained from (Eckmann et al. 2022).
Figure shows generated molecules with high binding affinities. See Appendix for more examples.
CONCLUSION AND DISCUSSION
We propose deep generative model, MolEBM, which models the joint distribution of molecules and molccular properties. It assumes an energy-based prior for low-dimensional continuous latent space, which effectively captures data regularities of the discrete molecule data. We then design distribution shifting method (SGDS) to shift the learned distribution to region with high property values. Molecule design can then be achieved by conditional sampling. Our experiments demon- strate that our method outperforms previous SOTA methods by a significant margin.
limiting factor is that the sampling with gradual distribution shifting (SGDS) requires that the properties of interest can be easily computed. 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.
Another limitation that, for our optimization method, good property scores may not translate to useful molecules for drug design. This problem may be partially addressed by multi-objective opti- mization as studied this paper. In our work, we focus on optimizing given objectives. Designing good objectives metries is an equally perhaps even more important problem that deserves careful investigation.
In particular, we exclude molecules with QED (f) smaller than 0.4, SA (L) larger than 5.5, and too small (less than 5 atoms) too large (more than 6 atoms) chemical rings.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Viraj Bagal, Rishal Aggarwal, PK Vinod, and U Deva Priyakumar. Liggpt: Molecular generation using a transformer-decoder model. 2021.
Nathan Brown, Marco Fiscato, Marwin H.S. Segler, and Alain C. Vaucher. Guacamol: Benchmark- ing models for de novo molecular design. Journal of Chemical Information and Modeling, 59 (3):109--1000 108, 2019. doi: 10.102 www.s...em....0....... URL https: http::::///cooo///./../oo/oo 1021/ acs. jeim. 8b00839
PMID: 30887799.
Binghong Chen, Tianzhe Wang, Chengtao Li, Hanjun Dai, and Le Song. Molecule optimization by explainable evolution. International Conference on Learning Representation (ICLR), 2021.
Hanjun Dai, Yingtao Tian, Bo Dai, Steven Skiena, and Le Song. Syntax-directee variational autoen- coder for structured data. International Conference on Learning Representations, 2018.
Nicola De Cao and Thomas Kipf. Molgan: An implicit generative model for small molecular graphs. arXiv preprint orXiv:/805 11973, 2018.
Yilun Du and Igor Mordatch. Implicit generation and generalization in energy-based models. CoRR, 0tss/903.08689) 2019, URL htt::tt/://axxx org abs/1903. 08689
Peter Eckmann, Kunyang Sun, Bo Zhao, Mudong Feng, Michael K Gilson, and Rose Yu. Limo: Latent inceptionism for targeted molecule generation. arXiv preprint arXiv:2206.0900 2022.
Chelsea Finn, Paul Christiano, Pieter Abbeel, and Sergey Levine. connection between gen- erative adversarial networks, inverse reinforcement learning, and energy- energy-based models. CoRR, abs//61 1. 1.03852, 2016. BRLhhtt:://
org/
www./r////.oooo..oo...o.eeeeeaee aba/ 03852
Tianfan Fu, Cao Xiao, and Jimeng Sun. Core: Automatic molecule optimization using copy & refine strategy. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 638-645, 2020.
Tianfan Fu, Wenhao Gao, Cao Xiao, Jacob Yasonik, Connor Coley, and Jimeng Sun. Differen- tiable scaffolding tree for molecular optimization. arXiv preprint arXiv:2109. 104.9 202
Tianfan Fu, Cao Xiao, Xinhao Li, Lucas M Glass, and Jimeng Sun. Mimosa: Multi-constrain molecule sampling for molecule optimization. In Proceedings of the AAAI Conference Artifi- cial Intelligece, volume 35, pp. 125-133, 2021b.
Rafael
Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato. Benjamín Beanchez-Lengelinng Dennis Sheberla, Jorge Aguilera-Iparraguine Timothy Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using data-driven contin- uous representation of molecules. ACS central science, 4(2):268-276, 2018.
Ian Goodfellow, Jean Pouget-Abadic, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neu- ral Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014. Montreal, Quebec, Canada, pp. 2672 2680, 2014. URL http:://appe www..nipc....c ec paper/ 5423 -e2---gnerattt.eeeaveeaaaa -
Sai Krishna Gottipati, Boris Sattarov, Sufeng Niu, Yashaswi Pathak, Haoran Wei, Shengchao Liu, Simon Blackburn, Karam Thomas, Connor Coley, Jian Tang, et al. Learning to navigate the synthetically accessible chemical space using reinforcement learning. In International Conference on Machine Learning, pp. 3668-3679 PMLR, 2020.
Will Grathwohl, Kuan-Chieh Wang, Jömn Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. arXiv preprint arXiv: 12.. 03263, 2019.
Shion Honda, Hirotaka Akita, Katsuhiko Ishiguro, Toshiki Nakanishi, and Kenta Oono. Graph residual flow for molecular graph generation. arXiv preprint arXiv: 1909. 13521, 2019.
10

---

Under review as a conference paper at ICLR 2023
Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development. arXiv preprint arXiv:2 2021.
John J Irwin, Teague Sterling, Michael M Mysinger, Erin Bolstad, and Ryan G Coleman. Zinc: free tool to discover chemistry for biology. Journal of chemical information and modeling, 52 (7):1757-1768, 2012.
Jan H Jensen. graph-based genetic algorithm and generative model/monte carlo tree search for the exploration of chemical space. Chemical science, 10(12):3567-3572. 2019.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning, pp. 2323-2332. PMLR, 2018.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Multi-objective molecule generation using interpretable substructures. International conference on machine learning, pp. 4849-488 PMLR, 2020.
Taesup Kim and Yoshua Bengio. Deep directed generative models with energy-hased probability estimation. CoRR, abs/1606.03433, 2016. URL httl:tpp//aazii..rrr///////1... 03439
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd Inter- national Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Truck Proceedings, 2015. URLhttp:://arxee org/ abs abs/aass/11 /122. 6980
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, hbsp//bs//// 1312. 6114
Conference Track Proceedings, 2014. http::////a
Panagiotis- Chrristos Kotsias, Josep Arús-Pous, Hongming Chen, Ola Engkvist, Christian Tyrchan, and Esben Jannik Bjerrum. Direct steering of de novo molecular generation with descriptor con- ditional recurrent neural networks. Nature Machine Intelligence, 2(5):254-265, 2020.
Mario Krenn, Florian Häse, AkshatKumar Nigam, Pascal Friederich, and Alan Aspuru-Guzik. Self- referencing embedded strings (selfies): A 100% robust molecular string representation. Machine Learning: Science and Technology, 1 14):0445224 2020.
Rithesh Kumar, Anirudh Goyal, Aaron C. Courville, and Yoshua Bengio. Maximum entropy gen- erators for energy- based models. CoRR, abs/1901.08508, 2019. URL /ttp:://axxxxx
abs/1901. 08508
Matt Kusner, Brooks Paige, and José Miguel Hemández-Looato. Grammar variational autoen- coder. In International Conference on Machine Learning, pp. 1945-1954, 2017.
Greg Landrum et al. Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling. Greg Landrum, 2013.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Huang. tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
Meng Liu, Keqiang Yan, Bora Oztekin, and Shuiwang Graphebm: Molecular graph generation with energy models. arXiv preprint arXiv:2102.00546, 2021.
Youzhi Luo, Keqiang Yan, and Shuiwang Ji. Graphdf: A discrete flow model for molecular graph generation. In International Conference Machine Learning, pp. 7192-7203 PMLR, 2021.
Kaushalya Madhawa, Katushiko Ishiguro, Kosuke Nakago, and Motoki Abe. Graphnvp: An invert- ible flow model for generating molecular graphs. arXiv preprint arXiv: 1600, 2019.
Radford M Neal. MCMC using hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2, 2011.
11

---

Under review as a conference paper at ICLR 2023
Jiquan Ngiam, Zhenghao Chen, Pang Wei Koh, and Andrew Y. Ng.. Learning deep energy models. In Proceedings of the 28th International Conference on Machine Learning, (CML2011) Bellevue, Washington. USA, June 28- July 2. 2011, pp. 1105-1112, 2011. URL https://.iiic ec/
2011 ht11/paper/5557 cm paper.ap...... pdf
Weili Nie, Arash Vahdat, and Anima Anandkumar. Controllable and compositional generation with latent-space energy-based models. Advances in Neural Information Processing Systems, 34:13497-13510, 2021.
AkshatKumar Nigam, Pascal Friederich, Mario Krenn, and Alán Aspuru-Guzik. Augmenting ge- netic algorithms with deep neural networks for exploring the chemical space. ICLR 2020, 2020.
Erik Nijkamp, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu. Learning non-convergent persistent short-run MCMC toward energy-based model. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, 8-14 December 2019, Vancouver, Canada, 2019.
Bo Pang. Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. arXiv preprini arXiv:2006.08203. 2020.
Pavel G Polishchuk, Timur Madzhidov, and Alexandre Varnek. Estimation of the size of drug- like chemical space based on gdb-17 data. Journal of computer-aided molecular design, 27(8): 675-679, 2013.
Daniil Polykovskiy. Alexander Zhebrak, Benjamin Sanchez-Lengeling, Sergey Golovanov, Oktai Tatanov, Stanislav Belyaev, Rauf Kurbanov, Aleksey Artamonov, Vladimir Aladinskiy, Mark Veselov, et Molecular sets (moses): a benchmarking platform for molecular generation models. Frontiers in pharmacology, 1 11:565644, 2020.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statistics, pp. 400-407, 1951.
Diogo Santos-Martins, Leonardo Koch, and Stefano Forli. Accelerating autodock4 with gpus and gradient-based local search. Journal of chemical theory and computation, 17(2):1060-1005 2021.
Andreas F Tillack, Michel F Sanner, Andreas
Marwin HS Segler, Thierry Kogej, Christian Tyrchan, and Mark P Waller. Generating focused molecule libraries for drug discovery with recurrent neural networks. ACS central science, 4(1): 120-131, 2018
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.09388 2020.
David Weininger. Smiles, chemical language and information system. 1. introduction to method- ology and encoding rules. Journal of chemical information and computer sciences, 28(1):31-4 1988.
Robin Winter, Floriane Montanari, Andreas Steffen, Hans Briem, Frank Noé, and Djork-Amé Clev- ert. Efficient multi-objective molecular optimization in continuous latent space. Chemical science, 10(3)):8016-804) 8024, 2019.
Jianwen Xie, Yang Lu, Song Zhu, and Ying Nian Wu. theory of generative convnet. In Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, pp. 2635 2644, 2016. URL URLhttp://proeeddng mlr. press/v4 hatps///vxxxiccie html
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 Confer- ence on Learning Representations, 2021. URL bttps://openzevieww het./foummiit
kHSu4ebbxXXY
12

---

Under review as a conference paper at ICLR 2023
Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec. Graph convolutional pol- icy network for goal-directed molecular graph generation. In Advances in neural information processing systems, pp. 6410-642) 2018.
Chengxi Zang and Fei Wang. Moflow: invertible flow model for generating molecular graphs. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. pp. 617-626, 2020.
Zhenpeng Zhou, Steven Keamne, Li Li, Richard N Zare, and Patrick Riley. Optimization of molecules via deep reinforcement learning. Scientific reports, 9(1):1-10, 2019.
13

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 DETAILS ABOUT MODEL AND LEARNING
Our model is the form Pa (a(2)p((z(2)p) (y/z). The marginal distribution of (a, y) is
po 3) = Po(a, y, 2)d2 (ol
We use po (2, y) to approximate the data distribution of (z, y).
For the data distribution of (I, y). y is deterministic function of However, machine learning method usually cannot learn the deterministic function exactly. Instead, we can only learn proba- bilistic po(y)z) x). Our model po (x.,)) y) seeks approximate the data distribution p(z, y) by maximum likelihood. A learnable and flexible prior model Pa (z) helps to make the approximate more accurate than fixed prior model such as that in VAE.
Let the training data be ((ri, y:), = 1,...n) The log-likelihood function is L(0) = Ei= log (Ti VK). The learning gradient is (0) = - Vo log (1 ya). In the following, provide details for calculating Ve log pa( y) for a single generic training example (æ, y) (where drop the subscript. for notation simplicity).
V log Po (3.
Vopo(z.1) v)
po(#, v)
po( v) / Vopo(T, V. 2,)
(x, v) / po(x; y, Volg y, d)
po (3x, U, 2)
V log po (2 y, z)d=
po( y)
= po(= xr v) Va po (I. z)dz
Epo(sle.y)
V lo a y, =)
Ep)() Ve Pa (2) log + lo (yl=))
For the prior model,
Va log Vafa(z) V log Z(a)
Vafalz) 1
Z(a) VaZ(a)
Vafa(z) Z(a) [// / (xx))(
Vafa(z) - Vafa(2) 1
exx((f(())))
Z(a)
Vafa(z) -pp--ppp--- [p) (a)]
Thus the learning gradient for given an example (z, y) is
(a v) Va log y) (p ( [2fa(2)] (x)
(17)
The above equation has an empirical Bayes nature. (0(2 y) is based on the empirical observation (a, y), while P. is the prior model. For the generation model.
81(2, y)
V log (I 4) ( VV (x)2)
(18)
Similarly, for the regression model,
(x y) = V., log v) (pp VV, (v12)
(19)
14

---

Under review as a conference paper at ICLR 2023
Estimating expectations in the above equations requires Monte Carlo sampling of the prior model Pz( and the posterior distribution (2 y). we can draw fair samples from the two dis- tributions, and use these Monte Carlo samples 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 tablished the convergence of such an algorithm to a local maximum of the log-likelihoo
For MCMC sampling using Langevin dynamics, the finite step short run Langevin dynamics may cause bias in Monte Carlo sampling. The bias was analyzed in Pang et al. (2020). The resulting algorithm is an approximate maximum likelihood learning algorithm.
A.2 ABLATION STUDIES
We conduct ablations on the key components of our method: 1) EBM Prior (EBM vs. standard Gaussian (V,O (2) SGDS (shift distributions with SGDS vs. shift), and (3) joint training (joint training of molecule and molecular property distribution vs. train molecule with LSEBM and learn property regression model in second step). The ablation results displayed in Table 6 It is clear that all the proposed components contribute significantly to the good performance of our method.
Penalized logP
EBM Prior
SGDS
Joint Training
1st
2nd
3rd
13.81
13.78
13.75
x
12.27
12.02
11.93
x
12.79
12.70
12.41
26.37
25.05
24.38
Table 6: Ablation studies.
A.3 ILLUSTRATION OF SAMPLING WITH GRADUAL DISTRIBUTION SHIFTING (SGDS)
Figure 3 shows property score densities of sampled molecules from MolEBM in the distribution shifting process. We can see the model distribution is gradually shifting towards the region supported by molecules with high property values.
15

---

Under review as a conference paper at ICLR 2023
10
20
Figure 3: Illustration of SGDS in a single-objective penalized logP optimization experiment.
A.4 UNCONDITIONAL GENERATION
For unconditional generation task, we use uniqueness, novelty and validity compare our gen- eration results with existing methods. Meanwhile, for SELFIES-basee ZINC dataset, we split this dataset into train split 240k) and test split (10k samples). Here, we randomly sample 10k molecules from the learned latent EBM and calculate their logP, QED and SA scores using RDKit. We com- pare these property densities with molecule property densities in test split. The results are shown in Figure We model match those of the data quite well, implying that our model indeed captures the chemical rules implicitly.
that the marginal distributions (property score densities) of the learned
0.5
Figure 4: Property score distributions.
16

---

Under review as a conference paper at ICLR 2023
A.5 GENERATION WITH DIFFERENT LENGTH OF MARKOV CHAIN
In our experiments, we use short-run MCMC (i.e. & 20) in Equation 13 for all experiments. From Figure 5 can with the increasing length of Markov chains, the molecules change accordingly, showing that Markov chain doesn't 't get stuck in the local mode.
Bno
appno
Pud
Sample 1;
xx
500
5-100
5.200
$-500
5000
Qu
Sample 2:
350 5-100
14000
$+500
140000
Figure Sampled molecules with the different length of Markov chain.
A.6 TRAINING TIME
The joint training of MolEBM takes 4 hours with 30 iterations a single Nvidia Titan XP GPU with batch size 2048. For non-biological single-objective property optimization, it takes around 0.5 hours do 50 distribution shifting (SGDS) iterations. For biological binding affinity maximization, the optimization time is mainly dependent on the number of queries of LutoDock-GPU We do 20 and 10 SGDS iterations for the single-objective and multi-objective tasks, respectively, which cost 10 hours and 5 hours. For biological property optimization tasks, we use two Nvidia Titan XP GPUs, one for running our code and another one for running AutoDock-GPU, We have added a table to compare with previous methods.
Model
Sanalize--logP/OE
JT-VAE
24
GCPN
8
MolDQN
24
GraphDF
8
Mars
12
LIMO
MolEBM
4.5
Table 7: Comparison of molecule generation time in (hrs). Results obtained from (eckmann al. 2022
Even if intensively MCMC sampling-basee methods, training speed is affordable com- paring to existing methods. That's due to our designed latent space EBM is low-dimensional (i.e. dim(=)=100) and we use short-run MCMC (i.e. with fixed iteration steps s 20) in our experi- ments. The sampling results with respect to the length of Markov chain is discussed in previous section.
A.7 OPTIMIZATION WITH CONSTRAINTS
In unconditional generation task, even with SMILES representation, our model can capture the valency constraints chemical space. This idea can further extended to constrained optimiza- tions, 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 We include logP targeting and
optimization experiments in the following.
For example, in logP targeting experiments, during each iteration, we only keep the molecules within the target logP range; in similarity constriiedd experiments, similarity constraints to update the model. By iteratively updating the model with those selected molecules, we shift the joint distribution towards the region that satisfies the constraints.
those molecules which satisfy
17

---

Under review as a conference paper at ICLR 2023
A.7.1 LOGP TARGETING
Comparing to previous methods, MolEBM able to get competitive diversity scores with signifi- cantly better success rate in both ranges. That's 's because after SGDS, our model is shifted towards the region that supported by molecules satisfying the logP constraints. 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.
2.5 S logp S -2
S logP S 5.5 Diversity
Method
Success
Diversity
Success
ZINC
0.4%
0.919
1.3%
0.901
JA-VAE
11.3%
0.846
7.6%
0.907
ORGAN
0
0.2%
0.909
GCPN
85.5%
0.392
54.7%
0.855
LIMO
10.4%
0.914
MolEBM
86.0%
0.874
62 2%
0.858
Table 8: logP targeting certain range (Eckmann al.
2022
You
2018
Luo al.
2021
Xie et al.
2021.
A.7.2 SIMILARITY-CONSTRAANE OPTIMIZATION
Following previous procedures in JT-VAE Jin al., 2018), we select 800 molecules with the lowest penalized-logP scores in ZINC250k dataset. This experiment aims generate novel molecules with high penalized-log while similarity to the target molecules. 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 6.
GCPN
GraphDF
LIMO
MolEBM
5
Improv.
% Succ.
Improv.
% Succ.
Improv.
% Succ.
Improv.
% Succ.
4.2 1.3
100
= 2.0
100
10.1 2.3
100
19.11 + 2.12
100
0.2
1.1 1.2
100
5.6 = 1.7
100
5.8 + 2.6
99.0
7.41 + 1.89
100
0.4
2.5 + 1.3
100
4.1 - 1.4
100
3.6 + 2.3
93.7
3.80 + 1.44
97.5
Table 9: Similarity constraty optimization results. Baseline results obtained from (Eckmann al.
2022
Luo et al. 2021
A.8 ADDITIONAL EXPERIMENTS ON GUACAMOL BENCHMARKS
We further evaluate our MolEBM on several distribution learning and goal-directed optimization tasks in GuacaMol benchmark (Brown et al. 2019).
To be specific, for distribution learning, we use the validity, uniqueness and novelty to evaluate our model. The results are shown in Table 10 Comparing existing methods, our MolEBM trained using SELFIES representations achieves highest scores among all three tasks.
For goal-directed benchmarks, we select five multiple property optimization tasks (MPO). 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. However, our MolEBM is still able to achieve comparable results to the strong baseline as Graph GA.
Again, GuacaMol an extremely useful benchmark since it provides the normalized scoring func- tions which are weighed sum of multiple diverse properties of interest. The normalized scoring functions enable the fast convergence of MolEBM. With GuacaMol, we can also investigate the design choice between single property regression network predict the pre-defined scores and mul- tiple property regression networks to predict original chemical properties individually. We leave this question for future studies.
18

---

Under review as a conference paper at ICLR 2023
A.8.1
DISTRIBUTION-LEARNNI BENCHMARKS
Benchmark
AAE
Graph MCTS
Random Sampler
SMILES LSTM
VAE
MolEBM
Validity
0.822
1.000
1.000
1.000
0.959
1.000
Uniqueness
1.000
1.000
0.997
1.000
0.999
1.000
Novelty 0.998
0.994
0.000
0.912
0.971
0.999
Table 10: Distribution learning results on GuacaMol benchmarks (Brown et
2019
A.8.2 GOAL-DIRECTED BENCHMARKS
Benchmark
Best of Dataset
SMILES GA
Graph MCTS
Graph GA
SMILES LSTM
MolEBM
Osimertinib MPO
0.839
0,886
0.784
0.953
0.907
0.933
Fexofenadine MPO
0.817
0,931
0.695
0.998
0.959
0.971
Ranolazinc MPO
0.792
0.881
0,616
0.920
0.855
0.924
Sitagliptin MPO
0.509
0.689
0,458
0.891
0.545
0.829
Table 11: Goal- directed optimization results on GuacaMol benchmarks (Brown et al. 2019). Top-2 results are highlighted as bold and italic respectively.
A.9 GENERATED SAMPLES
Figure 6 and figure 7 show generated molecules with high binding affinities towards ESR1 and ACAAT respectively single-objective property design experiments.
Figure 8 and Figure 9 show generated molecules with high binding affinities towards ESR1 and ACAAT respectively in multi-objective property design.
Comparing to the previous state-of-the-ar 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 the area supported by molecules with high binding affinities.
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). This observation is in accordance with (Eckmann et al., (002).. which is the case where the single-objective optimization not sufficient. Thus need multi-objective binding affinitics design settings because in contrast to non-bion non-biological properties, binding affinitics are hard to compute and optimize.
In multi-objective design settings, find that those issues mentioned above can be partially ad- dressed by optimizing binding affinities, QED and SA at the same time.
Meanwhile, compared to previous generative model based methods, to infer the posterior distribution p(z)a jr, y1...., Ba) without bothering to design different encoders when facing different combination of properties.
Langevin dynamics
19

---

Under review as a conference paper at ICLR 2023
0.54302
066776
106664
120041
120041
142111
165424
186170
186170
224151
231846
239806
2,40399
279146
Figure 6: Generated molecules in singe-objective esrl binding affinity maximization experiments with corre- sponding (D(1) amol/L.
arorg
Reoto
6.71418
6.94470
957027
10.06733
1484245
16A24277
23.41084
23.80933
26.79523
28.18692
28.18692
3119090
31.10900
3281089
35.10252
Figure 7: Generated molecules in singe-objective acaal binding affinity maximization experiments with cor- responding Kn()) mmol/L.
20

---

Under review as a conference paper at ICLR 2023
171104
174016
220399
235793
243886
256555
256555
260922
269880
274474
283897
298642
3.08895
353550
397889
Figure 8: Generated molecules in multi-objective esrl binding affinity maximization experiments with cor- responding Kp (4) mmol/L.
Suide
567143
6.60181
20.80208
2151626
o
2188250
21,88250
2263377
2263377
24.62676
25.47255
26.34677
2725132
28 7818892
30.68888
Figure 9: Generated molecules in multi-objective acaal binding affinity maximization experiments with corresponding (D()) in amol/L.
21

---

Under review as a conference paper at ICLR 2023
A.10 TOP-3 MOLECULES IN P-LOGP AND QED OPTIMIZATION
0.94119
0.98115
0.94809
Figure 10: Top-3 molecules in single-objective QED maximization.
26339733
25,05434
2437777
Figure 11: Top-3 molecules in single-objective p-logP maximization.
A.11 REPRODUCIBILIIT
Our code and saved checkpoints can be found here
22


Reviews:
[{"review_id": "dSCv2gJHqEp", "reviewer": "Reviewer_mgPa", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The aim of this paper is the generation of molecules with desired chemical and biological properties learning the joint distribution of molecules and properties . In order to achieve this the authors propose an energy based generative model that augments a top down generative model (conditional autoregressive model) with a latent space learnt via EBM and a property regression model for each of the properties. Then, to sample molecules with desired properties they design an algorithm to shift the learned distribution iteratively. \n", "summary_of_the_review": "The paper extends the LLSEBprior model with a property predictor on the latent space in order to generate molecules with desired chemical and biological properties. \nThe contribution of the paper is limited and more experiments are needed in order to illustrate the capability of the proposed model to generate molecules with specific properties (see Weaknesses).  \n", "strength_and_weaknesses": "Strength: \nThe paper seems to perform better than state of the art models for unconditional generations using smiles. Especially for unconditional generations it achieves 95.5% validity with 100% novelty using SMILES with the ZINC dataset. Moreover, the sampling with gradual distribution shifting (SGDS) is an interesting suggestion for out of distribution sampling. \n\nWeaknesses: \nThe paper is clearly written but is missing some more explanation about the sampling procedure and the theory behind it.  Even though to estimate the expirations MCMC sampling of the prior and the posterior is required the method is described as \"maximum likelihood”.  \nResults are not compelling enough. Event though the model is proposed for conditional generations given desired properties in the experimental results we see results only for property optimisation. So more experiments should be added to prove the capability for conditional molecule generation like\n- experiments with properties targeted to a predefined range, providing the percent of generated molecules within the target range and the diversity\n- experiments aiming to generate new molecules with a optimised property but similar to the original molecules (similarity-constrained Optimisation)\n\nMoreover, a comparison of the generation time of the proposed model with the other models would be useful. \n\nThe authors should not bold the best values only when they correspond to your model (Table3 GCPN, MolDQN, MARS, GraphDF achieve same QED 1st with the proposed model, MARS, GraphDF achieve same QED 2nd and 3rd with the proposed model) ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "ZdCiWkO9kwM", "reviewer": "Reviewer_gP5J", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a method for conditional generation of molecular structures, formulated as a joint latent variable generative model of the molecular structure and one or more properties. The latent variable are modelled using an energy based model, and the molecular structure is an LSTM neural network that generates a SELFIES string representation. Conditionaly sampling molecules with novel values of the desired properties is not likely to work well, since the generative model has not seen any training data in that region. To combat this issue, a procedure is proposed where additional data at the boundary is gradually added, and the model is finetuned.", "summary_of_the_review": "An interesting paper. Limited technical novelty. Fairly good results.\n", "strength_and_weaknesses": "Strenghts\n\nThe proposed method is fairly simple and seems to work well.\n\nWeaknesses\n\nThe sampling with gradual distribution shifting (SGDS) requires that the properties of interest can be easily computed on the fly. This could be a limiting factor, if the computing the properties cannot be computed quickly. It would be nice with a more detailed discussion of how this (a few comments are included in the appendix).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OVhJkhgruu", "reviewer": "Reviewer_vKUH", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a model and sampling method for de novo drug design. The model consists of\n-  multi-layer perceptron that outputs an 'energy' as a function of a latent vector\n- LSTM that generates SELFIES strings conditional on the latent vector\n- multi-layer perceptrons to predict molecule properties given latent vectors\n\nThe model is first trained on a large set of molecules and their properties. Then, it is fine-tuned to generate molecules with desired properties using a new method that the authors name 'sampling with gradual distributional shifting'.\n", "summary_of_the_review": "The paper introduces a new model architecture and sampling strategy for de novo drug design. It gets good results on some metrics, compared to the presented baselines, but I am not sure if the work represents a practically useful advance. The authors do not adequately reference and describe prior work on the same problem.", "strength_and_weaknesses": "- As far as I know, this model architecture has not previously been used for this problem, and it gets good results on the presented metrics compared to the presented baselines.\n-  Related work is not correctly described. In section 2 '... learn a surrogate function to predict properties, and then use Bayesian optimization' is not an accurate description of the molecular swarm optimization in Winter et al. (2019) nor the simulated design-synthesis-test cycles in Segler et al. (2018).  In 3.2 and 4.2 the authors fail to note that Segler et al. (2018) also generated a very high percentage of valid SMILES strings using an LSTM. It would be a good baseline to include.\n- Langevin dynamics is notoriously slow to sample multimodal distributions. Is this a problem for MolEBM? \n- What happens to validity, novelty, and uniqueness of generated molecules after SGDS?\n- Could the model be evaluated on the GuacaMol benchmarks (Brown et al., 2019)?\n- Good scores on the presented metrics may not indicate usefulness for de novo drug design. For example, Figure 2 shows that several of the molecules corresponding to bold numbers in table 4 are not reasonable drug candidates (large cycles etc.).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject