[논문 리뷰] Large Language Model Distilling Medication Recommendation Model

Qidong Liu, Xian Wu외 5명|arXiv (Cornell University)|2024. 02. 05.

Artificial Intelligence in Healthcare인용 수 5

한 줄 요약

LEADER는 분류 출력 계층을 가진 미세 조정된 LLM을 사용하여 약물 권고를 수행하고, 특징 수준의 지식 증류 및 프로필 정렬을 통해 의미론적 힘을 컴팩트한 학생 모델로 이전하여 단일 방문 및 다중 방문 환자에 대한 효과적인 권고를 가능하게 한다.

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.

연구 동기 및 목표

• 기존 약물 권고 시스템의 의미 이해 부족 및 단일 방문 한계 문제를 해결한다.
• 입력에 구애받지 않는 유연성으로 약물 권고를 향상시키기 위해 대형 언어 모델을 활용한다.
• LLM 지식을 더 작고 추론 친화적인 모델로 전이하기 위한 효율적인 증류 프레임워크를 개발한다.
• LEADER의 코드를 공개하여 재현성을 보장한다.

제안 방법

• EHR 데이터를 LLM 입력을 위한 자연어로 변환하기 위한 프롬프트 템플릿을 설계한다.
• 약물 확률을 산출하기 위해 분류 출력 계층과 감독 학습 미세 조정 손실(binary cross-entropy)로 LLM을 수정한다.
• 경량 매개변수만 업데이트하도록 LoRA 기반 미세 조정을 적용한다.
• 2단계 증류를 도입한다: (i) 교사 LLM(LEADER(T))를 학습하고 (ii) BCE 손실, 특징 수준 KD 손실, 프로필 정렬 손실로 학생 모델(LEADER(S))을 학습시킨다.
• 진단/시술/약물에 대한 세 개의 인코더, 공유 방문 인코더, 약물 예측을 위한 최종 두 층 프로젝션을 개발한다.
• 프로필 임베딩을 약물 표현과 정렬하기 위한 대비 손실을 사용하는 프로필 정렬 전략을 구현한다.

실험 결과

연구 질문

• RQ1 LEADER는 최첨단 약물 권고 모델 및 다른 LLM 기반 권고와 어떻게 비교되는가?
• RQ2 LEADER의 개별 설계 구성 요소가 성능 향상에 기여하는가?
• RQ3 지식 증류와 프로필 정렬이 LEADER의 성능에 미치는 영향은 무엇인가?
• RQ4 학생 모델이 정확도를 유지하면서 높은 효율을 달성할 수 있는가?

주요 결과

• LEADER(T)는 베이스라인과 비교했을 때 MIMIC-III 및 MIMIC-IV에서 overall, multi-visit, single-visit 모두에서 PRAUC, Jaccard, F1이 가장 높다.
• LEADER(S)는 여러 베이스라인을 능가하고 일부 LLM 기반 방법을 상회하며 특징 수준 지식 증류 및 프로필 정렬의 이점을 얻는다.
• 고찰 연구에서 KD 또는 정렬 제거가 성능 저하를 보이며 각 구성 요소의 효과를 확인한다.
• 단일 방문 성능은 LLM 의미 이해 및 프로필 정렬의 혜택으로 특히 크게 향상되어, 이전 처방 기록 없이도 견고한 권고를 가능하게 한다.