[論文レビュー] Large Language Model Distilling Medication Recommendation Model

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

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入力とするプロンプトテンプレートを設計する。
• LLMを分類出力層と監督付きファインチューニング損失（二値クロスエントロピー）で修正し、薬剤確率を生成する。
• 軽量パラメータのみを更新するLoRAベースのファインチューニングを適用する。
• 二段階蒸留を導入する：（i）教師LLM（LEADER(T)）を学習、（ii）生 BCE損失、特徴レベルKD損失、プロフィール整合損失を用いて生徒モデル（LEADER(S)）を学習。
• 診断・手技・薬剤の3つのエンコーダを開発; 共有の訪問エンコーダ; 薬剤を予測する最終の二層投影を設計。
• プロフィール埋め込みを薬物表現と整合させるために、対比損失を用いたプロフィール整合戦略を実装する。

実験結果

リサーチクエスチョン

• RQ1 LEADERは最先端の薬剤推奨モデルおよび他のLLMベースの推奨とどのように比較されるか？
• RQ2 LEADERの個々の設計要素は性能向上に寄与するか？
• RQ3 知識蒸留とプロファイル整合がLEADERの性能に与える影響は？
• RQ4 学生モデルは高い効率を達成しつつ精度を維持できるか？

主な発見

• LEADER(T) はベースラインと比較して、MIMIC-IIIおよびMIMIC-IVで総合・複数訪問・単一訪問の各評価指標PRAUC、Jaccard、F1で最高を達成。
• LEADER(S)は複数のベースラインを上回り、いくつかのLLMベース手法を凌駕し、特徴レベル知識蒸留とプロフィール整合の恩恵を受けている。
• アブレーション研究によりKDまたは整合を削除すると性能が低下することが示され、各コンポーネントの有効性が裏付けられている。
• 単一訪問の性能はLLMの意味理解とプロフィール整合性から顕著に恩恵を受け、処方履歴なしでも堅牢な推奨を可能にする。