Junction Tree Variational Autoencoder for Molecular Graph Generation

Wengong Jin 1 Regina Barzilay 1 Tommi Jaakkola 1

arXiv:1802.04364v2 [cs.LG] 19 Feb 2018

Abstract
We seek to automate the design of molecules
based on specific chemical properties. In computational terms, this task involves continuous
embedding and generation of molecular graphs.
Our primary contribution is the direct realization
of molecular graphs, a task previously approached
by generating linear SMILES strings instead of
graphs. Our junction tree variational autoencoder
generates molecular graphs in two phases, by first
generating a tree-structured scaffold over chemical substructures, and then combining them into a
molecule with a graph message passing network.
This approach allows us to incrementally expand
molecules while maintaining chemical validity
at every step. We evaluate our model on multiple tasks ranging from molecular generation to
optimization. Across these tasks, our model outperforms previous state-of-the-art baselines by a
significant margin.

1. Introduction
The key challenge of drug discovery is to find target
molecules with desired chemical properties. Currently, this
task takes years of development and exploration by expert
chemists and pharmacologists. Our ultimate goal is to automate this process. From a computational perspective, we
decompose the challenge into two complementary subtasks:
learning to represent molecules in a continuous manner that
facilitates the prediction and optimization of their properties
(encoding); and learning to map an optimized continuous
representation back into a molecular graph with improved
properties (decoding). While deep learning has been extensively investigated for molecular graph encoding (Duvenaud
et al., 2015; Kearnes et al., 2016; Gilmer et al., 2017), the
harder combinatorial task of molecular graph generation
from latent representation remains under-explored.
Prior work on drug design formulated the graph generation task as a string generation problem (Gómez-Bombarelli
et al., 2016; Kusner et al., 2017) in an attempt to side-step
1

MIT Computer Science & Artificial Intelligence Lab. Correspondence to: Wengong Jin <wengong@csail.mit.edu>.

Figure 1. Two almost identical molecules with markedly different
canonical SMILES in RDKit. The edit distance between two
strings is 22 (50.5% of the whole sequence).

direct generation of graphs. Specifically, these models start
by generating SMILES (Weininger, 1988), a linear string
notation used in chemistry to describe molecular structures.
SMILES strings can be translated into graphs via deterministic mappings (e.g., using RDKit (Landrum, 2006)).
However, this design has two critical limitations. First, the
SMILES representation is not designed to capture molecular similarity. For instance, two molecules with similar
chemical structures may be encoded into markedly different
SMILES strings (e.g., Figure 1). This prevents generative
models like variational autoencoders from learning smooth
molecular embeddings. Second, essential chemical properties such as molecule validity are easier to express on graphs
rather than linear SMILES representations. We hypothesize
that operating directly on graphs improves generative modeling of valid chemical structures.
Our primary contribution is a new generative model of
molecular graphs. While one could imagine solving the
problem in a standard manner – generating graphs node
by node – the approach is not ideal for molecules. This is
because creating molecules atom by atom would force the
model to generate chemically invalid intermediaries (see,
e.g., Figure 2), delaying validation until a complete graph
is generated. Instead, we propose to generate molecular
graphs in two phases by exploiting valid subgraphs as components. The overall generative approach, cast as a junction
tree variational autoencoder, first generates a tree structured object (a junction tree) whose role is to represent the
scaffold of subgraph components and their coarse relative
arrangements. The components are valid chemical substructures automatically extracted from the training set using tree
decomposition and are used as building blocks. In the second phase, the subgraphs (nodes in the tree) are assembled
together into a coherent molecular graph.

Junction Tree Variational Autoencoder for Molecular Graph Generation

Figure 2. Comparison of two graph generation schemes: Structure
by structure approach is preferred as it avoids invalid intermediate
states (marked in red) encountered in node by node approach.

We evaluate our model on multiple tasks ranging from
molecular generation to optimization of a given molecule
according to desired properties. As baselines, we utilize
state-of-the-art SMILES-based generation approaches (Kusner et al., 2017; Dai et al., 2018). We demonstrate that
our model produces 100% valid molecules when sampled
from a prior distribution, outperforming the top performing baseline by a significant margin. In addition, we show
that our model excels in discovering molecules with desired
properties, yielding a 30% relative gain over the baselines.

2. Junction Tree Variational Autoencoder
Our approach extends the variational autoencoder (Kingma
& Welling, 2013) to molecular graphs by introducing a suitable encoder and a matching decoder. Deviating from previous work (Gómez-Bombarelli et al., 2016; Kusner et al.,
2017), we interpret each molecule as having been built from
subgraphs chosen out of a vocabulary of valid components.
These components are used as building blocks both when
encoding a molecule into a vector representation as well
as when decoding latent vectors back into valid molecular
graphs. The key advantage of this view is that the decoder
can realize a valid molecule piece by piece by utilizing the
collection of valid components and how they interact, rather
than trying to build the molecule atom by atom through
chemically invalid intermediaries (Figure 2). An aromatic
bond, for example, is chemically invalid on its own unless
the entire aromatic ring is present. It would be therefore
challenging to learn to build rings atom by atom rather than
by introducing rings as part of the basic vocabulary.
Our vocabulary of components, such as rings, bonds and
individual atoms, is chosen to be large enough so that a
given molecule can be covered by overlapping components
or clusters of atoms. The clusters serve the role analogous to
cliques in graphical models, as they are expressive enough
that a molecule can be covered by overlapping clusters without forming cluster cycles. In this sense, the clusters serve
as cliques in a (non-optimal) triangulation of the molecular
graph. We form a junction tree of such clusters and use it
as the tree representation of the molecule. Since our choice
of cliques is constrained a priori, we cannot guarantee that
a junction tree exists with such clusters for an arbitrary
molecule. However, our clusters are built on the basis of the
molecules in the training set to ensure that a corresponding

Figure 3. Overview of our molecule generation paradigm: A
molecular graph G is first decomposed into its junction tree TG ,
where each colored node in the tree represents a substructure in
the molecule. We then encode both the tree and graph into their
latent embeddings zT and zG . To decode the molecule, we first
reconstruct junction tree from zT , and then assemble nodes in the
tree back to original molecule, guided by zG .

junction tree can be found. Empirically, our clusters cover
most of the molecules in the test set.
The original molecular graph and its associated junction tree
offer two complementary representations of a molecule. We
therefore encode the molecule into a two-part latent representation z = [zT , zG ] where zT encodes the tree structure
and what the clusters are in the tree without fully capturing how exactly the clusters are mutually connected. zG
encodes the graph to capture the fine-grained connectivity.
Both parts are created by tree and graph encoders q(zT |T )
and q(zG |G). The latent representation is then decoded
back into a molecular graph in two stages. As illustrated in
Figure 3, we first reproduce the junction tree using a tree
decoder p(T |zT ) based on the information in zT . Second,
we predict the fine grain connectivity between the clusters
in the junction tree using a graph decoder p(G|T , zG ) to
realize the full molecular graph. The junction tree approach
allows us to maintain chemical feasibility during generation.
Notation A molecular graph is defined as G = (V, E)
where V is the set of atoms (vertices) and E the set of bonds
(edges). Let N (x) be the neighbor of x. We denote sigmoid
function as σ(·) and ReLU function as τ (·). We use i, j, k
for nodes in the tree and u, v, w for nodes in the graph.

Junction Tree Variational Autoencoder for Molecular Graph Generation

2.1. Junction Tree
A tree decomposition maps a graph G into a junction tree
by contracting certain vertices into a single node so that G
becomes cycle-free. Formally, given a graph G, a junction
tree TG = (V, E, X ) is a connected labeled tree whose
node set is V = {C1 , · · · , Cn } and edge set is E. Each
node or cluster Ci = (Vi , Ei ) is an induced subgraph of G,
satisfying the following constraints:
S
1. The S
union of all clusters equals G. That is, i Vi = V
and i Ei = E.
2. Running intersection: For all clusters Ci , Cj and Ck ,
Vi ∩ Vj ⊆ Vk if Ck is on the path from Ci to Cj .
Viewing induced subgraphs as cluster labels, junction trees
are labeled trees with label vocabulary X . By our molecule
tree decomposition, X contains only cycles (rings) and single edges. Thus the vocabulary size is limited (|X | = 780
for a standard dataset with 250K molecules).
Tree Decomposition of Molecules Here we present our
tree decomposition algorithm tailored for molecules, which
finds its root in chemistry (Rarey & Dixon, 1998). Our
cluster vocabulary X includes chemical structures such as
bonds and rings (Figure 3). Given a graph G, we first find all
its simple cycles, and its edges not belonging to any cycles.
Two simple rings are merged together if they have more than
two overlapping atoms, as they constitute a specific structure
called bridged compounds (Clayden et al., 2001). Each of
those cycles or edges is considered as a cluster. Next, a
cluster graph is constructed by adding edges between all
intersecting clusters. Finally, we select one of its spanning
trees as the junction tree of G (Figure 3). As a result of ring
merging, any two clusters in the junction tree have at most
two atoms in common, facilitating efficient inference in the
graph decoding phase. The detailed procedure is described
in the supplementary.
2.2. Graph Encoder
We first encode the latent representation of G by a graph
message passing network (Dai et al., 2016; Gilmer et al.,
2017). Each vertex v has a feature vector xv indicating the
atom type, valence, and other properties. Similarly, each
edge (u, v) ∈ E has a feature vector xuv indicating its
bond type, and two hidden vectors ν uv and ν vu denoting
the message from u to v and vice versa. Due to the loopy
structure of the graph, messages are exchanged in a loopy
belief propagation fashion:
X
g
g
g
ν (t)
ν (t−1)
uv = τ (W1 xu + W2 xuv + W3
wu ) (1)
w∈N (u)\v
(t)

where ν uv is the message computed in t-th iteration, initial(0)
ized with ν uv = 0. After T steps of iteration, we aggregate

those messages as the latent vector of each vertex, which
captures its local graphical structure:
X
)
hu = τ (Ug1 xu +
Ug2 ν (T
(2)
vu )
v∈N (u)

P
The final graph representation is hG = i hi /|V |. The
mean µG and log variance log σ G of the variational posterior approximation are computed from hG with two separate
affine layers. zG is sampled from a Gaussian N (µG , σ G ).
2.3. Tree Encoder
We similarly encode TG with a tree message passing network. Each cluster Ci is represented by a one-hot encoding
xi representing its label type. Each edge (Ci , Cj ) is associated with two message vectors mij and mji . We pick an
arbitrary leaf node as the root and propagate messages in
two phases. In the first bottom-up phase, messages are initiated from the leaf nodes and propagated iteratively towards
root. In the top-down phase, messages are propagated from
the root to all the leaf nodes. Message mij is updated as:
mij = GRU(xi , {mki }k∈N (i)\j )

(3)

where GRU is a Gated Recurrent Unit (Chung et al., 2014;
Li et al., 2015) adapted for tree message passing:
X
sij =
mki
(4)
k∈N (i)\j

zij
rki
e ij
m

= σ(Wz xi + Uz sij + bz )
r

r

r

= σ(W xi + U mki + b )
X
= tanh(Wxi + U
rki

(5)
(6)
mki ) (7)

k∈N (i)\j

mij

=

(1 − zij )

sij + zij

e ij
m

(8)

The message passing follows the schedule where mij is
computed only when all its precursors {mki | k ∈ N (i)\j}
have been computed. This architectural design is motivated
by the belief propagation algorithm over trees and is thus
different from the graph encoder.
After the message passing, we obtain the latent representation of each node hi by aggregating its inward messages:
X
hi = τ (Wo xi +
Uo mki )
(9)
k∈N (i)

The final tree representation is hTG = hroot , which encodes
a rooted tree (T , root). Unlike the graph encoder, we do
not apply node average pooling because it confuses the tree
decoder which node to generate first. zTG is sampled in
a similar way as in the graph encoder. For simplicity, we
abbreviate zTG as zT from now on.
This tree encoder plays two roles in our framework. First, it
is used to compute zT , which only requires the bottom-up
phase of the network. Second, after a tree Tb is decoded

Junction Tree Variational Autoencoder for Molecular Graph Generation

Algorithm 1 Tree decoding at sampling time
Require: Latent representation zT
1: Initialize: Tree Tb ← ∅
2: function SampleTree(i, t)
Set Xi ← all cluster labels that are chemically com3:
patible with node i and its current neighbors.
4:
Set dt ← expand with probability pt .
. Eq.(11)
5:
if dt = expand and Xi 6= ∅ then
6:
Create a node j and add it to tree Tb .
7:
Sample the label of node j from Xi
.. Eq.(12)
8:
SampleTree(j, t + 1)
9:
end if
10: end function
Figure 4. Illustration of the tree decoding process. Nodes are labeled in the order in which they are generated. 1) Node 2 expands
child node 4 and predicts its label with message h24 . 2) As node 4
is a leaf node, decoder backtracks and computes message h42 . 3)
Decoder continues to backtrack as node 2 has no more children. 4)
Node 1 expands node 5 and predicts its label.

b ij over the enfrom zT , it is used to compute messages m
tire Tb , to provide essential contexts of every node during
graph decoding. This requires both top-down and bottom-up
phases. We will elaborate this in section 2.5.
2.4. Tree Decoder
We decode a junction tree T from its encoding zT with a tree
structured decoder. The tree is constructed in a top-down
fashion by generating one node at a time. As illustrated
in Figure 4, our tree decoder traverses the entire tree from
the root, and generates nodes in their depth-first order. For
every visited node, the decoder first makes a topological
prediction: whether this node has children to be generated.
When a new child node is created, we predict its label and
recurse this process. Recall that cluster labels represent
subgraphs in a molecule. The decoder backtracks when a
node has no more children to generate.
At each time step, a node receives information from other
nodes in the current tree for making those predictions. The
information is propagated through message vectors hij
when trees are incrementally constructed. Formally, let
Ẽ = {(i1 , j1 ), · · · , (im , jm )} be the edges traversed in a
depth first traversal over T = (V, E), where m = 2|E| as
each edge is traversed in both directions. The model visits node it at time t. Let Ẽt be the first t edges in Ẽ. The
message hit ,jt is updated through previous messages:
hit ,jt = GRU(xit , {hk,it }(k,it )∈Ẽt ,k6=jt )

(10)

where GRU is the same recurrent unit as in the tree encoder.
Topological Prediction When the model visits node it , it
makes a binary prediction on whether it still has children to
be generated. We compute this probability by combining

zT , node features xit and inward messages hk,it via a one
hidden layer network followed by a sigmoid function:
X
pt = σ(ud ·τ (W1d xit +W2d zT +W3d
hk,it ) (11)
(k,it )∈Ẽt

Label Prediction When a child node j is generated from
its parent i, we predict its node label with
qj = softmax(Ul τ (W1l zT + W2l hij ))

(12)

where qj is a distribution over label vocabulary X . When j
is a root node, its parent i is a virtual node and hij = 0.
Learning The tree decoder aims to maximize the likelihood p(T |zT ). Let p̂t ∈ {0, 1} and q̂j be the ground truth
topological and label values, the decoder minimizes the
following cross entropy loss:1
X
X
Lc (T ) =
Ld (pt , p̂t ) +
Ll (qj , q̂j )
(13)
t

j

Similar to sequence generation, during training we perform
teacher forcing: after topological and label prediction at
each step, we replace them with their ground truth so that
the model makes predictions given correct histories.
Decoding & Feasibility Check Algorithm 1 shows how a
tree is sampled from zT . The tree is constructed recursively
guided by topological predictions without any external guidance used in training. To ensure the sampled tree could be
realized into a valid molecule, we define set Xi to be cluster
labels that are chemically compatible with node i and its
current neighbors. When a child node j is generated from
node i, we sample its label from Xi with a renormalized
distribution qj over Xi by masking out invalid labels.
2.5. Graph Decoder
The final step of our model is to reproduce a molecular graph
b E).
b
G that underlies the predicted junction tree Tb = (V,
1

The node ordering is not unique as the order within sibling
nodes is ambiguous. In this paper we train our model with one
ordering and leave this issue for future work.

Junction Tree Variational Autoencoder for Molecular Graph Generation

the assembly of the root and its neighbors according to their
scores. Then we proceed to assemble the neighbors and
their associated clusters (removing the degrees of freedom
set by the root assembly), and so on.
It remains to be specified how each neighborhood realization is scored. Let Gi be the subgraph resulting from a
particular merging of cluster Ci in the tree with its neighbors Cj , j ∈ NTb (i). We score Gi as a candidate subgraph
by first deriving a vector representation hGi and then using
fia (Gi ) = hGi · zG as the subgraph score. To this end,
let u, v specify atoms in the candidate subgraph Gi and let
αv = i if v ∈ Ci and αv = j if v ∈ Cj \ Ci . The indices
αv are used to mark the position of the atoms in the junction
b i,j summarizing the subtree, and to retrieve messages m
tree under i along the edge (i, j) obtained by running the
tree encoding algorithm. The neural messages pertaining
to the atoms and bonds in subgraph Gi are obtained and
aggregated into hGi , similarly to the encoding step, but with
different (learned) parameters:
µ(t)
uv
Figure 5. Decode a molecule from a junction tree. 1) Ground truth
molecule G. 2) Predicted junction tree Tb . 3) We enumerate different combinations between red cluster C and its neighbors. Crossed
arrows indicate combinations that lead to chemically infeasible
molecules. Note that if we discard tree structure during enumeration (i.e., ignoring subtree A), the last two candidates will collapse
into the same molecule. 4) Rank subgraphs at each node. The final
graph is decoded by putting together all the predicted subgraphs
(dashed box).

Note that this step is not deterministic since there are potentially many molecules that correspond to the same junction
tree. The underlying degree of freedom pertains to how
neighboring clusters Ci and Cj are attached to each other
as subgraphs. Our goal here is to assemble the subgraphs
(nodes in the tree) together into the correct molecular graph.
Let G(T ) be the set of graphs whose junction tree is T . Deb E)
b is a structured prediction:
coding graph Ĝ from Tb = (V,
Ĝ = arg max f a (G0 )

(14)

G0 ∈G(Tb )

where f a is a scoring function over candidate graphs. We
only consider scoring functions that decompose across the
clusters and their neighbors. In other words, each term in
the scoring function depends only on how a cluster Ci is
attached to its neighboring clusters Cj , j ∈ NTb (i) in the
tree Tb . The problem of finding the highest scoring graph Ĝ –
the assembly task – could be cast as a graphical model inference task in a model induced by the junction tree. However,
for efficiency reasons, we will assemble the molecular graph
one neighborhood at a time, following the order in which the
tree itself was decoded. In other words, we start by sampling

e (t−1)
µ
uv

e (t−1)
= τ (W1a xu + W2a xuv + W3a µ
)
(15)
uv
(P
(t−1)
αu = αv
w∈N (u)\v µwu
=
P
(t−1)
b αu ,αv + w∈N (u)\v µwu
m
αu 6= αv

The major difference from Eq. (1) is that we augment the
b αu ,αv derived by running the
model with tree messages m
b αu ,αv provides a
tree encoder over the predicted tree Tb . m
tree dependent positional context for bond (u, v) (illustrated
as subtree A in Figure 5).
Learning The graph decoder parameters are learned to
maximize the log-likelihood of predicting correct subgraphs
Gi of the ground true graph G at each tree node:


X
X
f a (Gi ) − log
exp(f a (G0i )) (16)
Lg (G) =
i

G0i ∈Gi

where Gi is the set of possible candidate subgraphs at tree
node i. During training, we again apply teacher forcing, i.e.
we feed the graph decoder with ground truth trees as input.
Complexity By our tree decomposition, any two clusters
share at most two atoms, so we only need to merge at most
two atoms or one bond. By pruning chemically invalid
subgraphs and merging isomorphic graphs, |Gi | ≈ 4 on
average when tested on a standard ZINC drug dataset. The
computational complexity of JT-VAE is therefore linear in
the number of clusters, scaling nicely to large graphs.

3. Experiments
Our evaluation efforts measure various aspects of molecular
generation. The first two evaluations follow previously
proposed tasks (Kusner et al., 2017). We also introduce a
third task — constrained molecule optimization.

Junction Tree Variational Autoencoder for Molecular Graph Generation
O-

N
N

O

NH

NH

O

N

O

NH

N

N
H

H3N+

N

N

S

NH

N

O

N

NH

NH2+

N

N

N

NH

N

NH

NH

N
N

N

NH
N

N

N

N

N

N

NH

N

N
N

N

N

NH

N
H

O

N

H

NH

N
H

NH+

H

O

N
H

NH

O

S
NH

N

N

NH2+

N

NH2+

N

N

N

N

N

N

N

Cl
O

O

O

N

N

N

N

NH

S

N

N

N

N

NH

N

N

N

N

N

N

NH

NH

N
N

N

N

N
O

N

N

N

N

O
N

N
N

H

N

N

N

N

N

N

NH
S

O

NH

N

N
NH

N

S
NH

N
O

N

S
N
NH

S

N

NH

NH
N

O

N
S

N

O
O

O
O

NH2+

S
N

N

N

N

N

NH

N

N

S

NH
N

N

N

N

NH

N

N
NH

NH

N

N

N

NH

N

N

NH

N

N

N

N

O

N

NH

N

N

O

N

N

O

N

N

N

N
N

N

N

N

H2N

F

O

Cl

O

N
N

N

NH

S

N

NH
N

N

NH

N

N

N

N

O

NH+
O

N

N

NH
N

S

O

O

N
S
NH

N

NH

N

N

N

N

N

N

N

N

N

N

N

N
NH
N

NH

N

N

N

N

O

N

NH
N

O

NH

N

N

NH

N

N

NH

N

N

N

N

O

N

NH
N

N

N
N

N

O
NH3+

S
N

O

O
NH

N

O

O

N

H2N

N

N

N

N

N

N

N

N

N

N

N

N

N

N

N

NH

NH

NH

NH

N

NH

N

N

N

O

O

N

NH
S

S

N

NH

N

N

NH

N

N

NH

N

N

N

N

NH

NH

N

N

N

N

N

S

NH

N

N

O

N

N
N

NH

N

N

S

N

N

N

N

N

N

N

N

N

N

S

S

O

O

O

N

N

NH

S

S

NH

N

NH2

NH

NH2

N

NH

N

NH

N

N

NH

N

N

NH

N

N

N

N

NH

NH

N

N

N

N

NH

O

O
O

OH

O

O
O

H2N

O

NH

NH

NH2

N
N

N

N
N

NH

N

N

N

N

NH+

N

N

NH
N

NH 2

S

N

N

N
NH

N

S
NH2

N

S

N

S

S
NH2

N

NH

N

NH

S

S

NH

N

NH

N

N

NH

N

N

N

N

NH2

NH
N

N
N

O

N

OH

O

N
NH
N

N

NH+

NH2+
O

S

N

NH
O

N

N

N
O

N

N

NH2+

S

N

S

S

N
O

N
H

N

NH+

N

S
N

N

NH

NH2

NH

N

N

N

N

NH

N

N

N

N

N

O

H

O

N

S

N
NH
N

N

S
NH

O

S
N

N
NH+
N

N

NH2
S

N

O

NH2
O

NH
2

2

N

N

S

S

NH2

S

N

N
NH

N
NH
NH

O

N

NH

NH

N

O

N

N

N

O

NH2+

O

N

NH2+

H 2N

H 2N
S

N

O
NH

N
NH

OH

N

S

NH2

NH
N

S

S

S

N

S

NH
Cl

NH

N

S

NH

N

N
NH

NH

N

N

N
N
S
N

NH

NH

NH2

N
NH

2

NH+

N
S
N

N
N

S
N

NH+

N

N

S

NH2

NH2

NH
N

N

N

N

Figure 6. Left: Random molecules sampled from prior distribution N (0, I). Right: Visualization of the local neighborhood of a molecule
in the center. Three molecules highlighted in red dashed box have the same tree structure as the center molecule, but with different graph
structure as their clusters are combined differently. The same phenomenon emerges in another group of molecules (blue dashed box).

• Molecule reconstruction and validity We test the VAE
models on the task of reconstructing input molecules from
their latent representations, and decoding valid molecules
when sampling from prior distribution. (Section 3.1)
• Bayesian optimization Moving beyond generating valid
molecules, we test how the model can produce novel
molecules with desired properties. To this end, we perform Bayesian optimization in the latent space to search
molecules with specified properties. (Section 3.2)
• Constrained molecule optimization The task is to modify given molecules to improve specified properties, while
constraining the degree of deviation from the original
molecule. This is a more realistic scenario in drug discovery, where development of new drugs usually starts with
known molecules such as existing drugs (Besnard et al.,
2012). Since it is a new task, we cannot compare to any
existing baselines. (Section 3.3)
Below we describe the data, baselines and model configuration that are shared across the tasks. Additional setup details
are provided in the task-specific sections.
Data We use the ZINC molecule dataset from Kusner et al.
(2017) for our experiments. It contains about 250K drug
molecules extracted from the ZINC database (Sterling &
Irwin, 2015). We follow the same training/testing split as
the previous work.
Baselines We compare our approach with SMILES-based
baselines: 1) Character VAE (CVAE) (Gómez-Bombarelli
et al., 2016) which generates SMILES strings character by
character; 2) Grammar VAE (GVAE) (Kusner et al., 2017)

that generates SMILES following syntactic constraints given
by a context-free grammar; 3) Syntax-directed VAE (SDVAE) (Dai et al., 2018) that incorporates both syntactic
and semantic constraints of SMILES via attribute grammar. For molecule generation task, we also compare with
GraphVAE (Simonovsky & Komodakis, 2018) that directly
generates atom labels and adjacency matrices of graphs.
Model Configuration To be comparable with the above
baselines, we set the latent space dimension as 56, i.e., the
tree and graph representation hT and hG have 28 dimensions each. Full training details and model configurations
are provided in the appendix.
3.1. Molecule Reconstruction and Validity
Setup The first task is to reconstruct and sample molecules
from latent space. Since both encoding and decoding process are stochastic, we estimate reconstruction accuracy by
Monte Carlo method used in (Kusner et al., 2017): Each
molecule is encoded 10 times and each encoding is decoded 10 times. We report the portion of the 100 decoded
molecules that are identical to the input molecule. For a fair
comparison, we define two molecules as identical if they
have the same SMILES string. At testing time, we convert
all generated graphs to SMILES using RDKit.
To compute validity, we sample 1000 latent vectors from the
prior distribution N (0, I), and decode each of these vectors
100 times. We report the percentage of decoded molecules
that are chemically valid (checked by RDKit).

Junction Tree Variational Autoencoder for Molecular Graph Generation
Table 1. Reconstruction accuracy and prior validity results. Baseline results are copied from Kusner et al. (2017); Dai et al. (2018);
Simonovsky & Komodakis (2018).

Method
CVAE
GVAE
SD-VAE2
GraphVAE
JT-VAE

Reconstruction
44.6%
53.7%
76.2%
76.7%

Validity
0.7%
7.2%
43.5%
13.5%
100.0%

Results Table 1 shows that JT-VAE outperforms previous models in molecule reconstruction, and always produces valid molecules when sampled from prior distribution. These sampled molecules have non-trivial structures
such as simple chains (Figure 6). We further sampled 5000
molecules from prior and found they are all distinct from the
training set. Thus our model is not a simple memorization.
Analysis We qualitatively examine the latent space of JTVAE by visualizing the neighborhood of molecules. Given
a molecule, we follow the method in Kusner et al. (2017)
to construct a grid visualization of its neighborhood. For
comparison, we select the same molecule visualized in Dai
et al. (2018). Figure 6 shows the local neighborhood of
this molecule. Compared to the figure in Dai et al. (2018),
our neighborhood does not contain molecules with huge
rings (with more than 7 atoms), which rarely occur in the
dataset. We also highlight two groups of closely resembling
molecules that have identical tree structures but vary only
in how clusters are attached together. This demonstrates the
smoothness of learned molecular embeddings.
3.2. Bayesian Optimization
Setup The second task is to produce novel molecules with
desired properties. Following (Kusner et al., 2017), our
target chemical property y(·) is octanol-water partition coefficients (logP) penalized by the synthetic accessibility (SA)
score and number of long cycles.3 To perform Bayesian
optimization (BO), we first train a VAE and associate each
molecule with a latent vector, given by the mean of the variational encoding distribution. After the VAE is learned, we
train a sparse Gaussian process (SGP) to predict y(m) given
its latent representation. Then we perform five iterations of
batched BO using the expected improvement heuristic.
For comparison, we report 1) the predictive performance of
SGP trained on latent encodings learned by different VAEs,
measured by log-likelihood (LL) and root mean square error (RMSE) with 10-fold cross validation. 2) The top-3
molecules found by BO under different models.
2

The SD-VAE result is copied from Table 1 in Dai et al. (2018).
y(m) = logP (m) − SA(m) − cycle(m) where cycle(m)
counts the number of rings that have more than six atoms.
3

Table 2. Best molecule property scores found by each method.
Baseline results are from Kusner et al. (2017); Dai et al. (2018).

Method
CVAE
GVAE
SD-VAE
JT-VAE

1st
1.98
2.94
4.04
5.30

2nd
1.42
2.89
3.50
4.93

3rd
1.19
2.80
2.96
4.49

Figure 7. Best three molecules and their property scores found by
JT-VAE using Bayesian optimization.

Results As shown in Table 2, JT-VAE finds molecules with
significantly better scores than previous methods. Figure 7
lists the top-3 best molecules found by JT-VAE. In fact,
JT-VAE finds over 50 molecules with scores over 3.50 (the
second best molecule proposed by SD-VAE). Moreover, the
SGP yields better predictive performance when trained on
JT-VAE embeddings (Table 3).
3.3. Constrained Optimization
Setup The third task is to perform molecule optimization
in a constrained scenario. Given a molecule m, the task is
to find a different molecule m0 that has the highest property
value with the molecular similarity sim(m, m0 ) ≥ δ for
some threshold δ. We use Tanimoto similarity with Morgan
fingerprint (Rogers & Hahn, 2010) as the similarity metric,
and penalized logP coefficient as our target chemical property. For this task, we jointly train a property predictor F
(parameterized by a feed-forward network) with JT-VAE to
predict y(m) from the latent embedding of m. To optimize
a molecule m, we start from its latent representation, and
apply gradient ascent in the latent space to improve the predicted score F (·), similar to (Mueller et al., 2017). After
applying K = 80 gradient steps, K molecules are decoded
from resulting latent trajectories, and we report the molecule
with the highest F (·) that satisfies the similarity constraint.
A modification succeeds if one of the decoded molecules
satisfies the constraint and is distinct from the original.
To provide the greatest challenge, we selected 800 molecules
with the lowest property score y(·) from the test set. We
report the success rate (how often a modification succeeds),
and among success cases the average improvement y(m0 ) −
y(m) and molecular similarity sim(m, m0 ) between the
original and modified molecules m and m0 .

Junction Tree Variational Autoencoder for Molecular Graph Generation
Table 3. Predictive performance of sparse Gaussian Processes
trained on different VAEs. Baseline results are copied from Kusner
et al. (2017) and Dai et al. (2018).

Method
CVAE
GVAE
SD-VAE
JT-VAE

LL
−1.812 ± 0.004
−1.739 ± 0.004
−1.697 ± 0.015
−1.658 ± 0.023

RMSE
1.504 ± 0.006
1.404 ± 0.006
1.366 ± 0.023
1.290 ± 0.026

Table 4. Constrained optimization result of JT-VAE: mean and
standard deviation of property improvement, molecular similarity
and success rate under constraints sim(m, m0 ) ≥ δ with varied δ.

δ
0.0
0.2
0.4
0.6

Improvement
1.91 ± 2.04
1.68 ± 1.85
0.84 ± 1.45
0.21 ± 0.71

Similarity
0.28 ± 0.15
0.33 ± 0.13
0.51 ± 0.10
0.69 ± 0.06

Success
97.5%
97.1%
83.6%
46.4%

Results Our results are summarized in Table 4. The unconstrained scenario (δ = 0) has the best average improvement,
but often proposes dissimilar molecules. When we tighten
the constraint to δ = 0.4, about 80% of the time our model
finds similar molecules, with an average improvement 0.84.
This also demonstrates the smoothness of the learned latent
space. Figure 8 illustrates an effective modification resulting
in a similar molecule with great improvement.

4. Related Work
Molecule Generation Previous work on molecule generation mostly operates on SMILES strings. GómezBombarelli et al. (2016); Segler et al. (2017) built generative models of SMILES strings with recurrent decoders.
Unfortunately, these models could generate invalid SMILES
that do not result in any molecules. To remedy this issue,
Kusner et al. (2017); Dai et al. (2018) complemented the
decoder with syntactic and semantic constraints of SMILES
by context free and attribute grammars, but these grammars
do not fully capture chemical validity. Other techniques
such as active learning (Janz et al., 2017) and reinforcement
learning (Guimaraes et al., 2017) encourage the model to
generate valid SMILES through additional training signal.
Very recently, Simonovsky & Komodakis (2018) proposed
to generate molecular graphs by predicting their adjacency
matrices, and Li et al. (2018) generated molecules node by
node. In comparison, our method enforces chemical validity
and is more efficient due to the coarse-to-fine generation.
Graph-structured Encoders The neural network formulation on graphs was first proposed by Gori et al. (2005);
Scarselli et al. (2009), and later enhanced by Li et al. (2015)
with gated recurrent units. For recurrent architectures over

Figure 8. A molecule modification that yields an improvement of
4.0 with molecular similarity 0.617 (modified part is in red).

graphs, Lei et al. (2017) designed Weisfeiler-Lehman kernel
network inspired by graph kernels. Dai et al. (2016) considered a different architecture where graphs were viewed as latent variable graphical models, and derived their model from
message passing algorithms. Our tree and graph encoder are
closely related to this graphical model perspective, and to
neural message passing networks (Gilmer et al., 2017). For
convolutional architectures, Duvenaud et al. (2015) introduced a convolution-like propagation on molecular graphs,
which was generalized to other domains by Niepert et al.
(2016). Bruna et al. (2013); Henaff et al. (2015) developed
graph convolution in spectral domain via graph Laplacian.
For applications, graph neural networks are used in semisupervised classification (Kipf & Welling, 2016), computer
vision (Monti et al., 2016), and chemical domains (Kearnes
et al., 2016; Schütt et al., 2017; Jin et al., 2017).
Tree-structured Models Our tree encoder is related to recursive neural networks and tree-LSTM (Socher et al., 2013;
Tai et al., 2015; Zhu et al., 2015). These models encode
tree structures where nodes in the tree are bottom-up transformed into vector representations. In contrast, our model
propagates information both bottom-up and top-down.
On the decoding side, tree generation naturally arises in
natural language parsing (Dyer et al., 2016; Kiperwasser &
Goldberg, 2016). Different from our approach, natural language parsers have access to input words and only predict
the topology of the tree. For general purpose tree generation,
Vinyals et al. (2015); Aharoni & Goldberg (2017) applied recurrent networks to generate linearized version of trees, but
their architectures were entirely sequence-based. Dong &
Lapata (2016); Alvarez-Melis & Jaakkola (2016) proposed
tree-based architectures that construct trees top-down from
the root. Our model is most closely related to Alvarez-Melis
& Jaakkola (2016) that disentangles topological prediction
from label prediction, but we generate nodes in a depth-first
order and have additional steps that propagate information
bottom-up. This forward-backward propagation also appears in Parisotto et al. (2016), but their model is node
based whereas ours is based on message passing.

5. Conclusion
In this paper we present a junction tree variational autoencoder for generating molecular graphs. Our method significantly outperforms previous work in molecule generation
and optimization. For future work, we attempt to generalize
our method for general low-treewidth graphs.

Junction Tree Variational Autoencoder for Molecular Graph Generation

Acknowledgement
We thank Jonas Mueller, Chengtao Li, Tao Lei and MIT
NLP Group for their helpful comments. This work was
supported by the DARPA Make-It program under contract
ARO W911NF-16-2-0023.

References
Aharoni, Roee and Goldberg, Yoav.
to-tree neural machine translation.
arXiv:1704.04743, 2017.

Towards stringarXiv preprint

Alvarez-Melis, David and Jaakkola, Tommi S. Treestructured decoding with doubly-recurrent neural networks. 2016.
Besnard, Jérémy, Ruda, Gian Filippo, Setola, Vincent,
Abecassis, Keren, Rodriguiz, Ramona M, Huang, XiPing, Norval, Suzanne, Sassano, Maria F, Shin, Antony I,
Webster, Lauren A, et al. Automated design of ligands
to polypharmacological profiles. Nature, 492(7428):215–
220, 2012.
Bruna, Joan, Zaremba, Wojciech, Szlam, Arthur, and LeCun,
Yann. Spectral networks and locally connected networks
on graphs. arXiv preprint arXiv:1312.6203, 2013.
Chung, Junyoung, Gulcehre, Caglar, Cho, KyungHyun, and
Bengio, Yoshua. Empirical evaluation of gated recurrent
neural networks on sequence modeling. arXiv preprint
arXiv:1412.3555, 2014.
Clayden, Jonathan, Greeves, Nick, Warren, Stuart, and
Wothers, P. Organic Chemistry. Oxford University Press,
2001.
Dai, Hanjun, Dai, Bo, and Song, Le. Discriminative embeddings of latent variable models for structured data.
In International Conference on Machine Learning, pp.
2702–2711, 2016.
Dai, Hanjun, Tian, Yingtao, Dai, Bo, Skiena, Steven, and
Song, Le. Syntax-directed variational autoencoder for
structured data. International Conference on Learning
Representations, 2018. URL https://openreview.
net/forum?id=SyqShMZRb.
Dong, Li and Lapata, Mirella. Language to logical form
with neural attention. arXiv preprint arXiv:1601.01280,
2016.
Duvenaud, David K, Maclaurin, Dougal, Iparraguirre, Jorge,
Bombarell, Rafael, Hirzel, Timothy, Aspuru-Guzik, Alán,
and Adams, Ryan P. Convolutional networks on graphs
for learning molecular fingerprints. In Advances in neural
information processing systems, pp. 2224–2232, 2015.

Dyer, Chris, Kuncoro, Adhiguna, Ballesteros, Miguel, and
Smith, Noah A. Recurrent neural network grammars.
arXiv preprint arXiv:1602.07776, 2016.
Gilmer, Justin, Schoenholz, Samuel S, Riley, Patrick F,
Vinyals, Oriol, and Dahl, George E. Neural message passing for quantum chemistry. arXiv preprint
arXiv:1704.01212, 2017.
Gómez-Bombarelli, Rafael, Wei, Jennifer N, Duvenaud, David, Hernández-Lobato, José Miguel, SánchezLengeling, Benjamı́n, Sheberla, Dennis, AguileraIparraguirre, Jorge, Hirzel, Timothy D, Adams, Ryan P,
and Aspuru-Guzik, Alán. Automatic chemical design using a data-driven continuous representation of molecules.
ACS Central Science, 2016. doi: 10.1021/acscentsci.
7b00572.
Gori, Marco, Monfardini, Gabriele, and Scarselli, Franco.
A new model for learning in graph domains. In Neural
Networks, 2005. IJCNN’05. Proceedings. 2005 IEEE International Joint Conference on, volume 2, pp. 729–734.
IEEE, 2005.
Guimaraes, Gabriel Lima, Sanchez-Lengeling, Benjamin,
Farias, Pedro Luis Cunha, and Aspuru-Guzik, Alán.
Objective-reinforced generative adversarial networks (organ) for sequence generation models. arXiv preprint
arXiv:1705.10843, 2017.
Henaff, Mikael, Bruna, Joan, and LeCun, Yann. Deep
convolutional networks on graph-structured data. arXiv
preprint arXiv:1506.05163, 2015.
Janz, David, van der Westhuizen, Jos, and HernándezLobato, José Miguel. Actively learning what makes a
discrete sequence valid. arXiv preprint arXiv:1708.04465,
2017.
Jin, Wengong, Coley, Connor, Barzilay, Regina, and
Jaakkola, Tommi. Predicting organic reaction outcomes
with weisfeiler-lehman network. In Advances in Neural
Information Processing Systems, pp. 2604–2613, 2017.
Kearnes, Steven, McCloskey, Kevin, Berndl, Marc, Pande,
Vijay, and Riley, Patrick. Molecular graph convolutions:
moving beyond fingerprints. Journal of computer-aided
molecular design, 30(8):595–608, 2016.
Kingma, Diederik P and Welling, Max. Auto-encoding
variational bayes. arXiv preprint arXiv:1312.6114, 2013.
Kiperwasser, Eliyahu and Goldberg, Yoav. Easy-first dependency parsing with hierarchical tree lstms. arXiv preprint
arXiv:1603.00375, 2016.
Kipf, Thomas N and Welling, Max. Semi-supervised classification with graph convolutional networks. arXiv
preprint arXiv:1609.02907, 2016.

Junction Tree Variational Autoencoder for Molecular Graph Generation

Kusner, Matt J, Paige, Brooks, and Hernández-Lobato,
José Miguel. Grammar variational autoencoder. arXiv
preprint arXiv:1703.01925, 2017.

convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing
Systems, pp. 992–1002, 2017.

Landrum, Greg. Rdkit: Open-source cheminformatics. Online). http://www. rdkit. org. Accessed, 3(04):2012, 2006.

Segler, Marwin HS, Kogej, Thierry, Tyrchan, Christian, and
Waller, Mark P. Generating focussed molecule libraries
for drug discovery with recurrent neural networks. arXiv
preprint arXiv:1701.01329, 2017.

Lei, Tao, Jin, Wengong, Barzilay, Regina, and Jaakkola,
Tommi. Deriving neural architectures from sequence and
graph kernels. arXiv preprint arXiv:1705.09037, 2017.
Li, Yujia, Tarlow, Daniel, Brockschmidt, Marc, and Zemel,
Richard. Gated graph sequence neural networks. arXiv
preprint arXiv:1511.05493, 2015.
Li, Yujia, Vinyals, Oriol, Dyer, Chris, Pascanu, Razvan,
and Battaglia, Peter. Learning deep generative models of
graphs. 2018. URL https://openreview.net/
forum?id=Hy1d-ebAb.
Monti, Federico, Boscaini, Davide, Masci, Jonathan,
Rodolà, Emanuele, Svoboda, Jan, and Bronstein,
Michael M. Geometric deep learning on graphs and
manifolds using mixture model cnns. arXiv preprint
arXiv:1611.08402, 2016.
Mueller, Jonas, Gifford, David, and Jaakkola, Tommi. Sequence to better sequence: continuous revision of combinatorial structures. In International Conference on
Machine Learning, pp. 2536–2544, 2017.
Niepert, Mathias, Ahmed, Mohamed, and Kutzkov, Konstantin. Learning convolutional neural networks for
graphs. In International Conference on Machine Learning, pp. 2014–2023, 2016.
Parisotto, Emilio, Mohamed, Abdel-rahman, Singh,
Rishabh, Li, Lihong, Zhou, Dengyong, and Kohli, Pushmeet. Neuro-symbolic program synthesis. arXiv preprint
arXiv:1611.01855, 2016.
Rarey, Matthias and Dixon, J Scott. Feature trees: a new
molecular similarity measure based on tree matching.
Journal of computer-aided molecular design, 12(5):471–
490, 1998.
Rogers, David and Hahn, Mathew. Extended-connectivity
fingerprints. Journal of chemical information and modeling, 50(5):742–754, 2010.
Scarselli, Franco, Gori, Marco, Tsoi, Ah Chung, Hagenbuchner, Markus, and Monfardini, Gabriele. The graph
neural network model. IEEE Transactions on Neural
Networks, 20(1):61–80, 2009.
Schütt, Kristof, Kindermans, Pieter-Jan, Felix, Huziel
Enoc Sauceda, Chmiela, Stefan, Tkatchenko, Alexandre,
and Müller, Klaus-Robert. Schnet: A continuous-filter

Simonovsky, Martin and Komodakis, Nikos. Graphvae:
Towards generation of small graphs using variational autoencoders. arXiv preprint arXiv:1802.03480, 2018.
Socher, Richard, Perelygin, Alex, Wu, Jean, Chuang, Jason, Manning, Christopher D, Ng, Andrew, and Potts,
Christopher. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings
of the 2013 conference on empirical methods in natural
language processing, pp. 1631–1642, 2013.
Sterling, Teague and Irwin, John J. Zinc 15–ligand discovery
for everyone. J. Chem. Inf. Model, 55(11):2324–2337,
2015.
Tai, Kai Sheng, Socher, Richard, and Manning, Christopher D. Improved semantic representations from treestructured long short-term memory networks. arXiv
preprint arXiv:1503.00075, 2015.
Vinyals, Oriol, Kaiser, Łukasz, Koo, Terry, Petrov, Slav,
Sutskever, Ilya, and Hinton, Geoffrey. Grammar as a
foreign language. In Advances in Neural Information
Processing Systems, pp. 2773–2781, 2015.
Weininger, David. 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.
Zhu, Xiaodan, Sobihani, Parinaz, and Guo, Hongyu. Long
short-term memory over recursive structures. In International Conference on Machine Learning, pp. 1604–1612,
2015.

Supplementary Material

A. Tree Decomposition
Algorithm 2 presents our tree decomposition of molecules. V1 and V2 contain non-ring bonds and simple rings respectively.
Simple rings are extracted via RDKit’s GetSymmSSSR function. We then merge rings that share three or more atoms as
they form bridged compounds. We note that the junction tree of a molecule is not unique when its cluster graph contains
cycles. This introduces additional uncertainty for our probabilistic modeling. To reduce such variation, for any of the three
(or more) intersecting bonds, we add their intersecting atom as a cluster and remove the cycle connecting them in the cluster
graph. Finally, we construct a junction tree as the maximum spanning tree of a cluster graph (V, E). Note that we assign an
large weight over edges involving clusters in V0 to ensure no edges in any cycles will be selected into the junction tree.
Algorithm 2 Tree decomposition of molecule G = (V, E)
V1 ← the set of bonds (u, v) ∈ E that do not belong to any rings.
V2 ← the set of simple rings of G.
for r1 , r2 in V2 do
Merge rings r1 , r2 into one ring if they share more than two atoms (bridged rings).
end for
V0 ← atoms being the intersection of three or more clusters in V1 ∪ V2 .
V ← V0 ∪ V1 ∪ V2
E ← {(i, j, c) ∈ V × V × R | |i ∩ j| > 0}. Set c = ∞ if i ∈ V0 or j ∈ V0 , and c = 1 otherwise.
Return The maximum spanning tree over cluster graph (V, E).

Figure 9. Illustration of tree decomposition and sample of cluster label vocabulary.

B. Stereochemistry
Though usually presented as two-dimensional graphs, molecules are three-dimensional objects, i.e. molecules are defined
not only by its atom types and bond connections, but also the spatial configuration between atoms (chiral atoms and cis-trans
isomerism). Stereoisomers are molecules that have the same 2D structure, but differ in the 3D orientations of their atoms in
space. We note that stereochemical feasibility could not be simply encoded as context free or attribute grammars.
Empirically, we found it more efficient to predict the stereochemical configuration separately from the molecule generation.
Specifically, the JT-VAE first generates the 2D structure of a molecule m, following the same procedure described in
section 2. Then we generate all its stereoisomers Sm using RDKit’s EnumerateStereoisomers function, which
identifies atoms that could be chiral. For each isomer m0 ∈ Sm , we encode its graph representation hm0 with the graph
encoder and compute their cosine similarity f s (m0 ) = cos(hm0 , zm ) (note that zm is stochastic). We reconstruct the

Junction Tree Variational Autoencoder for Molecular Graph Generation

final 3D structure by picking the stereoisomer m
b = arg maxm0 f s (m0 ). Since on average only few atoms could have
stereochemical variations, this post ranking process is very efficient. Combining this with tree and graph generation, the
molecule reconstruction loss L becomes
X
L = Lc + Lg + Ls ;
Ls = f s (m) − log
exp(f s (m0 ))
(17)
m0 ∈Sm

C. Training Details
By applying tree decomposition over 240K molecules in ZINC dataset, we collected our vocabulary set X of size |X | = 780.
The hidden state dimension is 450 for all modules in JT-VAE and the latent bottleneck dimension is 56. For the graph
encoder, the initial atom features include its atom type, degree, its formal charge and its chiral configuration. Bond feature is
a concatenation of its bond type, whether the bond is in a ring, and its cis-trans configuration. For our tree encoder, we
represent each cluster with a neural embedding vector, similar to word embedding for words. The tree and graph decoder
use the same feature setting as encoders. The graph encoder and decoder runs three iterations of neural message passing.
For fair comparison to SMILES based method, we minimized feature engineering. We use PyTorch to implement all neural
components and RDKit to process molecules.

D. More Experimental Results
Sampled Molecules Note that a degenerate model could also achieve 100% prior validity by keep generating simple
structures like chains. To prove that our model does not converge to such trivial solutions, we randomly sample and plot 250
molecules from prior distribution N (0, I). As shown in Figure 10, our sampled molecules present rich variety and structural
complexity. This demonstrates the soundness of the prior validity improvement of our model.
Neighborhood Visualization Given a molecule, we follow Kusner et al. (2017) to construct a grid visualization of its
neighborhood. Specifically, we encode a molecule into the latent space and generate two random orthogonal unit vectors
as two axis of a grid. Moving in combinations of these directions yields a set of latent vectors and we decode them into
corresponding molecules. In Figure 11 and 12, we visualize the local neighborhood of two molecules presented in Dai et al.
(2018). Figure 11 visualizes the same molecule in Figure 6, but with wider neighborhood ranges.
Bayesian Optimization We directly used open sourced implementation in Kusner et al. (2017) for Bayesian optimization
(BO). Specifically, we train a sparse Gaussian process with 500 inducing points to predict properties of molecules. Five
iterations of batch BO with expected improvement heuristic is used to propose new latent vectors. In each iteration, 50 latent
vectors are proposed, from which molecules are decoded and added to the training set for next iteration. We perform 10
independent runs and aggregate results. In Figure 13, we present the top 50 molecules found among 10 runs using JT-VAE.
Following Kusner et al.’s implementation, the scores reported are normalized to zero mean and unit variance by the mean
and variance computed from training set.
Constrained Optimization For this task, a property predictor F is trained jointly with VAE to predict y(m) = logP (m) −
SA(m) from the latent embedding of m. F is a feed-forward network with one hidden layer of dimension 450 followed
by tanh activation. To optimize a molecule m, we start with its mean encoding z0m = µm and apply 80 gradient ascent
∂y
i
steps: ztm = zt−1
m + α ∂z with α = 2.0. 80 molecules are decoded from latent vectors {zm } and their property is calculated.
0
Molecular similarity sim(m, m ) is calculated via Morgan fingerprint of radius 2 with Tanimoto similarity. For each
molecule m, we report the best modified molecule m0 with sim(m, m0 ) > δ for some threshold δ. In Figure 14, we present
three groups of modification examples with δ = 0.2, 0.4, 0.6. For each group, we present top three pairs that leads to best
improvement y(m0 ) − y(m) as well as one pair decreased property (y(m0 ) < y(m)). This is caused by inaccurate property
prediction. From Figure 14, we can see that tighter similarity constraint forces the model to preserve the original structure.

Junction Tree Variational Autoencoder for Molecular Graph Generation

N

Cl

O
O

O

O

H

H

N

O

S

NH+

NH

H

O

NH

O

N

S

O

N

N

NH+
S

O

N

NH

OH

N

O

N

NH2+

NH2

O

N

O
NH

NH

N

N
N
O

F

O

NH
N

O

NH

NH

O

O

NH

OH

OH

O

F

O
N

O

O
O-

N

O

NH

NH

O

N

O

NH

H3N+

N

N

S

NH

N

H

N
NH

N

S

O

NH
O

NH

N

OH

O
NH

O

O

H

N

O

H
NH+

NH2+

N

S

N

NH

O
N

N

S

NH

O

N

N

N

N

O

H

O

N

O

N

N
NH

O
N

NH2+

N

S

O

N

O
O
NH

NH

S

NH

O

H2N

NH2

N
N

O
S

O

O

NH

O
S

NH2
O

O

OH

NH

NH
O

O

O
O
S

N
S
O

N

NH

N
O

O

N
NH

O

S
O

O
N
NH+
O

N

O
NH3+

O

N

N

O

NH
H2N

N

O

O
NH2+

O

N

N

O
O

Cl
O

N

N

N

NH
O

O

N

O

F

O

Cl

N

N
NH

NH+

H

O

N

S

O

S
NH

N

O

N
S

N

H

OH

Cl

NH
O

NH
O

O
O

O

NH

N

O

O
S

O

H2N

NH2

O

N

O

N+

N

N

NH2

N

NH2

H

O

OH

O
N

O

O

N

N

O
O

N

NH+

N

S

N
N

O

O

NH

N

NH
N

N
NH3+

O

Cl
NH+

O

N

S

NH

O-

NH
O

N

O

O

NH

N

H2 N

O

O
O

OH

N

N
Cl

F

O

S

N+

NH2

O-

NH

O

O

NH

O

S

NH2

NH

O
O

N

NH+

NH

N
NH

N

O

O

Cl

N

N

O

NH
O

O

H

S

F
O

NH

O
O

N+

NH
N

N-

S

O

NH2 N

NH

S

NH

O

N

O

O
O

O

HN

O

O

O

N
N

O
OH

H2N

N
O

O

N

NH+

N+

OH

N

NH
NH2

OH

O
NH+

OH

O

S

NH

S

N

N

N

O

S
N

O

O

O

O

O

N

O

NH

N

N
O

NH+

O

O

N
N

O

NH

O
NH

N

O

H
NH

SH

O

O

S

O

S

F
NH

N

O

Br

NH+

N

OH

F

F

HO

NH3+

N

H
S
NH

N

NH2

NH

O

O

NH2+
N

O

S

F
NH

N

O

O

O

N

N
N

F
N

O

N

O

NH2+

N

O

NH
NH+

S

H

NH

O

O

O

O
S

O

ON+

O

NH+

O

NH3+

H3N+

O

NH2

NH

H

Br

O

N

N

N

N

O

NH

N

O

O

F

S
N

O

N

O

S

O

F

NH

N

O

O

N

S

O

O
N

O

O
F

N

NH

O

NH2+

N

O

NH

N
S

N

Br

O

NH

S
H

O

O
O

NH

NH

N

NH

NH
O

O

NH

O

O

N
S

NH
N

N

F

N

F

N

N

HS

O

H

O

N

N

N

N

O

O

O

NH+

NH
N

O

S

NH

NH

H

N

NH

O

S

O

NH3+

NH

O
O

O

NH2

O

Cl

O

NH

O

NH2+

O
N

OH

O

NH2

OH

O
NH+

N

NH

N

NH2

N

N

O

O

N
N

Cl

O
H2N

O

NH3+

N

N

O

N
N

NH

F

NH

N

H

N

O

N

N

NH

OH

NH
O

NH

O
O

S

O

O

NH

Br
NH

O

O

O

O

S

N

N
O

O

NH

N

N

N

O
N

S

O

O

N

NH

S

N

NH2+

O

N

O

NH

H2N
O
NH2
NH

NH

O

S

NH2

O

NH
NH

H

N
O

O

OH

NH2

N

O

NH

NH
N

NH3+
NH

F

OH

O

N

N

NH

O

O

O

NH

O
H3N+

S

N

O

N
N
O

O

O

NH2+

N

N

NH
O

O

O

O

NH

NH
N

O

O

HO
O

O

H

NH+

NH

NH

N

H3N+

O

N

H

NH
NH

NH
O
N

N

O-

NH
H

NH
H
S

S

O

NH

NH2+

NH2+

O

O

N

NH

S

O

NH2+

OH

O

NH2

NH2+

N

O

NH
NH

O

O

N

O

O

O

N

N
N

NH+

O

O
F

S

NH
H2N

H3N+

O

NH2+

O-

Cl
NH

SH

O
N

O

S

H

S

NH

Cl

SH

H2N

O

NH
N

N

O

N

N

NH

NH

N

N

NH

O

O
O

N

H

O

F

O

NH
O

N
O

N

N

O

H2N

O

NH2+

HO

N
O

NH2+

O

N

O

S
N

N
NH2

N

O

O

NH

NH

O

N
O

OH

H3N

NH
O

+

NH3+

O
NH

N-

NH

NH3+
NH2+

NH

O

SH

O

O

NH

O

O

N

NH
S

N

NH2

O

N

O

O

N

NH

N

O

NH

N

N

N

O

N

NH

O

NH
O

O

N

NH

NH+
N

O

F

F

O

O

OH

S
O

NH2
O

O

O

O

O

Cl

O

F

NH
N
O

Br

N
O

F

NH+

S

O
O-

HO

F

S
N

N

N

Cl

O

N

N

O

N

H3N+

NH

O

O

N

H

S

O

N

O

N
O
NH

O

O

N

S

N
O

N
NH
O

NH2

H

N

NH

NH

S
NH

N

Cl

O
O

O

O

O

N

N

O

O

O

O

O
NH+
O

O

NH
NH

N

N
N

N

O

NH

N

N
NH

N

N
NH

O

N
N

O

O

OH

O

O

O

N

N-

N

N
O

N

N

NH

S

O

N

N

O

O

O

N

NH

N

O

O

NH3+
NH2
NH

NH2+

F

S

N

H2N

O

O

S

N
NH

NH
NH

N

NH

NH3+

OH

O

S

N

N

N

NH
O

O

NH2+

NH

H2N

O

O

N

O

O

O
N

NH

NH

O

O

O-

N

NH

F

OH

N

O

Br

O
N

NH

O

N

O

O

NH2+

OH

O

O-

N
O-

O

O

N

O

NH

N+

N

N

NH+

S

NH+

N

S

N

O

NH
N
NH

O

O

O

S

O

O
O

NH+

N

S
F

NH2

N

O

N

N

N

NH

NH

N

N

O

O
N

N

NH

O

S

O

O

O

O

NH2+

NH+

S
O

N

O

H2N

Cl
O

O

N

N

O

H

O

O

S

N

O

O

S
NH

NH

N

NH

NH+

NH

NH3+

HO

O

N

O

N

N

O

OH

O

N

F

O
N

F

O

N

O

S

O

N

N

NH

N

NH

NH

NH

N

NH

N

O

O
N

H2N

N

O

O

NH

O
N

O
NH2+

O

O

S

NH

N

F

O

O
N

N

NH

S

O

O

F

NH

NH

O

NH2

N

N

H
NH

NH2+

N

NH

Cl

O

O

N

O

H

NH+

NH2

O

N

N

O

N

O

N

S

Cl
O

N

N

H

H
NH

H2N

O

Cl

O
O

Br
NH2

F

O

N

O
O

O

N

N

O

O

O

NH

O

N

O

O

O

O

S

N

S

H2N
O
N

N

O

N

N

N

O

N

O

O

O

S

N

O O

N

N

NH

N
N

NH

N

N
N

N

N

O
N

NH

O

N

N
O

OH

O

NH

N

S

S

O

O

O

N

O

H2N

O

N-

O

O
H

NH

O

O

NH2+

O

O
NH3+

NH

N

N

S

N

NH

O

O

OH

O

NH

N

Cl

NH

S

O
N

O

NH

OH

N
N
S

OH

O

S

OH

O

O

O

O

O

NH2

H
O
O

HO

N

NH

O
N

N

O
S

S

Cl

O

O

S

S

O

N

N
H2N

O

O

NH

N

N
Br

H
O

N
O

N

O

NH

O
N

O
O

O-

N

N

H

O

O

N

NH

NH

N

S

N

NH+

O

NH

NH

N

O

O

O
S

H

S
O
S

HS
S

NH

O
NH2+

O
O

N

O

N
NH

NH

Cl

O

NH+

O

H
N

N

NH

NH

O

N

N

O
N

NH3+
O

NH

N
O

O

F
O

NH

O
N

NH
Cl

NH2

N

O

O
NH
O

O

O
N

Figure 10. 250 molecules sampled from prior distribution N (0, I).

O

O

Junction Tree Variational Autoencoder for Molecular Graph Generation

N

N

N

N

N

O

N

N
NH

O
N

N

NH

H

N

N

N

N
N

N
N

N

N

N

N

N
NH

N

N

N

N

N

NH
N
N

N

NH2+

N
N

N

N

NH
H

NH

NH

NH

H

N

N

N

N

H

N

N

NH

H

N

N

H

NH

H

NH
N

NH2

N

N

NH

N

H

O
N
N

NH

N

N

N
N

N

NH3+

N

N

NH

N
N

NH

N

NH
N
N

N

N

NH
N

NH
N

N

N

N
N

N

N

NH

S

N

N

N

N

NH

N

NH

N

N

N

N

N

N

NH

N

N

N

N

N

H

N

NH

NH

NH

H
N

N

N

N

N

N

N

N

NH+

N

N

N

N

N

N

NH2+

N

N

NH2+

N
N

N

N

N

N

N

N

N
NH

NH

NH

NH

N

NH

H

NH2+

N

N

H

H

N

H

H

N

H

N

N

N

N

N

N

N

NH

N

N

NH

N

N
N

NH

N

NH

N

NH+

N

NH+

N

N
N

S

N

N

N

N

S

S

N

NH

NH

N

N

NH

N

NH2+
N

N

N

NH

N

N

N

NH

N

N

N

NH

O

NH2

NH

N

N
NH

N

N

NH2+

N

N

NH2

S

N

N
NH

N

N
N

N

N

NH

H

N

N

NH

N

N

N

H

N

N

N

N

NH

H
N

N

N

NH

N

NH

NH

N

NH

N

N

N
N

N

N

N

N

N

N

N

N

N
NH

O

N

N

N

N

N

N

NH
N

O
N

N
N

N

N

N

NH2

N

N
NH2

S

N
NH2

N

N

NH2

N

N

N

NH

S

N

S

S

N
NH

NH

N

N

N

N
NH

NH2
N

NH

N

N

N

NH

N

N

N

N

N

N

N

O

N

NH

N

N

NH

N
N

N
NH
N

N

N

NH
NH

N
S
H
H
S

N

NH

NH

N

O

N

N

S

N

O

N

N

N

NH2

NH

N

N

N

N

N

N
N

N
S

N

NH2

N
NH

NH2

N

N

N

N

N

N

NH

NH

NH2

N

O

S

NH2

NH

N

N

N

N

N

N

N

NH

NH
NH2

N

N

NH

N

N

N

N

NH

N

N

NH

N

N

NH

N

N

NH

N
NH

N

N

O

N

NH

N

N

N

NH

N

S
N

N

N

N
N

S

N

NH2

N

N

N

N
H

N

H2N
N

N

N

N

N

N

N

N

N

NH

S

H2N

N

S

N

NH

N

N

N

N

N

N
S

NH
N

H

NH

N

N

N

S

NH

N

N

NH

N

N

NH

N

N

N

NH

N

NH2

NH

N

N

S

N

NH

N

N

N

N

NH2
N

N

S

N

N
NH

N

N

NH

H

N

NH2

S

N

NH

N

S
NH2

H

N

S

N
N

N

N

N

N

N

N

N

N

S

S

N

NH

S
N

N

NH

N

H

N

N

N

N

N

N

N

S

S

S

NH

NH2

NH

NH2

NH

N

N

NH

N

N

NH

N

N

NH

N

N

NH

N

N

N

N
NH2

NH2

N

S
NH2

N

H

N

S

N

S

S
H

N

N

NH

N

NH

N

NH

N

N

N

N

S

S

H

N

N

NH

NH2

NH

N

N

N

N

N

NH+

N

N

N

N

N

S

NH

N
NH

N

S
NH

NH2

NH

N

N

NH2

S

N

NH2

S
N

S

S
NH2

N

N

NH

N

NH

S

S
N

N

N
N

NH

NH

H2N

NH

NH

N

N

NH

N

N

NH2

NH

N

N
NH

S

N

N

NH

N

NH

N

N

N

N
NH2

S

S

NH2

H

N

N

NH2+

NH2

H
S

O

S

N

S

S

NH

NH

N

S
H

N

NH

N

S

S

S

NH2+
H

H2N

O

S

NH

N

NH+

N

NH

NH

N

N

S

NH2

NH

S

S

N

N
N

S

N

O
S

O

O

O

NH

H2N

S

S

NH

H2N
N

N

N

N

NH

NH2

NH2

N

N

NH+

S

N

S

S

N

NH

N

S

N

N

N
N

NH+
N

NH

NH2

N

NH+
N

NH

NH

NH
N

S

N

O

S
N
N

NH+

N

N

N

N

O

N
N

S

S

N

NH+

S

NH

H2N
N

N

N
S

N

N

N

N

S

NH

H2N

NH

N

N

N

NH+

NH

N
N
N
N

S
H

N

NH

N

N
N

N

NH2

H
S
S

H

N
N

S

N

N
NH+
NH2

N

N

S
NH2

H
S
S

O
NH

NH2

NH
N
N

NH2

S

S

N

NH

NH2

NH
N

N
N

H2N

N
NH

N
S
N

NH

N

NH2

NH+

N
S

N

NH2

NH
N

N

N

N

NH2

N

S

S

N
N

NH

NH+

NH

H2N
S

N
H

N

N

S

N

N

N

N

N

NH2+

O

S
S

H
NH

S
NH+

NH2

N
N

NH2

H
S
S

N

N

NH

N

NH

N

N

N

NH2

N
S

N
N
N

NH

N

N

NH+

N

NH

NH2
N

N

N
S

N

S

NH2

N

NH2

NH

S
S

NH+
N

N
S

S

NH2

NH
N

N

N

N

N
H

N

NH

NH2

N

NH2

NH2

O
N

O

N

NH+

N

S
S
H2N

NH

N

S

NH
N

N

N

N

NH

S

NH

N

N
S

NH
N

S
N

H2N

N
H2N+

H2N+

N

N

S

NH

S

N

N
N

N
NH2

S

S

N

N
NH+

N

O

S
N
N

N
S

O

NH

H2N
N
S

NH+

N
S

NH

H2N
NH2
N

N

NH

NH
N

S
H2N

NH

N

N

N
N

NH+
N

NH

N

NH+
N

NH

NH

N
N

N

N

S
NH2

N

N

NH2

NH2

Figure 11. Neighborhood visualization of molecule C[C@H]1CC(Nc2cncc(-c3nncn3C)c2)C[C@H](C)C1.

NH2

NH2

Junction Tree Variational Autoencoder for Molecular Graph Generation

O
N

NH3+

O

F

H2N

NH+

O

N

N

O

F

N

N

O

O

NH

N

O

O
O

NH

N

NH

N

O

NH

O
F
O

Cl

NH2+
F

OH

NH2+
F

O

NH2+

F

N

N

NH

F

N

NH+
F

O

NH2+

F

NH

O

F

F

O

F

Cl

N

O

NH

NH+

N
O
NH+

O

O

O

O

O

OH

O

O
O

O

NH3+
F

OH

OH

OH

NH
O

OH

F

OH

F

F
O

O

O

HO
F

NH+

O

NH+

F

F

NH+

O

N

F

N

F

N

N

O

O

NH

N

NH

N

O

NH

F

NH

F

N

N

F

O
NH+

F

O

NH2+

O

NH+

O

O

NH2+

F

F

F

O

NH2+

NH+

NH2+

F

F

O
O

O

F

O

O

O

O
O
NH2+

NH2+

NH2+

O-

O

NH2+

N+

N
O

O

NH2+

O
OH

F

N
O

O
F

O

O

OH

F

O

N

F

NH+

O

NH+

F

F

NH+

O

NH+

O

NH+

O

NH

F

NH

F

N

F

N
N

O

NH+

F

O

NH+

F

F

OH

NH2+

O
O

O

O

O
O

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

NH+

O

NH+

NH+

O

F

F

O

NH+

F

F

F

F
NH2+

O

F

F

O

O

O

F

O

O

F

O

F

F

F

F

O

O

O

NH+

O

NH+

O

NH+

O

NH+

F
F

F

O

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O
NH2+

O

F
H2 N

F

F

F

F

F

H3N+

O

NH+

F

F
F

O

O

O

O

O

NH+
O

O

S
O

F

F

O

N

O

NH+

F

F
F

NH

O

O

NH+

O

F

O

O

OH

F
NH+

O

NH2+

F
O

NH2+

O

N

O

NH2+

F

O

O

N

NH

F

N

F
O

N

O

F

F

S

O

F

O
F

F

F

O

O

F

O

F

O

F

O

F

O

F

O

F

O

F

O

O

F

O

O

F

F

O

O
O
F
O

O

NH+

O

NH+

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

NH2+

NH+

F

F

F

F

F
F

F

F

F
O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

F

O

H3N+

F

NH3+
O

F

O

O

NH2+

F

O

O
F
O

O

NH+

O

NH+

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

NH+

F

F

F

F
F

F

F

F
O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

F

O

NH3+
O

F

O

NH2+

O

H

O
NH2+
NH+

O

O

NH+

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

NH+

F

F

O
O

F

F

S
F

F

F

F

F

F

F

F

F

F

F
O

O

F

F

F

O

O

F

O

F

O

F

O

F

O

F

O

F

O

F

O

F

O

O

F

O

F

O

Br
O
O

OH

O

NH+

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

NH+

F

S

F

O

NH+
F

F

F

F

F

F

F

F

F

F

F

O

NH2+
F

F

F

F

F

F

F

F

F

F

O

F
N

O

O

O

O

O

O

O

O

O

O

O

O
N
O

O

O

NH+

NH+
NH+

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

NH2+

NH+

F

F
O

F

N

F
N

HO

O

F

F

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

F

F

O

O

O
N

NH+

NH+

O

NH+

O

NH+

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

O

NH+

F

NH+

F

N

F

F

F

F

F

F

O

F

F

F

F

F

F

O

O

O

O

O

O

O
O
NH+

NH+

O

NH+
O

O

NH+

O

NH+

F

O

NH+

F

F

N

F

NH+

N

F

F

F

F

O

O

O

F

O

O

N
NH+

NH+

F

N

O

O

O

O

NH+

NH+

F

NH+

O

NH+

NH+

F

F

N

O

S

F
F
F
O

F

F
O

F

F
O

F

F
O

F

NH2+

S

N

F

O
O

F
O

O

NH+

NH+
F

O

F
O

F

NH+

NH+

F
O

F

O

N

N

NH+

O

F

O

O

O

F
N

F

S

S

N

O

F

O

O

O
O

S

F
N

F
O

O

O

NH2+

NH+
N

F

F
F

F
O

O
N

NH+

F

O
N

F

O

O
NH+

NH+

F

O

F
N

O

F

O

O

O

F
O

O

F
O

O

O

NH+

F

F

S

S

N

O

N

O

O

F

HO
F

O

O

O
F

N
N

NH2+

NH+

F

OH
F

F

O

F

O

O

O

N

O

F

S

F

O

F

S
O

O

O

O

O

O

Figure 12. Neighborhood visualization of molecule COc1cc(OC)cc([C@H]2CC[NH+](CCC(F)(F)F)C2)c1.

N

F

Junction Tree Variational Autoencoder for Molecular Graph Generation

Cl

Cl
F
O
Cl

O

NH

Cl

Cl

Cl

Cl

NH
NH

NH

S

NH

NH

NH

NH

O

NH

NH

Cl
O

NH

NH

O

O

5.30

O
Cl

O

4.93

4.49

4.45

4.42

Cl
Br
O

NH

Cl

S

NH
Cl
NH

O

NH

S

O

Cl

Cl

Cl

NH

S
NH

NH
O

NH

NH

NH

O

S

O

O

Cl

Cl

NH

Cl

NH
NH

4.42

4.40

4.37

4.30

4.26

F
F

Cl
S
NH

NH
O
N

S

S

Cl

S

N

Cl

4.23

4.18

4.17

NH

4.08

Cl

Cl

NH
Cl

O

NH

4.07

Cl

O

Br

O

NH

O

NH

NH

Cl

O
O

O

NH
NH

NH

N

S

Cl
NH

O

NH

NH

S

Cl

S

Br

N

NH

S

O
NH

Cl

S

Cl

S

Cl

S
O
Cl
O

4.07

4.04

4.04

4.03

4.02
Cl

O
NH

NH
NH

O

NH

S

Cl

NH

Cl

NH

O
O

NH

N

NH

NH

NH

S

O
NH

Cl

S
O

NH

NH

Cl
O

O
N

O
Cl

Cl

S

4.01

4.00

3.99

3.98

3.96

Cl
Cl

Cl

O

Cl

NH

NH

S
Br

N

NH

O

O

N

Br

NH

NH

Cl
O

N

Cl

Cl

3.86

3.83

3.81

3.79

3.72

I
S

Cl

O

O

Cl

S

NH

S

O

O

S

N
N

Cl

Cl
Cl

O

S

NH

Cl

Br
N

O
N

3.69

3.68

3.64

Cl

Cl

O

3.61

3.61

Cl

Cl

S
Cl

S

N

NH
Cl

N

N

S
F

Cl
F

Br

O

Cl

O
Cl

3.61

3.60

3.58

3.58

3.58

F
OCl
N+

O

O

Cl

O

Cl

O
NH

S
N

S

S

Cl

NH

N
N

N

N

N

NH

O
N

S

N

O
Cl

3.57

3.56

Cl

3.53

O

F

Cl

S
O

Br

O

NH
O

Br

O

3.51

3.51

Br

NH
Cl

3.52

Br

3.50

3.50

3.50

Figure 13. Top 50 molecules found by Bayesian optimization using JT-VAE.

O

3.50

Junction Tree Variational Autoencoder for Molecular Graph Generation

O-

O
-S

N

+H

3N

N

H
N

O

S

+

+

N
H

O

O

NH+

S

H
N

N
OH

O

O

N
H

Cl

O
N+
O

+

6.80

6.69

6.51

-2.05

N
N

O

O

O

O
Br

N

HO

N
H

S

N

H
N

N

O
O-

OH

S

O

N

O
N+ O

N

-O

O
O
+

N
H

O

H
N

O

O-

O

O

O-

O
+

O

N

N
H

N
H 2N

N

O
5.69

4.94

4.52
O

Br

O

N

N

N

O

HN

O
N
S
O

NH
O O

O

NH2

-2.21

O

O

O

H
N

O
S
O

N

H
N

O

NH2
NH2

NH+
O

O

O
N
H

O
S
N
O H

N
N

O
S
N
O H

N
H

3.03

H
N
O

N
N
H

O

O
O-

NH2+

O

4.00

O

O

N

O
O-

N

N N

N+
O

OH

-1.92

O

N

O
Br

2.69

-O

N

N
H

O

O

N

+

H
N
N

N N
OH

Figure 14. Row 1-3: Molecule modification results with similarity constraint sim(m, m0 ) ≥ 0.2, 0.4, 0.6. For each group, we plot the top
three pairs that leads to actual property improvement, and one pair with decreased property. We can see that tighter similarity constraint
forces the model to preserve the original structure.

