H IERARCHICAL G RAPH - TO -G RAPH T RANSLATION FOR M OLECULES
Anonymous authors
Paper under double-blind review
A BSTRACT
The problem of accelerating drug discovery relies heavily on automatic tools to optimize precursor molecules to afford them with better biochemical properties.
Our work in this paper substantially extends prior state-of-the-art on graph-to-graph translation methods for molecular optimization.
In particular, we realize coherent multi-resolution representations by interweaving the encoding of substructure components with the atom-level encoding of the original molecular graph.
Moreover, our graph decoder is fully autoregressive, and interleaves each step of adding a new substructure with the process of resolving its attachment to the emerging molecule.
We evaluate our model on multiple molecular optimization tasks and show that our model signiﬁcantly outperforms previous state-of-the-art baselines.
I NTRODUCTION
Molecular optimization seeks to modify compounds in order to improve their biochemical properties.
This task can be formulated as a graph-to-graph translation problem analogous to machine translation.
Given a corpus of molecular pairs { ( X, Y ) } , where Y is a paraphrase of X with better chemical properties, the model is trained to translate an input molecular graph into its better form.
The task is difﬁcult since the space of potential candidates is vast, and molecular properties can be complex functions of structural features.
Moreover, graph generation is computationally challenging due to complex dependencies involved in the joint distribution over nodes and edges.
Similar to machine translation, success in this task is predicated on the inductive biases built into the encoder-decoder architecture, in particular the process of generating molecular graphs.
Prior work (Jin et al., 2019) proposed a junction tree encoder-decoder that utilized valid chemical substructures (e.g., aromatic rings) as building blocks to generate graphs.
Each molecule was represented as a junction tree over chemical substructures in addition to the original atom-level graph.
While successful, the approach remains limited in several ways.
The tree and graph encoding were carried out separately, and decoding proceeded in strictly successive steps: ﬁrst generating the junction tree for the new molecule, and then attaching its substructures together.
This means the predicted attachments do not impact the subsequent substructure choices (see Figure 1a).
Moreover, the attachment prediction process is non-autoregressive, thus it can predict inconsistent substructure attachments across different nodes in the junction tree (see Figure 1b).
We propose a multi-resolution, hierarchically coupled encoder-decoder for graph generation.
Our auto-regressive decoder interleaves the prediction of substructure components with their attachments to the molecule being generated.
In particular, a target graph is unraveled as a sequence of triplet predictions (where to expand the graph, new substructure type, its attachment).
This enables us to model strong dependencies between successive attachments and substructure choices.
The encoder is designed to represent molecules at different resolutions in order to match the proposed decoding process.
Speciﬁcally, the encoding of each molecule proceeds across three levels, with each layer capturing essential information for its corresponding decoding step.
The graph convolution of atoms at the lowest level supports the prediction of attachments and the convolution over substructures at the highest level supports the prediction of successive substructures.
Compared to prior work, our decoding process is much more efﬁcient because it decomposes each generation step into a hierarchy of smaller steps in order to avoid combinatorial explosion.
We also extend the method to handle conditional translation where desired criteria are fed as input to the translation process.
This enables our method to handle different combinations of criteria at test time.
isolated generate invalid cannot be assembled
This problem can be solved when we interleave the tree and graph decoding steps, allowing the predicted attachments to guide the substructure prediction; b) Their non-autoregressive graph decoder often predicts inconsistent local substructure attachments during training.
To this end, we propose an autoregressive decoder that interleaves the prediction of substructures with their attachments.
We evaluate our new model on multiple molecular optimization tasks.
Our baselines include previous state-of-the-art graph generation methods (You et al., 2018a; Liu et al., 2018; Jin et al., 2019) and an atom-based translation model we implemented for a more comprehensive comparison.
Our model signiﬁcantly outperforms these methods in discovering molecules with desired properties, yielding 3.3% and 8.1% improvement on QED and DRD2 optimization tasks.
During decoding, our model runs 6.3 times faster than previous substructure-based generation methods.
We further conduct ablation studies to validate the advantage of our hierarchical decoding and multi-resolution encoding.
Finally, we show that conditional translation can succeed (generalize) even when trained on molecular pairs with only 1.6% of them having desired target property combination.
R ELATED W ORK
Molecular Graph Generation Previous work have adopted various approaches for generating molecular graphs.
Methods (G ´ omez-Bombarelli et al., 2018; Segler et al., 2017; Kusner et al., 2017; Dai et al., 2018; Guimaraes et al., 2017; Olivecrona et al., 2017; Popova et al., 2018; Kang & Cho, 2018) generate molecules based on their SMILES strings (Weininger, 1988).
Simonovsky & Komodakis (2018); De Cao & Kipf (2018); Ma et al.
(2018) developed generative models which output the adjacency matrices and node labels of the graphs at once.
You et al.
(2018b)
; Li et al.
(2018); Samanta et al.
Liu et al.
(2018) proposed generative models decoding molecules sequentially node by node.
You et al.
(2018a)
; Zhou et al.
(2018) adopted similar node-by-node approaches in the context of reinforcement learning.
Kajino (2018) developed a hypergraph grammar based method for molecule generation.
Our work is most closely related to Jin et al.
(2018; 2019) that generate molecules based on substructures.
They adopted a two-stage procedure for realizing graphs.
The ﬁrst step generates a junction tree with substructures as nodes, capturing their coarse relative arrangements.
The second step resolves the full graph by specifying how the substructures should be attached to each other.
Their major drawbacks are 1)
The second step introduced local independence assumptions and therefore the decoder is not autoregressive.
These two steps are applied stage-wise during decoding – ﬁrst realizing the junction tree and then reconciling attachments without feedback.
In contrast, our method jointly predicts the substructures and their attachments with an autoregressive decoder.
Graph Encoders Graph neural networks have been extensively studied for graph encoding (Scarselli et al., 2009;
Bruna et al., 2013; Li et al., 2015; Niepert et al., 2016; Kipf & Welling, 2017; Hamilton et al., 2017; Lei et al., 2017; Velickovic et al., 2017; Xu et al., 2018).
Our method is related to graph encoders for molecules (Duvenaud et al., 2015; Kearnes et al., 2016;
Dai et al., 2016; Gilmer et al., 2017; Sch¨utt et al., 2017).
Different to these approaches, our method represents molecules as hierarchical graphs spanning from atom-level graphs to substructure-level trees.
Our work is most closely related to (Defferrard et al., 2016;
Ying et al., 2018; Gao & Ji, 2019) that learn to represent graphs in a hierarchical manner.
In particular, Defferrard et al.
(2016) utilized graph coarsening algorithms to construct multiple layers of graph hierarchy and Ying et al.
In adds a new attachment graph molecule
Gao & Ji (2019) proposed to learn the graph hierarchy jointly with the encoding process.
Despite some differences, all of these methods seek to represent graphs as a single vector for regression or classiﬁcation tasks.
In contrast, our focus is graph generation and a molecule is encoded into multiple sets of vectors, each representing the input at different resolutions.
Those vectors are dynamically aggregated by decoder attention modules in each graph generation step.
H IERARCHICAL G ENERATION OF M OLECULAR G RAPHS
The graph translation task seeks to learn a function F that maps a molecule X into another molecule
G with better chemical properties.
F is parameterized as an encoder-decoder with neural attention.
Both our encoder and decoder are illustrated in Figure 2.
In each generation step, our decoder adds a new substructure ( substructure prediction ) and decides how it should be attached to the current graph.
The attachment prediction proceeds in two steps: predicting attaching points in the new substructure and their corresponding attaching points in the current graph ( attachment prediction 1-2 ).
To support the above hierarchical generation, we need to design a matching encoder representing molecules at multiple resolutions in order to provide necessary information for each decoding step.
Therefore, we propose to represent a molecule X by a hierarchical graph H X with three components: 1) substructure layer representing how substructures are coarsely connected; 2) attachment layer showing the attachment conﬁguration of each substructure; 3) atom layer showing how atoms are connected in the graph.
Our model encodes nodes in H X into substructure vectors c S X , attachment vectors c A X and atom vectors c G X , which are fed to the decoder for corresponding prediction steps.
As our encoder is tailored for the decoder, we ﬁrst describe our decoder to clarify relevant concepts.
3. H IERARCHICAL G RAPH D ECODER
Notations We denote the sigmoid function as σ ( · ) .
MLP( a , b ) represents a multi-layer neural network whose input is the concatenation of a and b .
attention θ ( h ∗ , c X ) stands for a bilinear attention over vectors c X with query vector h ∗ .
Substructures We deﬁne a substructure S i
=
( V i , E i ) as subgraph of molecule G induced by atoms in V i and bonds in E i .
Given a molecule, we extract its substructures S 1 , · · · , S n such that their union covers the entire molecular graph: V = (cid:83) i V i and E =
(cid:83) i E i .
In this paper, we consider two types of substructures: rings and bonds.
We denote the vocabulary of substructures as S , which is constructed from the training set.
In our experiments, |S| < 500 and it has over 99.5% coverage on test sets.
Substructure Tree To characterize how substructures are connected in the molecule G , we construct its corresponding substructure tree T , whose nodes are substructures S 1 , · · · , S n .
Speciﬁcally, we construct the tree by ﬁrst drawing edges between S i and S j if they share common atoms, and then applying tree decomposition over T to ensure it is tree-structured.
Generation Our graph decoder generates a molecule G by incrementally expanding its substructure tree in its depth-ﬁrst order.
Suppose the model is currently visiting substructure node S k .
It makes the following predictions conditioned on encoding of input X (see Figure 3):
Topological Prediction : It ﬁrst predicts whether there will be a new substructure attached to S k .
If not, the model backtracks to its parent node S d k in the tree.
Let h S
k be the hidden representation of S k learned by decoder (which will be elaborated in § 3.2).
This probability is predicted by a MLP with attention over substructure vectors c S X of X :
[Equation]
Substructure Prediction :
If p k > 0 . the model decides to create a new substructure S t from S k and sets
its
parent d t = k .
It then predicts the substructure type of S t using another MLP
that outputs a
distribution
over the vocabulary S :
[Equation]
Attachment Prediction : Now the model needs to decide how S t should be attached to S k .
The attachment between S t and S k is deﬁned as atom pairs M t = { ( u j , v j ) | u j ∈ S t , v j ∈ S k } where atom u j and v j are attached together.
We predict those atom pairs in two steps: 1) We ﬁrst predict the atoms { v j } ⊂ S t that will be attached to S k .
Since the graph S t is always ﬁxed and the number of attaching atoms between two substructures is usually small,
we can
enumerate all possible conﬁgurations { v j } to form a vocabulary A ( S t ) for each substructure S t .
This allows us to formulate the prediction of { v j } as a classiﬁcation task – predicting the correct conﬁguration A t =
( S t , { v j } ) from the vocabulary A ( S t ) :
[Equation]
2) Given the predicted attaching points { v j } , we need to ﬁnd the corresponding atoms { u j } in the substructure S k .
As the attaching points are always consecutive, there exist at most 2 |S k | different attachments M = { ( u j , v j ) } .
The probability of a candidate attachment M is computed based on the atom representations h u j and h v j learned by the decoder:
[Equation]
The above three predictions together give an autoregressive factorization of the distribution over the next substructure and its attachment.
Each of the three decoding steps depends on the outcome of previous step, and predicted attachments will in turn affect the prediction of subsequent substructures.
During training, we apply teacher forcing to the above generation process, where the generation order is determined by a depth-ﬁrst traversal over the ground truth substructure tree.
The attachment enumeration is tractable because most of the substructures are small.
In our experiments, the average size of attachment vocabulary |A
( S t ) | < 5 and the number of candidate attachments is less than 20.
3. H IERARCHICAL G RAPH E NCODER
Our encoder represents a molecule X by a hierarchical graph H X in order to support the above decoding process.
The hierarchical graph has three components (see Figure 4):
Atom layer : The atom layer is the molecular graph of X representing how its atoms are connected.
Each atom node v is associated with a label a v indicating its atom type and charge.
Each edge ( u, v ) in the atom layer is labeled with b uv indicating its bond type.
Attachment layer : This layer is derived from the substructure tree of molecule X .
Each node A i in this layer represents a particular attachment conﬁguration of substructure S i in the vocabulary
A ( S i ) .
Speciﬁcally, A i =
( S i , { v j } ) where { v j } are the attaching atoms between S i and its parent S d i in the tree.
This layer provides necessary information for the attachment prediction (step 1).
Figure 4 illustrates how A i and the vocabulary A ( S i ) look like.
We further introduce edges that connect the atoms and substructures between different layers in order to propagate information in between.
In particular, we draw a directed edge from atom v in the atom layer to node A i in the attachment layer if v ∈ S i .
We also draw edges from node A
i
to node S i in the substructure layer.
This gives us the hierarchical graph H X for molecule X , which will be encoded by a hierarchical message passing network (MPN) (see Figure 4).
The encoder contains three MPNs that encode each of the three layer.
We use the MPN architecture from Jin et al.
For simplicity, we denote the MPN encoding process as MPN ψ ( · ) with parameter ψ .
Atom Layer MPN We ﬁrst encode the atom layer of H X (denoted as H gX ).
The inputs to this MPN are the embedding vectors { e ( a u ) } , { e ( b uv ) } of all the atoms and bonds in X .
During encoding, the network propagates the message vectors between different atoms for T iterations and then outputs the atom representation h v for each atom v :
[Equation]
Attachment Layer MPN
The input feature of each node
A i in the attachment layer H aX is an concatenation of the embedding e ( A i ) and the sum of its atom vectors { h v | v ∈ S i } :
[Equation]
The input feature for each edge ( A i , A j ) in this layer is an embedding vector e ( d ij ) , where d ij describes the relative ordering between node A
i and A j during decoding.
Speciﬁcally, we set d ij = k if node A i is the k -th child of node
A j and d ij =
if A i is the parent.
We then run T iterations of message passing over H aX to compute the substructure representations:
[Equation]
Substructure Layer MPN Similarly, the input feature of node S i in this layer is computed as the concatenation of embedding e ( S i ) and the node vector h
A i from the previous layer.
Finally, we run message passing over the substructure layer H sX to obtain the substructure representations:
[Equation]
In summary, the output of our hierarchical encoder is a set of vectors c X
= c S X ∪ c
A X ∪ c G X that represent a molecule X at multiple resolutions.
These vectors are input to the decoder attention.
Decoder MPN During decoding, we use the same hierarchical MPN architecture to encode the hierarchical graph H G at each step t .
This gives us the substructure vectors h S k and atom vectors h v j in § 3.1.
All future nodes and edges are masked to ensure the prediction of current substructure and attachment only depends on previously generated outputs.
3. T RAINING
Our training set contains molecular pairs ( X, Y ) where each compound X can be associated with multiple outputs Y since there are many ways to modify X to improve its properties.
In order to
1: for ( X, Y ) in the training set do 2:
Encode molecule X, Y into vectors c X , c Y 3: Compute µ X,Y , σ X,Y from δ X,Y 4: Sample latent code z ∼ Q ( z | X, Y ) 5: Generate molecule Y given c X and z 6: end for generate diverse outputs, we follow Jin et al.
and extend our method to a variational translation model F : ( X, z )
→ Y with an additional input z .
The latent vector z indicates the intended mode of translation which is sampled from a Gaussian prior P ( z ) during testing.
We train our model using variational inference (Kingma & Welling, 2013).
Given a training example ( X, Y ) , we sample z from the posterior Q ( z | X, Y ) =
N ( µ X,Y , σ X,Y ) .
To compute Q ( z | X, Y ) , we ﬁrst encode X and Y into their representations c X and c Y and then compute vector δ X,Y that summarizes the structural changes from molecule X to Y at both atom and substructure level:
[Equation]
Finally, we compute [ µ X,Y , σ X,Y ] =
MLP( δ S X,Y , δ G X,Y ) and sample z using reparameterization trick.
The latent code z is passed to the decoder along with the input representation
c X to reconstruct output Y .
The overall training objective follows a standard conditional VAE:
[Equation]
Conditional Translation In the above formulation, the model does not know what properties are being optimized during translation.
During testing, users cannot change the behavior of a trained model (i.e., what properties should be changed).
This may become a limitation of our method in a multi-property optimization setting.
Therefore, we extend our method to handle conditional translation where the desired criteria are also fed as input to the translation process.
In particular, let g X,
Y be a translation criteria indicating what properties should be changed.
During variational inference, we compute µ X,Y and σ X,Y with an additional input g X,Y :
[Equation]
We then augment the latent code as [ z , g X,Y ] and pass it to the decoder.
During testing, the user can specify their criteria in g X,Y to control the outcome (e.g., Y should be drug-like and bioactive).
E XPERIMENTS
We follow the experimental design by Jin et al.
and evaluate our translation model on their single-property optimization tasks.
As molecular optimization in the real-world often involves different property criteria, we further construct a novel conditional optimization task where the desired criteria is fed as input to the translation process.
To prevent the model from ignoring input X and translating it into arbitrary compound, we require the molecular similarity between X and output
Y to be above certain threshold sim( X, Y ) ≥ δ at test time.
The molecular similarity is deﬁned as the Tanimoto similarity over Morgan ﬁngerprints (Rogers & Hahn, 2010) of two molecules.
Single-property Optimization This dataset consists of four different tasks.
For each task, we train and evaluate our model on their provided training and test sets.
For these tasks, our model is trained under an unconditional setting (without g X,Y as input).
• LogP Optimization : The penalized logP score (Kusner et al., 2017) measures the solubility and synthetic accessibility of a compound.
In this task, the model needs to translate input X into output Y such that logP( Y ) >
logP( X ) .
We experiment with two similarity thresholds δ = { 0 . 4 , 0 . 6 } .
• QED Optimization : The QED score (Bickerton et al., 2012) quantiﬁes a compound’s druglikeness.
In this task, the model is required to translate molecules with QED scores from the lower range [0 . 7 , 0 . 8] into the higher range [0 . 9 , 1 . 0] .
The similarity constraint is sim( X, Y ) ≥ 0 . 4 .
[Table] • DRD2 Optimization : This task involves the optimization of a compound’s biological activity against dopamine type 2 receptor (DRD2).
The model needs to translate inactive compounds ( p < 0 . 05 ) into active compounds ( p ≥ 0 . 5 ), where the bioactivity is assessed by a property prediction model from Olivecrona et al.
The similarity constraint is sim( X, Y ) ≥ 0 . 4 .
Conditional Optimization This new task requires the model to translate input X into output Y to satisfy different combination of constraints over its QED and DRD2 scores.
We deﬁne a molecule Y as drug-like if QED( Y ) ≥ 0 .
9 and as DRD2-active if its predicted bioactivity DRD2( Y ) ≥ 0 .
At test time, our model needs to handle the following two criteria over output molecule Y :
Y is both drug-like and DRD2-active.
Here both properties need to be improved after translation.
Y is drug-like but DRD2-inactive.
In this case, DRD2 is an off-target that may cause side effects.
Therefore
only the drug- likeness should be improved after translation
As different users may be interested in different settings, we encode the desired criteria as vector g and train our model under the conditional translation setup in § 3.3.
Like single-property tasks, we impose a similarity constraint sim( X, Y ) ≥ 0 .
4 for both settings.
Our training set contains 120K molecular pairs and the test set has 780 compounds.
For each pair ( X, Y ) , we set g X,Y =
( I [ Y is drug-like ] , I
[ Y is DRD2-active ]) .
During testing, we translate each compound with g =
[1 , 1] , [1 , 0] for each setting.
We note that the ﬁrst criteria ( g = [1 , 1] ) is the most challenging because there are only 1.6% of the training pairs with target Y being both drug-like and DRD2-active.
To achieve good performance, the model must learn to transfer the knowledge from other pairs with g X,Y =
[1 , 0] , [0 , 1] ) that partially satisfy the criteria.
Evaluation Metrics Our evaluation metrics include translation accuracy and diversity.
Each test molecule X
i is translated K = 20 times with different latent codes sampled from the prior distribution.
On the logP optimization, we select compound Y i as the ﬁnal translation of X i that gives the highest property improvement and satisﬁes sim( X i , Y i ) ≥ δ .
We then report the average property improvement 1 D (cid:80) i
logP( Y i ) − logP
( X i ) over test set D .
For other tasks, we report the translation success rate.
A compound is successfully translated if one of its K translation candidates satisﬁes all the similarity and property constraints of the task.
To measure the diversity, for each molecule we compute the average pairwise Tanimoto distance between all its successfully translated compounds.
Here the Tanimoto distance is deﬁned as dist( X, Y ) = 1 − sim
( X, Y ) .
Baselines We compare our method (HierG2G) against the baselines including GCPN (You et al., 2018a), MMPA (Dalke et al., 2018) and translation based methods Seq2Seq and JTNN (Jin et al., 2019).
Seq2Seq is a sequence-to-sequence model that generates molecules by their SMILES strings.
JTNN is a graph-to-graph architecture that generates molecules structure by structure, but its decoder is not fully autoregressive.
We also compare with CG-VAE (Liu et al., 2018), a generative model that decodes molecules atom by atom and optimizes properties in the latent space using gradient ascent.
To make a direct comparison possible between our method and atom-based generation, we further developed an atom-based translation model (AtomG2G) as baseline.
It makes three predictions in each generation step.
First, it predicts whether the decoding process has completed (no more new [Table] [Table] atoms).
If not [Table] it creates a new atom a t and predicts its atom type.
Lastly [Table] it predicts the bond type between a t and other atoms autoregressively to fully capture edge dependencies (You et al., 2018b) [Table]
The encoder of AtomG2G encodes only the atom-layer graph and the decoder attention only sees the atom vectors c G X .
All translation models are trained under the same variational objective ( § 3.3).
Single-property Optimization Results As shown in Table 1, our model achieves the new state-ofthe-art on the four translation tasks.
In particular, our model signiﬁcantly outperforms JTNN in both translation accuracy (e.g., 76.9% versus 59.9% on the QED task) and output diversity (e.g., 0.
versus 0.480 on the logP task).
While both methods generate molecules by structures, our decoder is autoregressive which can learn more expressive mappings.
More importantly, our model runs 6.
times faster than JTNN during decoding.
Our model also outperforms AtomG2G on three datasets, with over 10% improvement on the DRD2 task.
This shows the advantage of our hierarchical model.
Conditional Optimization Results For this task, we compare our method with other translation methods: Seq2Seq, JTNN and AtomG2G. All these models are trained under the conditional translation setup where we feed the desired criteria g X,Y as input.
As shown in Table 2a, our model outperforms other models in both translation accuracy and output diversity.
Notably, all models achieved very low success rate on c = [1 , 1] because it has the strongest constraints and only 1.6K of the training pairs satisfy this criteria.
In fact, training our model on the 1.6K examples only gives 4.2% success rate as compared to 13.0% when trained with other pairs.
This shows our conditional translation setup can transfer the knowledge from other pairs with g X,Y =
Ablation Study To understand the importance of different architecture choices, we report ablation studies over the QED and DRD2 tasks in Table 2b.
We ﬁrst replace our hierarchical decoder with atom-based decoder of AtomG2G to see how much the structure-based decoding beneﬁts us.
We keep the same hierarchical encoder but modiﬁed the input of the decoder attention to include both atom and substructure vectors.
Using this setup, the model performance decreases by 0.8% and 10.9% on the two tasks.
We suspect the DRD2 task beneﬁts more from structure-based decoding because biological target binding often depends on the presence of speciﬁc functional groups.
Our second experiment reduces the number of hierarchies in our encoder and decoder MPN, while keeping the same hierarchical decoding process.
When the top substructure layer is removed, the translation accuracy drops slightly by 0.8% and 2.4%.
When we further remove the attachment layer, the performance degrades signiﬁcantly on both datasets.
This is because all the substructure information is lost and the model needs to infer what substructures are and how substructure layers are constructed for each molecule.
Implementation details of those ablations are shown in the appendix.
Lastly, we replaced our LSTM MPN with the original GRU MPN used in JTNN.
While the translation performance decreased by 4% and 2.2%, our method still outperforms JTNN by a wide margin.
Therefore we use the LSTM MPN architecture for both HierG2G and AtomG2G baseline.
C ONCLUSION
In this paper, we developed a hierarchical graph-to-graph translation model that generates molecular graphs using chemical substructures as building blocks.
In contrast to previous work, our model is fully autoregressive and learns coherent multi-resolution representations.
The experimental results show that our method outperforms previous models under various settings.
B N ETWORK A RCHITECTURE
LSTM MPN Architecture The LSTM MPN is a slight modiﬁcation from the MPN architecture used in Jin et al.
Let N ( v ) be the neighbors of node v , x v the node feature of v and x uv be the feature of edge ( u, v ) .
During encoding, each edge ( u, v ) is associated with two messages ν uv and ν vu , representing the message from u to v and vice versa.
The messages are updated by an
LSTM cell with parameters ψ
= { W zψ , W oψ , W rψ , W ψ } deﬁned as follows: Algorithm Passing function LSTM ψ (cid:16) x u , x uv , { ν ( t ) wu
[Equation]
c ( t ) wu
[Equation]
N ( u ) \
[Equation]
ν ( t +1) uv = o uv (cid:12) tanh (cid:16)
[Equation]
+
[Equation]
uv (cid:17)
Return ν ( t +1) uv , c ( t +1) uv end function
The message passing network MPN ψ ( H , { x u } , { x uv } ) over graph H is deﬁned as:
Algorithm 3 LSTM MPN with T message passing iterations function MPN ψ ( H , { x v } , { x uv } )
Initialize messages: ν 0 uv = 0 , c 0 uv =for t = 0 to T − 1 do
Compute messages ν
( t +1) uv , c ( t +1) uv = LSTM ψ (cid:16) x u , x uv , { ν ( t ) wu , c ( t ) wu } w ∈ N ( u ) \ v (cid:17) for all edges (
u, v ) ∈ H simultaneously.
end for
Return node representations h v =
MLP (cid:16) x v , (cid:80) [Equation] ∈ N ( v ) ν
( T ) uv (cid:17) end function
Attention Layer Our attention layer is a bilinear attention function with parameter θ =
{ A θ } : attention θ ( v , { h [Equation] [Equation] =
[Equation]
Then the model predicts the bond type between u t and other nodes in Q sequentially for |Q| steps (bond prediction, |Q| = 2 ).
Finally, it adds the new atom to the queue Q .
AtomG2G Architecture AtomG2G is an atom-based translation method that is directly comparable to HierG2G. Here molecules are represented solely as molecular graphs rather than a hierarchical graph with substructures.
The encoder of AtomG2G is the same LSTM MPN over molecular graph.
This gives us a set of atom vectors c G X representing molecule X only at the atom level.
The decoder of AtomG2G is illustrated in Figure 7.
Following You et al.
(2018b);
Liu et al.
(2018), the model generates molecule G atom by atom following their breath-ﬁrst order.
During generation, it maintains a FIFO queue Q that contains the frontier nodes in the graph (i.e., nodes who still have neighbors to be generated).
Let v t be the ﬁrst node in Q and G t be the current graph at step t .
In each step, the model makes three predictions to expand the graph G t :
It predicts whether there will be new atoms attached to v t .
If not, the model discards v and move on to the next node in Q .
The generation stops if Q is empty.
Otherwise, it creates a new atom u t and predicts its atom type.
Lastly, it predicts the bond type between u t and other frontier nodes in Q autoregressively to fully capture edge dependencies (You et al., 2018b).
Since nodes are generated in breath-ﬁrst order, there will be no edges between u t and nodes outside of Q .
To make those predictions, we use the same LSTM MPN to encode the current graph G t .
Let h v t be the atom representation of v t .
We represent G t as the sum of all its atom vectors h G t = (cid:80) v ∈G t h v .
In the ﬁrst step, we model the probability of expanding a new node from v t as:
[Equation]
In the second step, the atom type of the new node u t is predicted using another MLP:
[Equation]
In the last step, we predict the bonds between u t and nodes in Q = a 1 , · · · , a n sequentially starting with a 1 = v t .
Speciﬁcally, for each atom pair ( u t , a k ) , we predict their bond type (single, double, triple or none) as the following: [Equation]
[Equation]
where h a k is the atom representation of node a k and h ku t is the representation of node u t at the k th bond prediction.
Let N k ( u t ) be node u t ’s current neighbor predicted in the ﬁrst k steps.
h ku t is computed as follows to reﬂect its local graph structure after k th bond prediction:
[Equation]
where x u t is the atom feature of u t (i.e., predicted atom type) and x w,u t is the bond feature between w and u t (i.e., predicted bond type).
Intuitively, this can be viewed as running one-step message passing at each bond prediction step (i.e., passing the message ν w,u t from w to u t ).
AtomG2G is trained under the same variational objective as HierG2G, with the latent code z sampled from the posterior Q ( z | X, Y ) =
N ( µ X,Y , σ X,Y ) and [ µ X,Y , σ X,Y ] =
MLP( (cid:80) c
G Y − (cid:80) c G X ) .
[Table]
C E XPERIMENTAL D ETAILS
Data
The single-property optimization datasets are directly downloaded from the link provided in
Jin et al.
The training set and substructure vocabulary size for each dataset is listed in Table 3.
We constructed the multi-property optimization by combining the training set of QED and DRD optimization task.
The test set contains 780 compounds that are not drug-like and DRD2-inactive.
The training and test set is attached as part of the supplementary material.
Hyperparameters For HierG2G, we set the hidden layer dimension to be 270 and the embedding layer dimension 200.
We set the latent code dimension | z | = 8 and KL regularization weight λ KL = 0 .
We run T = 20 iterations of message passing in each layer of the encoder.
For AtomG2G, we set the hidden layer and embedding layer dimension to be 400 so that both models have roughly the same number of parameters.
We also set λ KL = 0 .
3 and number of message passing iterations to be T = 20 .
We train both models with Adam optimizer with default parameters.
For CG-VAE (Liu et al., 2018), we used their ofﬁcial implementation for our experiments.
Speciﬁcally, for each dataset, we trained a CG-VAE to generate molecules and predict property from the latent space.
This gives us three CG-VAE models for logP, QED and DRD2 optimization tasks, respectively.
At test time, each compound X is translated following the same procedure as in Jin et al.
First, we embed X into its latent representation z and perform gradient ascent over z to maximize the predicted property score.
This gives us z 1 , · · · , z K vectors for K gradient steps.
Then we decode K molecules from z 1 , · · · , z K and select the one with the best property improvement within similarity constraint.
We found that it is necessary to keep the KL regularization weight low ( λ KL = 0 . 005 ) to achieve meaningful results.
When λ KL = 1 . 0 , the above gradient ascent procedure always generate molecules very dissimilar to the input X .
Ablation Study Our ablation studies are illustrated in Figure 8.
In our ﬁrst experiment, we changed our decoder to the atom-based decoder of AtomG2G. As the encoder is still hierarchical, we modiﬁed the input of the decoder attention to include both atom and substructure vectors.
We set the hidden layer and embedding layer dimension to be 300 to match the original model size.
Our next two experiments reduces the number of hierarchies in both our encoder and decoder MPN.
In the two-layer model, molecules are represented by c X = c
G X ∪ c A X .
We make topological and substructure predictions based on hidden vector h
A k instead of h S k because the substructure layer is removed.
In the one-layer model, molecules are represented by c X = c
G X
and we make topological and substructure predictions based on atom vectors (cid:80) v ∈S k h v .
The hidden layer dimension is adjusted accordingly to match the original model size.