BERTKG-DDI: Towards Incorporating Entity-specific Knowledge Graph Information in Predicting Drug-Drug Interactions

Our model for extracting DDIs from texts is based on the pre-trained BERT-based relation classification model by (Wu and He 2019). Given a sentence $s$ with drugs $d_{1}$ and $d_{2}$, let the final hidden state output from BERT module is $H$. Let the vectors $H_{i}$ to $H_{j}$ are the final hidden state vectors from BERT for entity $d_{1}$, and $H_{k}$ to $H_{m}$ are the final hidden state vectors from BERT for entity $d_{2}$. An average operation is applied to obtain the vector representation for each of the drug entities. An activation operation tanh is applied followed by a fully connected layer to each of the two vectors, and the output for $d_{1}$ and $d_{2}$ are $H_{1}^{{}^{\prime}}$ and $H_{2}^{{}^{\prime}}$ respectively.

Matrices $W_{0}$, $W_{1}$, $W_{2}$ have the same dimensions, i.e. $W_{0}$ $\in$ $R^{d*d}$ ,$W_{1}$ $\in$ $R^{d*d}$, $W_{2}$ $\in$ $R^{d*d}$, where $d$ is the hidden state size from BERT. We concatenate $H_{0}^{{}^{\prime}}$, $H_{1}^{{}^{\prime}}$ and $H_{2}^{{}^{\prime}}$ and then add a fully connected layer and a softmax layer, which is expressed as :

To infuse external information of the entities in relation classification task, we obtain the representation of two Drug entities mentioned in each input instance of the relation classification task. We use an in-house heterogeneous biomedical Knowledge Graph (Bio-KG) consisting of the interactions of target-target, drug-drug, drug-disease, drug-target, disease-disease, disease-target interactions from a large number of ontologies such as : DrugBank¹¹1https://go.drugbank.com/, BioSNAP²²2http://snap.stanford.edu/biodata/, UniProt³³3https://www.uniprot.org/ (The UniProt Consortium 2016). The overall statistics of Bio-KG has been enumerated in table 1. The real-world information/facts observed in the Bio-KG are stored as a collection of triples in the form ($h$, $r$, t). Each triple is composed of a head entity $h$ $\in$ $E$, a tail entity $t$ $\in$ $E$, and a relation $r$ $\in$ $R$ between them, e.g., (paracetamol, treats, fever). The fact that paracetamol is effective in curing fever is being stored in Bio-KG. In this case, $E$ denotes set of entities, and $R$ denotes the set of relations. There are three different types of $E$ in Bio-KG such as drugs, diseases, targets and five different types of $R$ such as target-target, drug-disease, drug-target, disease-disease, disease-target interactions.

The aim of a Knowledge Graph embedding is to embed the entities and relations into a low-dimensional continuous vector space, so as to simplify the computations on the KG. They mostly use facts in the KG to perform the embedding task, enforcing embedding to be compatible with the facts. They provide a generalizable context about the overall Knowledge Graph (KG) that can be used to infer the relations. In this work, we employ some off-the-shelf KG embeddings to encode the representation of each of the drugs (in terms of their relationship with other entities). The knowledge graph embeddings are computed so that they satisfy certain properties; i.e., they follow a given KGE model. These KGE models define different score functions that measure the distance of two entities relative to its relation type in the low-dimensional embedding space. These score functions are used to train the KGE models so that the entities connected by relations are close to each other while the entities that are not connected are far away. Some of the KGEs used in our experiments as explained below:

• •

  TransE (Bordes et al. 2013): Given a fact ($h$, $r$, $t$), the relation in TransE is interpreted as a translation vector $r$ so that the embedded entities $h$ and $t$ can be connected by $r$, i.e., $h$ + $r$ $\approx$ $t$ when ($h$, $r$, $t$) holds. The scoring function is defined as (negative) distance between $h+r$ and $t$, i.e.,

  $$f_{r}(h,t)=\parallel h+r-t\parallel$$

  (6)

• •

  TransR (Lin et al. 2015): Given a fact ($h$, $r$, $t$), TransR first projects the entity representations $h$ and $t$ into the space specific to relation $r$, Here $M_{r}$ is a projection matrix from the entity space to the relation space of $r$, the scoring function is:

  $$h_{t}=M_{r}h,t_{t}=M_{r}t$$

  (7)

• •

  RESCAL (Nickel, Tresp, and Kriegel 2011): Each relation in RESCAL is represented as a matrix which models pairwise interactions between latent factors. The score of a fact ($h$, $r$, $t$) is defined by a bi-linear function where $h$, $t$ are vector representations of the entities, and $M_{r}$ is a matrix associated with the relation. This score captures pairwise interactions between all components of $h$ and $t$:

  $$f_{r}(h,t)=h^{T}M_{r}t=\sum_{i=0}^{d-1}\sum_{j=0}^{d-1}[M_{r}]_{ij}*[h]_{i}*[t]_{j}$$

  (8)

• •

  DistMult (Yang et al. 2015): DistMult simplifies RESCAL by restricting $M_{r}$ to diagonal matrices. For each relation $r$, it introduces a vector embedding $r$ and requires $M_{r}$ = $diag(r)$. The scoring function is defined as:

  $$f_{r}(h,t)=h^{T}diag(r)t=\sum_{i=0}^{d-1}r_{i}*[h]_{i}*[t]_{j}$$

  (9)

  This score captures pairwise interactions between only the components of $h$ and $t$ along the same dimension, and reduces the number of parameters to $O(d)$ per relation.

From the input instance $s$ with two tagged target drug entities $d_{1}$ and $d_{2}$, we obtain the KG embedding representation of two drugs ${kge}_{1}$ and ${kge}_{2}$ respectively using Bio-KG. We concatenate these two embeddings ${kge}_{1}$ and ${kge}_{2}$ and pass those through a fully connected layer as represented below:

We have followed the task setting of Task 9.2 in the DDIExtraction 2013 shared task (Herrero-Zazo et al. 2013) for evaluation. It consists of MEDLINE documents annotated with the drug mentions and five types of interactions: Mechanism, Effect, Advice, Interaction and Other. The task is a multi-class classification to classify each of the drug pairs in the sentences into one of the types and we evaluate using three standard evaluation metrics such as: Precision (P), Recall (R) and F1-score (F1).

During pre-processing, we obtain the DRUG mentions in the corpus and map those into unique DrugBank ⁴⁴4https://go.drugbank.com/ identifiers. This is a step for converting the drug mentions into their respective DrugBank ID, a step of entity linking (Mondal et al. 2019), (Leaman, Dogan, and lu 2013). This mention normalization has been performed based on the longest overlap of drug mentions in DrugBank and map the drugs to different Knowledge sources used to construct Bio-KG.

For the purpose of experiments, we use the initialization of various pre-trained contextual embeddings. For instance, we use the embeddings such as bert-base-cased ⁵⁵5https://huggingface.co/bert-base-cased, scibert-scivocab-uncased (Beltagy, Lo, and Cohan 2019) ⁶⁶6https://github.com/allenai/scibert and domain-specific biobert v1.0 pubmed pmc and biobert v1.0 pubmed⁷⁷7https://github.com/dmis-lab/biobert as the initialization of the transformer encoder in BERTKG-DDI. We uniformly keep the maximum sequence length as 300 for all the embedding ablations and trained for 5 epochs. For the KG embeddings, we use word embeddings dimensions to be 200. Stochastic Gradient Descent (SGD) was used for optimization with an initial learning rate of 0.0001 and the model is trained for 300 epochs. After training the embeddings, we obtain the final representation of each drug. For the drugs mentioned in the input instance, we make use of the obtained embeddings as shown in the equation 11. We initialize the non-normalized drugs using pre-trained word2vec (of dimension 200 same as the KG embedding) trained on PubMED ⁸⁸8http://evexdb.org/pmresources/ngrams/PubMed/.