[Paper Review] Large Language Model Distilling Medication Recommendation Model

Qidong Liu, Xian Wu+5 more|arXiv (Cornell University)|Feb 5, 2024

Artificial Intelligence in Healthcare5 citations

TL;DR

LEADER uses a fine-tuned LLM with a classification output layer for medication recommendation and transfers its semantic power to a compact student model via feature-level knowledge distillation and profile alignment, enabling effective recommendations for both single-visit and multi-visit patients.

ABSTRACT

The recommendation of medication is a vital aspect of intelligent healthcare systems, as it involves prescribing the most suitable drugs based on a patient's specific health needs. Unfortunately, many sophisticated models currently in use tend to overlook the nuanced semantics of medical data, while only relying heavily on identities. Furthermore, these models face significant challenges in handling cases involving patients who are visiting the hospital for the first time, as they lack prior prescription histories to draw upon. To tackle these issues, we harness the powerful semantic comprehension and input-agnostic characteristics of Large Language Models (LLMs). Our research aims to transform existing medication recommendation methodologies using LLMs. In this paper, we introduce a novel approach called Large Language Model Distilling Medication Recommendation (LEADER). We begin by creating appropriate prompt templates that enable LLMs to suggest medications effectively. However, the straightforward integration of LLMs into recommender systems leads to an out-of-corpus issue specific to drugs. We handle it by adapting the LLMs with a novel output layer and a refined tuning loss function. Although LLM-based models exhibit remarkable capabilities, they are plagued by high computational costs during inference, which is impractical for the healthcare sector. To mitigate this, we have developed a feature-level knowledge distillation technique, which transfers the LLM's proficiency to a more compact model. Extensive experiments conducted on two real-world datasets, MIMIC-III and MIMIC-IV, demonstrate that our proposed model not only delivers effective results but also is efficient. To ease the reproducibility of our experiments, we release the implementation code online.

Motivation & Objective

• Address lack of semantic understanding and single-visit limitations in existing medication recommender systems.
• Leverage large language models to enhance medication recommendation with input-agnostic flexibility.
• Develop an efficient distillation framework to transfer LLM knowledge to a smaller, inference-friendly model.
• Ensure reproducibility by releasing code for LEADER.

Proposed method

• Design prompt templates to convert EHR data into natural language for LLM input.
• Modify the LLM with a classification output layer and a supervised fine-tuning loss (binary cross-entropy) to produce medication probabilities.
• Apply LoRA-based fine-tuning to update only lightweight parameters.
• Introduce a two-stage distillation: (i) train a teacher LLM (LEADER(T)) and (ii) train a student model (LEADER(S)) with BCE loss, feature-level KD loss, and profile alignment loss.
• Develop three encoders for diagnosis, procedure, and medication; a shared visit encoder; and a final two-layer projection to predict medications.
• Implement a profile alignment strategy using contrastive losses to align profile embeddings with medication representations.

Experimental results

Research questions

• RQ1 How does LEADER compare with state-of-the-art medication recommendation models and other LLM-based recommendations?
• RQ2 Do the individual design components of LEADER contribute to performance gains?
• RQ3 What is the impact of knowledge distillation and profile alignment on LEADER’s performance?
• RQ4 Can the student model achieve high efficiency while maintaining accuracy?

Key findings

• LEADER(T) achieves the highest PRAUC, Jaccard, and F1 across overall, multi-visit, and single-visit on MIMIC-III and MIMIC-IV compared with baselines.
• LEADER(S) outperforms several baselines and surpasses some LLM-based methods, benefiting from feature-level knowledge distillation and profile alignment.
• Ablation studies show removing KD or alignment degrades performance, confirming the effectiveness of each component.
• Single-visit performance benefits notably from LLM semantic understanding and profile alignment, enabling robust recommendations without prior prescription history.