[논문 리뷰] MOTOR: A Time-To-Event Foundation Model For Structured Medical Records

Ethan Steinberg, Jason Fries외 2명|arXiv (Cornell University)|2023. 01. 09.

Machine Learning in Healthcare인용 수 7

한 줄 요약

MOTOR는 5,500만 건의 환자 기록(90억 건의 임상 이벤트)을 대규모로 사전학습한 자기지도 학습, 시간-이벤트(TTE) 기초 모델로, 수천 개의 TTE 과제를 활용해 강력한 시간적 표현을 학습하는 새로운 사전학습 목표를 사용한다. 이 모델은 최신 기술 대비 시간에 따라 변하는 C-통계치를 평균 4.6% 향상시키며, 레이블 효율성은 95% 향상시키고, 다양한 의료 데이터셋에서 뛰어난 전이 가능성과 강건성을 보이며, 의료 TTE 예측을 위한 첫 공개된 기초 모델로 기록된다.

ABSTRACT

We present a self-supervised, time-to-event (TTE) foundation model called MOTOR (Many Outcome Time Oriented Representations) which is pretrained on timestamped sequences of events in electronic health records (EHR) and health insurance claims. TTE models are used for estimating the probability distribution of the time until a specific event occurs, which is an important task in medical settings. TTE models provide many advantages over classification using fixed time horizons, including naturally handling censored observations, but are challenging to train with limited labeled data. MOTOR addresses this challenge by pretraining on up to 55M patient records (9B clinical events). We evaluate MOTOR's transfer learning performance on 19 tasks, across 3 patient databases (a private EHR system, MIMIC-IV, and Merative claims data). Task-specific models adapted from MOTOR improve time-dependent C statistics by 4.6% over state-of-the-art, improve label efficiency by up to 95% ,and are more robust to temporal distributional shifts. We further evaluate cross-site portability by adapting our MOTOR foundation model for six prediction tasks on the MIMIC-IV dataset, where it outperforms all baselines. MOTOR is the first foundation model for medical TTE predictions and we release a 143M parameter pretrained model for research use at [redacted URL].

연구 동기 및 목표

• 전자 건강 기록(EHR)의 시간-이벤트(TTE) 모델링에서 레이블이 부족한 문제를 해결함으로써, 모델의 일반화 능력과 정확도 향상을 도모한다.
• 대규모 비라벨 EHR 데이터로부터 광범위하게 유용한 시간적 표현을 학습하는 기초 모델을 개발하여, 다양한 의료 예측 과제에서 전이 학습 성능을 향상시킨다.
• 기존의 자동회귀 사전학습 목표의 한계를 극복함으로써, 의료 TTE 모델링에서 핵심적인 장기적 시간 간격을 포착하지 못하는 문제를 해결한다.
• 메모리 사용량을 최대 35% 감소시키면서도 수천 개의 TTE 사전학습 과제를 지원할 수 있는 확장성 있고 효율적인 아키텍처를 설계한다.
• 재현 가능한 연구를 촉진하고 의료 TTE 모델링 분야의 진전을 가속화하기 위해 공개된 143M 파라미터 사전학습 MOTOR 모델을 제공한다.

제안 방법

• 수천 개의 TTE 과제(예: 특정 진단 또는 절차까지의 시간 예측 등)를 사용하여 사전학습 신호로 삼는 새로운 자기지도 사전학습 목표를 제안함으로써, 다양한 임상 결과에 일반화 가능한 표현을 학습한다.
• 공유된 시간 의존 가중치와 희박성 구조를 갖춘 Transformer 기반의 조각별 지수 모델 아키텍처를 도입하여, 메모리 사용량을 줄이며 대규모 TTE 사전학습을 효율적으로 스케일링한다.
• 임상 코드 간 계층적 온톨로지 관계를 고려한 조건부 샤논 엔트로피를 사용하여 사전학습 과제를 선별함으로써, 부모 코드에 기반한 예측 불확실성이 높은 코드를 우선적으로 선별한다.
• 대조 학습 방식의 목표를 적용하여, 사전학습 중에 다양한 임상 이벤트에 대한 시간-이벤트 결과를 동시에 예측하도록 모델을 학습시킨다.
• 다음으로 적용되는 TTE 과제에 대해 사전학습된 MOTOR 인코더를 미세조정함으로써, 최소한의 레이블 데이터로도 전이 학습을 가능하게 한다.
• 조건부 엔트로피 기반의 선택 전략을 적용하여, 빈도와 온톨로지 구조를 균형 있게 고려한 8,192개의 가장 정보적인 임상 코드를 사전학습에 활용한다.

Figure 1: MOTOR architecture and TTE pretraining tasks, as applied to a patient. Note how tasks (survival curves) vary at each time step $j$ , e.g., atorvastatin risk increases given a history of cardiovascular disorders.

실험 결과

연구 질문

• RQ1 대규모 비라벨 EHR 데이터에서 자기지도 기초 모델을 사전학습하면, 다양한 의료 과제에서 시간-이벤트 예측 성능을 크게 향상시킬 수 있는가?
• RQ2 수천 개의 TTE 과제를 활용하는 사전학습 목표가, 의료 시퀀스의 유용한 시간적 표현을 학습하는 데 있어 표준 자동회귀 사전학습보다 우수한가?
• RQ3 MOTOR는 얼마나 높은 레이블 효율성을 제공하여, 최소한의 레이블 데이터로도 다음 과제에서 뛰어난 성능을 달성할 수 있는가?
• RQ4 MOTOR는 다양한 환자 집단과 의료 시스템 간의 시간적 분포 이동에 대해 얼마나 강건한가?
• RQ5 재학습 없이도 하나의 사전학습된 MOTOR 모델이 EHR과 청구 데이터를 포함한 여러 데이터셋에서 효과적으로 일반화될 수 있는가?

주요 결과

• MOTOR는 최신 기술 대비 19개의 다양한 의료 예측 과제에서 평균적으로 시간에 따라 변하는 C-통계치를 4.6% 향상시켰다.
• MOTOR는 레이블 효율성에서 최대 95% 향상을 기록하여, 다음 과제에서 최고 성능에 도달하기 위해 훨씬 적은 레이블 예제가 필요하다.
• 모델은 다양한 환자 집단과 데이터 소스 간의 시간적 분포 이동에 대해 강력한 강건성을 보이며 성능 유지 능력을 입증했다.
• 제안된 TTE 사전학습 목표는 표준 자동회귀 사전학습 대비 다음 과제 성능에서 2.0% 향상되어 효과성을 입증했다.
• MOTOR는 MIMIC-IV 데이터셋에서 6개의 예측 과제에서 최신 기술 수준의 성능을 달성했으며, 모든 베이스라인 대비 사이트 간 전이 학습에서 뛰어난 성능을 보였다.
• 공개된 143M 파라미터 MOTOR 모델은 정형 의료 기록에서 시간-이벤트 예측을 위한 첫 공개된 기초 모델이다.

Figure 2: Per-time horizon performance differences (in time-dependent C statistic) when replacing autoregressive pretraining with TTE pretraining. Blue indicates better performance with TTE pretraining and red indicates worse.