기술 요약: 질병 예측을 위한 지식 그래프 증강 대규모 언어 모델

문제 정의

전자 건강 기록(EHR)은 임상적 예측에 적합한 종단적 데이터를 포함하고 있으나, 기존 모델들은 종종 임상의가 실행 가능한 환자 수준의 추론을 제공하는 데 실패한다. 고전적 모델(예: 로지스틱 회귀, 트리 앙상블)은 견고성과 전역적 해석력을 제공하지만, 일반적으로 특정 임상 결정에 필요한 단계별 서사(narrative)가 결여된 거친 사후 설명(예: 특성 중요도 점수)만을 제공한다. 반대로, 대규모 언어 모델(LLM)은 사고의 사슬(Chain-of-Thought, CoT) 서사를 생성할 수 있지만, 의료 분야에서의 무분별한 적용은 고품질의 임상적으로 검증된 추론 데이터셋의 부족과 환각(hallucination) 또는 의학적으로 부정확한 중간 단계의 위험으로 인해 제약을 받는다. 핵심 과제는 LLM이 생물 의학적 지식에 근거하여 사실 관계가 명확하고 환자의 병력과 시간적으로 일관된 설명을 생성하면서, 정확한 방문 단위 질병 예측을 수행할 수 있도록 하는 것이다.

방법론

저자들은 PrimeKG라는 대규모 생물 의학 지식 그래프의 구조를 사용하여 LLM의 추론을 스캐폴딩(scaffolding)하도록 설계된 지식 그래프(KG) 유도 CoT 프레임워크를 제안한다. 파이프라인은 네 가지 주요 단계로 구성된다:

1. KG 엔티티 매핑:

   • MIMIC-III의 ICD-9 코드를 세 단계 절차를 통해 PrimeKG 노드와 정렬한다:
     • 정확 일치(Exact Match): ICD-9 설명과 KG 노드 레이블 간의 직접적인 문자열 매칭.
     • 유사도 매칭(Similarity Match): 비정확 일치에 대한 임베딩 기반 검색(코사인 유사도 > 0.85).
     • LLM 필터링: GPT-4o가 임상적 정확성을 보장하기 위해 잠정적 매핑을 검증하거나 수정함.
   • 이 과정은 1,513개의 매핑된 특성 노드와 10개의 매핑된 질병 노드를 생성한다.

2. KG 관련성 및 경로 마이닝:

   • 각 대상 질병에 대해, 시스템은 GPT-4o를 사용하여 관련 특성 노드의 압축된 집합($R_d$)을 식별한다.
   • 추론 경로(Reasoning Paths): PrimeKG 내에서 관련 특성 노드와 대상 질병 노드 사이의 최단 경로(홉 수 $L=5$로 제한)를 추출한다.
   • 경로 선택: GPT-4o가 이러한 경로들을 필터링하여 임상적으로 관련 있고 기전적으로 타당한 것들만 유지하며, 불필요한 연결 고리를 제거한다.

3. CoT 생성 및 필터링:

   • LLM은 관찰된 ICD-9 특성, 질병별 관련성 집합, 그리고 마이닝된 KG 경로를 포함하여 특정 방문 $t$에 대한 CoT 설명을 생성하도록 프롬프트된다.
   • 레이블 일관성 필터링(Label-Consistency Filtering): 생성된 CoT 궤적은 추론으로부터 도출된 명시적 이진 결론($\hat{y}$)이 실제 결과($y$)와 일치하는 경우에만 유지된다. 이를 통해 고품질의 KG 기반 감독 코퍼스를 구축한다.

4. 모델 미세 조정(Fine-Tuning):

   • 경량 오픈 웨이트 모델(LLaMA-3.1-Instruct-8B 및 Gemma-7B)을 이 정제된 코퍼스에 대해 미세 조정한다.
   • 추론 시, 이 모델들은 ICD-9 특성과 질병별 KG 증거를 입력으로 받아 질병 예측과 임상의 스타일의 추론 궤적을 생성한다.

주요 기여

• KG 유도 CoT 프레임워크: 생물 의학 KG 구조를 CoT 생성 과정에 직접 통합하여, 제약 없는 자유 형식의 설명을 넘어 임상적으로 근거 있는 추론으로 나아가는 새로운 방법을 제시한다.
• 데이터 효율적 감독: 레이블 일관성 필터링을 통한 압축된 고품질 감독 코퍼스 구축을 통해, 작은 규모의 LLM(8B 파라미터)이 제한된 데이터(400~1,000개 샘플)로 예후 추론을 학습할 수 있게 한다.
• 제로샷 전이(Zero-Shot Transfer): 본 프레임워크는 재학습 없이도 별도의 코호트(CRADLE)로 학습된 추론 패턴을 전달하는 능력을 입증하였으며, 당뇨 진단 후 심혈관 사건 발생을 예측하는 데 높은 정확도를 유지하였다.
• 임상의 중심 평가: 맹검 인간 연구를 통해 KG 유도 설명이 명확성, 관련성, 임상적 정확성 측면에서 베이스라인보다 선호됨을 검증하였다.

실험 결과

본 프레임워크는 MIMIC-III 코호트(PrimeKG에 매핑된 10가지 질병)와 CRADLE 코호트(제2형 당뇨병에서 심혈관 사건까지)에서 평가되었다.

• MIMIC-III 성능:

  • 저데이터 설정(400 및 1,000개의 인덱스 방문)에서, KG 유도 CoT 미세 조정은 LLaMA-3.1-8B 및 Gemma-7B의 성능을 향상시켜 강력한 고전적 베이스라인(XGBoost, Random Forest, Logistic Regression)과 대등하거나 더 우수한 결과를 달성했다.
  • LLaMA-3.1 (1,000개 샘플): AUROC 0.6995 및 macro-AUPR 0.4685를 달성하여, 최상의 고전적 베이스라인(Logistic Regression: AUROC 0.6874, AUPR 0.4178; XGBoost: AUROC 0.6799)을 약간 상회했다.
  • Gemma-7B: 유사한 개선을 보였으나, macro-AUPR이 원본 0.25에서 KG-CoT 튜닝 후 0.4363으로 증가한 것에 비해, AUROC(0.6609)는 최상위 고전적 베이스라인보다 약간 낮았다.
  • 절제 연구(Ablation): KG 가이드나 필터링을 제거하면 성능이 감소하였으며, 이는 구조적 근거와 레이블 일관성 필터 모두가 결정적임을 확인해 준다.

• CRADLE로의 제로샷 전이:

  • MIMIC-III에서 미세 조정된 모델을 재학습 없이 CRADLE에 적용하였다.
  • 정확도는 0.40–0.51(튜닝 전)에서 0.72–0.77(KG-CoT 튜닝 후)로 향상되었다.
  • AUROC와 AUPR에서도 완만한 상승을 보였으나, 상당한 정확도 점프는 모델이 데이터셋 특정 아티팩트가 아닌 근본적인 질병 기전을 포착하여 더 견고한 결정 경계를 학습했음을 시사한다.

• 인간 평가:

  • 115개 사례를 대상으로 한 맹검 연구에서, 임상의들은 명확성/일관성에서 96.5%, 범위/관련성에서 94.8%, 정확성/타당성에서 **98.3%**의 비율로 튜닝되지 않은 베이스라인보다 KG 유도 모델을 선호하였다.
  • 임상의들은 KG 유도 궤적이 시간적으로 더 잘 조직되어 있고, 의미 있는 위험 요인을 강조하며, 근거 없는 비약을 피한다고 언급하였다.

의의 및 주장

본 논문은 **KG-고정 CoT 감독(KG-anchored CoT supervision)**이 소규모 오픈 웨이트 LLM이 데이터 효율적이고 임상적으로 근거 있는 예후 추론을 학습할 수 있게 한다고 주장한다. 그 의의는 예측과 해석 가능한 정당화 사이의 간극을 메우는 데 있다:

1. 신뢰성: 추론을 PrimeKG에 고정하고 레이블 일관성을 필터링함으로써, 시스템은 환각을 완화하고 의학적 타당성을 보장한다.
2. 일반화 가능성: 학습된 추론 패턴은 외부 분포 코호트(CRADLE)로 효과적으로 전이되며, 이는 모델이 데이터셋 특정의 특징이 아닌 근본적인 질병 기전을 포착하고 있음을 시사한다.
3. 임상적 유용성: 프레임워크는 임상의들이 더 명확하고 신뢰할 수 있다고 느끼는 설명을 생성하며, 이는 투명성이 매우 중요한 임상 의사 결정 지원 시스템의 도입을 촉진할 수 있다.

저자들은 본 접근 방식이 PrimeKG의 커버리지에 의존하고 있으며, 향후 연구에서 더 풍부한 EHR 뷰(노트, 검사 결과)와 다중 방문 맥락을 탐구해야 한다고 겸허히 언급하였다.