D IFFERENTIABLE S CAFFOLDING T REE FOR M OLECULAR O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
The structural design of functional molecules, also called molecular optimization, is an essential chemical science and engineering task with important applications, such as drug discovery.
Deep generative models and combinatorial optimization methods achieve initial success but still struggle with directly modeling discrete chemical structures and often heavily rely on brute-force enumeration.
The challenge comes from the discrete and non-differentiable nature of molecule structures.
To address this, we propose differentiable scaffolding tree ( DST ) that utilizes a learned knowledge network to convert discrete chemical structures to locally differentiable ones.
DST enables a gradient-based optimization on a chemical graph structure by back-propagating the derivatives from the target properties through a graph neural network (GNN).
Our empirical studies show the gradient-based molecular optimizations are both effective and sample efﬁcient.
Furthermore, the learned graph parameters can also provide an explanation that helps domain experts understand the model output.
I NTRODUCTION
The structural design of new functional molecules, also called molecular optimization, is the key to many scientiﬁc and engineering challenges, such as ﬁnding energy storage materials (Hachmann et al., 2011; Janet et al., 2020), small molecule pharmaceutics (Kuntz, 1992; Zhavoronkov et al., 2019), and environment-friendly material (Zimmerman et al., 2020).
The objective is to identify novel molecular structures with desirable chemical or physical properties (Gómez-Bombarelli et al., 2018; Dai et al., 2018; Jin et al., 2018;
You et al., 2018; Jin et al., 2019; Shi et al., 2020; Zhou et al., 2019; Jin et al., 2020; Zang & Wang, 2020; Xie et al., 2021).
Recent advances in deep generative models (DGM) allow learning the distribution of molecules and optimizing the latent embedding vectors of molecules.
Models in this category are exempliﬁed by the variational autoencoder (VAE) (Gómez-Bombarelli et al., 2018; Dai et al., 2018; Jin et al., 2018; 2020) and generative adversarial network (GAN) (De Cao & Kipf, 2018).
On the other hand, because of the discrete and not explicitly combinatorial nature of the enormous chemical space, applying combinatorial optimization algorithms with some structure enumeration has been the predominant approach (You et al., 2018; Jensen, 2019;
Zhou et al., 2019;
Nigam et al., 2020; Xie et al., 2021).
Deep learning models have also been used to guide these combinatorial optimization algorithms.
For example, You et al.
; Zhou et al.
; Jin et al.
; Gottipati et al.
(2020) tried to solve the problem with deep reinforcement learning; Nigam et al.
(2020) enhanced a genetic algorithm with a neural network as a discriminator;
Xie et al.
Fu et al.
(2021) approached the problem with Markov Chain Monte Carlo (MCMC) to explore the target distribution guided by graph neural networks.
Despite the initial success of these previous attempts, the following limitations remain: (1) deep generative models optimize the molecular structures in a learned latent space, which requires the latent space to be smooth and discriminative.
Training such models needs carefully designed networks and well-distributed datasets.
(2) most combinatorial optimization algorithms, featured by evolutionary learning methods (Nigam et al., 2020; Jensen, 2019; Xie et al., 2021; Fu et al., 2021), exhibit randomwalk behavior, and leverage trial-and-error strategies to explore the discrete chemical space.
The recent deep reinforcement learning methods (You et al., 2018; Zhou et al., 2019; Jin et al., 2020;
Gottipati et al., 2020) aim to remove random-walk search using a deep neural network to guide the searching.
It is challenging to design the effective reward function into the objective (Jin et al., 2020).
(3) Most existing methods require a great number of oracle calls (a property evaluator; see
Def. 1) to proceed with an efﬁcient search.
However, realistic oracle functions, evaluating with either experiments or high-ﬁdelity computational simulations, are usually expensive.
Examples include using biological assays to determine the potency of drug candidates (Wang et al., 2017), or conducting electronic structure calculation to determine photoelectric properties (Long et al., 2011).
We propose differentiable scaffolding tree ( DST ) to address these challenges, where we deﬁne a differentiable scaffolding tree for molecular structure and utilize a trained GNN to obtain the local derivative that enables continuous optimization.
The main contributions are summarized as:
• We propose the differentiable scaffolding tree to deﬁne a local derivative of a chemical graph.
This concept enables a gradient-based optimization of a discrete graph structure.
• We present a general molecular optimization strategy utilizing the local derivative deﬁned by the differentiable scaffolding tree.
This strategy leverages the property landscape’s geometric structure and suppresses the random-walk behavior, exploring the chemical space more efﬁciently.
We also incorporate a determinantal point process (DPP) based selection strategy to enhance the diversity of generated molecules.
• We demonstrate encouraging preliminary results on de novo molecular optimization with multiple computational objective functions.
The local derivative shows consistency with chemical intuition, providing interpretability of the chemical structure-property relationship.
Our method also requires less oracle calls, maintaining good performance in limited oracle settings.
R ELATED W ORK
Existing molecular optimization methods can mainly be categorized as deep generative models and combinatorial optimization methods.
Deep generative models model a distribution of general molecular structure with a deep network model so that one can generate molecules by sampling from the learned distribution.
Typical algorithms include variational autoencoder (VAE), generative adversarial network (GAN), energybased models, ﬂow-based model (Gómez-Bombarelli et al., 2018; Jin et al., 2018; De Cao & Kipf, 2018;
Segler et al., 2018; Jin et al., 2019; Honda et al., 2019; Madhawa et al., 2019; Shi et al., 2020;
Jin et al., 2020; Zang & Wang, 2020; Kotsias et al., 2020; Chen et al., 2021; Liu et al., 2021; Bagal et al., 2021).
Shen et al.
(2021) also leverages inverse learning based on SELFIES representation.
However, its performance is not satisfactory, primarily due to the failure of training an adequate surrogate oracle.
In addition, DGMs can leverage Bayesian optimization in latent spaces to optimize latent vectors and reconstruct to obtain the optimized molecules (Jin et al., 2018).
However, such approaches usually require a smooth and discriminative latent space and thus an elaborate network architecture design and well-distributed data set.
Also, as they learn the reference data distribution, their ability to explore diverse chemical space is relatively limited, evidenced by the recent molecular optimization benchmarks (Brown et al., 2019; Huang et al., 2021).
Combinatorial optimization methods mainly include deep reinforcement learning (DRL) (You et al., 2018;
Zhou et al., 2019; Jin et al., 2020; Gottipati et al., 2020) and evolutionary learning methods (Nigam et al., 2020; Jensen, 2019; Xie et al., 2021; Fu et al., 2021).
They both formulate molecule optimization as a discrete optimization task.
Speciﬁcally, they modify molecule substructures (or tokens in a string representation (Weininger, 1988)) locally, with an oracle score or a policy/value network to tell if they keep it or not.
Due to the discrete nature of the formulation, most of them conduct an undirected search (random-walk behavior), while some recent ones like reinforcement learning try to guide the searching with a deep neural network, aiming to rid the random-walk nature.
However, it is challenging to incorporate the learning objective target into the guided search.
Those algorithms still require massive numbers of oracle calls, which is computationally inefﬁcient during the inference time (Korovina et al., 2020).
Our method, DST , falls into this category, explicitly leverages the objective function landscape and conducts an efﬁcient goal-oriented search.
Instead of operating on molecular substructure or tokens, we deﬁne the search space as a set of binary and multinomial variables to indicate the existence and identity of nodes respectively, and make it locally differentiable with a learned GNN as a surrogate of the oracle.
This problem formulation can ﬁnd its root in conventional computer-aided molecular design algorithms with branch-and-bound algorithms as solutions (Sinha et al., 1999; Sahinidis & Tawarmalani, 2000).
GNN Predict
Scaffolding Tree
Molecular Graph
Differentiable Scaffolding Tree 3
Molecular Graph
EXPAND
UNCHANGE
Abstract
GradientAscent
Assemble and Diversify
M ETHOD
We ﬁrst introduce the formulation of molecular optimization and differentiable scaffolding tree ( DST ) in Section 3.1, illustrate the pipeline in Figure 1, then describe the key steps following the order:
• Oracle GNN construction: We leverage GNNs to imitate property oracles, which are targets of molecular optimization (Section 3.2).
Oracle GNN is trained once and for all.
The training is separately from optimizing DST below.
• Optimizing differentiable scaffolding tree: We formulate the discrete molecule optimization into a locally differentiable problem with a differentiable scaffolding tree ( DST ).
Then a DST can be optimized by the gradient back-propagated from oracle GNN (Section 3.3).
• Molecule Diversiﬁcation After that, we describe how we design a determinantal point process (DPP) based method to output diverse molecules for iterative learning (Section 3.4).
3. P ROBLEM F ORMULATION AND N OTATIONS 3.1.
Molecular optimization problem Oracles are the objective functions for molecular optimization problems, e.g., QED quantifying a molecule’s drug-likeness (Bickerton et al., 2012).
Deﬁnition 1 (Oracle O ) .
Oracle O is a black-box function that evaluates certain chemical or biological properties of a molecule X and returns the ground truth property
O ( X ) .
In realistic discovery settings, the oracle acquisition cost is usually not negligible.
Suppose we want to optimize P molecular properties speciﬁed by oracle O 1 , · · · , O P , we can formulate a multi-objective molecule optimization problem through scalarization as represented in Eq.
[Equation]
where X is a molecule, Q denotes the set of valid molecules; f is the composite objective combining all the oracle scores, e.g., the mean value of P oracle scores.
3.1.2 Scaffolding Tree
The basic mathematical description of a molecule is molecular graph, which contains atoms as nodes and chemical bonds as edges.
However, molecular graphs are not easy to generate explicitly as graphs due to the presence of rings, relatively large size, and chemical validity constraints.
For ease of computation, we convert a molecular graph to a scaffolding tree as a higher-level representation, a tree of substructures, following Jin et al.
Deﬁnition 2 ( Substructure ) .
Substructures can be either an atom or a single ring.
The substructure set is denoted S (vocabulary set), which covers frequent atoms and single rings in drug-like molecules.
Deﬁnition 3 ( Scaffolding Tree T ) .
A scaffolding tree, T X , is a spanning tree whose nodes are substructures.
It is higher-level representation of molecular graph X .
T X is represented by (i) node indicator matrix, (ii) adjacency matrix, and (iii) node weight vector.
We distinguish leaf and non-leaf nodes in T X .
Among the K 1 nodes in T X , there are K leaf leaf nodes (nodes connecting to only one edge) and K − K leaf non-leaf nodes (otherwise).
The sets of leaf nodes and non-leaf nodes are denoted V leaf and V nonleaf correspondingly.
[Figure]
[Figure]
(cid:101) w and (cid:101)
A share the learnable parameters { (cid:98) w 3 , (cid:98) w 4 , (cid:98) w 5 | 3 , (cid:98) w 6 | 4 , (cid:98) w 7 | 1 , (cid:98) w 8 | 2 } .
Deﬁnition 4.
Node indicator matrix N is decomposed as N = (cid:18) N nonleaf
N leaf (cid:19) ∈ { 0 , 1 } K ×|S| , where
N nonleaf ∈ { 0 , 1 } ( K − K leaf )
×|S| corresponds to non-leaf nodes while N leaf ∈ { 0 , 1 } K leaf ×|S| corresponds to leaf nodes.
Each row of N is a one-hot vector, indicating which substructure the node belongs to.
Deﬁnition 5.
Adjacency matrix is denoted A ∈ { 0 , 1 } K × K .
A ij = 1 indicates the i -th node and j -th node are connected while 0 indicates unconnected.
Deﬁnition 6.
Node weight vector , w = [1 , · · · , 1] (cid:62) ∈ R K , indicates the K nodes in scaffolding tree are equally weighted.
3.1.3 Differentiable scaffolding tree Similar to a scaffolding tree, a differentiable scaffolding tree (DST) also contains (i) node indicator matrix, (ii) adjacency matrix, and (iii) node weight vector, but with additional expansion nodes.
Speciﬁcally, while inheriting leaf node set V leaf and non-leaf node set V nonleaf from the original scaffolding tree, we add expansion nodes and form expansion node set ,
V expand =
{ u v | v ∈ V leaf ∪ V nonleaf } , |V expand | =
K expand =
K , where u v is connected to v in the original scaffolding tree.
We also deﬁne differentiable edge set , Λ = { ( v, v (cid:48) )
| v ∈ V leaf OR v (cid:48) ∈
V expand ; v, v (cid:48) are connected } to incorporate all the edges involving leaf-nonleaf node and leaf/nonleafexpansion node connections.
To make it locally differentiable, we modify the tree parameters from two aspects: (A) node identity and (B) node existence .
Figure 2 shows an example to illustrate DST.
(A) We enable optimization on node identity by allowing the corresponding node indicator matrix learnable:
Deﬁnition 7.
Differentiable node indicator matrix (cid:101) N takes the form:
[Equation]
(cid:101) N nonleaf =
N nonleaf ∈ { 0 , 1 } ( K − K
[Equation]
) ×|S| are ﬁxed, equal to the part in the original scaffolding tree, each row is a one-hot vector, indicating that we ﬁx all the non-leaf nodes.
In contrast, both (cid:101) N expand and (cid:101) N leaf are learnable, we use softmax activation to implicitly encode the constraint (cid:80) j (cid:101) N ij =i.e., (cid:101) N ij =
[Equation]
j
[Equation]
N
[Equation]
j
[Equation]
) , (cid:98) N are the parameters to learn.
This constraint guarantee that each row of (cid:101)
N is a valid substructures’ distribution.
(B) We enable optimization on node
existence
by assigning
learnable weights for
the leaf
and
expansion
nodes
, construct adjacency
matrix and node weight vector
based on those values:
Deﬁnition 8.
Differentiable adjacency matrix (cid:101)
A ∈ R
( K + K expand ) ×
( K + K expand ) takes the form:
[Equation] [Equation]
[Equation]
where Λ is the differentiable edge set deﬁned above, Sigmoid function σ ( · ) imposes the constraint 0 ≤ (cid:101)
A ij ≤ 1 .
(cid:98) w ∈ R K leaf
+ K expand are the parameters, each leaf node and expansion node has one learnable parameter.
For connected i and j , when i ∈ V leaf , j ∈ V nonleaf , (cid:101)
A ij = σ ( (cid:98) w i ) measures the existence probability of leaf node i ; when i ∈ V expand , j ∈ V leaf ∪V nonleaf , (cid:101)
A ij = σ ( (cid:98) w i | j ) measures the conditional probability of the existence of expand node i given the original node j .
When j is a leaf node, it naturally embeds the inheritance relationship between the leaf node and the corresponding expansion node.
Deﬁnition 9.
Differentiable node weight vector (cid:101)
w
[Equation]
K + K expand takes the form:
[Equation]
where all the weights range from 0 to 1.
The weight of expansion node connecting to leaf node relies on the weight of corresponding leaf node.
(cid:101) w and (cid:101)
A (Def.
8) shares the learnable parameter (cid:98) w .
Figure 2 shows an example to illustrate DST.
3. T RAINING O RACLE G RAPH N EURAL N ETWORK
This section constructs a differentiable surrogate model to capture the knowledge from any oracle function.
We choose graph neural network architecture for its state-of-the-art performance in modeling structure-property relationships.
In particular, we imitate the objective function F with GNN:
[Equation]
where Θ represents the GNN’s parameters.
Concretely, we use a graph convolutional network (GCN) (Kipf & Welling, 2016).
Other GNN variants, such as Graph Attention Network (GAT) (Veli ˇ ckovi ´
c et al., 2017), Graph Isomorphism Network (GIN) (Xu et al., 2018), can also be used in our setting.
The initial node embeddings H (0)
=
NE ∈ R K × d stacks basic embeddings of all the nodes in the scaffolding tree, d is the GCN hidden dimension, N is the node indicator matrix (Def. 4).
E ∈ R |S|×
d is the embedding matrix of all the substructures in vocabulary set S , and is randomly initialized.
The updating rule of GCN for the l -th layer is
[Equation]
where L is GCN’s depth, A is the adjacency matrix (Def. 5),
H ( l ) ∈ R K
×
d is the nodes’ embedding of layer l , B ( l ) =
[ b ( l ) , b ( l ) , · · · , b ( l ) ] (cid:62) ∈
R K × d and U
( l ) ∈ R
d × d are bias and weight parameters of layer l , respectively.
We generalize the GNN from a discrete scaffolding tree to a differentiable one.
Based on learnable weights for each node, we leverage the weighted average as the readout function of the last layer’s
( L -th) node embeddings, followed by multilayer perceptron (MLP) to yield the prediction (cid:98) y , i.e., (cid:98) y
= MLP (cid:0)
[Equation]
=1 w k (cid:80) Kk
=
w k H ( L ) k (cid:1) , in discrete scaffolding tree, weights for all the nodes are equal to 1, H ( L ) k is the k -th row of H ( L ) .
In sum, the prediction can be written as
[Equation]
where Θ = { E } ∪ { B ( l ) , U ( l ) } Ll =1 are the GNN’s parameters.
We train the GNN by minimizing the discrepancy between GNN prediction (cid:98) y and the ground truth y .
[Equation]
where L is loss function, e.g., mean squared error; D is the training set.
After training, we have GNN parameterized by Θ ∗ to approximate the black-box objective function F (Eq. 1).
Worth to mention that Oracle GNN is trained once and for all.
The training is separately from optimizing DST below.
3. O PTIMIZING D IFFERENTIABLE S CAFFOLDING T REE
Overview With a little abuse of notations, via introducing DST , we approximate molecule optimization as a locally differentiable problem
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
where X ( t ) is the molecule at t -th iteration, N ( X ( t ) )
⊆
Q is the neighborhood set of X ( t )
(Def. 10).
Next, we explain the intuition behind these approximation steps.
Molecular optimization is generally a discrete optimization task, which is prohibitively expensive due to exhaustive search.
The ﬁrst approximation is to formulate the problem as an iterative local discrete search via introducing a neighborhood molecule set N ( X ( t ) ) .
Second, to enable differentiable learning, we use GNN to imitate black-box objective F (Section 3.2) and further reformulated it into a local differentiable optimization problem.
Then we can optimize DST ( T X = {
[Equation]
N X
[Equation]
t ) , (cid:101)
A X ( t ) , (cid:101) w X ( t ) } ) in a continuous domain for N ( X ( t ) ) using gradient-based optimization method.
3.3.1 Local Editing Operations For
a leaf node v in the scaffolding tree, we can perform three editing operations, (1) SHRINK : delete node v ; (2) REPLACE : replace a new substructure over v ; (3) EXPAND : add a new node u v that connects to node v .
For a nonleaf node v , we support (1)
EXPAND : add a new node u v connecting to v ; (2) do nothing .
If we EXPAND and REPLACE, the new substructures are sampled from the vocabulary S .
We deﬁne molecule neighborhood set:
Deﬁnition 10 (Neighborhood set) .
Neighborhood set of molecule X , denoted N ( X ) , is the set of all the possible molecules obtained by imposing one local editing operation to scaffolding tree T X and assembling the edited trees into molecules.
Optimizing DST Then within the domain of neighborhood molecule set N ( X ) , the objective function can be represented as a differentiable function of X’s DST ( (cid:101) N X , (cid:101)
A X , (cid:101) w X ).
We address the following optimization problem to get the best scaffolding tree within N ( X ) ,
[Equation] [Equation]
) where the GNN parameters Θ ∗ (Eq.
(8)) are ﬁxed.
Comparing with Eq.
(7), it is differentiable with regard to { (cid:101) N , (cid:101)
A , (cid:101) w } for all molecules in the neighborhood set N ( X ) .
Therefore, we can optimize the
DST using gradient-based optimization method, e.g., an Adam optimizer (Kingma & Ba, 2014).
The whole DST pipeline leverages iterative local discrete search (Equation 9, Algorithm 1), speciﬁcally, in t -th iteration, we optimize the DST of X ( t ) , i.e., X in Equation (10) is X ( t ) .
3.3.3 Sampling from DST
Then we sample the new scaffolding tree from the optimized DST .
Concretely, (i) for each leaf node v ∈ V leaf and the corresponding expansion node u v ∈ V expand , we select one of the following step with probabilities (w.p.) as follows,
[Equation] [Equation]
u
substructure
[Equation] [Equation]
[Equation] [Equation]
[Equation]
REPLACE: select substructure at v based on
[Equation]
(ii)
For each nonleaf node v , we expand a new node with probability ( (cid:101) w ∗ )
u v | v .
If
expanding,
we
select substructure at u v based on ( (cid:101) N ∗ ) u v .
Otherwise, we “ do nothing ”
3.3.4 Assemble Each scaffolding tree corresponds to multiple molecules due to the multiple ways substructures can be combined.
We enumerate all the possible molecules following Jin et al.
(See Section C.5 in Appendix for details) for the further selection as described below.
3. M OLECULE D IVERSIFICATION
In the current iteration, we have generated M molecules ( X 1 , · · · , X M ) and need to select C molecules for the next iteration.
We expect these molecules to have desirable chemical properties
(high F score) and simultaneously maintain higher structural diversity.
To do so, we resort to the determinantal point process (DPP) (Kulesza & Taskar, 2012), which models the repulsive correlation between data points.
Speciﬁcally, for M data points, whose indexes are { 1 , 2 , · · · , M } , S ∈ R M × M + denotes the similarity matrix between these data points.
To create a diverse subset (denoted R ) with ﬁxed size C , the sampling probability should be proportional to the determinant of the submatrix
S R ∈ R C × C , i.e., P ( R ) ∝ det( S R ) , where R ⊆ { 1 , 2 , · · · , M } , |R| = C .
Combining the objective ( F ) value and diversity, the composite objective is
[Equation]
where the hyperparamter λ > 0 balances the two terms, the diagonal scoring matrix V = diag ([exp( λF ( X 1 )) , · · · , exp( λF ( X M ))]) , V R ∈ R C × C is a sub-matrix of V indexed by R .
When λ goes to inﬁnity, it is equivalent to select top- C candidates with the highest F score, same as conventional evolutionary learning in (Jensen, 2019; Nigam et al., 2020).
Inspired by generalized
DPP methods (Chen et al., 2018), we further transform L DPP ( R ) ,
[Equation]
where V 12 S V 12 is symmetric positive semi-deﬁnite.
Then it can be solved by generalized DPP methods in O ( C 2 M )
(Chen et al., 2018)
(Section F in Appendix).
The computational complexity of
DST is O
( TMC 2 ) (Section C.8 in Appendix).
Algorithm 1 summarizes the entire algorithm.
13: end
E XPERIMENT 4. E XPERIMENTAL S ETUP
Molecular Properties contains QED ; LogP ; SA ; JNK3 ; GSK3 β , following (Jin et al., 2020;
Nigam et al., 2020; Moss et al., 2020;
Xie et al., 2021), where QED quantiﬁes drug-likeness; LogP indicates the water-octanol partition coefﬁcient; SA stands for synthetic accessibility and is used to prevents the formation of chemically unfeasible molecules; JNK3/GSK3 β measure inhibition against c-Jun
N-terminal kinase-3/Glycogen synthase kinase 3 beta.
For all the 5 scores (including normalized
SA), higher is better.
We conducted (1) single-objective generation that optimizes JNK3, GSK3 β and
LogP separately and (2) multi-objective generation that optimizes the mean value of “JNK3+GSK3 β ” and “QED+SA+JNK3+GSK3 β ” in the main text.
Details are in Section C.3.
Dataset : ZINC 250K contains around 250K druglike molecules (Sterling & Irwin, 2015).
We select the substructures that appear more than 1000 times in ZINC 250K as the vocabulary set S , which contains 82 most frequent substructures.
Details are in Section C.1.
Baselines .
(1) LigGPT (string-based distribution learning model with Transformer as a decoder) (Bagal et al., 2021); (2) GCPN (Graph Convolutional Policy Network)
(You et al., 2018); (3) MolDQN
(Molecule Deep Q-Network)
(Zhou et al., 2019); (4) GA+D
(Genetic Algorithm with Discriminator network) (Nigam et al., 2020); (5) MARS (Markov Molecular Sampling)
(Xie et al., 2021); (6)
RationaleRL (Jin et al., 2020); (7) ChemBO (Chemical Bayesian Optimization) (Korovina et al., 2020); (8) BOSS (Bayesian Optimization over String Space)
(Moss et al., 2020).
Among them,
LigGPT belongs to deep generative model, where all the oracle calls can be precomputed; GCPN,
MolDQN are deep reinforcement learning methods; GA+D, MARS are evolutionary learning methods; RationaleRL is deep generative model ﬁne-tuned with RL techniques.
ChemBO and BOSS are
Bayesian optimization methods.
We also consider a DST variant: DST -rand .
Instead of optimizing and sample from DST , DST -rand leverages random local search, i.e., randomly selecting basic operations (EXPAND, REPLACE, SHRINK) and substructure from vocabulary.
To improve efﬁciency, we also select a subset of all the random samples with high surrogate GNN prediction scores.
All the baselines except LigGPT require online oracle calls.
Details are in Section B.
Metrics .
For each method, we select top-100 molecules with highest property scores for evaluation, and consider the following metrics following Jin et al.
; You et al.
; Jin et al.
Xie et al.
(2021) (1) Novelty (Nov) (% of the generated molecules that are not in training set); (2)
Diversity (Div) (average pairwise Tanimoto distance between the Morgan ﬁngerprints); (3) Average
Property Score (APS) (average score of top-100 molecules); (4) # of oracle calls : DST needs to call oracle in labeling data for GNN ( precomputed ) and DST based de novo generation ( online ), we show the costs for both steps.
For each method in Table 1 and 2, we set the number of oracle calls so that the property score nearly converge w.r.t.
oracle call’s number.
Details are in Section C.4.
[Table]
[Figure]
[Table]
[Figure]
4. O PTIMIZATION P ERFORMANCE
The results of multi-objective and single-objective generation are shown in Table 1 and 2.
We ﬁnd that DGM (LigGPT) and RL based methods (GCPN and MolDQN) fails in some tasks, which is consistent with the results reported in RationaleRL (Jin et al., 2020) and MARS (Xie et al., 2021).
Overall, DST obtains the best results in most tasks.
In terms of success rate and diversity, DST outperformed all baselines in most tasks.
It also reached the highest scores within T = 50 iterations in most optimization tasks (see Table 5 and 6 in Appendix).
Especially in optimizing LogP, the model successfully learned to add a six-member ring (see Figure 9 in Appendix) each step, which is theoretically the optimal strategy under our setting.
Combined with the ablation study comparing with random selection (see Figure 12 in Appendix), our results show the local gradient deﬁned by DST is a useful direction indicator, consistent with the concept of gradient.
Further, achieving high diversity validates the effect of the DPP-based selection strategy.
Although the novelty is not the highest, it is still comparable to baseline methods.
These results show our gradient-based optimization strategy has a strong optimization ability to provide a diverse set of molecules with high objective functions.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
RationaleRL
[Table]
- - - - MARS
[Table]
4. O RACLE E FFICIENCY
As mentioned above, oracle calls for realistic optimization tasks can be time-consuming and expensive.
From Table 1 and 2, we can see that majority of de novo optimization methods require oracle calls online (instead of precomputation), including all of RL/evolutionary algorithm based baselines.
DST takes fewer oracle calls compared with baselines.
DST can leverage the precomputed oracle calls to label the molecules in an existing database (i.e., ZINC) for training the oracle GNN and dramatically saving the oracle calls during reference.
In the three tasks in Table 2, two-thirds of the oracle calls (10K) can be precomputed or collected from other sources.
To further verify the oracle efﬁciency, we explore a special setting of molecule optimization where the budget of oracle calls is limited to a ﬁxed number (2K, 5K, 10K, 20K, 50K) and compare the optimization performance.
For GCPN,
MolDQN, GA+D and MARS, the learning iteration number depends on the budget of oracle calls.
RationaleRL (Jin et al., 2020) is not included because it requires intensive oracle calls to collect enough reference data, exceeding the oracle budget in this scenario.
In DST , we use around 80% budget to label the dataset (i.e., training GNN) while the remaining budget to conduct de novo design.
Speciﬁcally, for 2K, 5K, 10K, 20K, 50K, we use 1.5K, 4K, 8K, 16K and 40K oracle calls to label the data for learning GNN, respectively.
We show the average objective values of top-100 molecules under different oracle budgets in Figure 3.
Our method shows a signiﬁcant advantage compared to all the baseline methods in all limited budget settings.
We conclude the reason as supervised learning is a well-studied and much easier task than generative modeling.
4. I NTERPRETABILITY A NALYSIS
To obtain more insights from the local gradient, We visualize two modiﬁcation steps in Figure 4.
The node weights and their gradient values interpret the property at the substructure level.
This is similar to most attribution-based interpretable ML methods, e.g., saliency map (Simonyan et al., 2013).
C ONCLUSION
This paper proposed differentiable scaffolding tree ( DST ) to make a molecular graph locally differentiable, allowing a continuous gradient-based optimization.
To the best of our knowledge, it is the ﬁrst attempt to make the molecular optimization problem differentiable at the substructure level, rather than resorting to latent spaces or using RL/evolutionary algorithms.
We constructed a general molecular optimization strategy based on DST , corroborated by thorough empirical studies.
R EPRODUCIBILITY S TATEMENT
The code repository is given in the supplementary material, including README ﬁle, all the codes for data preprocessing and model.
All the data we use are publicly available.
Section B describes experimental setup for all the baseline methods.
We elaborate the implementation details in Section C.
Concretely, Section C.1 describes the dataset we use; Section C.2 describes software and hardware conﬁguration; Section C.3 introduces the target molecular properties to optimize; Section C. describes all the metrics for evaluating the performance; Section C.6 describes the setup for model implementation, especially hyperparameter setup.
C ONTENTS 1 Introduction
2 Related Work 3 Method 3.
Problem Formulation and Notations . . . . . . . . . . . . . . . . . . . . . . . . .
Training Oracle Graph Neural Network . . . . . . . . . . . . . . . . . . . . . . .
Optimizing Differentiable Scaffolding Tree . .
Molecule Diversiﬁcation . .
4 Experiment 4.
Experimental Setup . .
Optimization Performance . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Oracle Efﬁciency . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Interpretability Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
5 Conclusion
A Complete Mathematical Notations.
B Baseline Setup
C Implementation Details
C.1 Dataset . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2 Software/Hardware Conﬁguration . . . . . . . . . . . . . . . . . . . . . . . . . .
C.3 Target molecular properties . . C
Evaluation
C Scaffolding
C
Details GNN C.7 Results
Different Random Seeds . . . . . . . . . . . . . . . . . . . . . . . . . .
C.8 Complexity Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D Additional Experimental Results D.1 Additional results of de novo molecular generation . . . . . . . . . . . . . . . . . D.2 De novo molecular optimization on QED (potential limitation of DST ) . . . . . . .
D.3 Results Analysis for Distribution Learning Methods (LigGPT) . . . . . . . . . . .
D.4 Ablation study . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.5 Chemical space visualization . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.6 Additional Interpretability Analysis . . . . . . . . . . . . . . . . . . . . . . . . .
E Theoretical Analysis
E.1 Convergence Analysis
F Extension of Molecule diversiﬁcation
G Proof of Theoretical Results
G.1 Proof of Lemma 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.2 Proof of Lemma 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.3 Proof of Lemma 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.
Proof
of Theorem
Section*Appendix to Differentiable Scaffolding Tree for Molecular Optimization
The appendix is organized as follows.
First, we list the complete mathematical notations for ease of exposition.
Then, we show additional experimental setup and empirical results, including baseline setup in Section B, implementation details of our method in Section C, additional experimental results in Section D.
Then, we provide theoretical analysis in Section E, extend molecule diversiﬁcation in
Section, and prove the theoretical results in the main paper in Section G.
[Table]
A C OMPLETE M ATHEMATICAL N OTATIONS .
In this Section, we show all the mathematical notations in Table 3 for completeness.
We also illustrate the difference between DST and existing methods in Figure 5.
Summary : DST makes the learning objective differentiable w.r t molecule graph structure, while prior works make the learning objective differentiable w r t. neural networks’ parameters.
Our approach directly optimizes molecular graph structures, while prior works indirectly search for the molecule graphs with the help of a neural network.
B B ASELINE S ETUP
In this section we describe the experimental setting for baseline methods.
Most of the settings follow the original papers.
• LigGPT (string-based distribution learning model with Transformer as a decoder) (Bagal et al., 2021) is trained for 10 epochs using the Adam optimizer with a learning rate of 6 e − 4 .
LigGPT comprises stacked decoder blocks, each of which, is composed of a masked self-attention layer and fully connected neural network.
Each self-attention layer returns a vector of size 256, that is taken as input by the fully connected network.
The hidden layer of the neural network outputs a vector of size 1024 and uses a GELU activation and the ﬁnal layer again returns a vector of size 256 to be 7 used as input for the next decoder block.
LigGPT consists of 8 such decoder blocks.
LigGPT has around 6M parameters.
• GCPN (Graph Convolutional Policy Network)
(You et al., 2018) leveraged graph convolutional network and policy gradient to optimize the reward function that incorporates target molecular properties and adversarial loss.
In each step, the allowable action to the current molecule could be either connecting a new substructure or an atom with an existing molecular graph or adding a bond to connect existing atoms.
GCPN predicts the actions and is trained via proximal policy optimization (PPO) to optimize an accumulative reward, including molecular property objectives and adversarial loss.
Both policy network and adversarial network (discriminative training) use the same neural architecture, which is a three-layer graph convolutional network (GCN) (Kipf & Welling, 2016) with 64 hidden nodes.
Batch normalization is adopted after each layer, and sum-pooling is used as the aggregation function.
Adam optimizer is used with 1e-3 initial learning rate, and batch size is 32.
• MolDQN (Molecule Deep Q-Networks)
(Zhou et al., 2019), same as GCPN, formulate the molecule generation procedure as a Markov Decision Process (MDP) and use Deep QNetwork to solve it.
The reward includes target property and similarity constraint.
Following
the original paper, the episode number is 5,000, maximal step in each episode is 40.
Each step calls oracle once; thus, 200K oracle calls are needed in one generation process.
The discount factor is 0.9.
Deep Q-network is a multilayer perceptron (MLP) whose hidden dimensions are 1024, 512, 128, 32, respectively.
The input of the Q-network is the concatenation of the molecule feature (2048-bit Morgan ﬁngerprint, with a radius of 3) and the number of left steps.
Adam is used as an optimizer with 1e-4 as the initial learning rate.
Only rings with a size of 5 and 6 are allowed.
It leverages (cid:15) -greedy together with randomized value functions (bootstrapped-DQN) as an exploration policy, (cid:15) is annealed from 1 to 0.01 in a piecewise linear way.
• GA+D
(Genetic Algorithm with Discriminator network) (Nigam et al., 2020) uses a deep neural network as a discriminator to enhance exploration in a genetic algorithm.
β is an important hyperparameter that weights the importance of the discriminator’s loss in the overall ﬁtness function, and we set it to 10.
The generator runs 100 generations with a population size of 100 for de novo molecular optimization and 50 generations with a population size of 50 for molecular modiﬁcation.
Following the original paper (Nigam et al., 2020), the architecture of the discriminator is a two-layer fully connected neural network with ReLU activation and a sigmoid output layer.
The hidden size is 100, while the size of the output layer is 1.
The input feature is a vector of chemical and geometrical properties characterizing the molecules.
We used Adam optimizer with 1e-3 as the initial learning rate.
• RationaleRL
(Jin et al., 2020) is a deep generative model that grows a molecule atom-byatom from an initial rationale (subgraph).
The architecture of the generator is a messagepassing network (MPN) followed by MLPs applied in breadth-ﬁrst order.
The generator is pre-trained on general molecules combined with an encoder and then ﬁne-tuned to maximize the reward function using policy gradient.
The encoder and decoder MPNs both have hidden dimensions of 400.
The dimension of the latent variable is 20.
Adam optimizer is used on both pre-training and ﬁne-tuning with initial learning rates of 1e-3, 5e-4, respectively.
The annealing rate is 0.9.
We pre-trained the model with 20 epochs.
• MARS (Xie et al., 2021) leverage Markov chain Monte Carlo sampling (MCMC) on molecules with an annealing scheme and an adaptive proposal.
The proposal is parameterized by a graph neural network, which is trained on MCMC samples.
We follow most of the settings in the original paper.
The message passing network has six layers, where the node embedding size is set to 64.
Adam is used as an optimizer with 3e-4 initial learning rate.
To generate a basic unit, top-1000 frequent fragments are drawn from ZINC database (Sterling & Irwin, 2015) by enumerating single bonds to break.
During the annealing process, the temperature T = 0 .
95 (cid:98) t/ 5 (cid:99) would gradually decrease to 0.
• ChemBO (chemical Bayesian optimization) (Korovina et al., 2020) leverage Bayesian optimization.
It also explores the synthesis graph in a sample-efﬁcient way and produces synthesizable candidates.
Following the default setting in the original paper, the number of steps of aquisition optimization is set to 20.
The initial pool size is set to 20, while the maximal pool size is set to 1000.
Regarding the kernel, we leveraged the kernel proposed in ChemBO, which is the optimal-transport-based distance and kernel that accounts for graphical information explicitly.
We did extra experiments to validate that it empirically outperforms the Tanimoto similarity kernel based on molecular ﬁngerprint. As for the acquisition, we followed the original paper and adopt the ensemble method using the EI (Expected Improvement), UCB (Upper Conﬁdence Bound), and TTEI (Top-Two Expected Improvement) acquisitions, we conducted extra experiment and found it empirically outperforms a single acquisition in our scenario, which is consistent to the ChemBO paper. GPs do not scale well with big data. To address this issue, we use Subset of DataChalupka et al.
(2013), which uses a subset of all the data points to reduce the size of the kernel matrix and learn the surrogate model efﬁciently. Following Chalupka et al. (2013), we selected a subset of data points using k-means clustering based on molecular ﬁngerprint, and sampled data points from each cluster. The size of the subset is set to m = 1000 .
The computationally complexity is O ( m 3 ) .
• BOSS (Bayesian Optimization over String Space)
(Moss et al., 2020) builds a Gaussian process surrogate model based on Sub-sequence String Kernel (SSK) 2 , which naturally supports SMILES strings with variable length, and maximizing acquisition function efﬁciently
for spaces with syntactical constraints.
Following their default setting, we only keep the SMILES string whose length is less than 50 (except LogP task).
At each BO step, BOSS samples 100 candidates, querying those that maximize the acquisition function predicted by SSK.
It utilized genetic algorithm (GA) to optimize acquisition function efﬁciently under syntactical constraints.
We empirically compared GA based acquisition optimizer with random search based acquisition optimizer (provided in the BOSS code repository) to validate its superiority.
GA requires intensive oracle calls and is leveraged in inner-loop maximization in BOSS, so we do not spend too many computational resources in a single GA step.
The population size is set to 100, the generation (evolution) number is set to 100.
These setups are nearly optimal and tuning these hyperparameters didn’t improve the performance.
• DST -rand a DST variant, which randomly selects basic operations (EXPAND, REPLACE, SHRINK) and substructure from vocabulary).
During each iteration, it randomly selects basic operations (EXPAND, SHRINK, REPLACE) and substructure from vocabulary to generate at most K 1 molecules, select at most top- K 2 molecules with highest GNN surrogate predictions.
Then we evaluate the K 2 candidates with the real oracle and select C molecules as starting points for the next iteration.
For fair comparison, C is set to 10, same as DST .
In order to get the near-optimal setup, we try different combinations of ( K 1 , K 2 ) , and ﬁnally ﬁnd that K 1 = 1000 , K 2 = 100 achieves the best performance in terms of average of top-100 molecules’ score.
Other setups of DST-rand follow DST .
C I MPLEMENTATION D ETAILS C. D ATASET
We use ZINC 250K dataset, which contains around 250K druglike molecules extracted from the ZINC database (Sterling & Irwin, 2015).
We ﬁrst clean the data by removing the molecules containing out-of-vocabulary substructure and having 195K molecules left.
Vocabulary S : set of substructure .
The substructure is the basic building block in our method, including frequent atoms and rings.
On the other hand, atom-wise molecule generation is difﬁcult due to the existence of rings.
To select the substructure set S , we break all the ZINC molecules into substructures (including single rings and single atoms), count their frequencies, and include the substructures whose frequencies are higher than 1000 into vocabulary set S .
The ﬁnal vocabulary contains 82 substructures, including the frequent atoms like carbon atom, oxygen atom, nitrogen atom, and frequent rings like benzene ring.
The vocabulary set covers over 80% molecules in ZINC databases.
After removing the molecules that contain out-of-vocabulary substructure, we use a random subset of the remaining molecules to train the GNNs, depending on the oracle budget.
The vocabulary size is big enough for this proof-of-concept study.
Other works also need to constrain their design space, such as MolDQN only allowing three types of atoms in a generation: “C”, “N”, “O” (Zhou et al., 2019); JTVAE (Jin et al., 2018), as well as RationaleRL (Jin et al., 2020) only using frequent substructures similar to our setting.
On the other hand, we may not want infrequent atoms or substructures because rare substructures in ZINC may have some undesired properties such as toxicity, may not be stable, may not be easily synthesizable (Gao & Coley, 2020).
Also, rare substructures may impede the learning of oracle GNN.
Note that users can enlarge the substructure space when they apply our method.
We show all the 82 substructures in S in Figure 6.
C. S OFTWARE /H ARDWARE C ONFIGURATION
We implemented DST using Pytorch 1.7.0, Python 3.7, RDKit v2020.09.1.0 on an Intel Xeon E5machine with 256G RAM and 8 NVIDIA Pascal Titan X GPUs.
C. T ARGET MOLECULAR PROPERTIES
Target molecular properties include
• QED represents a quantitative estimate of drug-likeness.
QED score ranges from 0 to 1.
It can be evaluated by the RDKit package ( https://www.rdkit.org/ ).
• LogP represents octanol-water partition coefﬁcient, measuring molecules’ solubility.
LogP score ranges from −∞ to + ∞ .
Thus, when optimizing LogP individually, we use the GNN model to do regression.
• SA (Synthetic Accessibility) score measures how hard it is to synthesize a given molecule, based on a combination of the molecule’s fragments contributions (Ertl & Schuffenhauer, 2009).
It is evaluated via RDKit (Landrum et al., 2006).
The raw SA score ranges from 1 to 10.
A higher SA score means the molecule is hard to be synthesized and is not desirable.
In the multiple-objective optimization, we normalize the SA score to [0 , 1] so that a higher normalized SA value mean easy to synthesize.
Following (Gao & Coley, 2020), we use the normalize function for raw SA score,
[Equation]
where µ = 2 . 230044 , σ = 0 .
• JNK3 (c-Jun N-terminal Kinases-3) belongs to the mitogen-activated protein kinase family and are responsive to stress stimuli, such as cytokines, ultraviolet irradiation, heat shock, and osmotic shock.
Similar to GSK3 β , JNK3 is also evaluated by well-trained 3 random forest classiﬁers using ECFP6 ﬁngerprints using ExCAPE-DB dataset (Li et al., 2018; Jin et al., 2020), and the range is also [0 , 1] .
• GSK3 β
(Glycogen synthase kinase 3 beta) is an enzyme that in humans is encoded by the GSK3 β gene.
Abnormal regulation and expression of GSK3 β is associated with an increased susceptibility towards bipolar disorder.
It is evaluated by well-trained 4 random forest classiﬁers using ECFP6 ﬁngerprints using ExCAPE-DB dataset (Li et al., 2018; Jin et al., 2020).
GSK3 β score of a molecule ranges from 0 to 1.
For QED, LogP, normalized SA, JNK3, and GSK3 β , higher scores are more desirable under our experimental setting.
C. E VALUATION METRICS
We leverage the following evaluation metrics to measure the optimization performance:
• Novelty is the fraction of the generated molecules that do not appear in the training set.
• Diversity of generated molecules is deﬁned as the average pairwise Tanimoto distance between the Morgan ﬁngerprints (You et al., 2018; Jin et al., 2020; Xie et al., 2021).
[Equation]
where Z is the set of generated molecules.
sim ( Z 1 , Z 2 ) is the Tanimoto similarity between molecule Z 1 and Z 2 .
• (Tanimoto) Similarity measures the similarity between the input molecule and generated molecules.
It is deﬁned as
[Equation]
b X is the binary Morgan ﬁngerprint vector for the molecule X .
In this paper, it is a 2048-bit binary vector.
• SR (Success Rate) is the percentage of the generated molecules that satisfy the property constraint measured by objective f deﬁned in Equation (1).
For single-objective de novo molecular generation, the objective f is the property score, the constraints for JNK3, GSK3 β and LogP are JNK3 ≥ 0 . 5 , GSK3 β ≥ 0 .
5 and LogP ≥ 5 . 0 respectively.
For multiobjective de novo molecular generation, the objective f is the average of all the normalized target property scores.
Concretely, when optimizing “JNK3+GSK3 β ”, both JNK3 and GSK3 β ranges from 0 to 1, f is average of JNK3 and GSK3 β scores; when optimizing “QED+SA+JNK3+GSK3 β ”, we ﬁrst normalized SA to 0 to 1.
f is average of QED, normalized SA, JNK3 and GSK3 β scores.
The constraint is the f score is greater than 0.4.
• # of oracle calls during the generation process.
DST needs to call oracle in labeling data for GNN and DST based de novo generation, thus we show the costs for both steps.
• chemical validities .
As we only enumerate valid chemical structures during the recovery from scaffolding trees (Section C.5), the chemical validities of the molecules produced by DST are always 100%.
Each scaffolding tree corresponds to multiple molecules due to rings’ multiple combination ways.
For each scaffolding tree, we enumerate all the possible molecules following Jin et al.
(2018) for further selection.
We provide two examples in Figure 7 to illustrate it.
Two examples are related to ring-atom combination and ring-ring combination, respectively.
For ring-ring combination, our current setting does not support the spiro compounds (contains rings sharing one atom but no bonds) or phenalenelike compounds (contains three rings sharing one atom, and each two of them sharing a bond).
These two cases are relatively rare chemical structures in the context of drug discovery (Supsana et al., 2005).
As we only enumerate valid chemical structures during the recovery from scaffolding trees, the chemical validities are always 100%.
C. D ETAILS ON GNN L EARNING AND DST O PTIMIZATION
Both the size of substructure embedding and hidden size of GCN (GNN) in Eq.
(6) are d = 100 .
The depth of GNN L is 3.
When training GNN, the training epoch number is 5, and we evaluate the loss function on the validation set every 20K data passes.
When the validation loss would not decrease, we terminate the training process.
When optimizing “JNK3”, “GSK3 β ”, “QED”, “JNK3+GSK3 β ” and “QED+SA+JNK3+GSK3 β ”, we use binary cross entropy as loss criterion.
When optimizing “LogP”, since LogP ranges from −∞ to + ∞ , we leverage GNN to conduct regression tasks and use mean square error (MSE) as loss criteria L .
In the de novo generation, in each generation, we keep
C = 10 molecules for the next iteration.
In most cases in experiment, the size of the neighborhood set (Deﬁnition. 10) is less than 100.
We use Adam optimizer with 1e-3 learning rate in training and inference procedure, optimizing the GNN and differentiable scaffolding tree, respectively.
When optimizing DST (Equation 10), our method processes one DST at a time, we conduct 1000 Adam steps (a) Ring-atom connection. When connecting atom and ring in a molecule, an atom can be connected to any possible atoms in the ring.
In the example, there are 4 possible ways to add a Chlorine atom (“Cl”) as an expansion node to the target ring, which is a leaf node in the scaffolding tree.
(b) Ring-ring connection.
When connecting ring and ring, there are two general ways, (1) one is to use a bond (single, double, or triple) to connect the atoms in the two rings.
(2) another is two rings share two atoms and one bond. In the example, there are 14 possible ways to add a Cyclohexane ring (SMILES is “C1CCCCC1”) and connect it to the target ring, which is a leaf node in the scaffolding tree. in every iteration, which is sufﬁcient to converge in almost all the cases. As a complete generation algorithm, we optimize a batch parallelly and select candidates based on DPP.
When we use up oracle budgets, we stop it.
All DST results in the tables takes at most T = 50 iterations.
C. R ESULTS OF D IFFERENT R ANDOM S EEDS
In this section, we present the empirical results that use different random seeds for multiple runs.
In our pipeline, the random error comes from in two steps: (1) Training oracle GNN: data selection/split, training process including data shufﬂe and GNN’s parameter initialization.
(2) Inference (Optimizing
DST ): before optimizing DST , we initialize the learnable parameter randomly, including (cid:101) N , (cid:101) w , (cid:101) A , which also brings randomness.
To measure the robustness of the proposed method, we use 5 different random seeds for the whole pipeline and compare the difference of 5 independent trials.
The results are reported in Table 4.
We ﬁnd that almost all the metrics would not changes signiﬁcantly among various trials, validating the robustness of the proposed method.
rate), we report the average value of 5 runs and their standard deviation
[Table]
C. C OMPLEXITY A NALYSIS
We did computational analysis in terms of oracle calls and computational complexity.
(1) oracle calls .
DST requires O ( TM ) oracle calls, where T is the number of iterations (Alg 1).
M is the number of generated molecules (Equation. 12), we have M ≤ KJ , K is the number of nodes in the scaffolding tree, for small molecule, K is very small.
J is the number of enumerated candidates in each node.
As shown in Figure 7, J is also upper-bounded ( J ≤ 4 + 14 for the example in Figure 7).
(2) computational complexity .
The computational complexity is O
( TMC 2 ) (the main bottleneck is
DPP method, Algorithm 2), where the size of selected molecules C = 10 for all the tasks (Section 3.
& C.6).
For all the tasks in Table 2 and 1, DST can be ﬁnished in 12 hours on an Intel Xeon E5562 machine with 256G RAM and 8 NVIDIA Pascal Titan X GPUs.
The complexity and runtime are acceptable for molecule optimization.
D A DDITIONAL E XPERIMENTAL R ESULTS
In this section, we present the additional empirical results, including additional results on de novo generation, ablation study, chemical space visualization, interpretability analysis (case study).
D. A DDITIONAL RESULTS OF de novo MOLECULAR GENERATION
In this section, we present some additional results of de novo molecular generation for completeness.
First, we present the optimization curve for all the optimization tasks in Figure 8.
We observe that our method is able to reach a high objective value efﬁciently within 10 iterations in all the optimization tasks.
Worth mentioning that when optimizing LogP, the model successfully learned to add a sixmember ring each step, as shown in Figure 9, and the objective ( F ) value grows linearly as a function of iteration number, which is theoretically the optimal strategy under our setting.
Then, in Figure 10, we show the molecules with the highest objective ( F ) scores generated by the proposed method on optimizing QED and “QED+SA+JNK3+GSK3 β ”.
Then we compare our method with baseline methods on 3 molecules with the highest objective ( F ) scores in Table 5 and 6 for single-objective and multi-objective generation, respectively. [Figure] [Figure]
We present the result of DST in the ﬁrst 50 iterations, but please note the setting of generation varies among the models, and a completely fair comparison is impossible.
[Table]
D. De novo MOLECULAR OPTIMIZATION ON QED ( POTENTIAL LIMITATION OF DST )
As we have touched in Section 4.2, the optimization on QED is not as satisfactory as other oracles.
We compare the performance of various methods on single-objective de novo molecular generation for optimizing QED score and show the result in Table 7.
Additional baseline methods include
JTVAE (junction tree variational autoencoder) (Jin et al., 2018) and GraphAF (Graph Flow-based
Autoregressive Model)
(Shi et al., 2020).
The main reason behind this result is that our GNN predicts the target property based on a scaffolding tree instead of a molecular graph, as shown in Equation (7).
A scaffolding tree omits rings’ assembling information, as shown in Figure 2.
Compared with other properties like LogP, JNK3, GSK3 β , drug-likeness is more sensitive to how substructures connect (Bickerton et al., 2012).
This behavior impedes the training of GNN and leads to the failure [Table] of optimization.
We report the learning curve in Figure [Table] where we plot the normalized loss on the validation set as a function of epoch numbers when learning GNN. For fairness of comparison
[Table] validation loss is normalized by dividing the validation loss at scratch (i [Table] e., 0-th epoch) so that all the validation losses are between 0 and 1.
For most of the target properties, the normalized loss value on the validation set would decrease signiﬁcantly [Table] and GNN can learn these properties well, except QED.
It veriﬁes the failure of training the GNN on optimizing QED.
A differentiable molecular graph at atom-wise resolution may potentially solve this problem.
[Table]
[Figure]
The key reason for the failure of GNN on optimizing QED is the limitation of the expressive power of scaffolding tree itself.
QED is a property that is highly dependent on how substructures connect (Bickerton et al., 2012), while our scaffolding tree currently ignores that information.
See
Section D.2 for more details and analysis.
D. R ESULTS A NALYSIS FOR D ISTRIBUTION L EARNING M ETHODS (L IG GPT)
As showed in Table 2 and 1, distribution learning methods (LigGPT) (Bagal et al., 2021) have much weaker optimization ability.
DST and all the other baselines fall into the category of goal-directed molecule generation, a.k.a., molecule optimization, which generates molecules with high scores for a given oracle.
In contrast, LigGPT belongs to distribution learning (a different category of method), which learns the distribution of the training set.
We refer to Brown et al.
for more description of two categories of methods.
Consequently, conditioned generation learns from the training set, is unable to generate molecules with property largely beyond the training set distribution and can not optimize a property directly, even though they claim to be able to solve the same problem.
Problem formulation of distribution learning methods leads to an inability to generate molecules with property largely beyond the training set distribution, which means they are much weaker in optimization.
D. A BLATION STUDY
As described in Section 3.3, during molecule sampling, we sample the new molecule from the differentiable scaffolding tree (Equation 11).
To verify the effectiveness of our strategy, we compare with a random-walk sampler , where the topological edition (i.e., expand, shrink or unchange) and substructure are both selected randomly.
We consider the following variants:
• “ DST + DPP”.
Both topology and substructure to ﬁll are sampled from the optimized differentiable scaffolding tree, as shown in Equation (11).
This is what we use in this paper.
• “random + DPP”.
Changing topology randomly, that is, at each leaf node, “expand”, “shrink” and “unchange” probabilities are ﬁxed to 0 .
Substructure selection is sampled from the substructures’ distribution in the optimized differentiable scaffolding tree.
Then it uses DPP (Section 3.4) to select diverse and desirable molecules for the next iteration.
• “ DST + top- K ”.
Same as “ DST + DPP”, it uses DST to sample new molecules.
The difference is when selecting molecules for the next iteration, it selects the top- K molecules with highest f score.
It is equivalent to λ −→ + ∞ in Equation (12).
We show the results in Figure 12.
We ﬁnd that both DST sampling and DPP-based diversiﬁcation play a critical role in performance.
We check the results for “ DST + top-K”, during some period, the objective does not grow, we ﬁnd it is trapped into local minimum, impeding its performance, especially convergence efﬁciency.
“random+DPP” exhibits the random-walk behaviour and it would not reach satisfactory performance.
When optimizing LogP, “ DST +DPP” and “ DST +top- K ” achieved similar performance, because logP score will prefer larger molecules with more carbon atoms, which is less sensitive to the diversity and relatively easier to optimize.
Overall, “ DST +DPP” is the best strategy compared with other variants.
[Figure]
D. C HEMICAL SPACE VISUALIZATION
We use principle component analysis (PCA) to visualize the distribution of explored chemical structures in optimizing “JNK3 & GSK3 β ”.
Speciﬁcally, we ﬁt a two-dimensional principal component analysis (PCA) (Bro & Smilde, 2014) for 2048-bit Morgan ﬁngerprint vectors of 20K ZINC molecules, which are randomly selected from ZINC database (Sterling & Irwin, 2015).
Then we use the PCA to project the ﬁngerprint of the generated molecule from various generations into a two-dimensional vector to observe their trajectories.
The results are reported in Figure 13, where the grey points represent the two-dimensional vector of ZINC molecules.
We ﬁnd that our method explores different parts of the 2D projection of the chemical space and covers a similar chemical space as the ZINC database after 20 to 30 iterations. [Figure]
D. A DDITIONAL I NTERPRETABILITY A NALYSIS
We provide an interpretability example in Figure 14.
At the leaf node (yellow), from the optimized differentiable scaffolding tree, we ﬁnd that the leaf weight and expand weight are both 0.99.
Thus we decide to EXPAND, the six-member ring is selected and ﬁlled in the expansion node (blue).
This is consistent with our intuition that logP score will prefer larger molecules with more carbon atoms.
E T HEORETICAL A NALYSIS
In this section, we present some theoretical results of the proposed method.
First, in Section E.1, we conduct a convergence analysis of DST under certain mild assumptions.
E. C ONVERGENCE A NALYSIS
In this section, we discuss the theoretical properties of DST in the context of de novo molecule design (learning from scratch).
We restrict our attention to a special variant of DST , named DST -greedy: at the t -th iteration, given one scaffolding tree Z ( t ) , DST -greedy pick up only one molecule with highest objective value from Z ( t ) ’s neighborhood set N ( Z ( t ) ) , i.e., Z ( t +1)
=
arg max Z ∈N ( Z ( t ) )
F ( Z ( t ) ) is exactly solved.
We theoretically guarantee the quality of the solution produced by DST -greedy.
First, we make some assumptions and explain why these assumptions hold.
Assumption 1 (Molecule Size Bound) .
The sizes (i.e., number of substructures) of all the scaffolding trees generated by DST are bound by N min and N max .
We focus on small molecule optimization; the target molecular properties would decrease signiﬁcantly when the molecule size is too large, e.g., QED (drug-likeness)
(Bickerton et al., 2012).
Thus it is reasonable to bound the size of the scaffolding tree.
In addition, we use submodularity and smoothness to characterize the geometry of the objective landscape.
Assumption 2 (Submodularity and Smoothness) .
Suppose X 1 , X 2 , X 3 are generated successively by DST -greedy via growing (i.e., EXPAND) a substructure on the corresponding scaffolding tree.
We assume that the objective gain (i.e., ∆ F ) brought by adding a single substructure would not increase as the molecule grows (EXPAND).
[Equation]
where X 2 = EXPAND ( X 1 , s 1 ) , X 3 = EXPAND ( X 2 , s 2 ) , s 1 , s 2 are substructures to add.
Submodularity plays the role of concavity/convexity in the discrete regime.
On the other hand, we specify the smoothness of the objective function F by assuming F
[Equation]
) holds for the X 1 , X 2 , X 3 described above, whose sizes are smaller than N min .
Then we theoretically guarantee the quality of the solution under these assumptions.
Theorem 1.
Suppose Assumption 1 and 2 hold; we have the following relative improvement bound with the optimum
[Equation]
where Z ∗ is the local optimum found by DST -greedy, X ∗ is the ideal optimal molecule, X 0 is an empty molecule, starting point of de novo molecule design.
In molecule generation setting, a molecule is a local optimum when its objective value is maximal within its neighbor molecule set, i.e.,
F ( Z ∗ ) ≥ F ( Z ) for ∀ Z ∈ N ( Z ∗ ) .
The proof is given in Section G.4.
Proof Sketch .
We ﬁrst show that DST -greedy can converge to local optimum within a ﬁnite step in
Lemma 1.
Then we decompose the successive generation path and leverage the geometric information of the objective landscape to analyze the quality of the local optimum.
Lemma 1 (Local optimum) .
DST -greedy would converge to local optimum within ﬁnite steps.
The proof is given in Section G.3.
F E XTENSION OF M OLECULE DIVERSIFICATION
In the current iteration, we have generated M molecules ( X 1 , · · · , X M ) and need to select C molecules for the next iteration.
We expect these molecules to have desirable chemical properties (high F score) and simultaneously maintain higher structural diversity.
To quantify diversity, we resort to the determinantal point process (DPP) (Kulesza & Taskar, 2012).
DPP models the repulsive correlation between data points (Kulesza & Taskar, 2012) and has been successfully applied to many applications such as text summarization (Cho et al., sampling (Zhang et al., 2017), and recommendation system (Chen et al., 2018).
Generally, we have
M data points, whose indexes are { 1 , 2 , · · · , M } , S ∈ R M × M + denotes the similarity kernel matrix between these data points, ( i, j ) -th element of S measures the Tanimoto similarity between i -th and j -th molecules.
We want to sample a subset (denoted R ) of M data, R is a subset of { 1 , 2 , · · · , M } with ﬁxed size C , it assigns the probability
P [Equation] , where R ⊆ { 1 , 2 , · · · , M } , |R| = C, (16)
where
S R ∈ R C
× C is the sub-matrix of S , det( S R ) is the determinant of the matrix S R .
For instance, if we want to sample a subset of size 2, i.e., R = { i, j } , then we have P ( R )
∝ det( S R ) =
S ii S jj − S ij S ji = 1 − S ij S ji , more similarity between i -th and j -th data points lower the probability of their co-occurrence.
DPP thus naturally diversiﬁes the selected subset.
DPP can be calculated efﬁciently using the following method.
Deﬁnition 11 (DPP-greedy (Chen et al., 2018)) .
For any symmetric positive semideﬁnite (PSD) matrix S ∈ R M × M + and ﬁxing the size of R to C , Problem (16) can be solved in a greedy manner by
DPP-greedy in polynomial time O
( C 2 M ) .
It is denoted R = DPP-greedy ( { X 1 , · · · , X M } , C ) .
We describe the DPP-greedy algorithm in Algorithm 2 for completeness.
Each iteration selects one data sample that maximizes the current objective, as described in Step 5 in Algorithm 2.
Algorithm 2 DPP-greedy (Chen et al., 2018)
∈ R M × M ,
number of selected data C ∈ N + ,
C < M .
: Output do = arg max k ∈W k } ) 8: end for As mentioned, our whole target is to select the molecules with desirable properties while maintaining the diversity between molecules.
The objective is formulated as
[Equation]
where the hyperparamter λ > 0 balances the two terms, the diagonal matrix V is
[Equation] , where v 1 = F ( X 1 ) , · · · , v M = F ( X M ) , (18) where v i is the F -score of the i -th molecule (Eq. 1), V R is a sub-matrix of V indexed by R .
For any square matrix M 1 , M 2 of the same shape, we have
[Equation]
we further transform L DPP ( R ) as below to construct symmetric matrix,
[Equation]
(19) where V 12 = diag (cid:0) [exp( λv 1 2 ) , · · · , exp( λv M 2 )] (cid:1) .
Then we present the following lemma for the usage of the DPP-greedy method.
Lemma 2.
Suppose S ∈ R M × M is the (Tanimoto) similarity kernal matrix of the M molecules, i.e., S ij =
[Equation]
2 , b i is the binary ﬁngerprint vector for the i -th molecule, V is diagonal matrix deﬁned in Eq.
(18) , then we have (1) V 12 S V 12 is positive semideﬁnite; (2) V 12 R S R V 12 R
=
( V 1 2 S V 1 2 ) R .
The proof is given in Section G.1.
Thus, Problem (19) can be transformed as
[Equation]
which means we can use DPP-greedy (Def. 11) to solve Problem (19) and obtain the optimal R .
Discussion .
In Eq.
(17), we have two terms to specify the constraints on the molecular property and structural diversity, respectively.
When we only consider the ﬁrst term ( λ (cid:80) t ∈R v t ), the selection strategy is to select C molecules with the highest F score for the next iteration, same as conventional evolutionary learning in Brown et al.
; Jensen (2019); Nigam et al.
On the other hand, if we only consider the second term in Eq.
(17), we show the effect of selection strategies under certain approximations.
Suppose we have C molecules X 1 , X 2 , · · · , X C with high diversity among them, then we leverage DST to optimize these C molecules respectively, and obtain C clusters of new molecules, i.e.,
[Equation]
X
[Equation]
X [Equation]
Then we present the following lemma to show that when only considering diversity, under certain assumptions, Problem (17) reduces to multiple chain
MCMC methods.
In Eq.
(17), λ is a key hyperparamter, a larger λ corresponds to more weights on objective function
F while smaller λ speciﬁes more diversity.
When λ goes to inﬁnity, i.e., only considering the ﬁrst term ( λ (cid:80) t ∈R v t ), it is equivalent to selecting C molecule candidates with the highest F score for the next iteration, same as conventional evolutionary learning in Jensen (2019); Nigam et al.
On the other hand, if we only consider the second term, we show the effect of selection strategies under certain approximations.
Suppose we have C molecules X 1 , X 2 , · · · , X C with high diversity among them, then we leverage DST to optimize these C molecules respectively, and obtain C clusters of new molecules, i.e.,
[Equation]
Then we present the following lemma to show that when only considering diversity, under certain assumptions, Problem (17) reduces to multiple independent Markov chains.
Lemma 3.
Assume (1) the inter-cluster similarity is upper-bounded, i.e., sim ( ˆ Z ip , ˆ Z jq ) ≤ (cid:15) 1 for any i (cid:54) =
j ; (2) the intra-cluster similarity is lower-bounded, i.e., sim ( ˆ Z ip , ˆ Z iq ) ≥ 1 − (cid:15) 2 for any i ∈ { 1 , 2 , · · · , M } and p (cid:54) =
q ; when both (cid:15) 1 , (cid:15) 2 approach to
+ , the optimal solution to Problem (17) is
[Equation]
where p c = arg max p F ( ˆ Z cp ) for c = 1 , · · · , C
The proof is given in Section G.2.
Remark .
When the inter-cluster similarity is low enough, and intra-cluster similarity is high enough, our molecule selection strategy reduces to multiple independent Markov chains.
However, these assumptions are usually too restrictive for small molecules.
G P ROOF OF T HEORETICAL R ESULTS
In this Section, we provide the proof of all the theoretical results in Section E and F.
G. P ROOF OF L EMMA
Proof.
(I) V 1 2 SV 1 2 is positive semideﬁnite.
First, let us prove similarity kernel matrix S ∈ RB M × M based on molecular Tanimoto similarity is positive semideﬁnite (PSD), we know that the ( i, j ) -th element of S measures the Tanimoto similarity between i -th and j -th molecules, i.e.,
[Equation]
where b i ∈ [0 , 1] P is the P -bit ﬁngerprint vector for the i -th molecule (in this paper, P = 2048 ).
S can be decomposed as
[Equation]
, where matrix B is the stack of all the normalized (divided by l 2 norm, (cid:107) · (cid:107) 2 ) ﬁngerprint vector, as
[Equation]
For ∀ x ∈ R M , we have
[Equation]
Thus, S is PSD.
Then, similarly, for ∀ x ∈ R M , we have
[Equation]
where V 12 is diagonal matrix, so
V 12 =
( V 12 ) (cid:62) .
Thus, V 12 SV 12 is symmetric and positive semideﬁnite. (II)
[Equation]
Without loss of generalization, we assume R = { t 1 , · · · , t C } , where t 1 < t 2 < · · · , t C .
V 12 is diagonal.
[Equation]
 , where v t i =
[Equation]
is the objective function of t i -th molecule X t i .
The i, j -th element of V 12 R S R V 12 R is
[Equation]
On the other hand, the i, j -th element of V 1 2 SV 1 2 is exp( λv i 2 )
S ij exp( λv j 2 ) .
Then the i, j -th element of
[Equation]
[Equation]
Combining Equation (21) and (22), we prove V 12 R S R V 12 R =
( V 12 SV 12 ) R .
G. P ROOF OF L EMMA
Proof.
We consider two cases in the solution R .
(A) one molecule for each input molecule
Z 1 , · · · , Z C .
(B) other cases.
Our solution belongs to Case (A).
(A) First, we prove for (A), our solution is optimal.
We consider the second term in Equation (17),
S R is diagonal dominant.
Also, a determinant function is a continuous function with regard to all the elements.
Thus, det ( S R ) =
(cid:81) Ci =1 ( S R ) ii goes to 1.
Intuitively, all the selected molecules are dissimilar to each other, and the diversity is maximized.
On the other hand, to maximizing the ﬁrst term in Equation (17), during each k ∈ { 1 , 2 , · · · , C } , we select molecule with highest F score from { ˆ Z k 1 , · · · , ˆ Z kl k } .
That is our solution.
(B) Then we prove all the possible combinations in (B) are worse than our solution.
In (B), based on pigeonhole principle, there are at least one input molecule Z k that corresponds to at least two selected molecules.
Without loss of generalization, we denoted them ˆ Z k 1 and ˆ Z k 2 .
Since S R is diagonal dominant, its determinant can be decomposed as
[Equation]
If there is at least one ˆ S k whose shape is greater than 1.
Based on the deﬁnition of determinant, for matrix A ∈ R M × M
[Equation]
where Perm ( M ) is the set of all permutations of the set { 1 , 2 , · · · , M } , sgn ( η ) denotes the signature of η , a value that is +1 whenever the reordering given by η can be achieved by successively interchanging two entries an even number of times, and -1 whenever it can be achieved by an odd number of such interchanges.
For exactly half of all η s, sgn ( η ) = 1 and the other half are equal to -1.
For the matrix A whose shape is greater than 1 and all the elements are equal to 1, the determinant is equal to (cid:80) [Equation]
The determinant function is a continuous function with regard to all the elements.
When (cid:15) 2 goes to 0 + , all the elements of ˆ S k approach to 1, the determinant goes to 0.
Thus, det( S R ) also goes to 0.
The objective in Equation (17) goes to negative inﬁnity.
Thus, it is worse than our solution.
G. P ROOF OF L EMMA
Proof.
For the de novo design, DST -greedy start from scratch (empty molecule).
First, we show in this setting, there is no “REPLACE” or “DELETE” by mathematical induction and contradiction.
Since we start from an empty molecule, at the 1-st step the action is “EXPAND”.
Then we show the ﬁrst t steps are “EXPAND”, the ( t +1) -th step is still “EXPAND”.
Now we have X ( t ) =
EXPAND ( X ( t − 1) , s t − 1 ) (where s t − 1 is a substructure).
Suppose the ( t + 1) -th
step DST ’s action is “REPLACE”, e.g.,
X ( t +1) =
REPLACE ( X ( t ) , s t ) (where s t is a substructure), based on deﬁnition of DST -greedy, we have F ( X ( t +1) ) ≥ F ( X ( t ) ) .
Since DST only REPLACE the leaf node, we ﬁnd that X ( t +1) and X ( t ) are both in neighbor molecule set of X ( t − 1) , i.e., N ( X ( t − 1) ) , which contradict with the fact that
X ( t ) =
arg max X ∈N ( X ( t − 1) )
F ( X ( t − 1) ) .
Similarly, we show that there would not exist “DELETE”.
Then based on Assummption 1, we ﬁnd that DST -greedy converges at most N max steps.
G. P ROOF OF T HEOREM
Proof.
Based on the Proof of Lemma 1, we ﬁnd that there is only “EXPAND” action, then we are able to decompose the generation path as follows.
Starting from scratch, i.e., X 0 , suppose the path to optimum X ∗ is
[Equation]
where each step one substructure is added.
The path produced by DST -greedy is
[Equation]
Based on the deﬁnition of in each step only one substructure is added.
Based on Assumption 1, we have N min ≤ k 1 , k 2 ≤ N max .
There might be some overlap within the ﬁrst several steps, without loss of generalization, we assume Z k
= X k and Z k +1 (cid:54) =
X k +1 , where k can be 0 , 1 , · · · , k 1 .
Based on Assumption 2, we have
[Equation]
Then, we have
[Equation]
Thus, we get
[Equation]
Since Z 0 = X 0 , according to the deﬁnition of greedy algorithm, we have F ( Z 1 ) ≥ F ( X 1 ) .
Based on Assumption 2, we have
[Equation]
Based on Assumption [Equation] we have F ( Z ∗ ) − F ( Z 0 ) ≥ F ( Z N min ) − F ( Z 0 ) .
Then we have
F ( Z ∗ ) − F ( Z 0 ) ≥ F
[Equation]
Z N min ) − F ( Z 0 ) =
[Equation]
Combining Equation (24) and (25), we have
[Equation]
We observe that objective F ’s improvement is relatively lower bounded.