하나의 손실로 모든 것을 지배한다 마크드 시간‑이벤트 사전학습으로 EHR 파운데이션 모델 강화

초록

본 논문은 전자건강기록(EHR)의 불규칙한 이벤트와 연속값을 동시에 모델링하는 새로운 사전학습 목표인 ORA(Observed‑Risk‑Association)를 제안한다. ORA는 마크드 포인트 프로세스를 기반으로 각 코드별 시간‑값 분포를 공동 최적화함으로써 기존의 다음 토큰 예측 방식보다 일반화 능력이 뛰어나며, 분류뿐 아니라 회귀·시간‑이벤트 예측에서도 성능 향상을 보인다. Transformer와 Mamba 두 가지 아키텍처에 적용했을 때 평균 10% 이상 성능 개선을 확인했다.

상세 분석

ORA는 EHR 데이터를 “시간, 마크(코드), 값”의 삼중 튜플(t, m, v)로 정의하고, 이를 마크드 포인트 프로세스(marked point process)로 수학화한다. 기존의 다음 토큰 예측(next‑token) 방식은 단일 다음 이벤트만을 예측하도록 설계돼, 환자 기록 전체에 걸친 다중 이벤트와 값(예: 실험실 수치)의 상관관계를 반영하지 못한다. ORA는 각 관찰 시점 j 에서 “다음에 관찰될 모든 코드 m 에 대한 첫 번째 발생 시간 Δtₘⱼ와 값 v, 그리고 관찰 여부 δ”를 동시에 예측하도록 손실을 구성한다. 이는 (1) 희소한 학습 신호 문제를 완화하고, (2) 코드 간 상호 배제 가정을 없애며, (3) 검열(censoring) 정보를 자연스럽게 포함한다는 장점을 제공한다.

구현 측면에서는 연속적인 시간·값 공간을 이산화(discretization)하여 각 코드 m 에 대해 T × V 크기의 확률 행렬 Pₘ(x) 를 출력한다. 여기서 행은 시간 구간, 열은 값 구간을 나타내며, 로그우도는 관측 이벤트에 대해 해당 셀의 로그 확률을, 검열된 경우에는 해당 시간 이전 전체 확률 질량의 로그를 최소화한다. 이 방식은 DeepHit과 유사한 비파라메트릭 접근법을 차용하면서도 마크드 포인트 프로세스의 복합 구조를 그대로 보존한다.

아키텍처 독립성을 강조하기 위해 Transformer와 최신 상태공간 모델인 Mamba에 동일 토크나이저와 ORA 손실을 적용했다. 실험 결과, 두 모델 모두 7개의 이진 분류, 3개의 회귀, 4개의 시간‑이벤트 예측 과제에서 평균 10.7% (Transformer) 및 11.4% (Mamba) 성능 향상을 기록했으며, 특히 희귀 이벤트와 연속값을 포함한 과제에서 기존 next‑token 기반 사전학습 대비 유의미한 개선을 보였다.

또한, 두 개의 대규모 EHR 데이터셋(대형 3차 및 4차 의료기관)에서 교차기관 검증을 수행해 ORA가 데이터 분포 변화에 강인함을 입증했다. 이는 사전학습 목표가 임상 데이터의 구조적 특성을 얼마나 반영하느냐가 downstream 일반화에 결정적 영향을 미친다는 중요한 교훈을 제공한다.

요약하면, ORA는 EHR의 불규칙성, 마크(코드) 다양성, 연속값을 동시에 고려한 통합 손실 함수로, 기존 토큰 기반 사전학습의 한계를 극복하고 다양한 임상 예측 작업에 폭넓게 적용 가능한 새로운 파운데이션 모델 학습 패러다임을 제시한다.