SparseEHR: Scalable Foundation Modeling for Structured EHR via Conditional Computation

Motivation

• Known: BEHRT, Med-BERT, CEHR-BERT 등 structured EHR 전용 transformer 기반 foundation model들이 표현 학습과 downstream prediction 성능을 개선해왔으며, MoE 아키텍처는 stored capacity와 token당 active computation을 분리하여 대규모 모델 확장을 가능하게 하는 것으로 알려져 있다.
• Gap: 기존 dense EHR transformer들은 긴 patient history와 이질적인 code vocabulary(흔한 만성질환 코드 vs 희귀 진단/시술)에 대해 확장 비용이 크며, 기존 medical MoE 연구는 주로 vision-language/multimodal imaging에 집중되어 longitudinal diagnosis/procedure 시퀀스에 대한 sparse routing 연구가 부족하다.
• Why: Conditional computation이 구조화된 EHR의 이질적 데이터 생성 구조(공통 backbone 패턴 vs 희귀하고 전문화된 궤적)에 자연스럽게 부합하는지, 그리고 대규모 사전학습이 서로 다른 의료 시스템(cross-system) 간에 zero-shot으로 전이 가능한지를 확인하는 것은 실용적 population-scale 임상 foundation model 구축에 중요하다.
• Approach: 4개의 dense warm-start layer 뒤에 12개의 MoE layer(4개 routed expert 중 top-1 선택 + 1개 always-active shared expert)를 쌓은 16-layer hybrid dense-to-sparse transformer를 설계하고, OLDW의 약 5천만 명 데이터로 autoregressive next-code objective로 사전학습한 뒤 MIMIC-IV에 fine-tuning 없이 zero-shot 전이 평가를 수행했다.

Achievement

Figure 1. A compact view of the SparseEHR design principle:

1. Zero-shot cross-system transfer 성능 입증: MIMIC-IV next-visit ICD-10 prediction에서 fine-tuning 없이 Recall@10 0.463, Recall@20 0.551을 달성하여, 고정된 preprocessing/code-mapping protocol 하에서도 공통 schema만으로 대규모 structured-EHR 모델의 전이 가능성을 확인했다.
2. Stored capacity와 active computation의 구현 수준 분리 회계: dense FFN, shared expert, routed expert, attention, embedding, router 등 컴포넌트별로 stored parameter와 token당 active parameter를 분리해 보고함으로써, 모델을 serving-speed 개선이 아닌 capacity-transfer trade-off로 명확히 위치시켰다.
3. Hybrid dense-MoE 설계의 학습 안정화: grouped-query attention, rotary positional embedding, expert-utilization 모니터링을 결합해 heterogeneous하고 long-tailed한 임상 코드 데이터에서 sparse routing의 학습 안정성을 확보했다.

How

Figure 1. A compact view of the SparseEHR design principle:

• 입력을 demographics(연령대, 성별, 지역)와 일별 diagnosis/procedure token bundle로 구성된 chronological token sequence로 표현하고, ICD 코드는 canonical prefix로, CPT 코드는 고정 numeric form으로 정규화
• 다음 token을 예측하는 autoregressive objective p(xt+1|x≤t)=softmax(Wht)를 사용하며 vocabulary projection은 입력 embedding matrix와 weight-tied
• 16-layer transformer(dmodel=2048, attention 12 head×170 dim, 2040-dim projected subspace)에서 앞의 D=4개 layer는 dense block, 나머지 12개 layer는 MoE block으로 구성
• MoE router는 4개 routed expert에 대해 logit을 계산하고 top-k(k=1) expert를 선택해 softmax 정규화된 가중치로 결합하며, 여기에 always-active shared expert의 출력을 더해 최종 출력 산출(식 2-4)
• inactive routed expert는 stored capacity로만 계산하고 active computation에는 포함하지 않는 component-level parameter accounting을 수행
• 약 5천만 명의 OLDW 데이터로 causal next-token loss 기반 사전학습(sequence length 2048로 truncate/pad) 후, MIMIC-IV 데이터를 동일 구조로 harmonize하여 라벨 튜닝 없이 zero-shot Recall@K 평가 수행

Originality

• 기존 medical MoE 연구가 주로 vision-language/multimodal imaging에 집중된 것과 달리, longitudinal structured diagnosis/procedure code trajectory에 특화된 sparse routing을 최초로 체계적으로 연구
• dense warm-start layer와 MoE layer를 결합한 hybrid 구조를 제안하여, 공통 backbone 패턴은 dense layer가 안정적으로 학습하고 희귀·전문화된 패턴은 routed expert가 담당하도록 설계
• stored capacity와 active per-token computation을 명확히 분리하는 component-level parameter accounting 프레임워크를 제시하여 MoE 모델의 효율성 주장을 과장하지 않고 명확히 경계짓는 방법론적 태도
• OLDW에서 MIMIC-IV로의 zero-shot cross-system transfer라는, 서로 다른 의료 데이터 시스템 간 실제 전이 가능성을 검증하는 실험 설계

Limitation & Further Study

• workshop paper 형식으로 지면이 짧아 ablation, baseline과의 protocol-matched 비교, 통계적 유의성 검증 등이 충분히 제시되지 않음
• 공개된 baseline들과 protocol이 다르다는 점을 인정하면서도 정량적 비교 결과를 state-of-the-art 주장 없이 "context"로만 제시해 실제 성능 우위 여부가 불명확
• expert-utilization term이 새로운 differentiable load-balancing 알고리즘이 아니라 구현 detail로 취급되어, load balancing의 이론적 안정성 보장이 부족
• 진단/시술 코드 기반 평가만 수행하고 임상 노트, 검사실 데이터 등 다른 모달리티나 실제 임상적 유용성(downstream clinical utility)에 대한 평가가 부재하여 후속 연구로 확장 필요
• 논문 발췌 상 전체 실험 결과(예: all-dense, all-MoE 비교, ablation 표)가 명확히 제공되지 않아 hybrid 설계의 실제 이점이 정량적으로 충분히 입증되었는지 추가 검증 필요

Evaluation

총평: Structured EHR에 대한 MoE 기반 conditional computation 적용과 cross-system zero-shot transfer 검증이라는 아이디어는 흥미롭고 실용적 의의가 있으나, workshop paper 특유의 제한된 지면으로 인해 정량적 비교와 ablation이 부족해 주장의 완전한 검증은 추후 full paper에서 이루어져야 할 것으로 보인다.