Background

Finding the most suitable drug for patients affected by heterogeneous diseases is a complex task. Different biological factors influence the individual drug response, while the prescription by general practitioners might be additionally influenced by drug availability, insurance coverage, or even personal preference. Computer-based recommendation systems can be utilized to support doctors in tackling this complex task.

To predict the drug best treating the patient, patterns from patient data present in electronic health records (EHRs) can be utilized. This patient information can be represented as a graph, which is a commonly used data structure in the field of computer science. A graph consists of nodes, representing entities like administered drugs or diagnoses, and edges, displaying their relationship (e.g., the drug treats the disease it is connected to).

Previous drug recommendation systems are usually trained on large datasets for specific diseases. In-context learning can be used to overcome the energy- and time-consuming fine-tuning of disease-specific drug recommendation models. Using this method, no task-specific training or fine-tuning has to be performed by pre-training a neural network.

Graph in-context learning on patient-specific graphs

GPT-3 is one of the first models that formally introduced in-context learning (ICL) [1]. This technique enables the model to answer questions it has not specifically seen before by using (prompt) examples to create a context for the question. Instead of additional training or gradient update, only pre-training is required to initially set the weights of the machine learning model and perform various downstream tasks. Since the first ICL models were only applicable to text data, newer approaches like Prodigy [2] construct a framework for applying in-context learning on graph-structured data. This is more complex due to the structural differences and unaligned semantic spaces. In the proposed framework, the weights of the graph neural networks are pre-trained to perform the predictions without requiring any fine-tuning, only using lightweight subgraphs as prompt examples.

To perform in-context-based personalized drug prediction for individual patients, we created patient-specific knowledge graphs representing the biomedical information of patients from the large cohort study UK Biobank [4]. By using patients with similar diseases and known prescriptions, the aim was to predict the best drug for a new patient based on his biomedical profile. In our case, each patient is represented by a graph (see Figure 1) incorporating his genomic information, the disease and treatment history, and demographic information (gender, age, education, …). To achieve further connections between the nodes of the graph, representing the different entities of the patient, additionally, a protein-protein interaction network was included.

For pre-training the Prodigy framework, we used the biomedical knowledge graph PrimeKG [3], which consists of drug-, disease-, and other information. The goal was to perform pre-training using a graph of a domain similar to the downstream task.

Results and Conclusion

Besides testing the performance of the model pre-trained on a graph from the biomedical domain, we further compared the prediction capability to the model pre-trained with a Wikipedia graph, representing a more general domain. This was done to validate the influence of the pre-trained weights. Furthermore, we compared the performance of the proposed model for the three psychiatric diseases schizophrenia, depression, and bipolar disorder to a fine-tuning-based baseline.

The outcome suggests that even though no downstream-specific fine-tuning has to be performed, the drug prediction capability of the presented model is comparatively high and outperformed a state-of-the-art baseline approach.

An additional major advantage of the proposed work is that since no fine-tuning is required, instead of training a model for one specific disease, only a subset of sample patients diagnosed with the disease of interest is required to perform the prediction for a new patient. This increases the speed and reduces the memory- and energy consumption.

By selecting random patients and comparing the recommendation results to real-world prescriptions, promising results were obtained. However, these results also show that further improvements have to be made for reliable drug recommendation support for doctors.

If interested, the master's thesis can be downloaded here.

Citations

[1] Wayne Xin Zhao et al. “A Survey of Large Language Models”. In: (Mar. 2023). url: https://arxiv.org/abs/2303.18223v13.

[2] Qian Huang et al. “PRODIGY:Enabling In-context Learning Over Graphs”. In: arXiv preprint arXiv:2305.12600 (2023).

[3] Payal Chandak, Kexin Huang, and Marinka Zitnik. “Building a knowledge graph to enable precision medicine”. In: Scientific Data 10.1 (2023). issn: 20524463. doi: 10.1038/s41597-023-01960-3.

[4] UK Biobank. 2023. url: https://www.ukbiobank.ac.uk/