[논문 정리] LLMs-based Few-Shot Disease Predictions using EHR: A Novel Approach Combining Predictive Agent Reasoning and Critical Agent Instruction

Introduction

이 논문에서는 Few-Shot Learning이 가능한 대규모 언어 모델(LLM)을 활용하여 질병의 발병률을 예측하고자 한다. 그러나 Few-Shot Learning은 데이터가 부족한 환경에서 적용되므로, 성능에 한계가 있을 수 있다는 문제가 존재한다. 이를 보완하기 위해 본 연구에서는 EHR-CoAgent 기법을 도입하였다. EHR-CoAgent 기법은 두 개의 LLM 에이전트를 사용하는 방식으로, Predictor LLM Agent와 Critic LLM Agent가 협업하여 예측 성능을 높이고자 한다.

먼저, Predictor LLM Agent는 질병 예측을 수행하고 예측 근거(이유)를 제시한다. 그런 다음, Critic LLM Agent는 Predictor Agent가 내린 예측과 실제 정답(Ground Truth)을 비교하여 만약 예측이 잘못되었다면, 그 원인을 분석하고 개선 방안을 제안한다. 이러한 구조를 통해 모델은 지속적으로 학습 환경을 개선하며, 소량의 데이터로도 질병 예측 성능을 높일 수 있다.

Method

본 연구에서 사용하는 EHR(Electronic Health Records) 데이터는 표(tabular) 형식으로 구성되어 있으며, 질병 코드(Cd), 약물 코드(Cm), 시술 코드(Cp)로 구분된다. 이와 같은 구조화된 데이터를 LLM이 처리하기에 적합하도록 자연어 서술 형태의 비구조화된 문장으로 변환하였다.

본 연구는 크게 두 가지 조사를 통해 진행되었다.

 

1. EHR 데이터를 이용한 질병 예측에 있어 Zero-Shot prompting, Few-Shot prompting의 성능 평가

• Zero-Shot Prompting: Chain of Thought 방식을 사용하여 모델이 단계별 설명을 생성하도록 유도함으로써, feature 간 상호작용을 고려하도록 한다. 또한 질병의 유병률 통계를 제공하여 추가적인 맥락을 제시하였다.
• Few-Shot Prompting: Zero-Shot Prompting에서 사용한 프롬프트에 추가로 3개의 양성(positive) 샘플과 3개의 음성(negative) 샘플을 제시한다.

 

2. coworking하는 llm agent를 활용하여 예측 성능을 향상

• Predictor Agent: Few-Shot Prompting을 사용하여 질병 예측을 수행한다.
• Critic Agent: 샘플별 피드백을 생성한 후, 여러 샘플(batch)의 피드백을 GPT-4를 활용해 통합하고 정제한다. 이 통합된 피드백(F_consolidated)을 다시 Predictor Agent의 프롬프트에 추가하여 m번 반복한다. 이를 통해 Instruction-Enhanced Prompting을 구현한다.

 

 

Predictor Agent의 프롬프트 예시

 

Critic Agent의 프롬프트 예시

 

Experimental Settings

1. Dataset : MIMIC-III와 CRADLE 데이터셋을 사용하였다.

• MIMIC-III: 본 연구에서는 MIMIC-III 데이터셋 중 지질 대사 장애(Disorders of Lipid Metabolism)라는 만성질환을 예측하는 데 초점을 두었다. Clinical Classifications Software를 통해 질환을 분류하였으며, 두 번 이상 병원에 방문한 환자만 포함하였다. 첫 번째 방문 기록은 모델에 입력되는 과거 진료 기록으로, 두 번째 방문 기록은 예측 대상의 결과로 사용하였다. 그 결과 12,353개의 데이터가 생성되었으나, 비용 문제로 1,000개의 데이터를 무작위로 선정하여 사용하였다.
• CRADLE (Emory Clinical Research Analytics Data Lake Environment): 제2형 당뇨병 환자가 1년 이내에 심혈관 질환(관상동맥심장질환, 울혈성심부전, 심근경색, 뇌졸중)에 걸릴지를 예측하는 데 초점을 두었다. ICD-9, ICD-10 코드를 통해 식별하였으며, 1년 내 심혈관질환이 발생한 양성 환자의 경우 발생 전 1년 내의 첫 번째 의료 기록을 입력 데이터로 사용하였다. 음성 환자의 경우 마지막 진료 기록을 기준으로 1년 전의 의료 기록을 선택하였다. 2회 미만으로 진료받은 사람, 첫 진료와 마지막 진료 간 간격이 1년 미만인 환자, 이미 심혈관 질환 경력이 있는 환자는 제외하였다. 그 결과 35,404명의 데이터가 생성되었으나, 동일한 비용 문제로 1,000개를 무작위로 추출하여 사용하였다.

 

두 데이터 모두 클래스 불균형 문제(양성 환자 비율이 적음)가 있어, Accuracy, Sensitivity, Specificity, F1 Score을 평가 지표로 사용하였다. 

 

2. Baseline 모델과 LLM 모델 

• Baseline 모델: 전통적인 머신러닝(ML) 알고리즘인 Decision Trees, Logistic Regression, Random Forests를 사용하였다. ML 모델은 두 가지 방식으로 훈련되었는데, 첫 번째는 대량의 레이블 데이터를 활용한 fully supervised 방식이고, 두 번째는 1,000개의 제한된 데이터를 사용하는 few-shot 방식이다.
• LLM 모델: GPT-4-0125-preview 모델과 GPT-3.5-turbo-16k-0613 모델을 활용하였다. 이를 Azure OpenAI Service를 통해 호출하고, 보안 규정을 준수하여 실험을 진행하였다.

3. 구현 환경

Python을 이용하였고, ML 모델은 sklearn 라이브러리를 활용하였으며, GPT 모델은 Azure OpenAI Service를 사용하여 호출하여 보안 규정을 준수하였다. 

Experimental Results

 

Fully-Supervised ML 모델이 충분한 양의 레이블 데이터를 사용했을 때 가장 높은 성능을 보였다. Few-Shot ML 모델은 데이터가 제한되어 있어 성능이 낮았다. 이는 전통적인 ML 모델이 대량의 레이블된 데이터를 필요로 함을 시사한다.

 

 

LLM들은 Few-Shot 환경에서 ML 모델보다 나은 성능을 기록하는 경우가 있었다. LLM 모델들은 Sensitivity는 높았지만, Specificity는 낮은 경향을 보였다. Zero-Shot+ 은 Zero-Shot에 추가적인 프롬프트 엔지니어링을 한 결과로, 성능이 전반적으로 개선될 수 있다. 그러나 Zero-Shot+ 모델보다 Few-Shot이 소량의 예제로도 향상된 Accuracy를 보였으며, EHR-CoAgent는 Fully-Supervised ML과 비슷하거나 Sensitivity 부분에서는 더 우수 성능을 보였다. 

 

Generated Instructions

 

Figure 2는 Critic Agent가 생성한 기준 및 지침에 대한 예시를 보여준다. 

예를 들어, MIMIC-III의 네번째 피드백의 경우, 처방된 약물의 역할을 반드시 고려하라는 지침을 제공한다. 

또한 CRADLE의 첫번째 피드백의 경우, LLM은 양성으로 진단을 내리는 보수적인 판단을 쉽게 하는 경향이 있음을 알려준다. 만약 환자가 적극적으로 관리를 받고 있다면, 반드시 심혈관질환이 발생한다고는 할 수 없다. 따라서 환자가 치료를 받고 있는지를 먼저 고려한 후 판단하라는 피드백을 주는 것이다. 

---

본 연구에서는 소량의 EHR 데이터로도 효과적으로 예측 성능을 높일 수 있는 EHR-CoAgent 구조를 제안하였다. 이를 통해 LLM 기반 Few-Shot Learning의 가능성과 한계를 탐색함과 동시에, Critic Agent의 피드백을 적극적으로 활용하여 모델의 예측 품질을 개선할 수 있음을 보였다.