Foundation Model for Biomedical Graphs: Integrating Knowledge Graphs and Protein Structures to Large Language Models (2024)

Abstract: Transformer model has been a de-facto standard in natural language processing. Its adaptations in other fields such as computer vision showed promising results that this architecture is a powerful neural network in representation learning regardless of the data type. This recent success has led to research in multimodal Large Language Model (LLM), which enabled us to new types of tasks and applications with multiple data types. However, multimodal LLM in the biomedical domain is primarily limited to images, text, and/or sequence data. Here I propose to work on multimodal LLM architecture for biomedical graphs such as protein structure and chemical molecules. The research hypothesis is based on the fact that clinicians and researchers in computational biology and clinical research take advantage of various information for their decision-making process. Therefore, an AI model being able to handle multiple data types should boost its ability to use diverse knowledge for improved performances in clinical applications.

heading: 1 Introduction
text: The foundation model revolutionized not only natural language processing (NLP) but also the humanAI interaction after the release of ChatGPT service by OpenAI (OpenAI, 2023a). ChatGPT enhanced the usability with a chat interface allowing users to instruct large language model (LLM) for any tasks even the ones requiring complex domain knowledge such as medical domain text (Savage et al., 2024). The emergence of open-source medical LLMs has further enhanced access to these technologies in healthcare settings, addressing privacy concerns associated with patient data (Toma et al., 2023; Kweon et al., 2023; Chen et al., 2023).
This success of the foundation model quickly extended to computer vision (CV), expanding the application of chat assistant tools to medical image analytics (OpenAI, 2023b; Li et al., 2023b;
Tu et al., 2023). Recently, visual instruction tuning was introduced to open the possibility of a visual assistant in medicine (Li et al., 2023a; Lee et al., 2023). Additionally, there has been notable progress in extending the model’s capabilities to handle biological sequences, including DNA sequences and chemical sequences represented by Simplified Molecular Input Line Entry Specification (SMILES) notation (Taylor et al., 2022; Consens et al., 2023; Zhang et al., 2024).
Despite these advancements, multimodal research in biomedicine has focused on integrating text, image, and sequence data. While these modalities have proven invaluable in capturing certain medical nuances, they often overlook the structural intricacies inherent in biomedical graph data, such as knowledge graphs and protein structures. Consequently, the full potential of multimodal learning remains largely unexplored in addressing the multifaceted challenges encountered in computational biology and clinical research.

heading: 1.1 Biomedical Graphs
text: Graph-based representations in biology and medicine are effective in elucidating the complex mechanisms of diseases and uncovering novel insights, such as biomarkers and therapeutic targets (Zhang et al., 2021; Chandak et al., 2023). Over the years, there has been a notable shift in graph representation learning methodologies, moving from traditional graph neural networks to transformer model architectures, mirroring advancements seen in other modalities. Notably, transformer models have shown considerable promise in graph representation learning, particularly for small biomedical graphs like chemical molecules. This approach has demonstrated the ability to overcome challenges such as over-smoothing observed in graph neural networks, while also exhibiting improved performance with deeper models (Ying et al., 2021).
346

heading: 1.2 Thesis Objective
text: As clinicians and researchers rely on multimodal data to make their decisions regarding patient care, there exists a pressing need to extend the scope of biomedical multimodal models to cover various modalities such as biomedical graphs (Soman et al., 2023; Lv et al., 2024). This extension holds the promise of significantly enhancing the capabilities of foundation models in biomedical research, thereby broadening the horizons for a myriad of biomedical tasks, including drug discovery, differential diagnoses, and treatment planning.
In light of these considerations, the proposed research aims to bridge the gap between foundation models and biomedical graph data, leveraging the rich structural information encoded in graphs to enhance the capabilities of multimodal learning in biomedical research. The overarching objective is to develop novel methodologies and frameworks that effectively harness the synergies between foundation models and biomedical graph data, enabling clinicians and researchers to derive deeper insights from complex biological networks.
Figure 1 shows how different modalities including the biomedical graphs such as protein and knowledge graph will be fused with the medical foundation model. With this foundation model, clinicians and researchers can use prompts with their expertise for clinical and biomedical applications for precision medicine and drug discovery.
For instance, the model can be queried to find a disease that can be cured with an existing drug.
The thesis proposes to explore the hypothesis that multimodal representation learning with biomedical graphs will improve the performance of drug discovery and precision medicine applications of foundation models. To achieve this objective, the research aims to:
1. Develop a novel state-of-the-art foundation model with genetics and pharmacology related biomedical guidelines to better understand human diseases. 2. Extend the modality of the developed LLM to interpret biomedical graphs as well as other modalities in the biomedical domain. 3. Compare the performance of the multimodal model with unimodal models and other stateof-the-art methods. 4. Use the foundation model for applications such as target identification and drug repurposing especially for neurodegenerative diseases such as Dementia.
Through these objectives, the research aims to contribute towards advancing precision medicine and healthcare innovation, paving the way for personalized and targeted approaches to healthcare using the foundation model. So far, I am working on the first objective, training a biomedical LLM and the preliminary result for this objective is included in this proposal.

heading: 2 Methods
text: The proposed foundation model for biomedical graphs, depicted in Figure 1, integrates multiple specialized encoders with a backbone LLM that is also trained for the biomedical domain to effectively process and understand diverse biomedical data types such as medical images and ontologies.

heading: 2.1 Medical Language Encoder
text: The following open-source LLMs were used to investigate their performance in handling medical domain text.
LLaMA2 (Touvron et al., 2023). LLaMA2, 7B, 13B, and 70B models without chat optimization were used in this work. These models were trained on 2 trillion (T) pretraining tokens in the general domain. There are several medical LLMs further trained from LLaMA2 model weights (Toma et al., 2023; Kweon et al., 2023; Chen et al., 2023). Among these medical LLMs, the Meditron 70B model claims to be the best-performing model (Chen et al., 2023). The recent version of the LLaMA family model, LLaMA-3 (Meta, 2024) 8B, was also used in this work. The pretraining corpus was increased to 15 T tokens.
Mistral (Jiang et al., 2023) Mistral-7B-v0.1 without chat optimization was used. While the details of the training dataset remain undisclosed, Mistral is known to utilize Grouped Query Attention, similar to Llama2-70B, along with Sliding Window Attention. For the biomedical LLM, BioMistral is one of the first models in the biomedicine domain based on the Mistral model (Labrak et al., 2024).
Phi-2 (Microsoft, 2023) Phi-2 model is the smallest model in this study. Phi-2 is 2.7B parameters and is trained on an augmented textbook corpus consisting of 1.4 T tokens. Other training details remain undisclosed. As far as my understanding, no Phi-2 model was trained for the biomedical domain at the time of conducting the research.
Phi-3 (Abdin et al., 2024) Lastly, I selected the Phi-3 model, which is slightly larger (3.8B parameters) and a recent version of the Phi model. The training corpus became larger as well (3.3 T tokens). Just like the Phi-2 model, the Phi-3 model trained for the biomedical domain did not exist.

heading: 2.2 Vision Encoder
text: SigLIP model trained at resolution 512X512 will be used for the vision encoder (Zhai et al., 2023).
It is a CLIP model with an improved loss, sigmoid loss. For medical image and text alignment training, the MIMIC-CXR dataset, which is made up of chest X-ray images and corresponding radiology reports, will be used (Johnson et al., 2019). Also, various types of clinical notes at University College London Hospitals will be used as well as national resources such as the Scottish Medical Imaging (SMI) archive (Baxter et al., 2023). It contains 54 million reports and medical images such as MRIs.

heading: 2.3 Sequence Encoder
text: For encoding biological sequences such as DNA sequences, protein sequences, and SMILES representations of chemical structures, I propose to use the Galactica mini and base models (Taylor et al., 2022). Galactica stands out as the only option specifically trained to handle a diverse range of biological sequence data types, for specialized embedding capturing the unique characteristics of DNA, protein, and chemical sequences.

heading: 2.4 Graph Encoder
text: Considering the absence of a single graph transformer model trained to handle knowledge graphs, protein structures, and chemical structures simultaneously, I plan to train a graph transformer model tailored for this purpose. However, one of the current limitations of existing graph transformer architectures lies in their constrained input size. To address this limitation, linear attention or any other efficient attention can enable the model to handle larger graphs effectively.
The training data for this encoder will be collected from previous works with protein structure and chemical molecule structure encoding (Hie et al., 2022; Ying et al., 2021). For the knowledge graph training dataset, I plan to construct the graph from biomedical entities recognized from clinical notes and biomedical papers. By leveraging these datasets, I aim to train a robust Graph Transformer model capable of effectively encoding diverse types of graph data.

heading: 2.5 Foundation Model for Biomedical Graphs
text: Once the encoder for each modality is trained, alignment using multi-layer perceptron adapters between the medical LLM and encoders will be implemented, an approach inspired by LLaVA family models (Li et al., 2023a; Lee et al., 2023). This will enable the foundation model to comprehend various modalities. Training data will be constructed
for this alignment, as well as for reinforcement learning to train the model for the expected output of various downstream tasks.

heading: 2.6 Downstream tasks
text: I aim to work on datasets for brain diseases such as dementia and multiple sclerosis.
Dementia is a syndrome caused by many diseases including Alzheimer’s disease. It affects memory and cognition, and symptoms become worse over time without cures. The foundation model will be used for the diagnosis and prognosis of dementia, aiding in precision medicine. For the diagnosis, the memory test report as well as the genetic expression profile will be used to diagnose the patient. The model will be also used to estimate biomarkers for Alzheimer’s disease prognosis such as brain volume from MRIs. The patient’s speech language ability is also another important data that the model will interpret for the prognosis.
Multiple sclerosis is a brain disease that changes our immune system to attack the myelin sheath. It can cause disability but has no cure. I aim to work on target identification for drug discovery. For target identification tasks, I propose to analyze single-cell disease-gene association networks sourced from the SC2disease dataset (Zhao et al., 2021). This dataset contains comparisons of gene expressions of different multiple sclerosis diseaserelated health status. It can thereby provide valuable insights into disease-gene associations at the single-cell level, and offer rich data for comprehensive analysis and interpretation.

heading: 3 Preliminary Experiment and Results

heading: 3.1 Medical LLM Training
text: The training dataset was collected from MedlinePlus1 which includes a medical encyclopedia and texts about drugs and genetics. The collected training dataset for continued pretraining was 2.2 million tokens based on Phi-2. Continued pretraining was done for all the models with an epoch of 3 and a learning rate of 5e-5.
Figure 2 illustrates the breakdown of the MedlinePlus corpus categories. The largest category, Health Conditions, comprises 26.1% of the corpus and includes information on the frequency, causes, synonyms, and inheritance patterns of various diseases. The Genes category, accounting for 20.3% of the corpus, describes the normal functions of
1https://medlineplus.gov/
human genes and the health implications of genetic modifications. The Medical Tests category, making up 19.9% of the corpus, covers tests such as allergy skin tests, detailing their purposes, procedures, and possible results. The Health Topics category constitutes 18.8% of the corpus and serves as an encyclopedia covering body parts, therapies, and wellness issues, with content regularly reviewed and updated daily. The Introduction to Genetics category, comprising 9.75% of the dataset, provides fundamental explanations of human genetics concepts. Finally, the Dietary Supplements category, representing 5.15% of the dataset, offers descriptions of the effectiveness, usual dosages, and potential drug interactions of various supplements.
The MedlinePlus training corpus is diverse and evenly distributed across various biomedical domains. For each category, one example is shown in Table 1. The examples highlight the diversity within the corpus which ensures a comprehensive representation of medical knowledge, which is crucial for training robust models capable of handling a wide range of medical and genetic information.

heading: 3.2 Medical LLM Evaluation
text: To evaluate the performance of the trained medical LLMs as well as the baseline models, this work uses the prevalent multiple choice question answering benchmarks in the medical language model domain, including MMLU medical subjects (MMLU_MED), MedQA, and MedMCQA (Jin et al., 2021; Pal et al., 2022; Hendrycks et al., 2020). The evaluation metric utilized is classification accuracy based on logits. As all the benchmarks are in MCQ format, the token with the highest logit value can be selected as the model’s predicted answer. The prompt used for evaluation as well as the example question and response are shown in Table 2. The models generate responses, and their accuracy is measured by comparing their responses to the expected correct answers.

heading: 3.2.1 MMLU_MED
text: MMLU (Massive Multitask Language Understanding) (Hendrycks et al., 2020) is a benchmark designed to measure the model’s ability in knowledgeintensive QA across 57 subjects. These subjects cover various levels of education: high school, college, and professional level. Questions in the dataset are structured as four-way multiple choice questions (MCQs), offering a standardized format for evaluation. Within the extensive list of subjects, there are nine healthcare-related subjects which are college medicine, professional medicine, clinical knowledge, anatomy, high school biology, college biology, medical genetics, nutrition, and virology. Collectively, these nine subjects comprise a total of 1,871 questions in the test set.

heading: 3.2.2 MedQA
text: MedQA (Jin et al., 2021) is an open-ended MCQ dataset made from professional medical doctor license exams. The dataset is available in three ver-
sions, one of which is an English version sourced from the United States Medical License Exams. While MMLU’s professional medicine subject also includes questions from USMLE practice examinations, MedQA’s English version sets itself apart by incorporating questions drawn from both real exams and mock tests for USMLE. 1,273 USMLEstyle questions are provided as the test dataset to benchmark the model’s ability to answer medical questions at the professional level. Each question is accompanied by four or five answer choices and corresponding relevant document collections, intended to help models in generating accurate responses.

heading: 3.2.3 MedMCQA
text: MedMCQA (Pal et al., 2022) is a benchmark with questions sourced from postgraduate-level Indian medical school entrance exams (AIIMS and NEET PG). Covering a breadth of medical specialties, the dataset has questions about 2,400 healthcare topics and 21 subjects within the medical domain. 4,183 MCQ, each offering four answer choices, are provided for evaluation.

heading: 3.3 Evaluation Results
text: The preliminary results in Table 3 for the medical large language training provide several notable trends. Firstly, there is a clear trend between
model size and performance, with larger models consistently achieving higher accuracy scores across all three benchmark datasets. For instance, the LLaMA-2 model, particularly in its larger 70- billion-parameter model, shows superior performance compared to smaller models. This underscores the importance of model scale in capturing the complexity of medical language and achieving better task performance. However, due to the constraints of the scarce computational resources at the hospital, smaller models with adequate performance can be preferred.
Additionally, the effect of continued training is observed. LLaMA-2-MedlinePlus and Phi-2MedlinePlus models demonstrate enhanced performance compared to their counterparts trained on general-domain data. However, it is worth noting that this trend is not universal, as observed with the Mistral-MedlinePlus model, which did not exhibit a significant increase in performance despite continued training. While the LLaMA-3-MedlinePlus model and Phi-3-MedlinePlus model showed improved performance in the MMLU_MED benchmark, these models showed a significant decrease in performance for the MedQA benchmark.
To ensure the integrity of the models regarding the pretraining corpus and evaluation benchmarks, a thorough analysis for data contamination was conducted using the recent method, MIN-K% PROB (Shi et al., 2023). The MIN-K% PROB score measures the average log-likelihood of the K% tokens with minimum probability, indicating how well a language model predicts the presence of tokens in
the given text. A higher log-likelihood might suggest that the model has been exposed to the evaluation data during its training phase, potentially leading to artificially inflated performance metrics.
Even for the pretraining corpus, a model with a higher score might have been trained with the same or a very similar corpus, making the gains from continued pretraining negligible. Ensuring the validity of our results is crucial to confirm that improvements in model performance are due to genuine learning and not the memorization of the evaluation dataset. However, due to computational limitations, contamination analysis of the pretraining data was not feasible for this preliminary work.
To test for contamination, 100 examples from each dataset were sampled. The more number of examples exhibit high log-likelihood and therefore low MIN-K% PROB scores (0.1 or lower), is more likely that the dataset has already been used for training. This helps in identifying and mitigating potential data leakage, ensuring fair and accurate model evaluations.
Table 4 shows the contamination result, and it shows that Phi-3 when compared to Phi-2 showed
the highest number of examples that were suspicious of the contamination. For MedQA, 67 out of 100 examples had a very high log likelihood value. This somewhat explains the performance drop with the continued pretraining as the baseline model already might have been already trained with the evaluation datasets used in this work.
To confirm Phi-3’s data contamination, especially for MedQA, the model responses were compared as shown in Table 5. Rather than giving the right cause for the symptom, Phi-3 models generated multiple choice options which did not have the desired answer. This hallucination effect was not seen in Phi-2 models which just generated a similar case of a patient rather than answering the cause. The example of response suggests that the
effect of continued pretraining was limited to logitbased classification as all the models did not give the desired answer.
Nevertheless, while these preliminary findings provide valuable insights, further in-depth analysis is warranted to explore the nuances of model performance in the medical domain fully. Future work will focus on leveraging other training methods and more comprehensive training data. Additionally, exploration of other evaluation methods for diverse tasks can contribute to more accurate and comprehensive assessments of LLM performance in real-world healthcare applications. Collaborations with healthcare professionals will ensure that the model is aligned with clinical needs and practices by evaluating and interpreting model outputs.

heading: 4 Conclusion
text: This proposal describes the plan to develop a foundation model architecture uniquely trained to understand the complexities of biomedical graphs. Unlike existing models that primarily focus on text, images, or sequences, the proposed model aims to bridge the gap by integrating information from diverse data types such as knowledge graphs, protein structures, and chemical molecules. By leveraging the strengths of large language models in capturing textual information and combining them with specialized encoders for biological sequences and graph structures, the foundation model holds immense potential to revolutionize various aspects of healthcare, including diagnosis, treatment planning, and drug discovery.
The model can be used to create an interactive agent that clinicians and researchers can utilize to help them navigate problems in biomedical research, thereby enhancing decision-making processes in clinical practice and computational biology research. For instance, the incorporation of knowledge graphs may allow for a more nuanced exploration of relationships between genes, drugs, and diseases, facilitating target identification for drug discovery as well as drug repurposing, which accelerates the clinical trial progress.
Moreover, the integration of protein structure and chemical molecule data should enable our model to delve deeper into molecular mechanisms underlying diseases and drug interactions. This deeper understanding opens the possibility of using an assistant tool for more effective protein-drug binding affinity prediction for drug discovery, as well as the identification of potential novel biomarkers for disease diagnosis and prognosis.
By leveraging the collective insights from diverse data modalities, the proposed foundation model has the potential to significantly improve performance across a spectrum of biomedical tasks. The development of a multimodal foundation model represents a pivotal step towards unlocking the full potential of artificial intelligence in biomedicine, thereby enhancing our understanding of complex biological systems and ultimately improving healthcare outcomes for patients.
Moving forward, future work will focus on developing the proposed foundation model to address specific challenges such as training with scarce data. Additionally, I will conduct the research with the help of the collective expertise of health infor-
matics researchers and clinicians in order to develop the foundation model with a focus on realworld biomedical applications, especially for neurodegenerative diseases.
Limitation
The limitation of this proposal is the lack of evaluation with clinicians and medical professionals. Incorporating feedback from domain experts could provide valuable insights into the practical utility and reliability of the models in real-world clinical settings. Additionally, while the study used several established medical benchmarks, these datasets may not fully capture the range of complexities and variances encountered in real-world medical data. Future research should focus on broader datasets, more diverse medical tasks, and extensive real-world evaluations to ensure the robustness and applicability of the proposed models in various clinical scenarios.
Broader Impacts and Ethics Statement
I fully comply with the copyright requirements of MedlinePlus. The content sourced from MedlinePlus for our pretraining corpus is used under their permissible use policy, ensuring that all derived data and models respect the original terms and conditions.
This work utilizes clinical data strictly for research purposes. All clinical data is or will be anonymized to protect patient privacy and confidentiality in accordance with ethical standards and regulatory requirements.
My work does not raise any major ethical concerns regarding the usage of LLMs as all LLMs tested were used for research purposes only. However, all LLMs even the ones further trained with the MedlinePlus pretraining corpus are not rigorously tested for use in real-world clinical applications or scenarios. Thus, they may not be suitable for use in the clinical decision making process.

heading: Acknowledgement
text: I would like to thank Professor Honghan Wu for his support as my supervisor. Additionally, I extend special thanks to Professor Ekapol Chuangsuwanich for his time and support as a mentor for this work. I also acknowledge Jinge Wu for her help with the training corpus data collection. Last but not least, I express my gratitude to the anonymous reviewers and for their valuable feedback.

heading: Future Work
text: Moving forward, future work will focus on developing the proposed foundation model to address specific challenges such as training with scarce data. Additionally, I will conduct the research with the help of the collective expertise of health infor-
matics researchers and clinicians in order to develop the foundation model with a focus on realworld biomedical applications, especially for neurodegenerative diseases. Limitation
The limitation of this proposal is the lack of evaluation with clinicians and medical professionals. Incorporating feedback from domain experts could provide valuable insights into the practical utility and reliability of the models in real-world clinical settings. Additionally, while the study used several established medical benchmarks, these datasets may not fully capture the range of complexities and variances encountered in real-world medical data. Future research should focus on broader datasets, more diverse medical tasks, and extensive real-world evaluations to ensure the robustness and applicability of the proposed models in various clinical scenarios