Retrieval-Augmented Foundation Model Enhances Risk Prediction Using Electronic Health Records

Motivation

• Known: 기존 EHR foundation model들은 예측 시점 이전의 특정 시간 윈도우 내 정보만 관련 있다고 가정하거나, lab·vital·diagnosis 등 일부 feature subset만으로 충분하다고 가정하여 환자 궤적을 모델링해왔다.
• Gap: 이러한 고정 윈도우 및 균일 집계 방식은 계산 제약과 long-range irregular trajectory 모델링의 어려움에서 비롯된 것으로, 만성 질환 패턴이나 여러 encounter 간 상호작용과 같이 임상적으로 중요한 장기 이력 정보를 누락시킬 수 있다는 한계가 있다.
• Why: 환자 상태는 누적된 longitudinal 정보에 의해 형성되므로, 전체 환자 이력을 효율적으로 통합하여 추론할 수 있는 모델은 EHR 기반 위험 예측의 정확도와 임상적 신뢰성을 크게 높일 수 있다.
• Approach: NLP 분야의 Retrieval-Augmented Generation(RAG) 성공에 착안하여, 과거 EHR 이벤트를 검색 가능한 지식으로 취급하고 longitudinal modeling 문제를 retrieval-augmented 문제로 재구성한다. Query segment와 전체 history를 clinically relevant chunking으로 vector database에 구성하고, prototype-guided retrieval module로 관련 chunk를 식별·가중 통합하는 EHR-RAGp를 제안한다.

Achievement

1. EHR vector database 구축 프레임워크: 임상적으로 의미 있는 chunking 전략을 지원하는 longitudinal EHR trajectory용 vector database 구축 방법을 제안했다.
2. Prototype-guided retrieval 메커니즘: 학습 가능한 prototype 공간에서의 정렬(alignment)을 기반으로 관련 이력 chunk를 통합하여 long-range dependency 모델링을 가능하게 하는 완전 미분 가능한 soft weighting 기법을 도입했다.
3. State-of-the-art 성능 달성: EHR-RAGp가 ICU-Readmit 30d, In-hospital Mortality, Long LOS 7d, 1YR Mortality 등 다중 임상 예측 task에서 DescEmb, GenHPF, Med-BERT, CEHR-BERT, BEHRT, Hi-BEHRT, EHRMamba, REMed 등 기존 EHR foundation model 대비 일관되게 AUROC/AUPRC에서 우수한 성능을 보였다.
4. 기존 모델과의 결합 가능성: EHR-RAGp를 Med-BERT, CEHR-BERT 등 기존 baseline에 결합했을 때도 모든 task에서 성능 향상이 관찰되어 프레임워크의 범용성을 입증했다.

How

Figure 1. Overview of EHR-RAGp.

• 문제 정의: 환자 p의 history Hp를 방문(visit) 시퀀스로 표현하고, 각 방문은 (clinical event, value, timestamp) 튜플로 구성. Index visit의 query segment Q(τ)_p와 이전 history H(<τ)_p를 이용해 outcome ŷ(τ)_p를 예측.
• 2단계 retrieval: (1) frozen pretrained retriever Rβ로 query와 history chunk를 임베딩해 cosine similarity 기반 top-M candidate chunk를 coarse하게 선택.
• Shared backbone encoding: 공유 encoder fα로 query와 top-M chunk를 동일 임베딩 공간에 재인코딩하여 semantic alignment 보장.
• Prototype-guided alignment: L개의 learnable prototype vector P를 두고, softmax(xP^T/Tπ)로 query와 chunk 각각의 prototype 분포 πq, πi를 계산 (Tq < Th로 query는 sharp, history는 broad하게 설정).
• 정렬 점수 및 가중치 계산: cross-entropy 기반 alignment score αi = -Σπq log πi를 계산하고, temperature Ts를 적용한 softmax로 정규화된 importance weight wi를 산출하여 모든 검색된 chunk를 미분 가능하게 soft weighting 후 fusion하여 zfused를 생성, end-to-end fine-tuning 수행.

Originality

• EHR longitudinal modeling을 retrieval-augmented 문제로 재구성한 최초의 시도 중 하나로, NLP의 RAG 패러다임을 임상 데이터의 이질성·불규칙성에 맞게 확장.
• 하드한 top-k selection 대신 learnable prototype space에서의 분포 정렬(cross-entropy alignment)을 이용한 soft, 완전 미분 가능한 weighting 메커니즘 도입.
• Query와 history에 서로 다른 temperature(Tq < Th)를 적용해 query는 sharp하게, history는 broad하게 prototype에 정렬시키는 설계적 차별성.
• 기존 EHR foundation model에 plug-in 형태로 결합 가능한 범용 retrieval augmentation 프레임워크 제시.

Limitation & Further Study

• Workshop 발췌 논문으로 전체 실험 세부사항(데이터셋 규모, ablation, hyperparameter 민감도 등)이 충분히 제시되지 않아 재현성 및 강건성 검증이 제한적이다.
• Coarse retrieval에 사용되는 frozen pretrained retriever Rβ의 선택이 성능에 미치는 영향과 다양한 chunking 전략 간 비교 분석이 부족하다.
• 단일 기관(예: MIMIC 계열로 추정) 데이터에 기반한 평가로 보이며, 다기관·이종 EHR 시스템에 대한 일반화 성능 검증이 필요하다.
• 계산 비용(vector database 구축 및 retrieval 오버헤드) 대비 성능 이득에 대한 정량적 분석이 추가되면 실제 임상 배포 가능성을 평가하는 데 도움이 될 것이다.

Evaluation

총평: 고정 윈도우와 균일 집계라는 EHR foundation model의 근본적 한계를 retrieval-augmented 패러다임으로 해결하려는 참신하고 실용적인 접근으로, 다중 임상 task에서 일관된 성능 향상을 보여 의미 있는 기여를 한다고 평가된다.