Pretraining EHR Foundation Models with Patient-Aware Sampling

Motivation

• Known: 언어모델링에서 유래한 표준 사전학습 방식은 문서(또는 환자 궤적)를 하나의 전역 토큰 스트림으로 concatenation한 뒤 고정 길이 윈도우를 샘플링하는 Global Stream 방식이며, ETHOS, Foresight, CoMET 등의 자기회귀 EHR 파운데이션 모델들이 이 방식을 그대로 차용해 왔다.
• Gap: EHR 데이터는 환자별 궤적 길이가 log-normal 분포를 따르며 매우 불균등한데, Global Stream 방식은 한 윈도우에 여러 환자의 토큰이 섞이거나 긴 기록을 가진 환자가 더 많은 최적화 업데이트에 기여하는 편향을 유발할 수 있음에도, 기존 연구들은 사전학습 시퀀스 구성 방식 자체를 주요 설계 변수로 명시적으로 연구하지 않았다.
• Why: 자기회귀 EHR 파운데이션 모델은 zero-shot 임상 예측을 지원하는 유망한 패러다임인데, 시퀀스 구성 방식에 내재된 편향을 이해하고 제어하는 것은 이러한 모델들의 임상 예측 신뢰성과 성능을 개선하는 데 근본적으로 중요하다.
• Approach: 토크나이제이션과 모델 아키텍처를 고정한 상태에서 Global Stream, Patient Chunks, Patient Sampling 세 가지 시퀀스 구성 방식을 비교하고, MIMIC-IV v2.2 및 v3.1의 다운스트림 임상 과제에서 성능을 평가한다.

Achievement

Figure 4. ∆positive score on ICU Mortality (MIMIC-IV v2.2)

1. Patient Sampling 기법 제안: 환자를 먼저 확률적으로 선택한 뒤 해당 환자 내에서 윈도우를 샘플링하는 방식으로, 파라미터 α를 통해 환자 균등 샘플링(α=0)과 궤적 길이 비례 샘플링(α=1) 사이를 매끄럽게 조절할 수 있는 방법을 설계함.
2. 성능 개선 입증: MIMIC-IV v2.2 및 v3.1 데이터셋에서 ICU Mortality, ICU Readmission, ICU Admission, Hospital Mortality 등 다운스트림 임상 과제에 대해 Patient Sampling(특히 α=0.6)이 Global Stream 베이스라인보다 Macro AUROC와 Macro AUPRC에서 우수한 성능을 보임을 실험적으로 확인함.
3. 설계 변수로서의 시퀀스 구성 중요성 규명: 자기회귀 EHR 파운데이션 모델 맥락에서 사전학습/검증 시퀀스 구성 방식을 최초로 체계적으로 탐구하여, 이 선택이 성능에 중요하지만 그동안 충분히 탐구되지 않은 설계 축임을 제시함.

How

• 데이터셋 D = {x(i)}를 N명의 환자 각각의 EHR 토큰 시퀀스로 정의하고, decoder-only transformer를 next-token prediction objective로 고정 길이 S의 윈도우에 대해 학습함.
• Global Stream: 모든 환자 시퀀스를 하나의 스트림 z로 concatenation하고 stride r로 시작 인덱스 집합 TGS를 정의하여 윈도우를 구성함(여러 환자 토큰이 섞일 수 있음).
• Patient Chunks: 환자별로 독립적으로 결정론적 청크 시작 인덱스 집합 Bi를 정의하고, 마지막 청크는 시퀀스 끝에 right-align하여 패딩 처리함(환자 경계는 보존하나 길이 불균형은 유지).
• Patient Sampling: 환자 i를 pα(i) = |W(i)|^α / Σ|W(j)|^α 확률로 샘플링(α∈[0,1]로 균등↔길이비례 조절)하고, 해당 환자 내에서 시작 인덱스를 균등하게 샘플링(음수 인덱스 클리핑으로 다양한 left-context 길이 확보)하여 윈도우를 구성함(학습 시에만 사용).
• MIMIC-IV v2.2/v3.1 + ED v2.2 데이터셋에서 세 방법과 α 값(0, 0.2, 0.4, 0.6, 0.8, 1)을 조합한 pretraining/validation 구성을 비교하고, ICU Mortality/Readmission/Admission, Hospital Mortality의 AUROC/AUPRC 및 Macro 지표로 평가함.

Originality

• 언어모델링에서 차용해 별다른 검토 없이 사용되던 Global Stream 방식의 EHR 데이터 특유의 문제(환자 경계 혼합, 긴 궤적 환자에 의한 편향)를 명시적으로 지적하고 정식화함.
• 환자 선택과 윈도우 선택을 분리하여 α라는 단일 파라미터로 학습 신호의 환자 간 분포를 연속적으로 제어할 수 있는 Patient Sampling이라는 새로운 확률적 시퀀스 구성 기법을 제안함.
• 자기회귀 EHR 파운데이션 모델링에서 사전학습 시퀀스 구성을 명시적 설계 변수로 다룬 최초의 연구라고 저자들이 주장함.

Limitation & Further Study

• 워크숍 논문(발췌본) 형태로 실험 설명과 모델/토크나이제이션 세부사항, ablation의 폭이 제한적이며, 성능 개선폭이 대체로 크지 않고 α에 따른 경향성이 비단조적(예: α=0.6에서 최고이나 α=0.8에서 하락)이어서 최적 α 선택에 대한 이론적 근거가 부족함.
• MIMIC-IV라는 단일 기관(단일 데이터 소스 계열)에서만 검증되어 다른 EHR 데이터셋이나 다른 모델 아키텍처(예: ETHOS, CoMET 등 실제 대규모 파운데이션 모델)로의 일반화 가능성이 명확하지 않음.
• 검증(validation) 시퀀스 구성에 대한 영향 분석이 제한적이며, Patient Sampling의 계산 비용(non-deterministic 샘플링에 따른 학습 효율성/재현성) 논의가 부족하여 후속 연구에서 다양한 데이터셋과 모델 스케일, 그리고 α 선택에 대한 원칙적 기준을 규명할 필요가 있음.

Evaluation

총평: EHR 파운데이션 모델 사전학습에서 간과되어 온 시퀀스 구성 문제를 명확히 짚어내고 실용적 해결책을 제시한 흥미로운 워크숍 논문이지만, 실험 규모와 분석 깊이가 제한적이어서 후속 연구를 통한 검증이 필요하다.