EveryQuery: Zero-Shot Clinical Prediction via Task-Conditioned Pretraining

Motivation

• Known: EHR 데이터에 대한 autoregressive foundation model들은 patient history로부터 다수의 synthetic future trajectory를 생성하고 그 통계를 집계함으로써 zero-shot 임상 예측을 수행할 수 있음이 알려져 있다.
• Gap: 기존 AR 기반 방법들은 (1) K개의 trajectory를 매 예측마다 시뮬레이션해야 하므로 연산 비용이 크고, (2) rare event에서는 기대 count Kp가 작아 통계적으로 노이즈가 심하며, (3) patient history만을 입력으로 받아 task-agnostic한 표현만 생성하므로 promptable하지 않다는 세 가지 근본적 한계가 있다.
• Why: 임상 현장에서는 다양한 예측 과제에 대해 즉각적이고 신뢰도 높은 zero-shot 추론이 필요한데, 특히 임상적으로 중요한 rare event에 대한 정확도와 실시간성(추론 속도)이 실용적 배포에 핵심적이기 때문이다.
• Approach: EveryQuery는 patient history x와 target code 및 prediction horizon으로 구성된 구조화된 query q=(c, Δt)를 함께 입력받아, bidirectional transformer(ModernBERT-base)를 통해 p(y|x,q)를 직접 예측하는 task-conditioned pretraining을 제안한다.

Achievement

Figure 2. Pareto frontier of zero-shot AUC vs. inference speed. EveryQuery offers significant improvements in zero-shot

1. Zero-shot 성능 우위: MIMIC-IV, AMC, NWICU 세 데이터셋 전반에서 AR baseline 대비 각각 84.6%, 58.0%, 61.3%의 win rate를 기록했으며(Wilcoxon signed-rank p<10^-4), MIMIC-IV에서는 mean ΔAUC +0.17을 달성했다.
2. Rare event에서의 강건성: AR 모델의 성능이 event prevalence와 강하게 결합되어 rare event에서 취약한 반면, EveryQuery는 prevalence-invariant한 성능을 보여 임상적으로 중요한 희귀 사건 예측에서 특히 우위를 보였다.
3. 추론 효율성: 단일 forward pass로 예측을 수행함으로써 AR baseline 대비 약 5,500~10,000배 빠른 추론 속도를 달성했다.
4. 일반화 능력: held-out code와 held-out duration(훈련 시 가장 긴 horizon을 넘어서는 extrapolation 포함)에 대해서도 일반화 성능을 보였다.
5. Prompt specificity: task와 patient의 모든 축에 걸쳐 풍부하게 구조화된 embedding 공간을 형성하여 query가 표현 기하에 직접 반영됨을 보였다.

How

Figure 2. Pareto frontier of zero-shot AUC vs. inference speed. EveryQuery offers significant improvements in zero-shot

• 환자 기록을 MEDS 표준에 따라 토큰화된 시퀀스 x=(x1,...,xL)로 표현하고, query를 q=(c, Δt)로 정의하여 이진 라벨 y(q,x)를 "코드 c가 tL부터 tL+Δt 이내에 발생하는가"로 자동 계산한다.
• 기록이 tL+Δt 이전에 끝나는 경우 censoring으로 처리하며, 이를 위한 전용 head(ŷcens)를 별도로 둔다.
• 아키텍처는 from scratch로 학습된 bidirectional transformer인 ModernBERT-base를 사용하며, query의 code는 공유 code embedding layer로, duration은 작은 MLP를 통해 연속적인 임베딩으로 인코딩되어 두 개의 query token으로 patient history 앞에 prepend된다.
• 최종 layer의 code-token 위치 hidden state를 task-conditioned representation으로 삼아 두 개의 MLP로 ŷoccurs와 ŷcens를 각각 산출한다.
• 학습 시 각 patient history에 독립적으로 샘플링된 query를 짝지어 사용하며, code는 어휘의 부분집합에서 균등 샘플링, duration은 1일~5년 범위의 log-uniform 분포에서 샘플링한다.
• 손실 함수는 L = Lcens + λLoccurs로, Lcens는 censoring 여부에 대한 BCE, Loccurs는 censoring되지 않은 샘플에 대해서만 마스킹되는 BCE이다.
• MIMIC-IV, NWICU, AMC 세 개의 MEDS 포맷 데이터셋에서 각 100개의 무작위 코드와 8개의 horizon(1,3,6개월; 1-5년)을 조합한 수백 개의 (code, duration) task로 MEDS-EIC-AR(K=20 trajectory 생성 AR baseline)과 비교 평가한다.

Originality

• Autoregressive 방식의 궤적 생성 및 집계 패러다임을 완전히 우회하고, 임상 예측 과제를 (code, duration) 구조화된 query language로 형식화하여 이를 직접 pretraining objective로 삼는 task-conditioned pretraining 패러다임을 제안한 점이 독창적이다.
• Query를 fixed output head가 아닌 model input으로 취급함으로써, instruction tuning 및 multitask EHR pretraining과 관련되면서도 이를 하나의 통합된 구조적 입력 포맷과 direct prediction objective로 일반화한 점이 신선하다.
• Bidirectional attention을 통해 query token과 patient token이 서로 주고받으며 attend하도록 하여 task-conditioned representation을 자연스럽게 형성하는 아키텍처 설계가 특징적이다.

Limitation & Further Study

• 논문에서 다루는 query language는 단일 code의 이진 발생 여부와 horizon만을 표현하는 매우 단순한 형태로, 복합 조건(다중 코드, 논리 연산, 값 기반 조건 등)을 포함하는 richer query에 대한 확장이 명시적으로 향후 과제로 언급된다.
• private AMC 데이터셋에 대한 세부 정보(전처리, 검증 재현성)가 제한적으로 공개되어 재현성 검증이 어려울 수 있다.
• 비교 대상이 단일 AR baseline(MEDS-EIC-AR)에 국한되어 있어, 다른 유형의 zero-shot 또는 supervised 접근법과의 폭넓은 비교가 부족하다.
• 발췌된 본문에서는 rare event에 대한 성능 우위의 통계적 near-term 메커니즘(즉 prevalence-invariance의 이론적 근거)이 실험적으로만 제시되고 이론적 분석은 제한적이다.

Evaluation

총평: AR 기반 궤적 샘플링의 근본적 한계를 구조화된 query 조건화 pretraining으로 우회한다는 아이디어가 명확하고 실험적으로 잘 뒷받침되어, EHR foundation model 연구에 실용적이고 의미 있는 기여를 하는 논문이다.