DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning

이 논문은 진단 중심의 상호작용과 비동기적 진행 패턴을 명시적으로 모델링하여 해석 가능한 환자 표현 학습을 가능하게 하는 '질병 궤적 인식 트랜스포머 (DT-BEHRT)'를 제안하고, 이를 통해 임상적 추론과 정렬된 강력한 예측 성능을 입증합니다.

핵심

🏥 1. 문제: 혼란스러운 병원 기록 도서관

환자의 병원 기록은 수천 개의 작은 메모 (진단명, 약, 검사 결과 등) 가 쌓인 거대한 도서관과 같습니다.
기존의 인공지능들은 이 도서관을 볼 때 두 가지 큰 실수를 저질렀습니다.

1. 순서를 잘못 이해함: 의사들이 기록한 순서가 실제 병의 발생 순서와 다를 수 있는데, 기계는 무조건 기록된 순서대로만 읽었습니다.
2. 모든 것을 똑같이 취급함: '심장마비'라는 진단명과 '진통제'라는 처방전을 똑같은 중요도로 다뤘습니다. 하지만 진단명은 병의 흐름을 이끄는 주인공이고, 약이나 검사는 그 흐름을 돕는 조연입니다. 이 차이를 무시하면 이야기의 핵심을 놓치게 됩니다.

🚀 2. 해결책: DT-BEHRT (질병 여정 탐정)

이 연구팀은 이 문제를 해결하기 위해 DT-BEHRT라는 새로운 '탐정'을 만들었습니다. 이 탐정은 두 가지 특별한 능력을 가지고 있습니다.

🧩 능력 1: 부서를 나누어 생각하는 '조직별 정리' (Disease Aggregation)

이 모델은 도서관의 책들을 무작위로 쌓아두지 않고, 부서 (장기 시스템) 별로 정리합니다.

• 예를 들어, '심장' 관련 책들은 한 책장에, '폐' 관련 책들은 다른 책장에 모읍니다.
• 환자가 여러 번 병원을 방문할 때, 단순히 "A 병원에서 B 약을 먹었다"가 아니라, **"심장 부서는 이렇게 발전했고, 폐 부서는 저렇게 변했다"**는 식으로 장기별 흐름을 파악합니다.
• 비유: 마치 사건을 해결할 때, '범인'과 '동행자'를 구분해서 조사하듯, 병의 원인 (진단) 과 치료 (약/검사) 를 구분하여 분석합니다.

⏳ 능력 2: 시간의 흐름을 따라가는 '여정 추적' (Disease Progression)

단순히 현재 상태만 보는 게 아니라, 시간에 따라 병이 어떻게 변해왔는지를 그래프로 연결합니다.

• "1 년 전에는 가벼운 기침이었는데, 6 개월 뒤엔 폐렴이 되고, 지금은 심부전으로 이어졌다"는 **시간의 흐름 (여정)**을 파악합니다.
• 비유: 추리 소설에서 범인이 처음에 작은 실수를 저지르고, 시간이 지나면서 범죄가 커지는 스토리 라인을 파악하는 것과 같습니다.

🎓 3. 학습 방법: 미리 공부하는 '수험생' 전략

이 모델은 실제 환자를 만나기 전에, 방대한 병원 기록으로 **미리 공부 (Pre-training)**를 합니다. 이때 두 가지 독특한 방법을 사용합니다.

1. 빈칸 채우기 (Global Code Masking): 기록에서 일부 단어를 가리고 "이게 무엇일까?"라고 맞추게 합니다. 하지만 단순히 단어 하나만 맞추는 게 아니라, 전체적인 흐름 (여정) 속에서 어떤 질병들이 함께 나타나는지 (예: 당뇨와 고혈압이 함께 오는 경우) 를 학습합니다.
2. 조상 찾기 (Ancestor Code Prediction): 특정 질병 코드가 나왔을 때, 그 질병이 속한 **큰 카테고리 (예: '심장병'이라는 큰 부류)**가 무엇인지 맞추게 합니다. 이렇게 하면 모델이 세부적인 증상뿐만 아니라 큰 그림 (의학적 지식) 을 이해하게 됩니다.

🌟 4. 결과: 왜 이 모델이 특별한가?

• 높은 정확도: 여러 병원 데이터 (MIMIC, eICU 등) 에서 환자의 사망 위험, 재입원 가능성, 장기 입원 여부 등을 예측할 때 기존 모델들보다 더 정확했습니다. 특히 재입원 예측처럼 복잡한 문제에서 빛을 발했습니다.
• 의사도 이해하는 설명 (해석 가능성): 이 모델의 가장 큰 장점은 **"왜 그렇게 예측했는지"**를 보여준다는 점입니다.
  • 기존 모델은 "이 환자는 위험합니다"라고만 말했지만, DT-BEHRT 는 **"환자의 폐 관련 기록 (DA 모듈) 이 나빠졌고, 시간이 지날수록 심장 질환 (DP 모듈) 이 진행되었기 때문에 위험합니다"**라고 설명합니다.
  • 이는 마치 탐정이 "범인은 A 시간대에 B 장소를 다녀왔고, C 물건을 소지했기 때문에 범인입니다"라고 논리적으로 추론 과정을 보여주는 것과 같습니다.

💡 요약

DT-BEHRT 는 환자의 병력을 단순한 데이터 나열이 아니라, 장기별 특성과 시간의 흐름을 고려한 하나의 완성된 이야기로 이해합니다. 덕분에 더 정확한 예측을 할 뿐만 아니라, 의사들이 그 결론을 신뢰하고 받아들일 수 있도록 투명한 이유를 제시해 줍니다.

이 기술은 앞으로 병원에서 의사를 돕는 '지능형 비서'로 활약하여, 환자에게 더 나은 치료 방향을 제시하는 데 기여할 것으로 기대됩니다.

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

전자의무기록 (EHR) 시스템의 보편화로 예측 모델링을 통한 임상 의사결정 지원의 기회가 확대되고 있습니다. 그러나 기존 EHR 기반 예측 모델들은 다음과 같은 한계를 가지고 있습니다.

• 이질적인 의료 코드의 균일한 처리: 기존 시계열 기반 (Sequence-based) 또는 그래프 기반 (Graph-based) 모델들은 진단 (Diagnosis), 약물 (Medication), 검사 (Lab), 시술 (Procedure) 등 서로 다른 임상적 역할과 특성을 가진 의료 코드를 동일하게 취급하는 경향이 있습니다.
  • 진단 코드: 질병의 진행을 주도하며, 장기 간 상호작용이 밀접함.
  • 약물/시술 코드: 치료 경로를 반영하며 시간적 상관관계는 강하지만, 단일 방문 내 상호작용은 제한적일 수 있음.
• 시퀀스 모델의 한계: 방문 내 코드 순서가 실제 임상 시간 순서를 반영하지 않을 수 있으며, 다중 핫 벡터 (Multi-hot vector) 표현으로는 코드 간의 공발생 (Co-occurrence) 및 의존성을 충분히 포착하지 못함.
• 그래프 모델의 한계: 방문 간 (Across visits) 순차적 의존성을 포착하는 데 어려움을 겪음.

따라서, 의료 코드의 이질적인 역할을 명시적으로 구분하고, **질병 진행 경로 (Disease Trajectory)**를 체계적으로 모델링할 수 있는 새로운 아키텍처가 필요합니다.

---

2. 제안 방법론: DT-BEHRT (Methodology)

저자들은 **DT-BEHRT (Disease Trajectory-aware Transformer for EHR)**를 제안했습니다. 이는 그래프 강화 시퀀스 아키텍처로, 환자의 건강 상태를 표현하기 위해 진단 중심의 상호작용과 비동기적 진행 패턴을 명시적으로 모델링합니다.

2.1 주요 구성 요소 (Architecture)

모델은 크게 4 가지 모듈로 구성됩니다 (그림 1 참조):

1. 시퀀스 표현 모듈 (Sequence Representation, SR):

   • BERT 스타일의 Transformer 레이어를 사용하여 의료 코드 시퀀스를 인코딩합니다.
   • 코드 유형 (진단, 약물 등) 임베딩과 방문 인덱스 임베딩을 추가하여 토큰 표현을 풍부하게 합니다.
   • 단일 방문 내 코드의 순서에 대한 가정을 배제합니다.

2. 질병 집계 모듈 (Disease Aggregation, DA):

   • 핵심 아이디어: ICD-9 분류 체계의 상위 수준 (19 개 장, Chapter) 을 기반으로 'DA 토큰'을 도입합니다.
   • 각 장 (예: 순환계, 호흡계) 에 해당하는 진단 코드가 일정 임계값 이상 존재할 때 해당 장을 대표하는 토큰을 생성합니다.
   • 목적: 동일한 장기/시스템 내에서의 질병 진행과 상호작용을 요약하여 고수준의 의미 패턴을 포착합니다.
   • 정규화: DA 토큰 간의 중복 정보를 줄이기 위해 공분산 정규화 (Covariance Regularization) 를 적용하여 토큰들이 서로 다른 장기 시스템을 독립적으로 표현하도록 유도합니다.

3. 질병 진행 모듈 (Disease Progression, DP):

   • 이질적 그래프 (Heterogeneous Graph) 구성: 각 방문을 나타내는 '가상 방문 노드 (DP 노드)'와 해당 방문의 '진단 노드'로 구성됩니다.
   • 연결 구조:
     • 시간 순서대로 방문 노드를 연결 (Temporal edges).
     • 방문 노드와 해당 방문의 진단 노드를 연결.
     • 자기 루프 (Self-loop) 를 추가하여 정보 보존.
   • GAT (Graph Attention Network): 그래프 어텐션을 통해 시간적 흐름에 따른 질병 발전 추세를 학습하고, 이전 방문의 정보를 현재 방문에 전달합니다.

4. 환자 표현 모듈 (Patient Representation, PR):

   • 최종 환자 표현 벡터 ($h_{[CLS]}$) 를 생성하기 위해 세 가지 소스 정보를 통합합니다:
     1. 전체 시퀀스를 요약한 [SEQ] 토큰.
     2. DA 모듈에서 추출된 장기별 질병 진행 토큰들.
     3. DP 모듈에서 추출된 시간적 진행 토큰들.
   • 어텐션 풀링 (Attention Pooling) 메커니즘을 사용하여 각 구성 요소의 상대적 중요도를 동적으로 가중치 부여합니다.

2.2 사전 학습 프레임워크 (Pre-training Framework)

모델의 강건성과 모듈 간 정렬 (Alignment) 을 향상시키기 위해 두 가지 태스크를 결합한 사전 학습을 수행합니다.

• 전역 코드 마스킹 예측 (Global Code Masking Prediction, GCMP):
  • 방문 내 시간 순서의 불확실성을 고려하여, 방문 단위가 아닌 환자 전체 경로 (Trajectory) 수준에서 고유 코드를 마스킹하고 예측합니다.
  • 진단, 약물, 검사, 시술 등 4 가지 코드 유형을 동시에 예측하도록 하여 다양한 다운스트림 작업에 일반화되도록 합니다.
• 조상 코드 예측 (Ancestor Code Prediction, ACP):
  • 마스킹된 진단 코드의 ICD-9 상위 분류 (Chapter) 를 예측하는 보조 태스크입니다.
  • SR 모듈과 DP 모듈이 ICD-9 온톨로지 구조를 이해하도록 유도하여, DA 모듈과의 의미적 정렬을 강화합니다.

---

3. 주요 기여 (Key Contributions)

1. 모델 아키텍처: 진단 코드를 중심으로 장기별 상호작용과 개인별 질병 진행 패턴을 분리하여 모델링하는 그래프 강화 시퀀스 모델 (DT-BEHRT) 을 제안했습니다.
2. 맞춤형 사전 학습: 경로 수준의 코드 마스킹과 온톨로지 기반 조상 코드 예측을 결합하여 모듈 간 정렬을 촉진하고 환자 표현의 강건성을 높였습니다.
3. 종합적 평가 및 해석 가능성: 다양한 임상 예측 작업에서 우수한 성능을 입증했을 뿐만 아니라, DA 및 DP 모듈의 어텐션 메커니즘이 의사의 진단 추론 (장기별 분류 및 시간적 진행) 과 일치함을 사례 연구를 통해 보여주었습니다.

---

4. 실험 결과 (Results)

MIMIC-III, MIMIC-IV, eICU 등 3 개의 공개 EHR 데이터셋에서 다양한 베이스라인 모델 (G-BERT, BEHRT, Med-BERT, HEART 등) 과 비교 평가되었습니다.

• 일반 결과 예측 (General Outcome Prediction):
  • 입원 사망률 (Mortality), 장기 입원 (PLOS), 재입원 (Readmission) 예측에서 대부분의 베이스라인을 능가했습니다.
  • 특히 재입원 예측에서 가장 큰 성능 향상을 보였으며, 이는 재입원이 다양한 요인과 장기적 질병 진행에 의존하기 때문에 DT-BEHRT 의 질병 진행 모델링이 효과적이었음을 시사합니다.
  • MIMIC-III(작은 데이터) 에서 더 두드러진 우위를 보였으며, 데이터 양이 많은 MIMIC-IV 에서는 베이스라인의 격차가 줄어들었으나 여전히 최상위 성능을 유지했습니다.
• 표현형 예측 (Phenotyping Prediction):
  • 다중 라벨 분류 작업에서 Macro-AUPRC 기준으로 전체 코호트 및 방문 횟수가 많은 환자군에서 가장 높은 성능을 기록했습니다.
  • 특히 빈번한 입원 이력을 가진 환자군에서 시간적 의존성을 잘 포착하는 능력이 입증되었습니다.
• Ablation Study:
  • DA 모듈 (장기별 집계) 은 사망률 예측에, DP 모듈 (시간적 진행) 은 재입원 예측에 각각 결정적인 기여를 했습니다.
  • ACP(조상 코드 예측) 태스크는 모델의 전반적인 강건성, 특히 사망률 예측 성능 향상에 크게 기여했습니다.
• 해석 가능성 (Interpretability):
  • 사례 연구에서 DA 모듈이 호흡기 질환 코드를, DP 모듈이 심혈관 질환의 시간적 진행을 의사의 추론과 일치하게 강조하는 것을 확인했습니다.

---

5. 의의 및 결론 (Significance)

DT-BEHRT 는 EHR 데이터 내 이질적인 의료 코드의 역할을 구분하고, 이를 **질병 중심 (Disease-centric)**으로 재구성하여 모델링했다는 점에서 의의가 큽니다.

• 임상적 신뢰성: 단순한 예측 정확도뿐만 아니라, 모델이 어떤 장기 시스템과 시간적 흐름에 주목했는지 해석 가능하게 제공하여 임상 의사들의 신뢰를 얻을 수 있습니다.
• 임상 의사결정 지원: 모델의 어텐션 메커니즘이 의사의 진단 추론 과정 (장기별 분류, 질병 진행 추적) 과 유사하게 작동하므로, 실제 임상 의사결정 지원 시스템 (CDSS) 으로 통합될 가능성이 높습니다.
• 향후 과제: 계산 비용 증가, 단일 방문 환자에서의 그래프 구조 한계, 그리고 약물/시술 코드에 대한 전용 모델링 구조 확장 등이 향후 연구 과제로 제시되었습니다.

요약하자면, DT-BEHRT 는 EHR 데이터의 복잡한 시공간적 특성을 효과적으로 포착하여 예측 성능과 해석 가능성을 동시에 극대화한 혁신적인 모델입니다.