11institutetext: 1 Université de Lyon, Lyon 2, ERIC UR 3083, 5 avenue Pierre
Mendès France, F69676 Bron Cedex, France
11email<EMAIL_ADDRESS>
11email<EMAIL_ADDRESS>
11email<EMAIL_ADDRESS>

# Dirichlet-Survival Process: Scalable Inference of Topic-Dependent Diffusion
Networks

Gaël Poux-Médard 11 0000-0002-0103-8778 Julien Velcin 11 0000-0002-2262-045X
Sabine Loudcher 11 0000-0002-0494-0169

###### Abstract

Information spread on networks can be efficiently modeled by considering three
features: documents’ content, time of publication relative to other
publications, and position of the spreader in the network. Most previous works
model up to two of those jointly, or rely on heavily parametric approaches.
Building on recent Dirichlet-Point processes literature, we introduce the
Houston (Hidden Online User-Topic Network) model, that jointly considers all
those features in a non-parametric unsupervised framework. It infers dynamic
topic-dependent underlying diffusion networks in a continuous-time setting
along with said topics. It is unsupervised; it considers an unlabeled stream
of triplets shaped as (time of publication, information’s content, spreading
entity) as input data. Online inference is conducted using a sequential Monte-
Carlo algorithm that scales linearly with the size of the dataset. Our
approach yields consequent improvements over existing baselines on both
cluster recovery and subnetworks inference tasks.

###### Keywords:

Spreading process Network Inference Clustering Bayesian Nonparametrics

##  1 Introduction

###  1.1 Overview of the contribution

Over the last decades, information spread patterns have become more and more
complicated. The volume of data that flows on social networks keeps increasing
every day that passes, and results in complex diffusion processes that can be
described by many factors. However, recent advances suggest that documents
complex diffusion processes can be efficiently modeled considering only three
variables: their publication date (when), the publisher (who) and their
semantic content (what). The idea of considering these three factors is not
novel. However, most of the models that tackle diffusion problems tend to
consider up to two of these, but seldom the three parameters.

We introduce the Houston model, that tackles the problem by jointly inferring
clusters of textual documents spreading online and the subnetworks they spread
on. Our method builds on recent Dirichlet-Point processes advances [9, 18, 23,
24]. To the best of our knowledge, it is the first model that considers
semantic content, publication dynamics and the network of spreading documents
in an online, non-parametric and unsupervised way.

Figure 1: From a stream of textual documents, we model the underlying topic-
dependent diffusion subnetworks. Inference is unsupervised, non-parametric and
conducted online, meaning data is processed sequentially. Results in the
bottom row come from the application of our method to the Memetracker dataset
[17]. Nodes colors represent traditional medias (red) and blog (blue).

###  1.2 Related works

It has been underlined on several occasions that efficiently modeling
information diffusion involves accounting for the network’s structure [16,
22], publication times [8, 12] and documents’ content [10, 15]. Some
approaches consider sequentially all three factors. Typically, they first
infer topics based on documents content, and only then they use this
information to infer the latent diffusion subnetworks [7, 10, 15, 26, 28, 29].
The work the closest to ours [4] is, to our knowledge, the only one that
jointly models documents’ content, dynamics and structure. It develops an
unsupervised topic-dependent network inference method. The approach breaks
down the topic-aware diffusion into two factors: each node is assumed to have
a given sensitivity to a topic, and a certain authority on them. Given this
assumption, the authors develop a parametric prior on the probability for a
diffusion cascade to belong to a given topic. The textual content (or side
information) is then accounted for using a homogeneous Poisson textual model
[19], combined with the above prior. The model is optimized using an EM
algorithm. However, the optimization algorithm is not designed for online
optimization –data cannot be added sequentially–, and topics optimization is
parametric –the number of topics must be provided.

##  2 Model

###  2.1 Background

To answer these limitations, we build a Dirichlet-Survival process that can be
used as a non-parametric Bayesian prior for online inference. The Dirichlet-
Survival prior is created by merging Dirichlet processes with Point processes.
The method has been explored by combining Hawkes processes to several variants
of Dirichlet processes (hierarchical [18], mixed membership [27], powered
[23], multivariate [24]). However, no work considered the combination with
other point processes than the Hawkes process. Our approach using Survival
analysis explores this new connection; it allows us to design an optimization
algorithm (Sequential Monte Carlo) for online non-parametric topics-aware
diffusion subnetworks inference (the number of topics/subnetworks does not
have to be chosen in advance).

In [14], the authors show that a large part of the literature on underlying
diffusion network inference [8, 12, 13, 14, 21, 28] can be expressed as
special cases of a counting point process. The method allows to infer dynamic
underlying diffusion networks using convex optimization tools.

###  2.2 Dirichlet process and Survival analysis

####  2.2.1 Dirichlet process

The Dirichlet process is used as a non-parametric prior distribution over
clusters in many clustering algorithms. It can be written as follows:

$P(s_{i}=k|\\{s_{m}\\}_{m=1,...,n-1},\alpha_{0})=\begin{cases}\frac{N_{k}}{\alpha_{0}+\sum_{k}^{K}N_{k}}\text{
if k = 1, ..., K}\\\ \frac{\alpha_{0}}{\alpha_{0}+\sum_{k}^{K}N_{k}}\text{ if
k = K+1}\end{cases}$ (1)

where $s_{i}$ is a variable that represents the cluster of the $i^{th}$
observation, $N_{k}=|\\{s_{i}|s_{i}=k\\}_{i=1,...,n-1}|$ the population of
cluster $k$, $K$ the total number of non-empty clusters and $\alpha_{0}$ a
concentration hyper-parameter. The choice of $K+1$ means a new cluster is
opened and $K$ in increased by 1. Note that references [23, 24] use the
powered version of this process [25].

####  2.2.2 Network inference model

The edges of topic-dependent networks are inferred using the NetRate model
[12], which is part of a broad literature on underlying spreading networks
inference [10, 12, 13, 14, 28]. In particular in [14], the authors demonstrate
that all these models can be expressed as special cases of a counting point
process. These processes take a collection of independent timestamped
diffusion cascades $\vec{c}=\\{(u_{i}^{c},t_{i}^{c})\\}_{i}$ as input, where
$u_{i}^{c}$ is the node on which the $i^{th}$ event occurred and $t_{i}^{c}$
the time at which it happened in cascade $c$. The process is entirely
characterized by a hazard function
$H(t_{i}^{c}|t_{j}^{c},\alpha_{u_{j}^{c},u_{i}^{c}})$, which is the
instantaneous infection rate of $u_{i}^{c}$ at time $t_{i}^{c}$ by $u_{j}^{c}$
previously infected at time $t_{j}^{c}$, given it infection did not happen
before $t_{i}$. In this paper, we express the hazard function as a constant
${H(t|t_{i},\alpha)=\alpha}$, implying by definition that the probability of
an event non happening before a time $t$ given $t_{i}$ decays exponentially as
$e^{-\alpha(t-t_{i})}$. The associated convex likelihood of $\alpha$ can be
found in [12] (Eq.7).

###  2.3 Dirichlet-Survival process

In [9] the authors define the Dirichlet-Hawkes process by replacing the
integer counts in Eq.1 by the intensity of a Hawkes process. It can be
interpreted as replacing integers counts in Dirichlet Processes by non-integer
time-dependent counts, encoded by the intensity of the point process. Here, we
consider the hazard rate of the NetRate model instead to account for networks
structure. Each node is associated to its own temporal point process, and
counts are replaced by the number of times any neighbour has been infected,
weighted according to time and to edges strength. Using the methodology
introduced in [9] and substituting the Hawkes process by the hazard rate of a
survival model [14], we make a yet unexplored bridge between Dirichlet
processes and Survival analysis. We remind that [14] reformulates the work of
[8, 12, 13, 28] in terms of Survival analysis and associated counting
processes; we settled on using NetRate here, but any of these models would fit
as well in our approach. The point process nature of survival analysis
discussed in [14] makes this extension sound with respect to previous works on
Dirichlet-Point processes [9, 18, 23, 27].

Let $\mathbf{A}^{(k)}$ be the adjacency matrix of the subnetwork associated to
cluster $k$, whose entries are $\alpha_{i,j}^{(k)}$. We define
$(u_{j}^{c},t_{j}^{c})^{(k)}$ as an event of cascade $c$ observed on node
$u_{j}$ at $t_{j}$ attributed to subnetwork $A^{(k)}$. We write the history of
events in cascade $c$ attributed to the subnetwork $k$ as
$\mathcal{H}_{i,c}^{(k)}=\\{(u_{j}^{c},t_{j}^{c})^{(k)}\\}_{j:t_{j}<t_{i}}$.
We note $\mathbf{\mathcal{H}_{i,c}}=\\{\mathcal{H}_{i,c}^{(k)}\\}_{k}$ and
$\mathbf{A}=\\{\mathbf{A}^{(k)}\\}_{k}$. We consider a new event from cascade
$c$ observed on node $u_{i}^{c}$ at time $t_{i}^{c}$. At this point, the new
event is not yet associated to any subnetwork. We write the Dirichlet-Survival
prior probability for the new event to belong to subnetwork $k$:

$P(s_{i}=k|\mathbf{\mathcal{H}_{i,c}},\mathbf{A},\mathbf{\lambda_{0}})=\begin{cases}\frac{\lambda_{0}^{(k)}+\sum_{\mathcal{H}_{i,c}^{(k)}}H(t_{i}^{c}|t_{j}^{c},\alpha_{u_{j},u_{i}}^{(k)})}{\lambda_{0}^{(K+1)}+\sum_{k}^{K}\lambda_{0}^{(k)}+\sum_{\mathcal{H}_{i,c}^{(k)}}H(t_{i}^{c}|t_{j}^{c},\alpha_{u_{j},u_{i}}^{(k)})}\text{
if k = 1, ..., K}\\\
\frac{\lambda_{0}^{(K+1)}}{\lambda_{0}^{(K+1)}+\sum_{k}^{K}\lambda_{0}^{(k)}+\sum_{\mathcal{H}_{i,c}^{(k)}}H(t_{i}^{c}|t_{j}^{c},\alpha_{u_{j},u_{i}}^{(k)})}\text{
if k = K+1}\end{cases}$ (2)

We introduced a new parameter
$\mathbf{\lambda_{0}}=\\{\lambda_{0}^{(k)}\\}_{k=1,...,K+1}$, which translates
the probability for a new observation not to have been triggered by any
neighbour. It represents the probability that an event of cluster $k$ is
exogenous [15, 20].

The Dirichlet-Survival prior is coupled to a sequential language model. For
simplicity, we consider the bag-of-words Dirichlet-Multinomial model, as in
[9, 18, 23]; note that more refined sequential language models are also fit to
our approach (Dynamic Topic Model [6], online LDA [3], online PLSA [5], etc.).

The input data is a stream of events. Each event takes the form of a triplet
$(u_{i}^{c},t_{i}^{c},v_{i}^{c})$, where $c$ is the cascade an event has been
observed in, $u_{i}^{c}$ is the node corresponding to the event, $t_{i}^{c}$
is its publication time, and $v_{i}^{c}$ represents its textual content (e.G.
words in a tweet or in a news article). By combining the Dirichlet-Survival
prior to the textual likelihood, we get the posterior distribution of the
$i^{th}$ observation belonging to cluster (or subnetwork) $k$ as:

$P(s_{i}|v_{i}^{c},\mathbf{N},\mathbf{\mathcal{H}_{i,c}},\mathbf{A},\theta_{0},\mathbf{\lambda_{0}})\propto\underbrace{P(v_{i}^{c}|s_{i},\mathbf{N},\theta_{0})}_{\text{Dirichlet-
Multinomial}}\times\underbrace{P(s_{i}|\mathbf{\mathcal{H}_{i,c}},\mathbf{A},\mathbf{\lambda_{0}})}_{\text{Dirichlet-
Survival prior (Eq.\ref{eq-DSP}})}$ (3)

where $\vec{N}$ contains the words counts within each cluster, $v_{i}^{c}$
contains the words count in document $i$, and $\theta_{0}$ the concentration
parameter of the model.

Finally, inference is conducted using a Sequential Monte Carlo algorithm
similar to [9, 18, 23]. We perform several parallel runs on the same data
stream. Within each run, each new observation in the stream is assigned to a
cluster according to Eq. 3. The adjacency matrix $\mathbf{A}$ is then updated
by optimizing the convex likelihood associated to the NetRate point process
(Eq. 7 in [12]). Finally, we compute the likelihood of the language model for
each run; runs that have a likelihood lesser than a threshold are discarded
and replaced by more likely ones. The process is repeated until the end of the
data stream. According to this algorithm, Eq. 1, and introducing a cutoff on
the exponential hazard function (observations older than a time $t_{old}$ are
ignored), the optimization runs in $\mathcal{O}(N_{obs}N_{runs}(N_{nodes}+K))$
where $N_{part}$ is the number of particles, $N_{nodes}$ is the maximum
network size and $K$ the number of clusters (typically $N_{runs}\ll K\ll
N_{nodes}$). Inference hence scales linearly with the size of the dataset.

We point out that the Dirichlet-Survival process is not about refining complex
diffusion models such as [4, 7, 26]. Instead, it introduces a different angle
for tackling content-aware diffusion problems. This new angle allows for
unsupervised, non-parametric and online inference.

##  3 Experiments

###  3.1 Data and experimental setup

All data, codes and results are available in open access
111https://github.com/GaelPouxMedard/HOUsToN. We consider 3 different network
types of 500 nodes each: power-law (PL) [2], random Erdös-Renye (ER) [11] and
a real network of hyperlinks between political blogs (Blogs) [1]. From each
network, we sample 5 subnetworks of 250 nodes and assign random weights
$\alpha$ between 0 and 1 to their edges. Each of the generated subnetworks is
used to propagate one given cluster of information. We then simulate infection
cascades on each subnetwork according to the exponential NetRate model.
Finally, we associate 5 words drawn from a vocabulary of size 100 to each so-
generated event according to its associated subnetwork (or cluster). We
generate a total of 55,000 events
$\\{(u_{i}^{c},t_{i}^{c},v_{i}^{c})\\}_{i,c}$ for each network.

Our hyperparameters are $\theta_{0}=0.1$ and $\lambda_{0}^{(k)}=0.001\ \forall
k$. The SMC algorithm considers 4 parallel runs. We consider a constant hazard
rate $H(t_{i}|t_{j},\alpha_{j,i})=\alpha_{j,i}$, so the probability of a new
event not happening decays exponentially with time.

Table 1: Results on clusters (NMI, ARI) and edges (AUC-ROC, F1, MAE) retrieval.  |  |  Houston |  NRxDM |  DHP |  NetRate |
---|---|---|---|---|---|---
PL |  NMI |  0.809 |  0.669 |  0.449 |  - |
ARI |  0.688 |  0.330 |  0.063 |  - |
2-6 |  AUC-ROC |  0.807 |  0.719 |  - |  0.731 |
|  F1 |  0.199 |  0.106 |  - |  0.005 |
|  MAE |  0.267 |  0.338 |  - |  0.460 |
ER |  NMI |  0.787 |  0.711 |  0.638 |  - |
ARI |  0.631 |  0.488 |  0.411 |  - |
2-6 |  AUC-ROC |  0.849 |  0.800 |  - |  0.659 |
|  F1 |  0.263 |  0.176 |  - |  0.005 |
|  MAE |  0.229 |  0.278 |  - |  0.481 |
Blogs |  NMI |  0.750 |  0.668 |  0.372 |  - |
ARI |  0.609 |  0.365 |  0.023 |  - |
2-6 |  AUC-ROC |  0.701 |  0.613 |  - |  0.710 |
|  F1 |  0.168 |  0.087 |  - |  0.005 |
|  MAE |  0.374 |  0.444 |  - |  0.499 |

###  3.2 Results

We compare to 3 similar baselines used as ablation tests: Dirichlet-Hawkes
process (DHP) [9] clusters textual data by using temporal dynamics, and does
not consider structure; NetRate [12] infers a dynamic network based on
observed cascades without considering their content; NetRate x Dirichlet-
Multinomial (NRxDM) first uses textual information to infer clusters, and only
then infers the underlying subnetwork for each cluster, in the same fashion as
[10, 15]. When applicable, we evaluate on a classification task (scores NMI
and ARI with respect to the clusters used for data generation) and a network
inference task (AUC-ROC, F1 and MAE on the true edges, same metrics as in
[12]).

We see in Table 1 that Houston consistently outperforms methods that do not
consider jointly text, time and structure of the network. To summarize, NRxDM
only considers textual information to build clusters, making the network
inference miss a great deal of temporal and structural information. DHP
considers textual information and temporal dynamics, but misses the structural
information. NetRate does not consider textual data and infers the network
based on temporal dynamics only. Houston bridges the gap between these models,
by making a joint use of textual, temporal and structural information.

As an illustration of what Dirichlet-Processes can yield on real-world data,
we draft its application to the Memetracker dataset [17] in Fig. 1 (bottom).
We retrieve the diffusion network associated to meme clusters and observe
diverse spreading dynamics. Topics spread in distinct parts of the global
network, and mostly do so through a reduced set of densely connected nodes, as
shown in [13].

##  4 Conclusion

In this paper, we propose the Dirichlet-Survival process as an alternative way
to jointly model textual, temporal and structural information in spreading
processes. Ablation tests demonstrate the relevance of the proposed approach.
As a prior, the Dirichlet-Survival process can add a dynamic network dimension
to any sequential Bayesian model; it could be coupled to models that account
for any type of clustering (e.g. images, time series, labels), or simply more
refined language models. Its introduction opens new perspectives on
traditional machine learning problems, including topic-dependent spreading
processes on networks.

## References

  * [1] Adamic, L.A., Glance, N.: The political blogosphere and the 2004 u.s. election: Divided they blog. In: Proceedings of the 3rd International Workshop on Link Discovery. p. 36–43. LinkKDD ’05, Association for Computing Machinery, New York, NY, USA (2005)
  * [2] Albert, R., Barabási, A.L.: Statistical mechanics of complex networks. Rev. Mod. Phys. 74, 47–97 (Jan 2002)
  * [3] AlSumait, L., Barbará, D., Domeniconi, C.: On-line lda: Adaptive topic models for mining text streams with applications to topic detection and tracking pp. 3–12 (2008). https://doi.org/10.1109/ICDM.2008.140
  * [4] Barbieri, N., Manco, G., Ritacco, E.: Survival factorization on diffusion networks. Machine Learning and Knowledge Discovery in Databases pp. 684–700 (2017). https://doi.org/10.1007/978-3-319-71249-9_41
  * [5] Bassiou, N.K., Kotropoulos, C.L.: Online plsa: Batch updating techniques including out-of-vocabulary words. IEEE Transactions on Neural Networks and Learning Systems 25(11), 1953–1966 (2014). https://doi.org/10.1109/TNNLS.2014.2299806
  * [6] Blei, D.M., Lafferty, J.D.: Dynamic topic models. In: Proceedings of the 23rd International Conference on Machine Learning. p. 113–120. ICML ’06, Association for Computing Machinery, New York, NY, USA (2006). https://doi.org/10.1145/1143844.1143859
  * [7] Choudhari, J., Dasgupta, A., Bhattacharya, I., Bedathur, S.: Discovering topical interactions in text-based cascades using hidden markov hawkes processes pp. 923–928 (2018). https://doi.org/10.1109/ICDM.2018.00112
  * [8] Du, N., Song, L., Smola, A., Yuan, M.: Learning networks of heterogeneous influence. NIPS 4, 2780–2788 (01 2012)
  * [9] Du, N., Farajtabar, M., Ahmed, A., Smola, A., Song, L.: Dirichlet-hawkes processes with applications to clustering continuous-time document streams. 21th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (2015). https://doi.org/10.1145/2783258.2783411
  * [10] Du, N., Song, L., Woo, H., Zha, H.: Uncover topic-sensitive information diffusion networks. In: Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statistics, AISTATS. JMLR Workshop and Conference Proceedings, vol. 31, pp. 229–237. JMLR.org (2013)
  * [11] Erdős, P., Rényi, A.: On the evolution of random graphs. In: PUBLICATION OF THE MATHEMATICAL INSTITUTE OF THE HUNGARIAN ACADEMY OF SCIENCES. pp. 17–61 (1960)
  * [12] Gomez-Rodriguez, M., Balduzzi, D., Schölkopf, B.: Uncovering the temporal dynamics of diffusion networks. In: ICML. p. 561–568 (2011)
  * [13] Gomez-Rodriguez, M., Leskovec, J., Schoelkopf, B.: Structure and dynamics of information pathways in online media. WSDM (2013)
  * [14] Gomez-Rodriguez, M., Leskovec, J., Schölkopf, B.: Modeling information propagation with survival theory. In: ICML. vol. 28, p. III–666–III–674 (2013)
  * [15] He, X., Rekatsinas, T., Foulds, J.R., Getoor, L., Liu, Y.: Hawkestopic: A joint model for network inference and topic modeling from text-based cascades. In: ICML (2015)
  * [16] Larremore, D., Carpenter, M., Ott, E., Restrepo, J.: Statistical properties of avalanches in networks. Physical Review E 85 (04 2012). https://doi.org/10.1103/PhysRevE.85.066131
  * [17] Leskovec, J., Backstrom, L., Kleinberg, J.: Meme-tracking and the dynamics of the news cycle. In: Proceedings of the 15th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. p. 497–506. KDD ’09, Association for Computing Machinery, New York, NY, USA (2009). https://doi.org/10.1145/1557019.1557077
  * [18] Mavroforakis, C., Valera, I., Gomez-Rodriguez, M.: Modeling the dynamics of learning activity on the web. In: Proceedings of the 26th International Conference on World Wide Web. p. 1421–1430. WWW ’17 (2017)
  * [19] Mei, Q., Fang, H., Zhai, C.: A study of poisson query generation model for information retrieval p. 319–326 (2007). https://doi.org/10.1145/1277741.1277797, https://doi.org/10.1145/1277741.1277797
  * [20] Myers, S.A., Zhu, C., Leskovec, J.: Information diffusion and external influence in networks. In: Proceedings of the 18th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. p. 33–41. KDD ’12, Association for Computing Machinery, New York, NY, USA (2012). https://doi.org/10.1145/2339530.2339540
  * [21] Nickel, M., Le, M.: Modeling sparse information diffusion at scale via lazy multivariate hawkes processes. In: Proceedings of the Web Conference 2021. p. 706–717. WWW ’21, Association for Computing Machinery, New York, NY, USA (2021). https://doi.org/10.1145/3442381.3450094, https://doi.org/10.1145/3442381.3450094
  * [22] Poux-Médard, G., Pastor-Satorras, R., Castellano, C.: Influential spreaders for recurrent epidemics on networks. Phys. Rev. Research 2, 023332 (Jun 2020). https://doi.org/10.1103/PhysRevResearch.2.023332
  * [23] Poux-Médard, G., Velcin, J., Loudcher, S.: Powered hawkes-dirichlet process: Challenging textual clustering using a flexible temporal prior. 2021 IEEE International Conference on Data Mining (ICDM) pp. 509–518 (2021)
  * [24] Poux-Médard, G., Velcin, J., Loudcher, S.: Multivariate powered dirichlet-hawkes process. ECIR (2023)
  * [25] Poux-Médard, G., Velcin, J., Loudcher, S.: Powered dirichlet process for controlling the importance of ”rich-get-richer” prior assumptions in bayesian clustering. ArXiv (2021)
  * [26] Suny, P., Li, J., Mao, Y., Zhang, R., Wang, L.: Inferring multiplex diffusion network via multivariate marked hawkes process. ArXiv abs/1809.07688 (2018)
  * [27] Tan, X., Rao, V.A., Neville, J.: The indian buffet hawkes process to model evolving latent influences. In: UAI (2018)
  * [28] Wang, L., Ermon, S., Hopcroft, J.E.: Feature-enhanced probabilistic models for diffusion network inference. In: Machine Learning and Knowledge Discovery in Databases. pp. 499–514. Springer Berlin Heidelberg, Berlin, Heidelberg (2012)
  * [29] Yang, S.H., Zha, H.: Mixture of mutually exciting processes for viral diffusion. In: Dasgupta, S., McAllester, D. (eds.) Proceedings of the 30th International Conference on Machine Learning. Proceedings of Machine Learning Research, vol. 28, pp. 1–9 (2013)


# Molecule Design by Latent Space Energy-Based Modeling and Gradual
Distribution Shifting

Deqian Kong Equal contribution Department of Statistics
University of California, Los Angeles  Bo Pang† Salesforce Research  Tian Han
Department of Computer Science
Stevens Institute of Technology  Ying Nian Wu  Department of Statistics
University of California, Los Angeles

###### Abstract

Generation of molecules with desired chemical and biological properties such
as high drug-likeness, high binding affinity to target proteins, is critical
for 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. Conditional on the
latent vector, the molecule and its properties are modeled by a molecule
generation model and a property regression model respectively. To search for
molecules with desired properties, we propose a sampling with gradual
distribution shifting (SGDS) algorithm, so that after learning the model
initially on the training data of existing molecules and their properties, the
proposed algorithm gradually shifts the model distribution towards the region
supported by molecules with desired values of properties. Our experiments show
that our method achieves very strong performances on various molecule design
tasks. The code and checkpoints are available at
https://github.com/deqiankong/SGDS.

##  1 Introduction

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. It is challenging to directly optimize
or search over the drug-like molecule space since it is discrete and enormous,
with an estimated size on the order of $10^{33}$ [Polishchuk et al., 2013].

Recently, a large body of work attempts to tackle this problem. 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 such as 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 and Kipf, 2018, Zhou et al.,
2019, Shi et al., 2020, Luo et al., 2021]. A number of other methods have been
proposed to optimize molecular properties with genetic algorithms [Nigam et
al., 2020], particle-swarm algorithms [Winter et al., 2019], and specialized
MCMC methods [Xie et al., 2021].

In this work, we propose a method along the first line mentioned above, by
learning a probabilistic latent space generative model of molecules and
optimizing molecular properties in the latent space. 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. Thus,
instead of assuming a simple Gaussian distribution in the latent space as in
prior work [Gómez-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 and Bengio, 2016, Xie et al., 2016, Kumar et al., 2019, Nijkamp
et al., 2019, Du and Mordatch, 2019, Grathwohl et al., 2019, Finn et al.,
2016] in latent space. This leads to a latent space energy-based model (LSEBM)
as studied in Pang et al. [2020], Nie et al. [2021], where LSEBM has been
shown to model the distributions of natural images and text well.

$x$$z$$y$$p_{\gamma}(y|z)$$p_{\beta}(x|z)$ConditionalGeneration$p_{\alpha}(z)$
(a) Single-Objective Optimization.

$x$$z$$\cdots$$y$$y_{1}$$y$$y_{m}$$p_{\beta}(x|z)$$p_{\gamma_{m}}(y_{m}|z)$$p_{\gamma_{1}}(y_{1}|z)$$p_{\alpha}(z)$
(b) Multi-Objective Optimization.

(c) Shift of distribution of $y$.

Figure 1: An illustration of the joint distribution of molecule with its
single property (a) or multiple properties (b). $x$ represents a molecule, $z$
is the latent vector, $y$ is a molecular property of interest,
$\\{y_{j}\\}_{j=1}^{m}$ indicates $m$ properties. (c) illustrates the shift of
the distribution of a single property in sampling with gradual distribution
shifting (SGDS).

Going beyond existing latent space energy-based models Pang et al. [2020], Nie
et al. [2021], our work makes two innovations:

First, given our goal of property optimization, we learn a joint distribution
of molecules and their properties. Our model consists of (1) an energy-based
model (EBM) in a low-dimensional continuous latent space, (2) a molecule
generation model that generates molecule given the latent vector, and (3) a
property regression model that predicts the value of the property given the
latent vector. See Figure 1(a) for an illustration of the model. We first
learn the initial model on the training data that consist of existing
molecules and their properties. All three components in our model are learned
jointly by an approximate maximum likelihood algorithm.

Second, and more importantly, we propose a sampling with gradual distribution
shifting (SGDS) method for molecule design. We first sample molecules and
their property values from the initial model learned on the training data
mentioned above. Then we gradually shift the joint distribution towards the
region supported by molecules with high property values. Specifically, our
method iterates the following steps. (1) Shift the sampled property values by
a small constant towards the desired target value. (2) Generate molecules
given the shifted property values. (3) Obtain the ground-truth property values
of the generated molecules by querying the software. (4) Update the model
parameters by learning from the generated molecules and their ground-truth
property values. Because of the flexibility of the latent space energy-based
model, the model can be updated to account for the change of the joint
distribution of the generated molecules and their ground-truth property values
in the gradual shifting process. Figure 1(c) illustrates the shifting of the
distribution of the property values of the generated molecules.

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 (see Figure 1(c)). We
can then simultaneously shift the values of the multiple properties for multi-
objective optimization.

We evaluate our method in various settings including single-objective and
multi-objective optimization. Our method outperforms prior methods by
significant margins.

In summary, our contributions are as follows:

  * •

We propose to learn a latent space energy-based model for the joint
distribution of molecules and their properties.

  * •

We develop a sampling with gradual distribution shifting method, which enables
us to extrapolate the data distribution and sample from the region supported
by molecules with high property values.

  * •

Our methods are versatile enough to be extended to optimizing multiple
properties simultaneously.

  * •

Our model achieves state-of-the-art performances on a range of molecule
optimization tasks.

Caveat. As in most existing work on molecule design, we assume that the value
of a property of interest of a given molecule can be obtained by querying an
existing software. There are two research problems in this endeavor. (1)
Developing software that can output biologically or chemically accurate value
of the property for an input molecule. (2) Developing method that can optimize
the property values output by a given software. While problem (1) is
critically important, our work is exclusively about problem (2). We duly
acknowledge that existing software may need much improvements. Meanwhile our
method can be readily applied to the improved versions of software.

##  2 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 and Welling, 2014] and generative
adversarial network (GAN) [Goodfellow et al., 2014] etc. to molecule data
[Gómez-Bombarelli et al., 2018, Jin et al., 2018, De Cao and Kipf, 2018, Honda
et al., 2019, Madhawa et al., 2019, Shi et al., 2020, Zang and 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]. After learning
continuous representations for molecules, they are further able to optimize
using different methods. [Segler et al., 2018] proposes to optimize by
simulating design-synthesis-test cycles. [Gómez-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. However, the performance of this latent optimization is not
satisfactory due to three major issues. First, it is difficult to 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 ability to explore the targets out of
the distribution is limited [Brown et al., 2019, Huang et al., 2021]. 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 the 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. In this paper,
we propose an encoder-free model in both training and optimization to learn
the joint distribution of molecules and properties. We then design an
efficient algorithm to shift the learned distribution gradually.

Optimization with Reinforcement Learning and Evolutionary Algorithms.
Reinforcement learning (RL) based methods directly optimize and generate
molecules in an explicit data space [You et al., 2018, Zhou et al., 2019, Jin
et al., 2020, Gottipati et al., 2020]. By formulating the property design as a
discrete optimization task, they can modify the molecular substructures guided
by an oracle 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-walk search behavior in the discrete space. Evolutionary algorithms
(EA) also formulate the optimization in a discrete manner [Nigam et al., 2020,
Jensen, 2019, Xie et al., 2021, Fu et al., 2021a, b]. By leveraging carefully-
crafted combinatorial search algorithms, they can search the molecule graph
space in a flexible and efficient way. However, the design of those algorithms
is non-trivial and domain specific.

##  3 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 $x=(x^{(1)},...,x^{(t)},...,x^{(T)})$ be a molecule
string encoded in SELFIES, where $x^{(t)}\in{\mathcal{V}}$ is the $t$-th
character and ${\mathcal{V}}$ is the vocabulary.

Suppose $y\in{\mathbb{R}}$ represents a molecular property of interest. Then
the problem we attempt to tackle is to optimize $x$ such that its property
$y=y^{*}$ where $y^{*}$ is some desirable value for $y$. We take a
probabilistic approach and treat the optimization problem as a sampling
problem, that is,

$x^{*}\sim p(x|y=y^{*}).$ (1)

This is a single-objective optimization problem since only one property is
targeted. In real-world drug design settings, we are more likely to optimize
multiple properties simultaneously, that is, multi-objective optimization.
Suppose we optimize for $\\{y_{j}\in{\mathbb{R}}\\}_{j=1}^{m}$, then our task
is to sample

$x^{*}\sim p(x|y_{1}=y_{1}^{*},...,y_{m}=y_{m}^{*}).$ (2)

To address these problems, we propose a solution within a unified
probabilistic framework. As a first step, we need to model or approximate the
data distribution of molecules and their properties, $p_{\rm data}(x,y)$. To
this end, we recruit latent space energy-based model (LSEBM) [Pang et al.,
2020, Nie et al., 2021] to model the molecule and properties. LSEBM assumes
that a latent vector $z\in{\mathbb{R}}^{d}$ in a low dimensional latent space
follows an energy-based prior model $p(z)$. Conditional on $z$, the molecule
$x$ and the property $y$ are independent, so that the joint distribution
$p(x,y,z)$ can be factorized as $p(z)p(x|z)p(y|z)$, which leads to
$p(x,y)=\int p(z)p(x|z)p(y|z)dz$ as an approximation to $p_{\rm data}(x,y)$.
The latent space energy-based prior model $p(z)$, the molecule generation
model $p(x|z)$, and the property regression model $p(y|z)$ can be jointly
learned by an approximate maximum likelihood algorithm (see §3.3 and Algorithm
1). LSEBM within the context of molecule data is presented in §3.2.

For the purpose of property optimization, we are required to generate
molecules $x$ with some desirable property $y^{*}$. Rather than direct
optimization in the molecule space, we choose to optimize $z$ in the latent
space. We first consider the single-objective optimization problem (Equation
1). With the learned model, we propose to optimize $x$ given $y=y^{*}$ by
ancestral sampling,

$\displaystyle z^{*}\sim p(z|y=y^{*}),\quad x^{*}\sim p(x|z=z^{*}).$ (3)

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. To address this problem, we propose a Sampling with
Gradual Distribution Shifting (SGDS) approach where we gradually shift the
learned distribution to a region where it is supported by high property values
(see §3.4 and Algorithm 2).

Our model is designed to be versatile such that it admits straightforward
extension to multi-objective optimization. To optimize $x$ given
$\\{y_{j}=y_{j}^{*}\\}_{j=1}^{m}$, we can simply augment the joint
distribution with more regression models, i.e.,
$p(x,z,y_{1},...,y_{m})=p(z)p(x|z)\prod_{j=1}^{m}p(y_{j}|z)$. The optimization
procedure follows the same SGDS approach. See §3.5 for more details on multi-
objective optimization.

###  3.2 Joint Distribution of Molecule and Molecular Property

Suppose $x=(x^{(1)},...,x^{(t)},...,x^{(T)})$ is a molecule string in SELFIES,
$y\in\mathbb{R}$ is the target property of interest, and $z\in\mathbb{R}^{d}$
is the latent vector. Consider the following model,

$\displaystyle z\sim p_{\alpha}(z),\quad[x\mid z]\sim
p_{\beta}(x|z),\quad[y\mid z]\sim p_{\gamma}(y|z),$ (4)

where $p_{\alpha}(z)$ is a prior model with parameters $\alpha$,
$p_{\beta}(x|z)$ is a molecule generation model with parameters $\beta$, and
$p_{\gamma}(y|z)$ is a property regression model with parameter $\gamma$. In
VAE [Kingma and Welling, 2014], the prior is simply assumed to be an isotropic
Gaussian distribution. In our model, $p_{\alpha}(z)$ is formulated as a
learnable energy-based model,

$\displaystyle p_{\alpha}(z)=\frac{1}{Z(\alpha)}\exp(f_{\alpha}(z))p_{0}(z),$
(5)

where $p_{0}(z)$ is a reference distribution, assumed to be isotropic Gaussian
as in VAE. $f_{\alpha}:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ is a scalar-
valued negative energy function and is parameterized by a small multi-layer
perceptron (MLP) with parameters $\alpha$.
$Z(\alpha)=\int\exp(f_{\alpha}(z))p_{0}(z)dz=\mathbb{E}_{p_{0}}[\exp(f_{\alpha}(z))]$
is the normalizing constant or partition function.

The molecule generation model $p_{\beta}(x|z)$ is a conditional autoregressive
model,

$\displaystyle
p_{\beta}(x|z)=\prod_{t=1}^{T}p_{\beta}(x^{(t)}|x^{(1)},...,x^{(t-1)},z)$ (6)

which is parameterized by a one-layer LSTM Hochreiter and Schmidhuber [1997]
with parameters $\beta$. Note that the latent vector $z$ controls every step
of the autoregressive model. It is worth pointing out the simplicity of the
molecule generation model of our method considering that those in prior work
involve complicated graph search algorithm or alternating generation of atoms
and bonds with multiple networks.

Given a molecule $x$, suppose $y$ is the chemical property of interest, such
as drug likeliness or protein binding affinity. The ground-truth property
value can be computed for an input $x$ via open-sourced software such as RDKit
[Landrum et al., 2013] and AutoDock-GPU [Santos-Martins et al., 2021]. We
assume that given $z$, $x$ and $y$ are conditionally independent, so that

$\displaystyle p_{\theta}(x,y,z)$
$\displaystyle=p_{\alpha}(z)p_{\beta}(x|z)p_{\gamma}(y|z),$ (7)

where $\theta=(\alpha,\beta,\gamma)$. We use the model $p_{\theta}(x,y)=\int
p_{\theta}(x,y,z)dz$ to approximate the data distribution
$p_{\mathrm{data}}(x,y)$. See Supplement for a detailed discussion.

The property regression model can be written as

$\displaystyle
p_{\gamma}(y|z)=\frac{1}{\sqrt{2\pi\sigma^{2}}}\exp\left(-\frac{1}{2\sigma^{2}}(y-s_{\gamma}(z))^{2}\right),$
(8)

where $s_{\gamma}(z)$ is a small MLP, with parameters $\gamma$, predicting $y$
based on the latent $z$. The variance $\sigma^{2}$ is set as a constant or a
hyperparameter in our work.

###  3.3 Learning Joint Distribution

Suppose we observe training examples $\\{(x_{i},y_{i}),i=1,...,n\\}$. The log-
likelihood function is $L(\theta)=\sum_{i=1}^{n}\log p_{\theta}(x_{i},y_{i})$.
The learning gradient can be calculated according to

$\displaystyle\nabla_{\theta}\log
p_{\theta}(x,y)=\mathbb{E}_{p_{\theta}(z|x,y)}\left[\nabla_{\theta}\log
p_{\theta}(x,y,z)\right]$
$\displaystyle=\mathbb{E}_{p_{\theta}(z|x,y)}\left[\nabla_{\theta}(\log
p_{\alpha}(z)+\log p_{\beta}(x|z)+\log p_{\gamma}(y|z))\right].$ (9)

For the prior model,

$\displaystyle\nabla_{\alpha}\log
p_{\alpha}(z)=\nabla_{\alpha}f_{\alpha}(z)-\mathbb{E}_{p_{\alpha}(z)}[\nabla_{\alpha}f_{\alpha}(z)].$
(10)

The learning gradient given an example $(x,y)$ is

$\displaystyle\delta_{\alpha}(x,y)=\nabla_{\alpha}\log p_{\theta}(x,y)$
$\displaystyle=\mathbb{E}_{p_{\theta}(z|x,y)}[\nabla_{\alpha}f_{\alpha}(z)]-\mathbb{E}_{p_{\alpha}(z)}[\nabla_{\alpha}f_{\alpha}(z)].$
(11)

Thus $\alpha$ is updated based on the difference between $z$ inferred from
empirical observation $(x,y)$, and $z$ sampled from the current prior.

For the molecule generation model,

$\displaystyle\begin{split}\delta_{\beta}(x,y)=\nabla_{\beta}\log
p_{\theta}(x,y)=\mathbb{E}_{p_{\theta}(z|x,y)}[\nabla_{\beta}\log
p_{\beta}(x|z)].\end{split}$ (12)

Similarly, for the property regression model,

$\displaystyle\begin{split}\delta_{\gamma}(x,y)=\nabla_{\gamma}\log
p_{\theta}(x,y)=\mathbb{E}_{p_{\theta}(z|x,y)}[\nabla_{\gamma}\log
p_{\gamma}(y|z)].\end{split}$ (13)

Estimating expectations in Equations 11, 12, and 13 requires MCMC sampling of
the prior model $p_{\alpha}(z)$ and the posterior distribution
$p_{\theta}(z|x,y)$. We recruit Langevin dynamics [Neal, 2011, Han et al.,
2017]. For a target distribution $\pi(z)$, the dynamics iterates

$\displaystyle
z_{\tau+1}=z_{\tau}+s\nabla_{z}\log\pi(z_{\tau})+\sqrt{2s}\epsilon_{\tau},$
(14)

where $\tau$ indexes the time step of the Langevin dynamics, $s$ is step size,
and $\epsilon_{\tau}\sim{\mathcal{N}}(0,I_{d})$ is the Gaussian white noise.
$\pi(z)$ can be either the prior $p_{\alpha}(z)$ or the posterior
$p_{\theta}(z|x,y)$. In either case, $\nabla_{z}\log\pi(z)$ can be efficiently
computed by back-propagation.

We initialize $z_{0}\sim{\mathcal{N}}(0,I_{d})$, and we run $\Gamma$ steps of
Langevin dynamics (e.g. $\Gamma=20$) to approximately sample from the prior
and the posterior distributions. The resulting learning algorithm is an
approximate maximum likelihood learning algorithm. See [Pang et al., 2020,
Nijkamp et al., 2020] for a theoretical understanding of the learning
algorithm based on the finite-step MCMC. See also [Gao et al., 2020, Yu et
al., 2022] for learning EBMs at multiple noise levels for effective modeling
and sampling of multimodal density.

The learning algorithm is summarized in Algorithm 1.

input : Learning iterations $T$, learning rates for the prior, generation, and
regression models $\\{\eta_{0},\eta_{1},\eta_{2}\\}$, initial parameters
$\theta_{0}=(\alpha_{0},\beta_{0},\gamma_{0})$, observed examples
$\\{(x_{i},y_{i})\\}_{i=1}^{n}$, batch size $m$, number of prior and posterior
sampling steps $\\{\Gamma_{0},\Gamma_{1}\\}$, and prior and posterior sampling
step sizes $\\{s_{0},s_{1}\\}$.

output : $\theta_{T}=(\alpha_{T},\beta_{T},\gamma_{T})$.

for  _$t=0:T-1$_ do

1\. Mini-batch: Sample observed examples $\\{(x_{i},y_{i})\\}_{i=1}^{m}$.

2\. Prior sampling: For each $i$, sample $z_{i}^{-}\sim{p}_{\alpha_{t}}(z)$
using Equation 14, where the target distribution $\pi(z)={p}_{\alpha_{t}}(z)$,
and $s=s_{0}$, $\Gamma=\Gamma_{0}$.

3\. Posterior sampling: For each $(x_{i},y_{i})$, sample
$z_{i}^{+}\sim{p}_{\theta_{t}}(z|x_{i},y_{i})$ using Equation 14, where the
target distribution $\pi(z)={p}_{\theta_{t}}(z|x_{i},y_{i})$, and $s=s_{1}$,
$\Gamma=\Gamma_{1}$.

4\. Update prior model:
$\alpha_{t+1}=\alpha_{t}+\eta_{0}\frac{1}{m}\sum_{i=1}^{m}[\nabla_{\alpha}f_{\alpha_{t}}(z_{i}^{+})-\nabla_{\alpha}f_{\alpha_{t}}(z_{i}^{-})]$.

5\. Update generation model:
$\beta_{t+1}=\beta_{t}+\eta_{1}\frac{1}{m}\sum_{i=1}^{m}\nabla_{\beta}\log
p_{\beta_{t}}(x_{i}|z_{i}^{+})$.

6\. Update regression model:
$\gamma_{t+1}=\gamma_{t}+\eta_{2}\frac{1}{m}\sum_{i=1}^{m}\nabla_{\gamma}\log
p_{\gamma_{t}}(y_{i}|z_{i}^{+})$.

Algorithm 1 Learning joint distribution.

###  3.4 Sampling with Gradual Distribution Shifting (SGDS)

To tackle the single-objective optimization problem (Equation 1), one naive
approach is to perform ancestral sampling with two steps, given some desirable
property value $y^{*}$,

$\displaystyle(\mathrm{i})\ z^{*}\sim p_{\theta}(z|y=y^{*})\propto
p_{\alpha}(z)p_{\gamma}(y=y^{*}|z),$ (15) $\displaystyle(\mathrm{ii})\
x^{*}\sim p_{\beta}(x|z=z^{*}),$ (16)

where $(\mathrm{i})$ is an application of Bayes rule, with $p_{\alpha}(z)$ as
the prior and $p_{\gamma}(y|z)$ as the likelihood. Sampling from
$p_{\theta}(z|y)$ can be carried out by Langevin dynamics in Equation 14 by
replacing the target distribution $\pi(z)$ with $p_{\theta}(z|y)$.

Our model $p_{\theta}(x,y,z)$ 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 15 does not work well
since it involves extrapolating the learned distribution.

We propose an iterative updating method called sampling with gradual
distribution shifting (SGDS) to address this issue. In particular, we first
leverage the $n$ samples collected from the common dataset
$\\{(x_{i}^{0},y_{i}^{0})\\}_{i=1}^{n}$ (e.g. ZINC, $n=250,000$) to learn the
initial joint distribution $p_{\theta_{0}}(x,y)$ as a valid starting point.
Then we shift the joint distribution progressively using a smaller number $k$
(e.g., $k=10,000$) of synthesized samples
$\\{(x_{i}^{t},y_{i}^{t})\\}_{i=1}^{k}$ from distribution $p_{\theta_{t-1}}$
at the previous iteration, where $k\ll n$. Therefore, by progressively
learning the joint distribution with $T$ (e.g., $T=30$) shift iterations,
$p_{\theta_{1}},\dots,p_{\theta_{T}}$, at the last several iterations, we
expect to generate molecules with desirable properties that are significantly
distant from the initial distribution as shown in Figure 1(c).

Next, we shall explain in detail the steps to generate
$\\{(x_{i}^{t},y_{i}^{t})\\}_{i=1}^{k}$ given $p_{\theta_{t-1}}$. In a
property maximization task, we shift the support slightly by adding a small
$\Delta_{y}$ to all $y$’s,

$\displaystyle\tilde{y}^{t}=y^{t-1}+\Delta_{y},$ (17)

and generate $x^{t}$ conditional on shifted $\tilde{y}^{t}$, following
Equation 15,

$\displaystyle(\mathrm{i})\ $ $\displaystyle z^{t}\sim
p_{\theta_{t-1}}(z|y=\tilde{y}^{t}),$ (18) $\displaystyle(\mathrm{ii})\ $
$\displaystyle x^{t}\sim p_{\beta_{t-1}}(x|z=z^{t}).$ (19)

$\Delta_{y}$ can be chosen as a fixed small value. After generating $x^{t}$,
its ground-truth property value $y^{t}$ can be computed by calling the
corresponding engines such as RDKit and AutoDock-GPU. In Equation 18, sampling
can be achieved by langevin dynamics as in Equation 14. For the sake of
efficiency, we propose to run persistent chain by initializing the Langevin
dynamics from the latent vectors generated in the previous iteration Han et
al. [2017]. This is also called warm start. Specifically, we have

$\displaystyle z_{0}^{t}=z^{t-1}_{\Gamma},$ $\displaystyle
z^{t}_{\tau+1}=z^{t}_{\tau}+s\nabla_{z}\log
p_{\theta_{t-1}}(z|\tilde{y}^{t})+\sqrt{2s}\epsilon_{\tau},$ (20)

for $\tau=1,...,\Gamma$, where $\Gamma$ is the length of Markov chain in each
iteration. With warm start, we use $\Gamma=2$ in our experiments.

For more efficient optimization via distribution shifting, we further
introduce a rank-and-select scheme by maintaining a buffer of top-$k$ samples
of $(z,x,y)$ (where $z$ is the sampled latent vector, $x$ is the generated
molecule, and $y$ is the ground-truth property value of $x$). Specifically, we
maintain a buffer which consists of $k$ samples of $(z,x,y)$ with the highest
values of $y$ in the past shifting iterations. In each shift iteration,
conditioned on the shifted property values, with warm start, initialized from
these $k$ vectors $z$ in the buffer, a new batch of $k$ latent vectors $z$,
molecules $x$, and their ground-truth values $y$ can be produced in the new
shift iteration. We rank all the $2k$ samples of $(z,x,y)$ (including $k$
newly generated ones and $k$ samples in the buffer) and select the top-$k$
samples of $(z,x,y)$ based on the ground-truth values $y$. The $k$ samples of
$(z,x,y)$ in the buffer are then updated by those newly selected $k$ samples
of $(z,x,y)$. We call this procedure as rank-and-select. This rank-and-select
procedure can also be applied to constrained optimization tasks, where we
select those sampled molecules that satisfy the given constraints. With the
selected samples, we then shift the model distribution by learning from these
samples with several learning iterations. The SGDS algorithm is summarized in
Algorithm 2.

input : Shift iterations $T$, initial pretrained parameters
$\theta_{0}=(\alpha_{0},\beta_{0},\gamma_{0})$, initial examples
$\\{(x_{i}^{0},y_{i}^{0})\\}_{i=1}^{k}$ from the data distribution boundary,
shift magnitude $\Delta_{y}$, ${\rm PropertyComputeEngine}={\rm RDKit}$ or
${\rm AutoDock}$-${\rm GPU}$, ${\rm LearningAlgorithm}={\rm Algorithm}$ 1.

output : $\\{(x_{i}^{T},y_{i}^{T})\\}_{i=1}^{k}$.

for  _$t=1:T$_ do

1\. Property shift: For each $y_{i}^{t-1}$,
$\tilde{y}_{i}^{t}=y_{i}^{t-1}+\Delta_{y}$.

2\. Latent sampling with warm start: For each $\tilde{y}_{i}^{t}$, sample
$z_{i}^{t}\sim p_{\theta_{t-1}}(z|\tilde{y}_{i}^{t})$ using Equation 20.

3\. Molecule generation: For each $z_{i}^{t}$, sample $x_{i}^{t}\sim
p_{\theta_{t-1}}(x|z_{i}^{t})$.

4\. Property computation: For each $x_{i}^{t}$, compute $y_{i}^{t}={\rm
PropertyComputeEngine}(x_{i}^{t})$.

5\. Rank-and-select: Update the buffer of top-$k$ samples
$\\{z_{i}^{t},x_{i}^{t},y_{i}^{t}\\}_{i=1}^{k}$ by rank-and-select.

6\. Distribution shift: $\theta_{t}={\rm
LearningAlgorithm}(\\{(x_{i}^{t},y_{i}^{t})\\}_{i=1}^{k},\theta_{t-1})$.

Algorithm 2 SGDS for single property optimization.

###  3.5 Multi-Objective Optimization

We next consider the multi-objective optimization problem. Suppose we optimize
for a set of properties $\\{y_{j}\\}_{j=1}^{m}$, then we learn a property
regression model for each property $y_{j}$,

$p_{\gamma_{j}}(y_{j}|z)=\frac{1}{\sqrt{2\pi\sigma_{j}^{2}}}\exp\left(-\frac{1}{2\sigma_{j}^{2}}(y_{j}-s_{\gamma_{j}}(z))^{2}\right),$
(21)

where each $s_{\gamma_{j}}$ is a small MLP with parameters $\gamma_{j}$. We
assume that given $z$ the properties are conditionally independent, so the
joint distribution is

$p_{\theta}(x,z,y_{1},...,y_{m})=p_{\alpha}(z)p_{\beta}(x|z)\prod_{j=1}^{m}p_{\gamma_{j}}(y_{i}|z).$
(22)

Under our framework, both the learning and the sampling algorithm for the
single-objective problem can be straightforwardly extended to the multi-
objective setting. In SGDS, we shift the values of the multiple properties
simultaneously, and generate molecules conditional on the multiple properties.

##  4 Experiments

To demonstrate the effectiveness of our proposed method, SGDS, we compare our
model with previous SOTA methods for molecule design including single-
objective optimization (§4.2), multi-objective optimization (§4.3) and
constrained optimization (§4.4). In molecule design experiments, we consider
both non-biological and biological properties. Finally, we add ablation
studies to analyze the effects of different components in SGDS. We also
conduct unconditional molecule generation experiments for sanity check of the
model and discuss the mode traversing in the latent space at the end of this
section.

###  4.1 Experimental Setup

Datasets. For the molecule property optimization task, we report results on
ZINC [Irwin et al., 2012] and MOSES [Polykovskiy et al., 2020], which consist
of around $250$k and $2$M molecules respectively.

Encoding systems in molecular studies typically include SMILES [Weininger,
1988], SELFIES [Krenn et al., 2020], and graph representations. SMILES and
SELFIES linearize a molecular graph into character strings. SMILES has
historically faced challenges regarding validity (the percentage of molecules
that satisfy the chemical valency rules). Recently, SELFIES was introduced,
offering an encoding system where each string inherently corresponds to a
valid molecule. We use SELFIES representation in our work.

The non-biological properties (such as penalized logP, QED, etc.) can be
computed using RDKit [Landrum et al., 2013]. Following [Eckmann et al., 2022],
we use the docking scores from AutoDock-GPU [Santos-Martins et al., 2021] to
approximate the binding affinity to two protein targets, human estrogen
receptor (ESR1) and human peroxisomal acetyl-CoA acyl transferase 1 (ACAA1).

Training Details. There are three modules in our method, the molecule
generation model $p_{\beta}(x|z)$, the energy-based prior model
$p_{\alpha}(z)$, and property regression model
$\\{p_{\gamma_{j}}(y|z)\\}_{j=1}^{m}$, where $m$ is the total number of
properties we aim to optimize. The generation model $p_{\beta}(x|z)$ is
parameterized by a single-layer LSTM with $1024$ hidden units where the
dimension of latent vector $z$ is $100$. The energy-based prior model
$p_{\alpha}(z)$ is a $3$-layer MLP. Each of the property regression model
$p_{\gamma_{j}}(y|z)$ is a 3-layer MLP. It is worth mentioning that compared
to most previous models, SGDS is characterized by its simplicity without
adding inference networks for sampling, or RL-related modules for
optimization. In order to get valid initial distribution $\theta_{0}$ for
SGDS, we first train our model for $30$ epochs on ZINC. We use Adam optimizer
[Kingma and Ba, 2015] to train our models with learning rates $10^{-4}$ for
energy-based prior model, and $10^{-3}$ for the molecule generation model and
property regression model. During SGDS, we use $30$ shifting iterations for
single-objective optimization and $20$ for multi-objective optimization. For
each iteration of distribution shifting, we sample $10^{4}$ boundary examples
except binding affinity, where $2\times 10^{3}$ examples are used to speed up
calculation, and then we update the model parameters $\theta$ for $10$
iterations using Algorithm 1 with Adam optimizer and the same learning rates
mentioned above. All experiments are conducted on Nvidia Titan XP GPU.

###  4.2 Single-Objective Optimization

Penalized logP and QED Maximization. For non-biological properties, we are
interested in Penalized logP and QED, both of which can be calculated by RDKit
[Landrum et al., 2013].

Since we know Penalized logP scores have a positive relationship with the
lengths of molecules, we maximize Penalized logP either with or without
maximum length limit. Following [Eckmann et al., 2022], the maximum length is
set to be the maximum length of molecules in ZINC using SELFIES. From Table 1,
we can see that with length limit, SGDS outperforms previous methods by a
large margin. 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 Figure 1(c). One can notice that, the
distribution of the property is gradually shifted towards the region with
higher values, and the final distribution is significantly distant from the
initial one.

Table 1: Non-biological single-objective optimization. Report top-3 highest
scores found by each model. LL (Length Limit) denotes whether it has the
maximum length limit. Baseline results obtained from [Eckmann et al., 2022,
You et al., 2018, Luo et al., 2021, Xie et al., 2021].

Method | LL | Penalized logP $(\uparrow)$ | QED $(\uparrow)$
---|---|---|---
1st | 2rd | 3rd | 1st | 2rd | 3rd
JT-VAE | ✗ | 5.30 | 4.93 | 4.49 | 0.925 | 0.911 | 0.910
GCPN | ✓ | 7.98 | 7.85 | 7.80 | 0.948 | 0.947 | 0.946
MolDQN | ✓ | 11.8 | 11.8 | 11.8 | 0.948 | 0.943 | 0.943
MARS | ✗ | 45.0 | 44.3 | 43.8 | 0.948 | 0.948 | 0.948
GraphDF | ✗ | 13.7 | 13.2 | 13.2 | 0.948 | 0.948 | 0.948
LIMO | ✓ | 10.5 | 9.69 | 9.60 | 0.947 | 0.946 | 0.945
SGDS | ✓ | 26.4 | 25.8 | 25.5 | 0.948 | 0.948 | 0.948
SGDS | ✗ | 158.0 | 157.8 | 157.5 | 0.948 | 0.948 | 0.948

Biological Property Optimization. ESR1 and ACAA1 are two human proteins. We
aim to design ligands (molecules) that have the maximum binding affinities
towards those target proteins. ESR1 is well-studied, which has many existing
binders. However, we do not use any binder-related information in SGDS.
Binding affinity is measured by the estimated dissociation constants
$\mathrm{K_{D}}$, which can be approximated by docking scores from AutoDock-
GPU [Santos-Martins et al., 2021]. Large binding affinities corresponds to
small $\mathrm{K_{D}}$. That is, we aim to minimize $\mathrm{K_{D}}$. Table 2
shows that our model outperforms previous methods on both ESR1 and ACAA1
binding affinity maximization tasks by large margins. Comparing to existing
methods, much more molecules with high binding affinities can be directly
sampled from the last several shifting iterations. See Supplement for more
examples. Producing those ligands with high binding affinity plays a vital
role in the early stage of drug discovery.

Table 2: Biological single-objective optimization. Report top-3 lowest
$\mathrm{K_{D}}$ (in nanomoles/liter) found by each model. Baseline results
obtained from [Eckmann et al., 2022].

Method | ESR1 $\mathrm{K_{D}}$ $(\downarrow)$ | ACAA1 $\mathrm{K_{D}}$ $(\downarrow)$
---|---|---
1st | 2rd | 3rd | 1st | 2rd | 3rd
GCPN | 6.4 | 6.6 | 8.5 | 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 | 41
SGDS | 0.03 | 0.03 | 0.04 | 0.11 | 0.11 | 0.12

###  4.3 Multi-objective Optimization

#### Multi-objective Binding Affinity Optimization.

We consider maximizing binding affinity, QED and minimizing synthetic
accessibility score (SA) simultaneously. Following Eckmann et al. [2022], we
exclude molecules with abnormal behaviors 111 In particular, we exclude
molecules with QED$\uparrow$ smaller than 0.4, SA$\downarrow$ larger than 5.5,
and too small (less than 5 atoms) or too large (more than 6 atoms) chemical
rings. to encourage the joint distribution shifts towards a desirable region
in terms of pharmacologic and synthetic properties. Those heuristics can be
conveniently added in our rank-and-select step. Table 3 shows our multi-
objective results compared to LIMO and GCPN. From the results, we can see that
SGDS is able to find the ligands with desired properties while keeping the
pharmacologic structures. For ESR1, we have two existing binders on the
market, Tamoxifen and Raloxifene. Our designed ligands have similar QED and
SA, with very low $\mathrm{K_{D}}$. Compared to existing methods, SGDS obtains
better results in overall adjustments. For ACAA1, we do not have any existing
binders. Compared with prior SOTA methods, our optimized ligands outperform
those by a large margin in terms of $\mathrm{K_{D}}$. When comparing with
single-objective optimization, we find that multi-objective optimization is
more complicated, but it may be more useful in real world molecule design.
While we still need domain expertise to determine the effectiveness of those
ligands discovered by SGDS, we believe the ability of SGDS to generate many
high quality molecules given multiple metrics is extremely useful in the early
stage of drug discovery.

Table 3: Muli-objective optimization for both ESR1 and ACAA1. Report Top-2
average scores related to $\mathrm{K_{D}}$ (in nmol/L), QED and SA. Baseline
results obtained from [Eckmann et al., 2022].

Ligand | ESR1 | ACAA1
---|---|---
$\mathrm{K_{D}}$ $\downarrow$ | QED $\uparrow$ | SA $\downarrow$ |  $\mathrm{K_{D}}$ $\downarrow$ | QED $\uparrow$ | SA $\downarrow$
Tamoxifen | 87 | 0.45 | 2.0 | $-$ | $-$ | $-$
Raloxifene | $7.9\times 10^{6}$ | 0.32 | 2.4 | $-$ | $-$ | $-$
GCPN $\texttt{1}^{\texttt{st}}$ | 810 | 0.43 | 4.2 | 8500 | 0.69 | 4.2
GCPN $\texttt{2}^{\texttt{nd}}$ | $27000$ | 0.80 | 3.7 | 8500 | 0.54 | 4.3
LIMO $\texttt{1}^{\texttt{st}}$ | 4.6 | 0.43 | 4.8 | 28 | 0.57 | 5.5
LIMO $\texttt{2}^{\texttt{nd}}$ | 2.8 | 0.64 | 4.9 | 31 | 0.44 | 4.9
SGDS $\texttt{1}^{\texttt{st}}$ | 0.36 | 0.44 | 3.99 | 4.55 | 0.56 | 4.07
SGDS $\texttt{2}^{\texttt{nd}}$ | 1.28 | 0.44 | 3.86 | 5.67 | 0.60 | 4.58

###  4.4 Constrained Optimization

To optimize single-objective under some constrains, we use the original SGDS
steps and in rank-and-select, we only keep the molecules that satisfy the
constraints.

#### Similarity-constrained Penalized logP Maximization

Following JT-VAE [Jin et al., 2018], this experiment aims to generate
molecules with high penalized logP while being similar to the target
molecules. Similarity is measured by Tanimoto similarity between Morgan
fingerprints with a cutoff value $\delta$. We compare our results with
previous SOTA method in Table 4. The results show that SGDS tends to obtain
better results with weak constraints (i.e. $\delta=0,0.2$) with $100\%$
success rate, since different from optimized property, the constraints are
added implicitly.

Table 4: Similarity-constrained optimization results. LIMO results obtained
from [Eckmann et al., 2022, Luo et al., 2021].

$\delta$ | GraphDF | LIMO | SGDS
---|---|---|---
Improv. |  $\%$ Succ. | Improv. |  $\%$ Succ. | Improv. |  $\%$ Succ.
0.0 | $5.9\pm 2.0$ | 100 | $10.1\pm 2.3$ | 100 | $\mathbf{19.1}\pm 2.1$ | 100
0.2 | $5.6\pm 1.7$ | 100 | $5.8\pm 2.6$ | 99.0 | $\mathbf{7.4}\pm 1.9$ | 100
0.4 | $\mathbf{4.1}\pm 1.4$ | 100 | $3.6\pm 2.3$ | 93.7 | $3.8\pm 1.4$ | 97.5
0.6 | $1.7\pm 1.2$ | 93.0 | $1.8\pm 2.0$ | 85.5 | $\mathbf{2.6}\pm 2.0$ | 95.6

#### logP targeting

In Table 5, comparing to previous methods, SGDS is able to get competitive
diversity scores with significantly better success rate in both ranges. That
is because after SGDS, our model is shifted towards the region that is
supported by molecules satisfying the logP constraints. Due to the flexibility
of our EBM prior, SGDS achieves high diversity scores while keeping most of
the sampled molecules within the logP range.

Table 5: logP targeting to a certain range [Eckmann et al., 2022, You et al.,
2018, Luo et al., 2021, Xie et al., 2021].

Method |  $-2.5\leq$ logP $\leq-2$ |  $5\leq$ logP $\leq 5.5$
---|---|---
Success | Diversity | Success | Diversity
ZINC | 0.4$\%$ | 0.919 | $1.3\%$ | 0.901
JT-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 | $-$ | $-$
SGDS | $\mathbf{86.0\%}$ | 0.874 | $\mathbf{62.2\%}$ | 0.858

###  4.5 Ablation Studies

SGDS outperforms previous methods by a significant margin, especially on
binding affinity related experiments. Hence, we conduct ablations on the key
components of our method on a challenging single-objective ACAA1 maximization
experiment. Since SGDS is optimized based on shifting the joint distribution
rather than per-molecule based optimization method (such as [Eckmann et al.,
2022]), we use the summarized statistics (i.e. the mean and standard deviation
of $100$ lowest $\mathrm{K_{d}}$ from uniquely generated molecules) of last
three shifted distributions as our metric to compare the key components rather
than top-3 optimized $\mathrm{K_{D}}$ in §4.2. Ablation studies are discussed
as follows.

(1) Without EBM Prior: for joint distribution, we replace the learnable EBM
prior by a fixed $\mathcal{N}(0,I_{d})$. (2)Without Property Regression
$p_{\gamma}(y|z)$: we only learn the distribution of molecules as
$p_{\alpha}(z)p_{\beta}(x|z)$. For each iteration of distribution shifting, we
only use rank-and-select and update model parameters based on those molecules
with high values. The molecule can be generated by first sampling $z\sim
p_{\alpha}(z)$ and then $x\sim p_{\beta}(x|z)$. (3) Without Gradual Shifting:
rather than iterative distribution shifting in SGDS, we directly sample $z\sim
p_{\theta}(z|y=y^{\star})$, where $y^{\star}$ is set to be the minimal value
we can get in §4.2. (4) Without Rank-and-Select: we skip the rank-and-select
step in Algorithm 2. (5) Without Warm Start: when sampling $z\sim
p_{\theta}(z|y)$ in current iteration, we replace the warm start algorithm in
Equation 20 by 20-step langevin dynamics in Equation 14 with the same step
size.

Table 6: Ablation Studies. Report the mean and standard deviation of $100$
uniquely generated molecules with the lowest $\mathrm{K_{D}}$(in
$10^{-9}$mol/L) from last three shifted iterations (i.e. the 28th, 29th, 30th
iterations of total 30 iterations).

Method | 28th | 29th | 30th
---|---|---|---
SGDS | $0.74\pm 0.04$ | $0.61\pm 0.03$ | $0.59\pm 0.03$
Without EBM Prior | $47.8\pm 26.9$ | $38.8\pm 21.1$ | $35.1\pm 20.2$
Without Property Regression | $140\pm 74.8$ | $114\pm 67.0$ | $103\pm 56.3$
Without Gradual Shifting | $211\pm 125$ | $166\pm 97.9$ | $137\pm 74.8$
Without Rank-and-Select | $9.71\pm 5.52$ | $5.75\pm 3.39$ | $3.91\pm 2.17$
Without Warm Start | $6.27\pm 3.92$ | $3.27\pm 2.99$ | $2.37\pm 1.35$

The ablation studies are displayed in Table 6. It is clear that all the
proposed components contribute significantly to the good performance of our
method.

###  4.6 Unconditional Generation

We employ unconditional molecule generation tasks as a sanity check of the
latent space EBM. The goal is to model the molecules in the training dataset
and generate similar molecules. We evaluate the model based on validity (the
percentage of molecules that satisfy the chemical valency rules), uniqueness
(the percentage of unique molecules in all generated samples) and novelty (the
percentage of generated molecules that are not in the training set) of
generated molecules. Note that we are not concerned with optimization of
molecular properties in this subsection.

Following previous work, we randomly sample $10$k molecules for ZINC and $30$k
for MOSES, comparing the results based on the aforementioned metrics.
Generation results are shown in Table 7 for ZINC and Table 8 for MOSES. In
Table 7, we present generation results for both SMILES and SELFIES. Despite
lacking a validity constraint during generation, our model attains $95.5\%$
validity using SMILES, outperforming other SMILES-based methods and rivaling
those with valency checks. This demonstrates our model’s effective and
implicit capture of valency rules. Furthermore, our model’s samples exhibit
perfect uniqueness and novelty.

Table 7: Unconditional generation on ZINC. ⋆ denotes valency check. [Jin et
al., 2018, You et al., 2018, Madhawa et al., 2019, Shi et al., 2020, Luo et
al., 2021, Gómez-Bombarelli et al., 2018, Kusner et al., 2017]

Model | Representation | Validity | Novelty | Uniqueness
---|---|---|---|---
JT-VAE | Graph | 1.000⋆ | 1.000 | 1.000
GCPN | Graph | 1.000⋆ | 1.000 | 1.000
GraphNVP | Graph | 0.426 | 1.000 | 0.948
GraphAF | Graph | 1.000⋆ | 1.000 | 0.991
GraphDF | Graph | 1.000⋆ | 1.000 | 1.000
ChemVAE | SMILES | 0.170 | 0.980 | 0.310
GrammarVAE | SMILES | 0.310 | 1.000 | 0.108
Ours | SMILES | 0.955 | 1.000 | 1.000
Ours | SELFIES | 1.000 | 1.000 | 1.000

Table 8: Unconditional generation on MOSES. ⋆ denotes valency check. Results
obtained from [Polykovskiy et al., 2020, Eckmann et al., 2022].

Model | Representation | Validity | Novelty | Uniqueness
---|---|---|---|---
JT-VAE | Graph | 1.000⋆ | 0.914 | 1.000
GraphAF | Graph | 1.000⋆ | 1.000 | 0.991
GraphDF | Graph | 1.000⋆ | 1.000 | 1.000
LIMO | SELFIES | 1.000 | 1.000 | 0.976
Ours | SELFIES | 1.000 | 1.000 | 1.000

Then we randomly sample $10$k molecules from the learned latent space EBM and
compute their PlogP and QED using RDKit. Their empirical densities are then
compared not only with the molecule property densities from the test split but
also with the predictions made by the regression model $p_{\gamma}(y|z)$. As
shown in Figure 2, the property densities from both our learned model and
predicted values from regression model align closely with those of the data,
suggesting that our model effectively captures regularities in the data.

Figure 2: Property distributions of PlogP (left) and QED (right).

In our experiments, we employ short-run MCMC [Nijkamp et al., 2019] with a
Markov chain length of $K=20$ and step size $s=0.1$ for all tests. As shown in
Figure 3, with increasing Markov chain length, the molecules evolve
correspondingly, suggesting that the Markov chain is not trapped in local
modes.

Figure 3: Two sequences of sampled molecules with different lengths of Markov
chain.

##  5 Conclusion and discussion

We propose a deep generative model for the joint distribution of molecules and
their properties. It assumes an energy-based prior model in a low-dimensional
continuous latent space, and the latent vector can generate the molecule and
predict its property value. We then design a sampling with gradual
distribution shifting method to shift the learned distribution to a region
with high property values. Molecule design can then be achieved by conditional
sampling. Our experiments demonstrate that our method outperforms previous
SOTA methods on some tasks by significant margins.

###### Acknowledgements.

Y. N. Wu was partially supported by NSF DMS-2015577 and a gift fund from
Amazon.

## References

  * Bagal et al. [2021] Viraj Bagal, Rishal Aggarwal, PK Vinod, and U Deva Priyakumar.  Liggpt: Molecular generation using a transformer-decoder model.  2021\.
  * Brown et al. [2019] Nathan Brown, Marco Fiscato, Marwin H.S. Segler, and Alain C. Vaucher.  Guacamol: Benchmarking models for de novo molecular design.  _Journal of Chemical Information and Modeling_ , 59(3):1096–1108, 2019.  10.1021/acs.jcim.8b00839.  URL https://doi.org/10.1021/acs.jcim.8b00839.  PMID: 30887799.
  * Chen et al. [2021] Binghong Chen, Tianzhe Wang, Chengtao Li, Hanjun Dai, and Le Song.  Molecule optimization by explainable evolution.  In _International Conference on Learning Representation (ICLR)_ , 2021\.
  * De Cao and Kipf [2018] Nicola De Cao and Thomas Kipf.  Molgan: An implicit generative model for small molecular graphs.  _arXiv preprint arXiv:1805.11973_ , 2018.
  * Du and Mordatch [2019] Yilun Du and Igor Mordatch.  Implicit generation and generalization in energy-based models.  _CoRR_ , abs/1903.08689, 2019.  URL http://arxiv.org/abs/1903.08689.
  * Eckmann et al. [2022] 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.09010_ , 2022.
  * Finn et al. [2016] Chelsea Finn, Paul F. Christiano, Pieter Abbeel, and Sergey Levine.  A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models.  _CoRR_ , abs/1611.03852, 2016.  URL http://arxiv.org/abs/1611.03852.
  * Fu et al. [2020] 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, pages 638–645, 2020.
  * Fu et al. [2021a] Tianfan Fu, Wenhao Gao, Cao Xiao, Jacob Yasonik, Connor W Coley, and Jimeng Sun.  Differentiable scaffolding tree for molecular optimization.  _arXiv preprint arXiv:2109.10469_ , 2021a.
  * Fu et al. [2021b] Tianfan Fu, Cao Xiao, Xinhao Li, Lucas M Glass, and Jimeng Sun.  Mimosa: Multi-constraint molecule sampling for molecule optimization.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 125–133, 2021b.
  * Gao et al. [2020] Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma.  Learning energy-based models by diffusion recovery likelihood.  _arXiv preprint arXiv:2012.08125_ , 2020.
  * Gómez-Bombarelli et al. [2018] Rafael Gómez-Bombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik.  Automatic chemical design using a data-driven continuous representation of molecules.  _ACS central science_ , 4(2):268–276, 2018.
  * Goodfellow et al. [2014] Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio.  Generative adversarial nets.  In _Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada_ , pages 2672–2680, 2014.  URL http://papers.nips.cc/paper/5423-generative-adversarial-nets.
  * Gottipati et al. [2020] 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_ , pages 3668–3679. PMLR, 2020.
  * Grathwohl et al. [2019] Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik 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:1912.03263_ , 2019.
  * Han et al. [2017] Tian Han, Yang Lu, Song-Chun Zhu, and Ying Nian Wu.  Alternating back-propagation for generator network.  In _Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence, February 4-9, 2017, San Francisco, California, USA._ , pages 1976–1984, 2017.  URL http://aaai.org/ocs/index.php/AAAI/AAAI17/paper/view/14784.
  * Hochreiter and Schmidhuber [1997] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Honda et al. [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.
  * Huang et al. [2021] 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:2102.09548_ , 2021.
  * Irwin et al. [2012] John J Irwin, Teague Sterling, Michael M Mysinger, Erin S Bolstad, and Ryan G Coleman.  Zinc: a free tool to discover chemistry for biology.  _Journal of chemical information and modeling_ , 52(7):1757–1768, 2012.
  * Jensen [2019] Jan H Jensen.  A graph-based genetic algorithm and generative model/monte carlo tree search for the exploration of chemical space.  _Chemical science_ , 10(12):3567–3572, 2019.
  * Jin et al. [2018] Wengong Jin, Regina Barzilay, and Tommi Jaakkola.  Junction tree variational autoencoder for molecular graph generation.  In _International conference on machine learning_ , pages 2323–2332. PMLR, 2018.
  * Jin et al. [2020] Wengong Jin, Regina Barzilay, and Tommi Jaakkola.  Multi-objective molecule generation using interpretable substructures.  In _International conference on machine learning_ , pages 4849–4859. PMLR, 2020.
  * Kim and Bengio [2016] Taesup Kim and Yoshua Bengio.  Deep directed generative models with energy-based probability estimation.  _CoRR_ , abs/1606.03439, 2016.  URL http://arxiv.org/abs/1606.03439.
  * Kingma and Ba [2015] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ , 2015.  URL http://arxiv.org/abs/1412.6980.
  * Kingma and Welling [2014] 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, Conference Track Proceedings_ , 2014.  URL http://arxiv.org/abs/1312.6114.
  * Kotsias et al. [2020] Panagiotis-Christos Kotsias, Josep Arús-Pous, Hongming Chen, Ola Engkvist, Christian Tyrchan, and Esben Jannik Bjerrum.  Direct steering of de novo molecular generation with descriptor conditional recurrent neural networks.  _Nature Machine Intelligence_ , 2(5):254–265, 2020.
  * Krenn et al. [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(4):045024, 2020.
  * Kumar et al. [2019] Rithesh Kumar, Anirudh Goyal, Aaron C. Courville, and Yoshua Bengio.  Maximum entropy generators for energy-based models.  _CoRR_ , abs/1901.08508, 2019.  URL http://arxiv.org/abs/1901.08508.
  * Kusner et al. [2017] Matt J Kusner, Brooks Paige, and José Miguel Hernández-Lobato.  Grammar variational autoencoder.  In _International Conference on Machine Learning_ , pages 1945–1954, 2017.
  * Landrum et al. [2013] Greg Landrum et al.  Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling.  _Greg Landrum_ , 2013.
  * LeCun et al. [2006] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang.  A tutorial on energy-based learning.  _Predicting structured data_ , 1(0), 2006.
  * Liu et al. [2021] Meng Liu, Keqiang Yan, Bora Oztekin, and Shuiwang Ji.  Graphebm: Molecular graph generation with energy-based models.  _arXiv preprint arXiv:2102.00546_ , 2021.
  * Luo et al. [2021] Youzhi Luo, Keqiang Yan, and Shuiwang Ji.  Graphdf: A discrete flow model for molecular graph generation.  In _International Conference on Machine Learning_ , pages 7192–7203. PMLR, 2021.
  * Madhawa et al. [2019] Kaushalya Madhawa, Katushiko Ishiguro, Kosuke Nakago, and Motoki Abe.  Graphnvp: An invertible flow model for generating molecular graphs.  _arXiv preprint arXiv:1905.11600_ , 2019.
  * Neal [2011] Radford M Neal.  MCMC using hamiltonian dynamics.  _Handbook of Markov Chain Monte Carlo_ , 2, 2011.
  * Ngiam et al. [2011] 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, ICML 2011, Bellevue, Washington, USA, June 28 - July 2, 2011_ , pages 1105–1112, 2011.  URL https://icml.cc/2011/papers/557_icmlpaper.pdf.
  * Nie et al. [2021] 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.
  * Nigam et al. [2020] AkshatKumar Nigam, Pascal Friederich, Mario Krenn, and Alán Aspuru-Guzik.  Augmenting genetic algorithms with deep neural networks for exploring the chemical space.  _ICLR 2020_ , 2020.
  * Nijkamp et al. [2019] Erik Nijkamp, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu.  Learning non-convergent non-persistent short-run MCMC toward energy-based model.  _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, Canada_ , 2019.
  * Nijkamp et al. [2020] Erik Nijkamp, Bo Pang, Tian Han, Linqi Zhou, Song-Chun Zhu, and Ying Nian Wu.  Learning multi-layer latent variable model via variational optimization of short run mcmc for approximate inference.  In _Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part VI 16_ , pages 361–378. Springer, 2020.
  * Pang et al. [2020] Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu.  Learning latent space energy-based prior model.  _arXiv preprint arXiv:2006.08205_ , 2020.
  * Polishchuk et al. [2013] Pavel G Polishchuk, Timur I 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.
  * Polykovskiy et al. [2020] Daniil Polykovskiy, Alexander Zhebrak, Benjamin Sanchez-Lengeling, Sergey Golovanov, Oktai Tatanov, Stanislav Belyaev, Rauf Kurbanov, Aleksey Artamonov, Vladimir Aladinskiy, Mark Veselov, et al.  Molecular sets (moses): a benchmarking platform for molecular generation models.  _Frontiers in pharmacology_ , 11:565644, 2020.
  * Robbins and Monro [1951] Herbert Robbins and Sutton Monro.  A stochastic approximation method.  _The annals of mathematical statistics_ , pages 400–407, 1951.
  * Santos-Martins et al. [2021] Diogo Santos-Martins, Leonardo Solis-Vasquez, Andreas F Tillack, Michel F Sanner, Andreas Koch, and Stefano Forli.  Accelerating autodock4 with gpus and gradient-based local search.  _Journal of chemical theory and computation_ , 17(2):1060–1073, 2021.
  * Segler et al. [2018] 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.
  * Shi et al. [2020] Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang.  Graphaf: a flow-based autoregressive model for molecular graph generation.  _arXiv preprint arXiv:2001.09382_ , 2020.
  * Weininger [1988] David Weininger.  Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules.  _Journal of chemical information and computer sciences_ , 28(1):31–36, 1988.
  * Winter et al. [2019] Robin Winter, Floriane Montanari, Andreas Steffen, Hans Briem, Frank Noé, and Djork-Arné Clevert.  Efficient multi-objective molecular optimization in a continuous latent space.  _Chemical science_ , 10(34):8016–8024, 2019.
  * Xie et al. [2016] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Ying Nian Wu.  A 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_ , pages 2635–2644, 2016.  URL http://proceedings.mlr.press/v48/xiec16.html.
  * Xie et al. [2021] Yutong Xie, Chence Shi, Hao Zhou, Yuwei Yang, Weinan Zhang, Yong Yu, and Lei Li.  Mars: Markov molecular sampling for multi-objective drug discovery.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=kHSu4ebxFXY.
  * You et al. [2018] Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec.  Graph convolutional policy network for goal-directed molecular graph generation.  In _Advances in neural information processing systems_ , pages 6410–6421, 2018.
  * Yu et al. [2022] Peiyu Yu, Sirui Xie, Xiaojian Ma, Baoxiong Jia, Bo Pang, Ruigi Gao, Yixin Zhu, Song-Chun Zhu, and Ying Nian Wu.  Latent diffusion energy-based model for interpretable text modeling.  _arXiv preprint arXiv:2206.05895_ , 2022.
  * Zang and Wang [2020] Chengxi Zang and Fei Wang.  Moflow: an invertible flow model for generating molecular graphs.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pages 617–626, 2020.
  * Zhou et al. [2019] Zhenpeng Zhou, Steven Kearnes, Li Li, Richard N Zare, and Patrick Riley.  Optimization of molecules via deep reinforcement learning.  _Scientific reports_ , 9(1):1–10, 2019.

##  Appendix A Details about model and learning

Our model is of the form $p_{\alpha}(z)p_{\beta}(x|z)p_{\gamma}(y|z)$. The
marginal distribution of $(x,y)$ is

$\displaystyle p_{\theta}(x,y)=\int p_{\theta}(x,y,z)dz=\int
p_{\alpha}(z)p_{\beta}(x|z)p_{\gamma}(y|z)dz.$

We use $p_{\theta}(x,y)$ to approximate the data distribution of $(x,y)$.

For the data distribution of $(x,y)$, $y$ is a deterministic function of $x$.
However, a machine learning method usually cannot learn the deterministic
function exactly. Instead, we can only learn a probabilistic
$p_{\theta}(y|x)$. Our model $p_{\theta}(x,y)$ seeks to approximate the data
distribution $p(x,y)$ by maximum likelihood. A learnable and flexible prior
model $p_{\alpha}(z)$ helps to make the approximation more accurate than a
fixed prior model such as that in VAE.

Let the training data be $\\{(x_{i},y_{i}),i=1,...,n\\}$. The log-likelihood
function is $L(\theta)=\sum_{i=1}^{n}\log p_{\theta}(x_{i},y_{i})$. The
learning gradient is $L^{\prime}(\theta)=\sum_{i=1}^{n}\nabla_{\theta}\log
p_{\theta}(x_{i},y_{i})$. In the following, we provide details for calculating
$\nabla_{\theta}\log p_{\theta}(x,y)$ for a single generic training example
$(x,y)$ (where we drop the subscript i for notation simplicity).

$\displaystyle\nabla_{\theta}\log p_{\theta}(x,y)$
$\displaystyle=\frac{1}{p_{\theta}(x,y)}\nabla_{\theta}p_{\theta}(x,y)$
$\displaystyle=\frac{1}{p_{\theta}(x,y)}\int\nabla_{\theta}p_{\theta}(x,y,z)dz$
$\displaystyle=\frac{1}{p_{\theta}(x,y)}\int
p_{\theta}(x,y,z)\nabla_{\theta}\log p_{\theta}(x,y,z)dz$
$\displaystyle=\int\frac{p_{\theta}(x,y,z)}{p_{\theta}(x,y)}\nabla_{\theta}\log
p_{\theta}(x,y,z)dz$ $\displaystyle=\int p_{\theta}(z\mid
x,y)\nabla_{\theta}\log p_{\theta}(x,y,z)dz$
$\displaystyle=\mathbb{E}_{p_{\theta}(z|x,y)}\left[\nabla_{\theta}\log
p_{\theta}(x,y,z)\right]$
$\displaystyle=\mathbb{E}_{p_{\theta}(z|x,y)}\left[\nabla_{\theta}(\log
p_{\alpha}(z)+\log p_{\beta}(x|z)+\log p_{\gamma}(y|z))\right].$

For the prior model,

$\displaystyle\nabla_{\alpha}\log p_{\alpha}(z)$
$\displaystyle=\nabla_{\alpha}f_{\alpha}(z)-\nabla_{\alpha}\log Z(\alpha)$
$\displaystyle=\nabla_{\alpha}f_{\alpha}(z)-\frac{1}{Z(\alpha)}\nabla_{\alpha}Z(\alpha)$
$\displaystyle=\nabla_{\alpha}f_{\alpha}(z)-\frac{1}{Z(\alpha)}\int\nabla_{\alpha}\exp(f_{\alpha}(z))p_{0}(z)dz$
$\displaystyle=\nabla_{\alpha}f_{\alpha}(z)-\int\nabla_{\alpha}f_{\alpha}(z)\frac{1}{Z(\alpha)}\exp(f_{\alpha}(z))p_{0}(z)dz$
$\displaystyle=\nabla_{\alpha}f_{\alpha}(z)-\mathbb{E}_{p_{\alpha}(z)}[\nabla_{\alpha}f_{\alpha}(z)].$

Thus the learning gradient for $\alpha$ given an example $(x,y)$ is The above
equation has an empirical Bayes nature. $p_{\theta}(z|x,y)$ is based on the
empirical observation $(x,y)$, while $p_{\alpha}$ is the prior model. For the
generation model,

$\displaystyle\begin{split}\delta_{\beta}(x,y)=\nabla_{\beta}\log
p_{\theta}(x,y)=\mathbb{E}_{p_{\theta}(z|x,y)}[\nabla_{\beta}\log
p_{\beta}(x|z)].\end{split}$ (24)

Similarly, for the regression model,

$\displaystyle\begin{split}\delta_{\gamma}(x,y)=\nabla_{\gamma}\log
p_{\theta}(x,y)=\mathbb{E}_{p_{\theta}(z|x,y)}[\nabla_{\gamma}\log
p_{\gamma}(y|z)].\end{split}$ (25)

Estimating expectations in the above equations requires Monte Carlo sampling
of the prior model $p_{\alpha}(z)$ and the posterior distribution
$p_{\theta}(z|x,y)$. If we can draw fair samples from the two distributions,
and use these Monte Carlo samples to approximate the expectations, then 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 and Monro, 1951], who established the convergence of such
an algorithm to a local maximum of the log-likelihood.

For MCMC sampling using Langevin dynamics, the finite step or 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.

##  Appendix B Training Time

The training of joint distribution of molecule and its properties takes around
4 hours with $25$ iterations on a single Nvidia Titan XP GPU with batch size
$2048$. For non-biological single-objective property optimization, it takes
around $20$ minutes to do $30$ distribution shifting (SGDS) iterations. If we
use SGDS without warm start, it takes around half an hour. For biological
binding affinity maximization, the optimization time is mainly dependent on
the number of queries of AutoDock-GPU. We do $30$ and $20$ SGDS iterations for
the single-objective and multi-objective tasks, respectively, which cost $10$
hours and $8$ hours with warm start, and cost $14$ hours and $10$ hours
without warm start. For biological property optimization tasks, we use two
Nvidia Titan XP GPUs, one for running our code and the other for running
AutoDock-GPU. We have added a table to compare with previous methods.

Model | Penalized-logP/QED | Single binding affinity
---|---|---
JT-VAE | 24 | $-$
GCPN | 8 | 6
MolDQN | 24 | 6
GraphDF | 8 | 12
Mars | 12 | 6
LIMO | 1 | 1
SGDS without warm start | $4.5$ | $18$
SGDS with warm start | $4.3$ | $14$
Table 9: Comparison of molecule generation time in (hrs). Results obtained
from [Eckmann et al., 2022].

Even if we use MCMC sampling-based methods, our training speed is affordable
comparing to existing methods. That is because 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 $20$) in our experiments. The major bottleneck of the
training speed is the time of querying the property compute engines.

##  Appendix C Generated Samples

###  C.1 Biological Property Optimization

Figure 4 and Figure 5 show generated molecules with high binding affinities
towards ESR1 and ACAA1 respectively in single-objective property design
experiments.

Figure 6 and Figure 7 show generated molecules with high binding affinities
towards ESR1 and ACAA1 respectively in multi-objective property design.

Comparing to the previous state-of-the-art methods, SGDS is able to produce
more high quality molecules than top-3 molecules because after gradual
distribution shifting, the joint distribution locates at the area supported by
molecules with high binding affinities.

Meanwhile, compared to previous generative model-based methods, we use
Langevin dynamics to infer the posterior distribution
$p(z|x,y_{1},\dots,y_{n})$ without bothering to design different encoders when
facing different combination of properties.

Figure 4: Generated molecules in singe-objective esr1 binding affinity
maximization experiments with corresponding $\mathrm{K_{D}}(\downarrow)$ in
nmol/L. Figure 5: Generated molecules in singe-objective acaa1 binding
affinity maximization experiments with corresponding
$\mathrm{K_{D}}(\downarrow)$ in nmol/L. Figure 6: Generated molecules in
multi-objective esr1 binding affinity maximization experiments with
corresponding $\mathrm{K_{D}}(\downarrow)$ in nmol/L, SA$(\downarrow)$ and
QED$(\uparrow)$ respectively. Figure 7: Generated molecules in multi-objective
acaa1 binding affinity maximization experiments with corresponding
$\mathrm{K_{D}}(\downarrow)$ in nmol/L, SA$(\downarrow)$ and QED$(\uparrow)$
respectively.

###  C.2 P-logP and QED Optimization

Figure 8: Top-3 molecules in single-objective QED maximization. Figure 9:
Top-3 molecules in single-objective p-logP maximization.

##  Appendix D Illustration of Sampling with Gradual Distribution Shifting
(SGDS)

Figures 10, 11 and 12 show property densities of sampled molecules of the
distribution shifting process in single-objective penalized logP, esr1 and
acaa1 optimization respectively. SGDS is implemented with warm start. We can
see the model distribution is gradually shifting towards the region supported
by molecules with high property values. To better visualize the shifting
process, we plot the docking scores rather than $\mathrm{K_{D}}$. The increase
in docking scores corresponds to the exponential decrease in $\mathrm{K_{D}}$.

Figure 10: Illustration of SGDS in a single-objective penalized logP
optimization experiment. Figure 11: Illustration of SGDS in a single-objective
esr1 optimization experiment. Figure 12: Illustration of SGDS in a single-
objective acaa1 optimization experiment.