[Paper Review] Health-LLM: Personalized Retrieval-Augmented Disease Prediction System

Qinkai Yu, Mingyu Jin+9 more|arXiv (Cornell University)|Feb 1, 2024

Artificial Intelligence in Healthcare11 citations

TL;DR

Health-LLM combines large language models, retrieval-augmented generation, and automated feature engineering to predict diseases from health reports, outperforming several baselines.

ABSTRACT

Recent advancements in artificial intelligence (AI), especially large language models (LLMs), have significantly advanced healthcare applications and demonstrated potentials in intelligent medical treatment. However, there are conspicuous challenges such as vast data volumes and inconsistent symptom characterization standards, preventing full integration of healthcare AI systems with individual patients' needs. To promote professional and personalized healthcare, we propose an innovative framework, Heath-LLM, which combines large-scale feature extraction and medical knowledge trade-off scoring. Compared to traditional health management applications, our system has three main advantages: (1) It integrates health reports and medical knowledge into a large model to ask relevant questions to large language model for disease prediction; (2) It leverages a retrieval augmented generation (RAG) mechanism to enhance feature extraction; (3) It incorporates a semi-automated feature updating framework that can merge and delete features to improve accuracy of disease prediction. We experiment on a large number of health reports to assess the effectiveness of Health-LLM system. The results indicate that the proposed system surpasses the existing ones and has the potential to significantly advance disease prediction and personalized health management.

Motivation & Objective

• Motivate personalized health management by predicting potential diseases from patient health reports.
• Integrate health reports into a large model and adjust feature weights with professional medical knowledge.
• Develop a semi- to fully automated feature extraction framework to improve disease prediction accuracy.

Proposed method

• Use in-context learning to generate symptom features for various diseases.
• Apply Retrieval Augmented Generation (RAG) with a medical knowledge base to enrich prompts for LLMs.
• Score features using Llama Index with professional medical information and embed scores as features for downstream modeling.
• Train an XGBoost classifier on the engineered feature scores to predict disease labels.
• Incorporate Context-Aware Automated Feature Engineering (CAAFE) to iteratively generate semantically relevant features.
• Evaluate against multiple baselines on the IMCS-21 dataset with ACC and macro F1 as metrics.

Experimental results

Research questions

• RQ1 Can Health-LLM improve disease prediction accuracy over GPT-4 with information retrieval and other baselines on the IMCS-21 dataset?
• RQ2 What is the impact of external medical knowledge retrieval and feature preprocessing (CAAFE) on prediction performance?
• RQ3 Does a retrieval-augmented architecture combined with automated feature engineering improve personalized health predictions?
• RQ4 How does Health-LLM perform relative to fine-tuned Llama2 and other large models in a few-shot setting?

Key findings

• Health-LLM achieves Accuracy 0.833 and F1 0.762 on IMCS-21, outperforming GPT-4 with retrieval (Accuracy 0.680, F1 0.718).
• Health-LLM outperforms several baselines including RoBERTa, TextBiLSTM+Attention, and Hierarchical Attention on the same task.
• Ablation shows that removing external medical knowledge retrieval drops Accuracy to 0.78 and F1 to 0.714, underscoring the value of professional knowledge indexing.
• Removing CAAFE feature preprocessing reduces Accuracy to 0.77 and F1 to 0.721, indicating benefits from context-aware automated feature engineering.
• Compared to fine-tuned Llama2 models, Health-LLM remains superior in the reported setup (Health-LLM 0.833 ACC, 0.762 F1 vs. Fine-tuned-Llama2-13B 0.730 ACC, 0.671 F1).