[논문리뷰] - Zero shot health trajectory prediction using transformer

Abstract 

연구 배경 및 목적 

ETHOS(Enhanced Transformer for Health Outcome Simulation)는 의료 데이터를 효과적으로 분석하여 환자 맞춤형 예측과 시뮬레이션을 수행하는 데 중점을 둔 Transformer 기반 모델입니다. 기존 의료 인공 지능 적용의 주요 장애물인 대규모 레이블 데이터 부족 문제를 해결하고, 다중 태스크를 하나의 모델에서 수행할 수 있는 Zero-shot learning 을 도입하였습니다.

 

본 논문의 초반에는 미국의 의료 시스템에 대해 가장 많은 비용이 들지만(most expensive), 의료의 품질과 안전성은 다른 선진국과 비교했을 때 뒤떨어진다고 언급합니다. 전자의료기록(EHR)이 사용되고 있고, 의사결정 지원 기술도 많이 도입되고 있지만, 대부분 규칙 기반(Rule-based)으로 운영되고 있어 효과가 제한적이라고 주장합니다.

 

인공지능은 의료의 퀄리티를 개선할 잠재력을 가진 기술로 떠오르고 있지만, 대규모 레이블링된 데이터셋의 부족으로 인해 개발에 많은 비용과 시간이 소요되며, 적절한 시점에 적합한 의료진에게 권장 사항을 전달할 시스템 역량의 제한이 현재 문제라고 합니다.

본 논문에서는 ETHOS이라는 새로운 방법론을 소개하며, 이 방법이 AI의 광범위한 채택을 제한하는 여러 요소들은 해결할 수 있다고 주장합니다.

 

데이터

ETHOS는 MIMIC-IV(https://physionet.org/content/mimiciv/3.1/) 데이터셋을 활용했으며, 아래의 특징이 있습니다.

• Tokenization : 환자의 건강 기록(Patient Health Timelines, PHT)을 시간 순서에 따라 토큰화하여 구성
• 다양한 데이터 유형 처리 : PHT는 입원, 약물 처방, 실험실 결과 등 다양한 정보를 포함하며 모든 데이터를 정량화하여 토큰화
• 데이터 전처리 최소화 : 데이터 노이즈(입퇴원 날짜 불일치 등)를 정제하지 않고 모델이 자체적으로 처리하도록 설계

모델 구조 

Figure1. ETHOS 모델의 구조 및 흐름

 

a. Extraction of raw patient data from the MIMIC-IV database, encompassing tables of admissions, patient demographics, medical procedures, among others.

 

원시 데이터를 추출하는 단계 a에 대한 설명입니다.

우선 MIMIC-IV 데이터베이스에서 환자 정보를 추출합니다. ( 입원 기록, 처방, 실험실 데이터 등)

해당 모델에서 사용된 주요 데이터 테이블은 아래와 같습니다.

• 환자 기본 정보(Patients)
• 입원 및 퇴원 기록(Admissions)
• 약물 처방(Prescriptions)
• 실험실 결과(Labevents)
• 치료 및 절차 기록(Procedures)
• 진단 정보(Diagnoses)

b The tokenization process, utilizing data from 90% of patients for model training and the remaining 10% for testing, transforms complex medical records into structured PHT for efficient model processing.

 

PHT는 환자의 건강 기록을 시간 순서대로 정리하여 이벤트 단위로 나누고 이를 토큰(Token)으로 변환합니다.

각 토큰은 의료 이벤트를 나타내며 시간 간격 토큰을 삽입하여 이벤트 간의 시간을 표현합니다.

토큰화 과정 :

• 정량 데이터(ex.혈압, SOFA 점수)는 분위수(Quantiles)로 변환하여 범주화합니다.
• 텍스트 기반 데이터(ICD-10, ATC 코드 등)는 계층적으로 토큰화 합니다.

토큰화 된 상태를 보면 환자 의료 데이터를 시간 순서에 따라 배열하고 각 데이터 항목을 토큰으로 변환한 것을 볼 수 있습니다. 그리고 시간 간격을 나타내는 시간 토큰을 삽입합니다.(ex. 5분, 1시간, 6개월 등)

해당 Figure를 통해 ETHOS 모델이 원시 데이터를 토큰화하고, Transformer를 사용해 훈련 및 예측을 수행하는 전체 과정을 명확히 보여줍니다. 특히 제로-샷 학습을 통해 다중 태스크를 처리하는 효율적인 워크플로를 강조합니다.

 

Figure 2는 ETHOS 모델이 MIMIC-IV 데이터를 어떻게 토큰화하고, 이를 임베딩 공간으로 변환하는지 시각적으로 보여줍니다.

 

1. 토큰화(Tokenization)

• PHT 구성:
  • 원시 데이터를 시간 순서에 따라 정렬하고, 이벤트 단위로 토큰화.
  • 각 이벤트(예: 입원, 약물 처방)는 1~7개의 토큰으로 변환.
  • 이벤트 간 시간 간격을 나타내는 시간 토큰 추가.
• 정량 데이터 처리:
  • 혈압, SOFA 점수와 같은 연속형 데이터를 분위수(quantiles)로 변환.
  • 10개 분위수로 나눠 각 범위에 해당하는 값을 특정 토큰으로 표현.

2. 임베딩(Embedding)

• 양적 데이터 임베딩:
  • Fig. 2b에서, 정량 데이터를 나타내는 토큰의 768차원 임베딩을 2D로 축소(PCA 사용).
  • 분위수 토큰 사이의 거리와 패턴은 임상적 데이터의 순서를 학습했음을 보여줌.
• 시간 간격 임베딩:
  • Fig. 2c에서 시간 간격을 나타내는 토큰 임베딩의 분포를 시각화.
  • 짧은 시간 간격(5분1시간)과 긴 간격(36개월)이 모델에서 유사하게 처리되는 경향을 보임.

3. 데이터 분포 및 관계 학습

• 높은 값(예: Q9, Q10)의 임베딩 간 거리가 더 큰 이유:
  • 혈압과 같은 임상적 데이터에서 높은 값 간의 차이가 더 중요한 임상적 의미를 가짐.
• 시간 간격 토큰에서도 유사한 패턴이 관찰되며, 이는 모델이 시간 데이터를 효과적으로 학습했음을 의미.

모델 Summary

ETHOS 연구는 제로-샷 학습(Zero-shot learning)을 통해 다음과 같은 예측 작업에서의 모델 성능을 보여줍니다.

• 입원 및 ICU 사망률 예측
• ICU 체류 기간(LOS) 추정
• 재입원 가능성 예측

또한, 회귀 작업을 수행하여 ICU 입원 전 정보를 사용하여 첫날의 SOFA(Sequential Organ Failure Assessment) 점수를 추정합니다.

• SOFA 점수는 ICU에서 환자의 상태를 모니터링하는 데 사용되며 호흡기, 심혈관, 간, 응고, 신장, 신경계 6개 시스템을 각각 평가하여 24점 사이에 분포하게 됩니다.

ETHOS는 또한 병원 퇴원 시 진단 관련 그룹(DRG) 분류를 예측하며, 이는 771개의 카테고리로 구성됩니다. DRG는 병원 사례를 표준화된 복잡성 기준으로 분류하여 효율적인 환자 관리를 장려하고 품질을 유지합니다.

ETHOS의 제로-샷 학습 접근 방식은 다양한 테스크에서 세부 훈련 없이 정확한 추론을 가능하게 하며, 미래 건강 기록(PHT)을 예측할 수 있는 기능은 의료의 다양한 상황에서 혁신적이고 맞춤화된 응용 프로그램 개발을 가능하게 합니다.

ETHOS와 기존 모델의 차이

ETHOS는 기존의 두 가지 주요 의료 AI 모델 범주와 차별화됩니다:

1. 임상 언어 모델(CLaMs):
   • 임상 텍스트 데이터를 처리하도록 설계된 대형 언어 모델(LLM)로, 의료 노트, 생의학 문헌 등 방대한 데이터셋으로 훈련.
   • 약물명 추출, 의료 대화 요약, 임상 결과 예측 등 다양한 임상 태스크 수행.
2. 전자 의료 기록(EMR) 기반 모델(FEMRs):
   • 구조화된 데이터(예: 인구통계 및 검사 결과)와 비구조화된 데이터(예: 임상 노트 및 영상 보고서)를 분석.
   • 환자 표현 및 결과 예측을 지원하지만 텍스트 생성 기능은 없음.

ETHOS는 라벨링된 데이터나 태스크별 미세 조정 없이도 정확한 예측을 제공하며, 훈련 및 테스트 데이터 선택에 특정 기준을 적용하지 않아, 수백만에서 수억 명의 환자 데이터를 포함하는 대규모 데이터셋에도 확장 가능성을 보여줍니다.

 

Fig. 4: PHT(Patient Health Timeline) 구성 및 토큰화 과정

 

Fig. 4는 ETHOS 모델이 데이터를 정리하고 처리하는 방식을 명확히 보여줍니다. 이 과정의 장점은 다음과 같습니다:

• 시간 순서 보존: 데이터의 시간적 흐름이 유지되어, 모델이 이벤트 간의 시간적 관계를 학습할 수 있습니다.
• 다양한 데이터 처리: 정량 데이터, 텍스트 데이터, 시간 간격 등을 모두 통합하여 처리 가능.
• 효율적 학습: 정제된 PHT는 ETHOS 모델의 Transformer 아키텍처가 효율적으로 학습할 수 있는 기반을 제공합니다.

 

Fig. 5: SOFA 점수 추정 및 DRG 분류에서 ETHOS 모델의 성능

 

a) SOFA 점수 추정의 프로세스(Fig. 5a)

• 입력 데이터:
  • 모델은 환자가 ICU에 입원한 후 첫날의 SOFA 점수를 추정합니다.
  • 입력 데이터는 ICU 입원 시점까지의 환자 건강 이력을 포함합니다.
• 토큰 시퀀스 생성:
  • ETHOS는 세 가지 주요 토큰을 생성합니다:
    1. Admission type 토큰: 환자의 ICU 입원 유형을 나타냄(오렌지색 토큰).
    2. SOFA 토큰: SOFA 점수 예측이 이어질 것을 나타냄.
    3. 분위수 토큰(Quantile Token): SOFA 점수의 분위수에 대한 확률 분포를 예측.
• 점수 계산 방식:
  • 모델은 각 분위수의 확률 분포를 생성하며, 각 분위수의 평균 값(1.0, 3.5, ...)을 사용하여 SOFA 점수를 계산합니다.
  • 데이터셋에 SOFA 값 24가 없기 때문에 점수 범위는 0~23으로 제한됩니다.

b) 실제 값과 예측 값의 상관관계(Fig. 5b)

• 결과 시각화:
  • 실제 SOFA 점수와 ETHOS가 예측한 점수 간의 상관 관계를 그래프로 표시.
  • 모델이 높은 예측 정확도를 보여주며, 실제 값과 예측 값이 강한 양의 상관관계를 가짐.
• 임상적 유용성:
  • SOFA 점수는 중환자실(ICU)에서 환자의 장기 기능 상태를 모니터링하는 데 중요한 지표로, ETHOS가 정확한 점수를 추정할 수 있는 능력을 보여줍니다.

c) DRG 분류 방법

• 입력 데이터:
  • DRG 분류는 환자의 퇴원 시점을 기준으로 수행되며, 환자의 PHT(건강 기록) 데이터를 사용합니다.
• 토큰화:
  • ETHOS는 퇴원 이벤트 관련 토큰(예: 퇴원 유형, 병원 체류 기간, 퇴원 목적) 뒤에 DRG 토큰을 삽입하여 DRG 분류를 수행.
  • DRG 데이터가 훈련 세트에서 제공되지 않을 경우, "DRG_UNKNOWN"이라는 자리표시자 토큰을 사용합니다.
• 예측 방식:
  • 모델은 DRG 카테고리(771개 중 하나)에 대한 확률 분포를 생성합니다.
  • Top-{1, 2, 3, 5} 정확도를 사용하여 모델 성능을 평가.

d) DRG 분류 정확도

• 결과 시각화:
  • Top-1, Top-2, Top-3, Top-5 정확도를 그래프로 표시하여 모델의 예측 성능을 강조.
• 모델의 강점:
  • ETHOS는 퇴원 이벤트 이전의 환자 데이터를 사용하여 높은 정확도로 DRG를 예측.
  • 이는 DRG 분류를 통해 병원의 자원 사용과 효율성을 최적화하는 데 유용합니다.

Fig. 5는 ETHOS가 SOFA 점수 추정과 DRG 분류라는 두 가지 중요한 의료 작업에서 뛰어난 성능을 발휘함을 보여줍니다. 특히, 모델이 시간적 데이터를 활용하여 복잡한 예측 작업을 수행할 수 있음을 강조하며, 이는 ETHOS가 의료 데이터 분석과 자원 관리에 중요한 도구로 활용될 수 있음을 시사합니다.

 

성능 평가

모델은 다음과 같은 임상 작업에서 성과를 보였습니다:

• 중환자실(ICU) 및 입원 사망 예측: AUC 0.921 (입원 사망), AUC 0.927 (ICU 사망).
• ICU 체류 기간 추정: 평균 절대 오차(MAE) 2.262일.
• SOFA 점수 예측: ICU 입원 후 첫날 SOFA 점수 MAE 1.502.
• DRG(Diagnosis-Related Group) 분류: Top-1 정확도 84.8%.
• ICU 재입원 예측: AUC 0.807

Summary

ETHOS의 장점은 시간 기반 시뮬레이션을 통해 미래 건강 경로를 예측하며 다양한 임상 작업을 수행합니다.

SOFA 점수와 DRG 분류에서 높은 정확도를 보이며 다양한 임상 데이터에 적용 가능합니다.

전통적인 머신러닝 모델이나 GPT-4o 보다 뛰어난 성능과 효율성을 보입니다.

 

ETHOS는 입원 사망률, ICU 체류 기간, SOFA 점수 및 DRG 분류와 같은 다양한 임상 태스크에서 뛰어난 예측 성능을 보여줍니다. 특히, 제로-샷 학습 접근 방식을 통해 추가적인 데이터 라벨링이나 미세 조정 없이 높은 효율성과 확장성을 제공합니다. ETHOS는 기존 모델에 비해 임상적 유용성과 정확도를 개선한 강력한 도구임을 입증합니다.

 

 

Insight

저는 해당 논문에서 EHR 데이터 토큰화에 대해 흥미롭게 봤습니다. 

대규모 전자의료기록을 트랜스포머 모델에 학습 가능한 형태로 변환하기 위해 토큰화하는 과정은 중요한 단계입니다. 

단순 토큰화가 아닌 시간적 특성을 반영하여 모델 학습에 최적화된 구조를 제공하는 부분이 인상깊었습니다. 

 

 

EHR 데이터는 정량적 데이터(혈압, 검사 수치)와 비정형 데이터(진단 코드, 약물 코드)를 포함하고 있고, 의료 이벤트는 시간적 맥락을 유지해야 의미있는 예측이 가능합니다. 그런데 해당 데이터는 잘 정리되지 않은 경우가 많습니다. 

 

ETHOS는 데이터 전처리를 최소화하여 노이즈 데이터도 처리 가능한 모델로 설계하고 시간적 맥락을 유지하기 위해 데이터를 PHT로 변환하였죠. 

 

이벤트 기반 토큰화는 해당 논문의 토큰화 과정의 핵심입니다. 

모든 의료 이벤트(입원, 검사, 약물 처방 등) 를 토큰 단위로 변환하고 각 이벤트는 하나 이상의 토큰으로 표현되며, 세부 정보까지 포함합니다. ( ex. 약물 처방 이벤트는 약물 코드, 용량, 투여 방식 등을 각각의 토큰으로 변환 ) 

 

혈압, SOFA 점수와 같은 연속형 데이터를 10개의 분위수로 나눠 범주화하는 방법을 사용하는데 연속형 데이터의 세부 값을 유지하면서도 모델 학습 효율성을 높입니다. 

 

13개의 시간 토큰으로 시간 간격(5분, 1시간, 6개월)을 표현하여 긴 시간 간격은 반복된 시간 토큰으로 세분화 합니다. 

이러한  방식의 장점은 시간적 맥락을 학습할 수 있으며, 긴 간격과 짧은 간격을 모두 적절히 처리할 수 있습니다. 

 

고정 정보(Static Data)는 환자의 성별, 연령, 인종 등 시간에 따라 변하지 않는 정보입니다. 이런 정보는 PHT의 시작에 고정 토큰으로 삽입합니다. 예를 들어 '45-50 years', 'female' 과 같은 환자의 고유한 특성이죠. 

 

 정리하자면 시간 토큰의 기술적 설계는 해당 논문에서 얻을 수 있는 가장 좋은 인사이트 같습니다. 

범주화의 이점이 있고, 긴 시간 간격을 세분화 하고, 시간 토큰의 일관성을 적용하는 것은 추후 활용성이 높을 것 같습니다.