# Improving Molecule Generation and Drug Discovery with a Knowledge-enhanced
Generative Model

Aditya Malusare  Vaneet Aggarwal

###### Abstract

Recent advancements in generative models have established state-of-the-art
benchmarks in the generation of molecules and novel drug candidates. Despite
these successes, a significant gap persists between generative models and the
utilization of extensive biomedical knowledge, often systematized within
knowledge graphs, whose potential to inform and enhance generative processes
has not been realized. In this paper, we present a novel approach that bridges
this divide by developing a framework for knowledge-enhanced generative models
called K-DReAM. We develop a scalable methodology to extend the functionality
of knowledge graphs while preserving semantic integrity, and incorporate this
contextual information into a generative framework to guide a diffusion-based
model. The integration of knowledge graph embeddings with our generative model
furnishes a robust mechanism for producing novel drug candidates possessing
specific characteristics while ensuring validity and synthesizability. K-DReAM
outperforms state-of-the-art generative models on both unconditional and
targeted generation tasks.

Drug discovery, Knowledge graphs, Generative models

Figure 1: K-DReAM a. A biomedical knowledge graph showing the links between
drugs, target proteins, diseases, genes, biological pathways, etc. The edges
represent different types of relationships like targeting, synergistic or
unwanted interactions, or side effects. b. Given a relation and a subject, we
can complete the triple and obtain Knowledge Graph Embeddings (KGEs) (Sec 4).
c. These embeddings are used to guide a generative process (Sec 5). d. The end
result of the generative process is a novel drug candidate.

##  1 Introduction

Drug discovery is an expensive endeavor, with costs often surging beyond the
billion-dollar mark, primarily due to the extensive stages involved in the
development pipeline, ranging from target identification to clinical
evaluations (Morgan et al., 2011). Nevertheless, the advent of computational
techniques and machine learning has ushered in a paradigm shift in
pharmaceutical research, streamlining drug development, mitigating expenses,
and enhancing the discovery of bioactive compounds (Rudrapal & Chetia, 2020).
Predominantly, the emergence of sophisticated machine learning applications
has substantially bolstered predictive capabilities, accelerating the pace of
in-silico drug design (Rifaioglu et al., 2018). However, this rapid
progression has not been devoid of challenges. Present-day generative models
in machine learning showcase impressive results but often miss out on tapping
into the extensive biological knowledge available in the domain (Tian et al.,
2023). These models frequently grapple with issues like overfitting, lack of
generalizability to new data (Ma et al., 2015), and difficulties in handling
novel drug or target interactions (Nguyen et al., 2022). Although there is a
vast amount of biomedical knowledge, there is a missing link between these
collective datasets and the efforts to build generative models for drug
discovery.

Addressing these limitations, we introduce an innovative methodology for
generating drugs that not only harnesses structured biological knowledge but
ensures semantic coherence when generating new insights from existing
knowledge graphs. Our scalable model employs knowledge graph embeddings to
direct a diffusion-based generative model, fine-tuned by an original
reinforcement learning reward strategy. This symbiosis between comprehensive
bio-knowledge and generative models paves a pathway for the interpretative,
scalable, and controllable generation of drug molecules.

In order to improve the process of computational drug discovery, we focus on
two key aspects: (i) Improved generative models: These models are capable of
learning complex distributions of molecular structures from vast datasets and
generating new compounds with desired properties. By automating the molecular
design process, generative models can quickly propose viable drug candidates
that match specific efficacy and safety profiles, potentially leading to
breakthroughs in the treatment of diseases, and (ii) Knowledge Graphs (KGs) of
biomedical data:  KGs in drug discovery constitute a structured representation
of vast amounts of heterogeneous data, encompassing the relations between
different entities such as genes, proteins, drugs, and diseases. By leveraging
knowledge graphs, we can create interpretable and generalizable models that
encapsulate both the biomedical knowledge encoded within the relationships and
the data-driven insights garnered from machine learning models. These graphs
also enable a deeper analysis of biological pathways, off-target effects, and
drug repurposing opportunities.

Molecular generation is a critical aspect of drug discovery, material science,
and chemical exploration. Generative models have demonstrated the ability to
generate novel structures across a variety of representations, from 1-D SMILES
strings to 2-D and 3-D molecular graphs. However, while these models have
shown great results in producing drug-like molecules, they are typically
evaluated based on their unconditioned generation capabilities. Generative
models have also achieved significant improvements in targeted generation with
modalities like text, images (Black et al., 2023), sound, and videos. The
development of a targeted generative model for graphs is of particular
interest in the problem of drug discovery.

Biomedical knowledge can be organized and represented by a powerful tool
called a Knowledge Graph (KG). A good way of guiding the generative process is
by using the KG as a structural and semantic scaffold, which can provide
context-relevant constraints and objectives for the generative models. This
integration of KGs with generative models can facilitate the targeted
generation of molecules by driving the model to consider biological relevance
and plausibility in the generation process. By doing so, the models can
leverage both the relational knowledge embedded within the KG, such as the
interactions between proteins, genes, drugs, and diseases, and the generative
capabilities to propose novel compounds that are more likely to exhibit
desired therapeutic effects and pharmacological profiles.

Incorporating KGs into generative modeling allows for a more informed
exploration of the chemical space, focusing on areas with higher potential for
successful drug development. The properties and relationships within KGs can
guide the generative model in synthesizing molecules that not only are
structurally novel but also align with known biological pathways and
mechanisms of action. This coupling can lead to the development of generative
models that produce candidates which are not only chemically valid but also
biologically relevant. Moreover, such directed generation can help in
hypothesis generation for drug repurposing and identifying previously
unrecognized drug-target interactions, greatly enriching the pipeline of drug
discovery and potentially reducing the timeline for the development of new
therapies.

We propose an end-to-end framework 1 for Knowledge-enhanced Drug discovery
with a Reinforcement learning-Augmented Model (K-DReAM), with the following
key contributions:

  1. 1.

A diffusion-based generative model for graphs steered with Discrete Diffusion
Policy Optimization (DDPO) for improved molecule generation. The generative
model, initially trained on extensive chemical structure datasets, is fine-
tuned with the DDPO approach that rewards the model for producing chemically
valid, pharmacologically active structures that have the ability to satisfy
user-defined properties. (Section 3)

  2. 2.

A Knowledge Graph Embedding (KGE) model that generates contextual embeddings
from large-scale biomedical databases that are used for targeted extrapolation
and synthesis of drug candidates. By reformulating models that generate KGEs
from energy-based models to ones that minimize a Maximum Likelihood Estimation
(MLE) objective, we demonstrate the ability to generate more meaningful
embeddings that also incorporate domain constraints, which ensure that
predictions follow user-specified logical formulae. (Section 4)

  3. 3.

The development of a training workflow that is used by the model to learn
conditional generation, with the help of a property inference network that
integrates the knowledge represented by the KGEs into the training of the
diffusion model to create a unified process for knowledge-enhanced novel drug
discovery (Section 5)

The novelty of our framework lies in: 1) the extension of the DDPO procedure
from the fixed topology of images to the variable message-passing framework of
graphs, 2) the creation of a KGE model with domain constraints specific to
drug-based knowledge, and 3) the development of a mapping from molecule-space
to KGE-space that is used for score-based conditional generation.

##  2 Related Work

Molecular Generation.  The use of SMILES strings (Weininger, 1988), a one-
dimensional notation that encodes molecular structures into strings of ASCII
characters, was popularized for generative modeling with the advent of
recurrent neural networks (RNNs) and variational autoencoders (VAEs). (Gómez-
Bombarelli et al., 2018) demonstrated that VAEs could be trained on SMILES
strings to interpolate in latent space and generate molecules with desired
chemical properties. Subsequent work by (Neil et al., 2018) employed RNNs for
a similar purpose and introduced reinforcement learning to bias the generation
process towards certain properties. While the SMILES-based approach remains
popular due to its simplicity and the wealth of available chemical data, it is
not without limitations, as SMILES strings do not explicitly capture the
geometrical and topological features of molecules, and small changes in SMILES
strings can lead to vastly different molecular structures (Lavecchia, 2019).

To address the limitations of SMILES, researchers have increasingly turned to
graph-based representations, where molecules are represented as graphs with
atoms as nodes and bonds as edges. This representation aligns with the
intrinsic structure of molecules. Generative models have been devised using
GANs (Martinkus et al., 2022), diffusion (Jo et al., 2022a; Vignac et al.,
2023), autoregressive methods (Kong et al., 2023), and normalizing flows (Shi
et al., 2019; Zang & Wang, 2020).

Knowledge-enhanced Models.  Knowledge graphs have gained traction in drug
discovery (James & Hennig, 2023) by structuring diverse biological data into
interconnected frameworks that facilitate the identification of new
therapeutic targets and drug repurposing opportunities (Santos et al., 2022;
Zheng et al., 2020; Chandak et al., 2023). These graphs represent entities
such as molecules, proteins, and diseases as nodes, with edges capturing their
intricate relationships. (Yu et al., 2021) showcase the use of entity
embeddings from knowledge graphs for predicting drug-disease associations,
while (Zitnik et al., 2018) employed graph neural networks on similar
structures to uncover drug-target interactions, surpassing traditional
predictive models. By incorporating relational information from knowledge
graphs, molecular generative models can impose biologically relevant
constraints during the generation process (Bilodeau et al., 2022). Knowledge
graphs also aid in systematically identifying adverse drug reactions (Nováček
& Mohamed, 2020) and repurposing existing drugs by analyzing the network’s
topology to reveal hidden biological pathways (Pan et al., 2022).

Diffusion Models.  Diffusion models have emerged as a powerful class of
generative models with applications in drug discovery, offering an alternative
to traditional variational autoencoders and generative adversarial networks.
They work by gradually introducing noise to a data distribution and learning
to reverse this process, thereby generating new data points. In the context of
drug discovery, diffusion models can generate novel molecular structures by
learning the distribution of drug-like molecules, with Ho et al. (2020)
demonstrating their potential through the generation of high-fidelity images.
Specifically, these models can potentially be adapted to generate 2D molecular
structures that could lead to novel compounds with desired properties.

##  3 Molecular Generative Model

Figure 2: Generative process.  The random initialization of the molecular
graph G converges to a valid non-benzeneoid aromatic compound called tropone
through the diffusion model.

Molecular structures can be represented using a planar graph
$\textbf{G}=(\textbf{X},\textbf{A})$ where $\textbf{X}\in\mathbb{R}^{N\times
M}$ is a feature matrix for $N$ nodes (heavy atoms) described by
$M$-dimensional vectors encoding atom information, and
$\textbf{A}\in\mathbb{R}^{N\times N}$ is the adjacency matrix indicating the
presence of single, double or triple bonds between the nodes.

Our generative model is built on the foundations of the GDSS (Jo et al.,
2022a) and MOOD (Lee et al., 2022) diffusion models. Graph Diffusion via the
System of SDEs (GDSS) defines the forward diffusion $q$ of a graph
$\textbf{G}_{t}=(\textbf{X}_{t},\textbf{A}_{t})$ with a Stochastic
Differential Equation (SDE):

$\mathrm{d}\textbf{G}_{t}=\textbf{f}_{t}(\textbf{G}_{t})+g_{t}\mathrm{d}\textbf{w}$
(1)

where w is the standard Wiener process and $\textbf{f}_{t}$ and $g_{t}$ are
the coefficients of linear drift and scalar diffusion, respectively. GDSS
performed well at generating molecular graphs, but its distributional learning
led to generated molecules closely resembling the training dataset. The
Molecular Out-of-distribution (MOOD) framework overcame the restricted
exploration space of the training process by modifying the above equation into
a conditional SDE, with the marginal distribution of the forward process
becoming $p_{\theta}(\textbf{G}_{t}|\textbf{y}_{o}=\lambda)$. Here,
$\textbf{y}_{o}$ represents the OOD condition and the hyperparameter
$\lambda\in[0,1)$ tunes the “OOD-ness” of the generative process. Controlling
$\lambda$ allows MOOD to explore areas outside the training space and generate
novel molecules.

Since we aim to use contextual information from KGs for our generative model,
we formulate an extension that samples from a more general marginal
distribution $p_{\theta}(\textbf{G}_{t}|\textbf{c})$ conditioned on contextual
data c. We then introduce a RL-based objective to maximize a reward $r$:

$\max\mathbb{E}_{\textbf{c}\sim p(\textbf{c}),\textbf{G}_{0}\sim
p(\textbf{G}_{0}|\textbf{c})}\left[r(\textbf{G}_{0},\textbf{c})\right]$ (2)

(Black et al., 2023) introduce a fine-tuning technique known as Denoising
Diffusion Policy Optimization (DDPO) for image-based models. This method
excels in handling images by incorporating human feedback to ensure metrics
like aesthetic qualities—a factor typically difficult to quantify in
computational models. Despite its success in the domain of images, extending
these techniques to the realm of molecule optimization presents significant
challenges. The intricate dependencies between nodes and edges in molecular
structures demand careful consideration, as these relationships are critical
to determining the validity and inherent properties of the molecules.

We introduce a reward function $r$ that simultaneously prioritizes multiple
molecular properties: (i) The drug-likeness $Q$ of a molecule, quantified by
its QED score (Bickerton et al., 2012), (ii) the synthesizability $S$
calculated using the SAScore (Ertl & Schuffenhauer, 2009) which is a rule-
based determination of synthetic accessibility, and (iii) a provision for
defining a property $C$ in the chemical space like novelty/molecular
similarity, docking scores, structure validity, etc.

We define our overall reward function as:

$r(\textbf{G},\textbf{c})=\begin{cases}\kappa_{1}Q(\textbf{G})+\kappa_{2}S(\textbf{G})+\kappa_{3}C(\textbf{G},\textbf{c});&\\\
-\kappa_{4},\hskip 30.0pt\text{if {G} is invalid}&\end{cases}$ (3)

We elaborate on the training process of the model for conditional generation
in Section 5.

##  4 Knowledge Graph Embeddings

A knowledge graph $\mathcal{G}$, involving a set of entities $\mathcal{E}$ and
relations $\mathcal{R}$, is a directed multigraph composed of triples
$(s,r,o)\in\mathcal{E}\times\mathcal{R}\times\mathcal{E}$ that represent the
relation $r$ between subject $s$ and object $o$ nodes. In biomedical knowledge
graphs, $\mathcal{E}$ contains entities like drugs, diseases, genes,
phenotypes, and biological pathways while $\mathcal{R}$ contains relations
like interactions, drug targets and side effects. Some examples of public
domain knowledge graphs in this field like Hetionet (Himmelstein et al.,
2017), CKG (Santos et al., 2022), BioKG (Walsh et al., 2020), PharmKG (Zheng
et al., 2020), and PrimeKG (Chandak et al., 2023), which range from $\sim$10K
to 10M nodes and up to $\sim$200M relations.

Knowledge Graph Embeddings (KGEs) are low dimensional representations of
relations and entities that are used for various tasks like link prediction,
graph completion and attribute inference. KGEs are constructed with a scoring
function
$\phi_{\mathcal{G}}:\mathcal{E}\times\mathcal{R}\times\mathcal{E}\to\mathbb{R}$
that measures the plausibility of any given triple.

KGEs help preserve the structure and information of neighbors, enabling the
efficient encoding of the topology and semantic relationships inherent in
knowledge graphs. Particularly in the biomedical domain, where the complexity
of interactions is high, the capacity of KGEs to facilitate the extraction of
latent relationships between disparate entities is crucial for advancements in
drug repurposing, disease gene prioritization, and patient outcome prediction.

KGE models involve various translational distance-based metrics like TransE
(Lin et al., 2015), tensor factorization methods such as RESCAL (Nickel et
al., 2011), and neural network-based models like ConvE (Dettmers et al.,
2018). These models differ in how they interpret the relationships and
interactions between entities, and each comes with relative strengths and
weaknesses with regard to specific types of relational data and structures.

Traditionally, score-based KGE models have been interpreted as energy-based
models, where the score is seen as a measure of the negative energy of an
$(s,r,o)$ triplet. Recasting these models into a probabilistic interpretation
would enable exact training by Maximum Likelihood Estimation (MLE), as well as
the ability to encode domain constraints into the learning process. In order
to transform these negative energies into probabilities over the
$\mathcal{E}\times\mathcal{R}\times\mathcal{E}$ space, we need to calculate
the partition function, which is infeasible for large-scale biomedical KGs.
(Loconte et al., 2023) demonstrate that KGE models can be represented by
structured computational graphs, called circuits, which are expressive
probabilistic functions over the triplet space and can be efficiently trained
with MLE. This enables us to use KGE models as efficient generative models of
new $(s,r,o)$ triples, consistent with the statistics of existing the KG while
guaranteeing the satisfaction of constraints that would be crucial to
pharmacological applications.

The MLE objective is evaluated as:

$\displaystyle\mathcal{L}_{MLE}$
$\displaystyle=\sum_{(s,r,o)\in\mathcal{G}}\log{p(s,r,o)}$ (4)
$\displaystyle=-|\mathcal{G}|\log{Z}+\sum_{(s,r,o)\in\mathcal{G}}\phi_{pc}(s,r,o)$

Circuit-based score functions bring down the complexity of evaluating $Z$ from
$\mathcal{O}(|\mathcal{E}|^{2}\cdot|\mathcal{R}|\cdot\text{cost}(\phi))$ to
$\mathcal{O}(|\mathcal{E}+\mathcal{R}|\cdot\text{cost}(\phi))$ (Loconte et
al., 2023).

The generation of new $(s,r,o)$ triplets is ensured to be semantically
coherent within the rules of the KG by introducing domain constraints (Ahmed
et al., 2022) that modify the score functions with indicator functions for
valid triples. Given a relation $r\in\mathcal{R}$ the subsets
$S_{r},O_{r}\subset\mathcal{E}$ represent the sets of entities between which
$r$ is a valid relation. This domain $K_{r}=S_{r}\wedge\\{r\\}\wedge O_{r}$
can be extended across all $r$ to form $K=\vee_{r}K_{r}$. We restrict the
score function to only valid triples $(s,r,o)\in K$ by introducing an
indicator function $c_{K}$ in the score calculation $\phi(s,r,o)=p(s,r,o)\cdot
c_{K}(s,r,o)$. For example, in our experiments, we are particularly interested
in drug-protein interaction. For this type of relation $r=\texttt{targets}$,
we restrict $s\in\texttt{drugs}$ and $o\in\texttt{proteins}$. Once again, the
circuit-based score functions help us train these models with a complexity of
$\mathcal{O}(|\mathcal{E}+\mathcal{R}|\cdot\text{cost}(\phi)\cdot\text{cost}(c_{K}))$

In order to generate KGEs, K-DReAM uses the RotatE (Sun et al., 2019)
algorithm with the above modifications to the training process in order to
restrict the model to generate coherent triples. The embeddings generated from
the knowledge graph are used to guide the generative process as described
below.

##  5 Implementation of K-DReAM

An essential aspect of drug design is the ability to specify the target
properties in our generative process. This section describes our
implementation of the guidance scheme for our diffusion model, which we refer
to as the conditional diffusion model. We first describe the creation of a
regressor that predicts properties based on graph structure, and then
formulate the conditional process that guides the diffusion model to push it
to generate molecules with the desired properties.

###  5.1 Property Inference Network $P_{\phi}(\textbf{G})$

To guide the conditional generation process, we create a neural network to
estimate knowledge-based embeddings c from a noised version of an input
molecular graph $\textbf{G}_{T}$. $P_{\phi}(\textbf{G}_{T})\approx\textbf{c}$
is used to implement a modified version of the classifier guidance algorithm
by (Sohl-Dickstein et al., 2015). While previous work by (Lee et al., 2022)
and (Vignac et al., 2023) uses a similar algorithm to guide conditional
generative processes, ours is a novel approach that utilizes a combination of
graph attention and convolutional layers to estimate knowledge-based
embeddings, effectively creating a map between chemical space and KGE space.

$P{{}_{\phi}}(\textbf{G})=P_{\phi}(\textbf{X},\textbf{A})$ is constructed by
first passing the feature X and adjacency matrices A through an aggregation
operation:

$H^{1}=\sigma(\textbf{A}\textbf{X}W^{0}_{\phi})$ (5)

We then use a stack of self-attention layers:

$h_{i}^{l+1}=\sum_{j}\alpha_{ij}~{}\cdot~{}\mathbf{W^{l}}h_{j}^{l}$ (6)

where

$\alpha_{ij}=\text{SoftMax}(\sigma(\mathbf{W_{a}}^{l}\cdot[\mathbf{W}^{l}h^{l}_{i}~{}\oplus~{}\mathbf{W}^{l}h^{l}_{j}]))$
(7)

Here, $\textbf{W}^{l},\textbf{W}_{a}^{l}$ are learnable parameters at the
$l$-th layer. The stack of attention layers produces a final output of
dimension $|\textbf{c}|$. (Training details in Appendix A)

###  5.2 Conditional Diffusion Training

The stochastic forward process described in Eq. (1) can be used for generation
by solving it’s reverse-time version:

$\mathrm{d}\textbf{G}_{t}=[\textbf{f}_{t}(\textbf{G}_{t})-g_{t}^{2}\nabla_{\textbf{G}_{t}}\log{p_{\theta}(\textbf{G}_{t}})]\mathrm{d}\overline{t}+g_{t}\mathrm{d}\overline{\textbf{w}}$
(8)

where $\overline{t}$ and $\overline{\textbf{w}}$ represent a reverse time-step
and stochastic process. A score network $s_{\theta}$ is used to approximate
$\nabla_{\textbf{G}_{t}}\log{p_{t}(\textbf{G}_{t}})$ and simulate the reverse
process in time, to generate $\textbf{G}_{t-1}$.

Conditioning this process can be achieved by adding the conditioning
information c at each diffusion step: $p_{\theta}(\textbf{G}_{t}|\textbf{c})$.
The score network is then used to approximate the modified gradient
$\nabla_{\textbf{G}_{t}}\log{p_{\theta}(\textbf{G}_{t}}|\textbf{c})$. The
conditional distribution is rearranged to give

$\nabla_{\textbf{G}_{t}}\log{p_{\theta}(\textbf{G}_{t}}|\textbf{c})\propto\nabla_{\textbf{G}_{t}}\log
p_{\theta}(\textbf{G}_{t})+\nabla_{\textbf{G}_{t}}\log
p_{\theta}(\textbf{c}|\textbf{G}_{t})$ (9)

The term $\nabla_{\textbf{G}_{t}}\log p_{\theta}(\textbf{c}|\textbf{G}_{t})$
in the above equation steers the model towards optimizing for the condition,
while the first term $\nabla_{\textbf{G}_{t}}\log p_{\theta}(\textbf{G}_{t})$
introduces variation into the trajectory and helps explore newer regions.

We model the conditional probability using a distribution of the form:

$p_{\theta}(\textbf{c}|\textbf{G})=\dfrac{1}{Z_{\theta}}\exp({-\alpha_{\theta}||\textbf{c}-P_{\phi}(\textbf{G})||^{2}})$
(10)

where $\alpha_{\theta}$ is a scaling factor and $Z_{\theta}$ is the partition
function. This procedure has a tractable complexity, supported by the analysis
in Section 4.

##  6 Experiments

We perform individual evaluations of each component, followed by an end-to-end
study of K-DReAM in synthesizing novel drug candidates. We compare
unconditional generation against variants of generative models that use GANs
(Martinkus et al., 2022), diffusion (Jo et al., 2022a; Vignac et al., 2023),
autoregressive methods (Kong et al., 2023), and normalizing flows (Shi et al.,
2019; Zang & Wang, 2020).

###  6.1 Evaluating Unconditional Molecular Generation

Experimental Setup.  The Quantum Machines 9 (QM9) (Ramakrishnan et al., 2014)
and ZINC (Irwin et al., 2012) datasets contain 134k and 250k chemical
structures and are designed to aid in the exploration of chemical space. We
follow previous work (Kong et al., 2023) and perform unconditional generation
of 10,000 molecular structures. We report the percentage of valid and unique
molecules, along with the novelty as defined by (Jin et al., 2020b) and the
Frechet ChemNet Distance (FCD) (Preuer et al., 2018), which evaluates the
similarity between the generated and training datasets using the activations
of the penultimate layer of the ChemNet. The reward function (Eq. (3)) is used
to fine-tune K-DReAM for maximizing the validity and novelty of the molecules.

Results.  Tables 1 and 2 demonstrate that K-DReAM achieves state-of-the-art
results on all metrics except the FCD on the QM9 dataset. We omit reporting
uniqueness on ZINC since all models were $\gtrapprox 99.99\%$ unique. Previous
results sourced from (Kong et al., 2023). K-DReAM generates valid molecules at
rates of 99.25% and 98.29% on QM9 and ZINC, respectively as a result of the
fine-tuning procedure.

Model |  Validity$\uparrow$ |  FCD$\downarrow$ |  Unique$\uparrow$ |  Novelty$\uparrow$
---|---|---|---|---
GraphAF | 74.43 | 5.27 | 88.64 | 86.59
MoFlow | 91.36 | 4.47 | 98.65 | 94.72
SPECTRE | 87.3 | 47.96 | 35.7 | 97.28
GDSS | 95.72 | 2.9 | 98.46 | 86.27
DiGress | 99.0 | 0.36 | 96.66 | 33.4
GraphARM | 90.25 | 1.22 | 95.62 | 70.39
K-DReAM (w/o RL) | 94.20 | 2.21 | 97.64 | 95.63
K-DReAM | 99.25 | 1.54 | 98.62 | 99.74
Table 1: Generation results on the QM9 dataset. (Best, Second) Model |  Validity$\uparrow$ |  FCD$\downarrow$ |  Unique$\uparrow$
---|---|---|---
GraphAF | 68.47 | 16.02 | 98.64
MoFlow | 63.11 | 20.93 | 99.99
SPECTRE | 90.2 | 18.44 | 67.05
GDSS | 97.01 | 14.66 | 99.64
DiGress | 91.02 | 23.06 | 81.23
GraphARM | 88.23 | 16.26 | 99.46
K-DReAM (w/o RL) | 96.17 | 17.77 | 97.64
K-DReAM | 98.29 | 14.56 | 99.95
Table 2: Generation results on ZINC. (Best, Second)

###  6.2 Drug-based Knowledge Graph Embeddings

Experimental Setup.  (Bonner et al., 2022) devise a procedure to assess how
well biomedical KGE models can predict missing links in knowledge graphs,
standard evaluation employs knowledge graph metrics like Mean Reciprocal Rank
and Hits@k. These metrics measure the model’s ability to rank true triples
higher than corrupted ones in link prediction tasks. Models are evaluated on
predicting both head and tail entities by corrupting each in turn. Robustness
is examined by testing across multiple random parameter initializations and
dataset splits. Performance is evaluated on real-world KGs - BioKG (Walsh et
al., 2020) and Hetionet (Himmelstein et al., 2017). The following metrics are
used for comparison: (i) Mean Reciprocal Rank (MRR) (ii) Hits@1 and Hits@10
(iii) Adjusted Mean Rank (Berrendorf et al., 2020). Hits@k is computed using
the following procedure: For each corrupted triple (with head or tail entity
removed) in the test set, the model ranks all possible candidate entities to
complete the triple. A “hit” is when the true missing entity is ranked within
the top $k$ predictions. Hits@k calculates the proportion of test triples
where the true entity was ranked in the top $k$.

Results.  Table 3 presents the mean performance over 10 random seeds of all
models and both datasets as measured by the above metrics. Only the random
seed used to initialize the model parameters is changed. The mean score of
K-DReAM outperforms all metrics except the AMR score on BioKG, showing the
improved quality of embeddings due to the introduction of domain constraints.
We compare against ComplEx (Trouillon et al., 2017), DistMult (Zhang et al.,
2018), RotatE (Sun et al., 2019), and TransE (Lin et al., 2015).

Dataset | Approach | Metric
---|---|---
|  | AMR $\downarrow$ | MRR $\uparrow$ | Hits@1 $\uparrow$ | Hits@10 $\uparrow$
Hetionet | ComplEx | 0.167$\pm$0.009 | 0.026$\pm$0.009 | 0.008$\pm$0.003 | 0.064$\pm$0.024
DistMult | 0.201$\pm$0.303 | 0.036$\pm$0.019 | 0.012$\pm$0.007 | 0.079$\pm$0.045
RotatE | 0.035$\pm$0.000 |  0.127$\pm$0.000 | 0.063$\pm$0.000 | 0.262$\pm$0.001
TransE | 0.053$\pm$0.000 | 0.079$\pm$0.001 | 0.034$\pm$0.001 | 0.117$\pm$0.002
TransH | 0.126$\pm$0.000 | 0.033$\pm$0.001 | 0.007$\pm$0.000 | 0.088$\pm$0.002
K-DReAM |  0.030$\pm$0.002 |  0.125$\pm$0.002 |  0.071$\pm$0.000 |  0.33$\pm$0.010
BioKG | ComplEx | 0.213$\pm$0.011 | 0.008$\pm$0.001 | 0.003$\pm$0.000 | 0.008$\pm$0.003
DistMult | 0.560$\pm$0.339 | 0.015$\pm$0.003 | 0.007$\pm$0.001 | 0.027$\pm$0.006
RotatE |  0.022$\pm$0.000 | 0.123$\pm$0.000 | 0.059$\pm$0.000 | 0.240$\pm$0.001
TransE | 0.021$\pm$0.000 | 0.062$\pm$0.000 | 0.019$\pm$0.000 | 0.134$\pm$0.001
TransH | 0.078$\pm$0.001 | 0.022$\pm$0.000 | 0.008$\pm$0.000 | 0.042$\pm$0.001
K-DReAM | 0.025$\pm$0.010 |  0.136$\pm$0.010 |  0.065$\pm$0.000 |  0.28$\pm$0.001
Table 3: Evaluation of KGE Model.  Previous results sourced from (Bonner et
al., 2022).

Method | Target protein
---|---
parp1 | fa7 | 5ht1b | braf | jak2
REINVENT (Olivecrona et al., 2017) |  0-8.702 ($\pm$ 0.523) | -7.205 ($\pm$ 0.264) |  0-8.770 ($\pm$ 0.316) |  0-8.392 ($\pm$ 0.400) |  0-8.165 ($\pm$ 0.277)
MORLD (Jeon & Kim, 2020) |  0-7.532 ($\pm$ 0.260) | -6.263 ($\pm$ 0.165) |  0-7.869 ($\pm$ 0.650) |  0-8.040 ($\pm$ 0.337) |  0-7.816 ($\pm$ 0.133)
HierVAE (Jin et al., 2020a) |  0-9.487 ($\pm$ 0.278) | -6.812 ($\pm$ 0.274) |  0-8.081 ($\pm$ 0.252) |  0-8.978 ($\pm$ 0.525) |  0-8.285 ($\pm$ 0.370)
FREED (Yang et al., 2021) | -10.427 ($\pm$ 0.177) |  -8.297 ($\pm$ 0.094) | -10.425 ($\pm$ 0.331) | -10.325 ($\pm$ 0.164) |  0-9.624 ($\pm$ 0.102)
FREED-QS | -10.579 ($\pm$ 0.104) |  -8.378 ($\pm$ 0.044) | -10.714 ($\pm$ 0.183) | -10.561 ($\pm$ 0.080) |  0-9.735 ($\pm$ 0.022)
GDSS (Jo et al., 2022b) |  0-9.967 ($\pm$ 0.028) | -7.775 ($\pm$ 0.039) |  0-9.459 ($\pm$ 0.101) |  0-9.224 ($\pm$ 0.068) |  0-8.926 ($\pm$ 0.089)
MOOD |  -10.865 ($\pm$ 0.113) | -8.160 ($\pm$ 0.071) |  -11.145 ($\pm$ 0.042) |  -11.063 ($\pm$ 0.034) |  -10.147 ($\pm$ 0.060)
K-DReAM (w/o KGE conditioning) | -9.416 ($\pm$ 0.463) | -7.920 ($\pm$ 0.321) | -9.679 ($\pm$ 0.345) | -9.465 ($\pm$ 0.098) | -9.782 ($\pm$ 0.303)
K-DReAM |  -11.475 ($\pm$ 0.410) |  -8.270 ($\pm$ 0.211) |  -11.419 ($\pm$ 0.072) |  -11.364 ($\pm$ 0.035) |  -10.636 ($\pm$ 0.211)

Table 4:  Top 5% docking score (kcal/mol).  The results are the means and the
standard deviations of 5 runs. Previous results sourced from (Lee et al.,
2022).

###  6.3 Knowledge-enhanced Drug Discovery: Novelty and Protein Targeting

Experimental Setup.  K-DReAM generates 3,000 drug candidates with the reward
function (Eq. (3)) calibrated to optimize synthesizability, drug-likeness and
binding affinity to a target protein. Following (Lee et al., 2022), we choose
the targets: parp1, fa7, 5ht1b, braf, jak2. We set the chemical property $C$
to emphasize novelty, introducing a penalty for a similarity score $>0.4$ with
the training dataset. The knowledge embeddings are generated by asking the
model to complete the triple (_____, targets, protein) for all five targets,
searching over the molecule space for the highest scoring embeddings. These
embeddings are used for conditional generation as desctribed in Section 5.

Results.  Table 4 reports the average Docking Score (DS) of the top 5% of
generated molecules. K-DReAM is able to generate novel molecules with the best
docking scores among the models we compare against, with previous results
sourced from (Lee et al., 2022). The procedure for calculating these scores is
described in Appendix A.

###  6.4 Ablation Study: Generation without Reinforcement Learning

We perform an ablation study to examine the effect of fine-tuning K-DReAM with
the reward function (Eq. (3)) by comparing the performance before the fine-
tuning procedure in Tables 1 and 2. We find that even without DDPO, K-DReAM
has a similar performance as other diffusion- and SDE-based systems. DDPO
improves the validity and novelty to $\geq 99\%$, an improvement of 4-5%
points over the non fine-tuned performance.

###  6.5 Ablation Study: Generation without Knowledge Enhancement

We validate the impact of integrating KGEs into the model by running the
target protein experiments with a version of K-DReAM that is fine-tuned using
DDPO but not trained for conditional generation. Table 4 shows that K-DReAM
without KGE information obtains molecules that do not score as well in the
docking metric. In most cases, the docking score is similar to that of GDSS,
which is in line with our expectations, since our base model is derived from a
similar framework. The DDPO fine-tuning is identical in both cases. We thus
demonstrate that the incorporation of knowledge through embeddings leads to a
measurable improvement in targeted drug parameters. We can also consider a
comparison between K-DReAM and RotatE in Table 3 as an ablation study, since
our model is an improved version of RotatE with the modifications in Section
4.

##  7 Conclusion

We enhance molecular generation and drug discovery models by incorporating
information from biomedical knowledge graphs into the generative process. This
framework, called K-DReAM, demonstrates: 1) Performance on unconditional
molecular generation on par with state-of-the-art generative models, 2)
Improvements in Knowledge Graph Embedding (KGE) metrics by introducing domain
constraints and reformulating the scoring process, and 3) Enhancement with
KGEs improves targeted generation, evaluated with docking scores for five
proteins.

Through multiple experiments and ablation studies, we demonstrate that the
extraction of embeddings from knowledge graphs and their use in guiding
generative models leads to a measurable improvement in the quality of
generated drug candidates.

##  8 Acknowledgements

This work uses the Anvil supercomputer (Song et al., 2022) at Purdue
University through allocation CIS230228 from the Advanced Cyberinfrastructure
Coordination Ecosystem: Services & Support (ACCESS) program, which is
supported by National Science Foundation grants #2138259, #2138286, #2138307,
#2137603, and #2138296. The authors gratefully acknowledge the Walther Cancer
Foundation and support from the Purdue University Institute for Cancer
Research, P30CA023168

## References

  * Ahmed et al. (2022) Ahmed, K., Teso, S., Chang, K.-W., Van den Broeck, G., and Vergari, A.  Semantic probabilistic layers for neuro-symbolic learning.  _Advances in Neural Information Processing Systems_ , 35:29944–29959, 2022.
  * Ali et al. (2021) Ali, M., Berrendorf, M., Hoyt, C. T., Vermue, L., Sharifzadeh, S., Tresp, V., and Lehmann, J.  PyKEEN 1.0: A Python Library for Training and Evaluating Knowledge Graph Embeddings.  _Journal of Machine Learning Research_ , 22(82):1–6, 2021.  URL http://jmlr.org/papers/v22/20-825.html.
  * Berrendorf et al. (2020) Berrendorf, M., Faerman, E., Vermue, L., and Tresp, V.  On the ambiguity of rank-based evaluation of entity alignment or link prediction methods.  _arXiv preprint arXiv:2002.06914_ , 2020.
  * Bickerton et al. (2012) Bickerton, G. R., Paolini, G. V., Besnard, J., Muresan, S., and Hopkins, A. L.  Quantifying the chemical beauty of drugs.  _Nature chemistry_ , 4(2):90–98, 2012.
  * Bilodeau et al. (2022) Bilodeau, C., Jin, W., Jaakkola, T., Barzilay, R., and Jensen, K. F.  Generative models for molecular discovery: Recent advances and challenges.  _Wiley Interdisciplinary Reviews: Computational Molecular Science_ , 12(5):e1608, 2022.
  * Black et al. (2023) Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S.  Training diffusion models with reinforcement learning.  _arXiv preprint arXiv:2305.13301_ , 2023.
  * Bonner et al. (2022) Bonner, S., Barrett, I. P., Ye, C., Swiers, R., Engkvist, O., Hoyt, C. T., and Hamilton, W. L.  Understanding the performance of knowledge graph embeddings in drug discovery.  _Artificial Intelligence in the Life Sciences_ , 2:100036, December 2022.  ISSN 2667-3185.  doi: 10.1016/j.ailsci.2022.100036.  URL http://dx.doi.org/10.1016/j.ailsci.2022.100036.
  * Chandak et al. (2023) Chandak, P., Huang, K., and Zitnik, M.  Building a knowledge graph to enable precision medicine.  _Scientific Data_ , 10(1), February 2023.  ISSN 2052-4463.  doi: 10.1038/s41597-023-01960-3.  URL http://dx.doi.org/10.1038/s41597-023-01960-3.
  * Dettmers et al. (2018) Dettmers, T., Minervini, P., Stenetorp, P., and Riedel, S.  Convolutional 2d knowledge graph embeddings.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 32, 2018.
  * Ertl & Schuffenhauer (2009) Ertl, P. and Schuffenhauer, A.  Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions.  _Journal of cheminformatics_ , 1:1–11, 2009.
  * Gómez-Bombarelli et al. (2018) Gómez-Bombarelli, R., Wei, J. N., Duvenaud, D., Hernández-Lobato, J. M., Sánchez-Lengeling, B., Sheberla, D., Aguilera-Iparraguirre, J., Hirzel, T. D., Adams, R. P., and Aspuru-Guzik, A.  Automatic chemical design using a data-driven continuous representation of molecules.  _ACS central science_ , 4(2):268–276, 2018.
  * Himmelstein et al. (2017) Himmelstein, D. S., Lizee, A., Hessler, C., Brueggeman, L., Chen, S. L., Hadley, D., Green, A., Khankhanian, P., and Baranzini, S. E.  Systematic integration of biomedical knowledge prioritizes drugs for repurposing.  _Elife_ , 6:e26726, 2017.
  * Irwin et al. (2012) Irwin, J. J., Sterling, T., Mysinger, M. M., Bolstad, E. S., and Coleman, R. G.  Zinc: a free tool to discover chemistry for biology.  _Journal of chemical information and modeling_ , 52(7):1757–1768, 2012.
  * James & Hennig (2023) James, T. and Hennig, H.  _Knowledge Graphs and Their Applications in Drug Discovery_ , pp. 203–221.  Springer US, September 2023.  ISBN 9781071634493.  doi: 10.1007/978-1-0716-3449-3˙9.  URL http://dx.doi.org/10.1007/978-1-0716-3449-3_9.
  * Jeon & Kim (2020) Jeon, W. and Kim, D.  Autonomous molecule generation using reinforcement learning and docking to develop potential novel inhibitors.  _Scientific Reports_ , 10, 12 2020.
  * Jin et al. (2020a) Jin, W., Barzilay, R., and Jaakkola, T.  Hierarchical generation of molecular graphs using structural motifs.  In _International Conference on Machine Learning_ , pp. 4839–4848. PMLR, 2020a.
  * Jin et al. (2020b) Jin, W., Barzilay, R., and Jaakkola, T.  Multi-objective molecule generation using interpretable substructures.  In _International conference on machine learning_ , pp. 4849–4859. PMLR, 2020b.
  * Jo et al. (2022a) Jo, J., Lee, S., and Hwang, S. J.  Score-based generative modeling of graphs via the system of stochastic differential equations.  In _International Conference on Machine Learning_ , pp. 10362–10383. PMLR, 2022a.
  * Jo et al. (2022b) Jo, J., Lee, S., and Hwang, S. J.  Score-based generative modeling of graphs via the system of stochastic differential equations.  In _Proceedings of the 39th International Conference on Machine Learning_ , 2022b.
  * Kong et al. (2023) Kong, L., Cui, J., Sun, H., Zhuang, Y., Prakash, B. A., and Zhang, C.  Autoregressive diffusion model for graph generation, 2023.  URL https://openreview.net/forum?id=98J48HZXxd5.
  * Lavecchia (2019) Lavecchia, A.  Deep learning in drug discovery: opportunities, challenges and future prospects.  _Drug discovery today_ , 24(10):2017–2032, 2019.
  * Lee et al. (2022) Lee, S., Jo, J., and Hwang, S. J.  Exploring chemical space with score-based out-of-distribution generation.  _ArXiv_ , abs/2206.07632, 2022.  URL https://api.semanticscholar.org/CorpusID:249674702.
  * Lin et al. (2015) Lin, Y., Liu, Z., Sun, M., Liu, Y., and Zhu, X.  Learning entity and relation embeddings for knowledge graph completion.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 29, 2015.
  * Loconte et al. (2023) Loconte, L., Mauro, N. D., Peharz, R., and Vergari, A.  How to turn your knowledge graph embeddings into generative models.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=RSGNGiB1q4.
  * Ma et al. (2015) Ma, J., Sheridan, R., Liaw, A., Dahl, G., and Svetnik, V.  Deep neural nets as a method for quantitative structure–activity relationships.  _Journal of Chemical Information and Modeling_ , 55:263–274, 2015.  doi: 10.1021/ci500747n.
  * Martinkus et al. (2022) Martinkus, K., Loukas, A., Perraudin, N., and Wattenhofer, R.  Spectre: Spectral conditioning helps to overcome the expressivity limits of one-shot graph generators.  In _International Conference on Machine Learning_ , pp. 15159–15179. PMLR, 2022.
  * Morgan et al. (2011) Morgan, S., Grootendorst, P., Lexchin, J., Cunningham, C., and Greyson, D.  The cost of drug development: a systematic review.  _Health Policy_ , 100:4–17, 2011.  doi: 10.1016/j.healthpol.2010.12.002.
  * Neil et al. (2018) Neil, D., Segler, M., Guasch, L., Ahmed, M., Plumbley, D., Sellwood, M., and Brown, N.  Exploring deep recurrent models with reinforcement learning for molecule design.  2018\.
  * Nguyen et al. (2022) Nguyen, T., Nguyen, T., and Tran, T.  Mitigating cold-start problems in drug-target affinity prediction with interaction knowledge transferring.  _Briefings in Bioinformatics_ , 23, 2022.  doi: 10.1093/bib/bbac269.
  * Nickel et al. (2011) Nickel, M., Tresp, V., Kriegel, H.-P., et al.  A three-way model for collective learning on multi-relational data.  In _Icml_ , volume 11, pp. 3104482–3104584, 2011.
  * Nováček & Mohamed (2020) Nováček, V. and Mohamed, S. K.  Predicting polypharmacy side-effects using knowledge graph embeddings.  _AMIA Summits on Translational Science Proceedings_ , 2020:449, 2020.
  * Olivecrona et al. (2017) Olivecrona, M., Blaschke, T., Engkvist, O., and Chen, H.  Molecular de-novo design through deep reinforcement learning.  _Journal of cheminformatics_ , 9(1):1–14, 2017.
  * Pan et al. (2022) Pan, X., Lin, X., Cao, D., Zeng, X., Yu, P. S., He, L., Nussinov, R., and Cheng, F.  Deep learning for drug repurposing: Methods, databases, and applications.  _Wiley interdisciplinary reviews: Computational molecular science_ , 12(4):e1597, 2022.
  * Preuer et al. (2018) Preuer, K., Renz, P., Unterthiner, T., Hochreiter, S., and Klambauer, G.  Fréchet chemnet distance: a metric for generative models for molecules in drug discovery.  _Journal of chemical information and modeling_ , 58(9):1736–1741, 2018.
  * Ramakrishnan et al. (2014) Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A.  Quantum chemistry structures and properties of 134 kilo molecules.  _Scientific data_ , 1(1):1–7, 2014.
  * Rifaioglu et al. (2018) Rifaioglu, A., Ataş, H., Martin, M., Cetin-Atalay, R., Atalay, V., and Doğan, T.  Recent applications of deep learning and machine intelligence on in silico drug discovery: methods, tools and databases.  _Briefings in Bioinformatics_ , 20:1878–1912, 2018.  doi: 10.1093/bib/bby061.
  * Rudrapal & Chetia (2020) Rudrapal, M. and Chetia, D.  Virtual screening, molecular docking and qsar studies in drug discovery and development programme.  _Journal of Drug Delivery and Therapeutics_ , 10:225–233, 2020.  doi: 10.22270/jddt.v10i4.4218.
  * Santos et al. (2022) Santos, A., Colaço, A. R., Nielsen, A. B., Niu, L., Strauss, M., Geyer, P. E., Coscia, F., Albrechtsen, N. J. W., Mundt, F., Jensen, L. J., and Mann, M.  A knowledge graph to interpret clinical proteomics data.  _Nature Biotechnology_ , 40(5):692–702, January 2022.  ISSN 1546-1696.  doi: 10.1038/s41587-021-01145-6.  URL http://dx.doi.org/10.1038/s41587-021-01145-6.
  * Shi et al. (2019) Shi, C., Xu, M., Zhu, Z., Zhang, W., Zhang, M., and Tang, J.  Graphaf: a flow-based autoregressive model for molecular graph generation.  In _International Conference on Learning Representations_ , 2019.
  * Sohl-Dickstein et al. (2015) Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International conference on machine learning_ , pp. 2256–2265. PMLR, 2015.
  * Song et al. (2022) Song, X. C., Smith, P., Kalyanam, R., Zhu, X., Adams, E., Colby, K., Finnegan, P., Gough, E., Hillery, E., Irvine, R., Maji, A., and St. John, J.  Anvil - system architecture and experiences from deployment and early user operations.  In _Practice and Experience in Advanced Research Computing_ , PEARC ’22, New York, NY, USA, 2022. Association for Computing Machinery.  ISBN 9781450391610.  doi: 10.1145/3491418.3530766.  URL https://doi.org/10.1145/3491418.3530766.
  * Sun et al. (2019) Sun, Z., Deng, Z.-H., Nie, J.-Y., and Tang, J.  Rotate: Knowledge graph embedding by relational rotation in complex space.  _arXiv preprint arXiv:1902.10197_ , 2019.
  * Tian et al. (2023) Tian, T., Li, S., Fang, M., Zhao, D., and Zeng, J.  Molshap: interpreting quantitative structure–activity relationships using shapley values of r-groups.  _Journal of Chemical Information and Modeling_ , 2023.  doi: 10.1021/acs.jcim.3c00465.
  * Trouillon et al. (2017) Trouillon, T., Dance, C. R., Gaussier, É., Welbl, J., Riedel, S., and Bouchard, G.  Knowledge graph completion via complex tensor factorization.  _Journal of Machine Learning Research_ , 18(130):1–38, 2017.
  * Vignac et al. (2023) Vignac, C., Krawczuk, I., Siraudin, A., Wang, B., Cevher, V., and Frossard, P.  Digress: Discrete denoising diffusion for graph generation.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=UaAD-Nu86WX.
  * Walsh et al. (2020) Walsh, B., Mohamed, S. K., and Nováček, V.  Biokg: A knowledge graph for relational learning on biological data.  In _Proceedings of the 29th ACM International Conference on Information & Knowledge Management_, pp. 3173–3180, 2020.
  * Weininger (1988) Weininger, D.  Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules.  _Journal of chemical information and computer sciences_ , 28(1):31–36, 1988.
  * Yang et al. (2021) Yang, S., Hwang, D., Lee, S., Ryu, S., and Hwang, S. J.  Hit and lead discovery with explorative RL and fragment-based molecule generation.  _arXiv:2110.01219_ , 2021.
  * Yu et al. (2021) Yu, Z., Huang, F., Zhao, X., Xiao, W., and Zhang, W.  Predicting drug–disease associations through layer attention graph convolutional network.  _Briefings in bioinformatics_ , 22(4):bbaa243, 2021.
  * Zang & Wang (2020) Zang, C. and Wang, F.  Moflow: an invertible flow model for generating molecular graphs.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pp. 617–626, 2020.
  * Zhang et al. (2018) Zhang, Z., Zhuang, F., Qu, M., Lin, F., and He, Q.  Knowledge graph embedding with hierarchical relation structure.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pp. 3198–3207, 2018.
  * Zheng et al. (2020) Zheng, S., Rao, J., Song, Y., Zhang, J., Xiao, X., Fang, E. F., Yang, Y., and Niu, Z.  Pharmkg: a dedicated knowledge graph benchmark for bomedical data mining.  _Briefings in Bioinformatics_ , 22(4), December 2020.  ISSN 1477-4054.  doi: 10.1093/bib/bbaa344.  URL http://dx.doi.org/10.1093/bib/bbaa344.
  * Zitnik et al. (2018) Zitnik, M., Agrawal, M., and Leskovec, J.  Modeling polypharmacy side effects with graph convolutional networks.  _Bioinformatics_ , 34(13):i457–i466, 2018.

##  Appendix A Implementation Details

Our models are written in PyTorch 2.0, with the training workflow implemented
using the Lightning framework. Hyperparameters are managed using Hydra. All
experiments are performed on machines with an Intel (R) Core i9-9980XE CPU and
four NVIDIA (R) RTX 2080 Ti (12 GB) GPUs.

###  A.1 Knowledge graphs

PyKEEN (Ali et al., 2021) (built using Pytorch) is used to implement the
Knowledge Graph Embedding Model (KGEM). The knowledge graph dataloaders for
BioKG and Hetionet in PyKEEN are used with a custom defined model based on the
formulation of (Loconte et al., 2023).

###  A.2 Docking scores

We use AutoDock Vina to calculate docking scores, which reflect the predicted
binding affinity between a receptor (protein) and a ligand (such as a drug
molecule). We first prepare the receptor (protein) and ligand (drug) structure
file in PDBQT format, which requires the removal of water molecules, addition
of hydrogen atoms, assignment of charges, and determination of torsion angles.
Once the PDBQT files are ready, the Python script defines the search space and
creates a configuration file for Vina. The script then invokes the Vina
command-line interface with the appropriate parameters, captures output, and
saves the docking results, typically a PDBQT file containing the ranked
binding poses according to their scores.