Enhancing Biomedical Knowledge Discovery for Diseases: An End-To-End Open-Source Framework

Leveraging the extracted co-occurrence graph, we define two distinct probability distributions to select sentences for manual annotation. The first distribution $\mathcal{P}$ focuses on common pairs of co-occurred entities, with higher frequency in the co-occurrence graph resulting in a higher likelihood of sampling. The second distribution $\mathcal{IP}$ prioritizes novel/rare pairs of co-occurred entities, selecting sentences where the entities have a lower frequency in the co-occurrence graph. We sample 50% of sentences using $\mathcal{P}$ and 50% using $\mathcal{IP}$ to ensure a balance of common and potentially novel pairs of co-occurring entities in the datasets.⁶⁶6Sentence sampling algorithm details in Appendix A.

Then, we develop an annotation portal using the streamlit⁷⁷7https://streamlit.io/ library, providing a user-friendly interface for annotators. Annotators are presented with a sentence containing two highlighted entities and are prompted to categorize the semantic relation between them. Options include positive (direct semantic connection), negative (negative semantic connection where negative words like "no" and "absence" are present), complex (semantic connection with complex reasoning), and no relation. The annotation portal offers additional functions such as sentence removal (for non-informative sentences), entity removal (for incorrect entity types or spans), and context addition (for providing additional text to aid in relation type determination). We enlist the expertise of three medical experts to ensure the accuracy and reliability of the annotation process.

The result of the expert annotation yields two curated datasets. The Relation Detection dataset for Rett Syndrome (ReDReS) contains 601 sentences with 5,259 instances and 1,148 unique CUIs (Tab. 1). The inter-annotator agreement is measured using the Fleiss kappa score (McHugh, 2012), resulting in 0.6143 in the multi-class setup (4 classes) and indicating substantial agreement among annotators (Landis and Koch, 1977). In the binary setup (relation or no relation), the Fleiss kappa score is 0.7139. The Relation Detection dataset for Alzheimer’s Disease (ReDAD) comprises 641 sentences with 8,565 instances and 1,480 unique CUIs (Tab. 1). The Fleiss kappa score is 0.6403 in the multi-class setup and 0.7064 in the binary setup, showing substantial consensus among annotators. The final labels are determined through majority voting, leveraging the labels provided by each expert. While the label distribution across classes is relatively balanced, the negative class is under-represented with 97 and 125 instances in ReDReS and ReDAD respectively (Tab. 1). Each dataset is randomly split into train, development, and test sets.

LaMEL learns an embedding space optimized for relation detection (Fig. 3). As the backbone language model (LM), we opt for PubMedBERT (Gu et al., 2021; Tinn et al., 2023) (MIT License), available in both uncased base and uncased large versions.⁸⁸8HuggingFace’s Transformers library (Wolf et al., 2019) PubMedBERT is pretrained on the PubMed corpus, making it well-suited for our task as the curated datasets consist of sentences of abstracts from PubMed papers. Leveraging PubMedBERT ensures that the model can capture the language patterns prevalent in biomedical text. Following the LM encoding, we construct the representation of each entity by extracting its contextualized embedding $E_{i}$ corresponding to each entity $e_{i}$ from the encoded sequence. Subsequently, the entity representations are projected to the embedding space using a linear layer without changing the embedding dimension. The final prediction is based on cosine similarity between the two projected entity representations. If the cosine similarity exceeds a predefined threshold, the model predicts that there is a semantic relation between the two entities. We experiment with diverse strategies for learning entity representations (Fig. 3), aiming to optimize the effectiveness of the embedding space for the relation detection task. The explored types of entity representation $E$ are:

• •

  A, B, C - Special Tokens:

  $$E_{A}=t_{[ent]},$$

  (1)

  $$E_{B}=t_{[/ent]},$$

  (2)

  $$E_{C}=t_{[ent]};t_{[/ent]},$$

  (3)

• •

  D - Entity Pool:

  $$E_{D}=[t_{E}],$$

  (4)

• •

  E - Entity & Middle Pool:

  $$E_{E}=[t_{E}]*[t_{Inter}],$$

  (5)

• •

  F, G, H - Special Tokens & Middle Pool:

  $$E_{F}=t_{[ent]}*[t_{Inter}],$$

  (6)

  $$E_{G}=t_{[/ent]}*[t_{Inter}],$$

  (7)

  $$E_{H}=t_{[ent]}*t_{[/ent]}*[t_{Inter}],$$

  (8)

LaMReD provides two variations that differ in information synthesis (Fig. 3), aiming to explore the potential effect of different aggregations (Theodoropoulos and Moens, 2023). LaMReDA utilizes element-wise addition to aggregate the entities’ representations, while LaMReDM employs element-wise multiplication. The input text is encoded using PubMedBERT (base or large). Following LM encoding, we construct the relation representation by sampling and aggregating tokens from the input sequence. This step enables the model to capture essential features and contextual information relevant to semantic relation classification. To mitigate the risk of overfitting and enhance model generalization, we incorporate a dropout layer (Srivastava et al., 2014) with a probability of 0.3. The linear classification layer takes the aggregated representation and outputs the predicted label.

Following the paradigm proposed by Baldini Soares et al. (2019) and Hogan et al. (2021), we experiment with various approaches for learning relation representations tailored to the relation detection task to empirically ascertain the effectiveness of each strategy (Fig. 3). The explored types of relation representation R are the following:

• •

  A, B, C - Special Tokens:

  $$R_{A}=f(l(t_{[ent]_{1}}),l(t_{[ent]_{2}})),$$

  (9)

  $$R_{B}=f(l(t_{[/ent]_{1}}),l(t_{[/ent]_{2}})),$$

  (10)

  	$\displaystyle R_{C}=f(l(t_{[ent]_{1}}),l(t_{[/ent]_{1}}),$		(11)
  	$\displaystyle l(t_{[ent]_{2}}),l(t_{[/ent]_{2}})),$

• •

  D - Entity Pool:

  $$R_{D}=f(l([t_{E1}]),l([t_{E2}])),$$

  (12)

• •

  E - Middle Pool:

  $$R_{E}=l([t_{Inter}]),$$

  (13)

• •

  F - [CLS] token & Entity Pool:

  $$R_{F}=f(l(t_{[CLS]}),l([t_{E1}]),l([t_{E2}])),$$

  (14)

• •

  G, H, I - [CLS] token & Special Tokens:

  $$R_{G}=f(l(t_{[CLS]}),l(t_{[ent]_{1}}),l(t_{[ent]_{2}})),$$

  (15)

  $$R_{H}=f(l(t_{[CLS]}),l(t_{[/ent]_{1}}),l(t_{[/ent]_{2}})),$$

  (16)

  	$\displaystyle R_{I}=f(l(t_{[CLS]}),l(t_{[ent]_{1}}),l(t_{[/ent]_{1}}),$		(17)
  	$\displaystyle\vspace{-3mm}l(t_{[ent]_{2}}),l(t_{[/ent]_{2}})),$

• •

  J - [CLS] token & Middle Pool:

  $$R_{J}=f(l(t_{[CLS]}),l([t_{Inter}])),$$

  (18)

• •

  K, L, M - Special tokens & Middle Pool:

  $$R_{K}=f(l(t_{[ent]_{1}}),l([t_{Inter}]),l(t_{[ent]_{2}})),$$

  (19)

  $$R_{L}=f(l(t_{[/ent]_{1}}),l(t_{Inter}]),l(t_{[/ent]_{2}})),$$

  (20)

  	$\displaystyle R_{M}=f(l(t_{[ent]_{1}}),l(t_{[/ent]_{1}}),l([t_{Inter}]),$		(21)
  	$\displaystyle\vspace{-3mm}l(t_{[ent]_{2}}),l(t_{[/ent]_{2}})),$

• •

  N - Entity & Middle Pool:

  $$R_{N}=f(l([t_{E1}]),l([t_{Inter}]),l([t_{E2}])),$$

  (22)

• •

  O, P - Context Vector & Entity Pool:

  $$R_{O}=l(cv),$$

  (23)

  $$R_{P}=f(l([t_{E1}]),l([t_{E2}]),l(cv)),$$

  (24)

The models are trained for 50 epochs and the best checkpoints are retained based on the performance on the development set, measured using the F1-score. We utilize the Adam (Kingma and Ba, 2014) optimizer with a learning rate of 10^-5. The batch size is set to 16. We conduct experiments in two distinct setups. In the multi-class setup, we evaluate performance using micro and macro F1-score, considering four relation types: positive, negative, complex, and no relation. In the binary setup, the objective is the prediction of the presence of relation. LaMEL is specifically designed for the binary setup. We utilize the official splits of ReDReS and ReDAD (Tab. 1) and repeat the experiments 10 times with different seeds. To ensure robustness of results, we also employ a 5-fold cross-validation approach. To explore the cross-disease capabilities of our approach, we train the models using one dataset (e.g., ReDReS) and evaluate on the other (e.g., ReDAD), and vice versa. We utilize the relation representation $R_{A}$ (Eq. 9) for LaMReDA and LaMReDM and the entity representation $E_{A}$ (Eq. 1) for LaMEL. These experiments are repeated 10 times with different seeds, and 15% of the training data is excluded to define the development set¹⁰¹⁰10Hardware: single NVIDIA RTX 3090 GPU 24GB..

The cross-entropy loss function is used to train LaMReDA and LaMReDM. For LaMEL, the following cosine embedding loss function is used: