기술 요약: CLIN-JEPA

문제 정의

중환자실 (ICU) 의 전자 건강 기록 (EHR) 모델링은 근본적인 한계에 직면해 있습니다. 기존 접근법은 환자 데이터를 정적 텍스트 시퀀스 (토큰 자기회귀 모델) 로 취급하여 연속적인 생리학적 상태를 모델링하지 못하거나, 각 하위 작업마다 특정 특징 공학과 미세 조정이 필요한 작업별 기반 모델로 취급합니다. 어느 접근법도 환자 궤적을 예측하고 작업별 적응 없이 다양한 위험 예측 작업을 지원할 수 있는 통합된 잠재 표현을 제공하지 못합니다.

로보틱스와 비전에서 잠재 공간 계획을 가능하게 한 Joint-Embedding Predictive Architectures(JEPA) 가 성공적으로 적용되었음에도 불구하고, 이 패러다임을 임상 EHR 에 적용하는 것은 여전히 열린 과제로 남아 있습니다. 이전 JEPA 설계는 임상 시뮬레이션에 적용될 때 다음과 같은 특정 아키텍처적 결함을 겪습니다.

• I-JEPA 및 V-JEPA: 사전 학습 후 예측기를 폐기하여 하위 작업에서 궤적 시뮬레이션이 불가능합니다.
• V-JEPA 2-AC: 예측기를 유지하지만 고정된 인코더에서 훈련합니다. 결과적으로 인코더는 추론에 필요한 "롤아웃 신호"를 인지하지 못해 표현 공간과 예측 동역학 간의 단절이 발생합니다.

공통 JEPA 목적 하에 인코더와 예측기를 단순하게 공동 훈련하는 것은 불안정하며, 표현 붕괴(인코더가 상수 출력으로 퇴화하는 현상) 와 온라인/타겟 드리프트(이동하는 잠재 공간에서 예측기가 오차를 누적하며 발산하는 현상) 에 취약합니다.

방법론: CLIN-JEPA 프레임워크

저자들은 단일 백본을 안정적으로 훈련하여 환자 궤적을 동시에 예측하고 하위 위험 예측 작업을 지원할 수 있는 다단계 공동 훈련 프레임워크인 CLIN-JEPA를 제안합니다.

1. 데이터 표현 및 인코더

• 입력 형식: 환자 입원 기간은 1 시간 단위 (최대 72 시간) 로 상태 - 행동 쌍의 시퀀스로 모델링되며, 정적 인구통계학적 정보가 선행됩니다. 수치 벡터화와 달리, 이 모델은 상태 (생체 징후, 검사 결과, 중증도 점수) 와 행동 (개입, 투여량) 을 나타내기 위해 구조화된 자연어 텍스트를 사용합니다. 이는 인간 가독성을 유지하고, 누락 값을 자연스럽게 처리하며(제거를 통해), 모달리티별 특징 공학을 피합니다.
• 인코더 아키텍처: 시스템은 기본 인코더로 Qwen3-8B언어 모델을 활용합니다. 기본 가중치는 고정되며, 경량 LoRA 어댑터를 훈련하여 텍스트 입력 (인구통계학적 정보 $d$, 상태 $s_t$, 행동 $a_t$) 을 4096 차원 잠재 공간에 매핑합니다.

2. 잠재 궤적 예측기

• 아키텍처: 92M 파라미터의 Transformer 인코더가 4096 차원 잠재 공간 내에서 완전히 작동합니다.
• 메커니즘: 인구통계학적 정보, 상태, 행동이 교차된 시퀀스를 소비합니다. 블록 인과적 어텐션(시간 단계 내에서는 양방향, 시간 단계 간에는 엄격한 인과적) 을 사용하여 다음 잠재 상태 $\hat{z}_{t+1}$를 예측합니다.
• 추론: 예측기는 추론 시 유지됩니다. 예측기는 자기회귀 롤아웃을 수행하여 자신의 연속적 임베딩 예측을 다음 단계의 입력으로 피드백하여 $H$시간 범위의 환자 궤적을 시뮬레이션합니다.

3. 5 단계 공동 훈련 커리큘럼

단순한 공동 훈련의 불안정성을 해결하기 위해 저자들은 5 단계 사전 학습 일정을 도입합니다.

1. 워밍업: 예측기를 고정된 감독 미세 조정 (SFT) 인코더에 대해 단독으로 훈련합니다. 이는 "차가운" 예측기가 즉시 표현 붕괴로 인코더를 끌어당기는 것을 방지합니다.
2. 공동 훈련: 인코더의 LoRA 어댑터를 잠금 해제하고 공유 JEPA 목적 하에 예측기와 함께 훈련합니다. 인코더가 예측기를 위해 역동적으로 기반을 둔 표현을 학습하도록 하기 위해 교사 강제 롤아웃을 사용합니다.
3. 정렬: 온라인 인코더를 다시 고정하고, 지수 이동 평균 (EMA) 타겟 인코더가 온라인 인코더의 매개변수를 부드럽게 따라잡도록 합니다.
4. 하드 동기화: 즉각적인 매개변수 복사로 타겟 및 온라인 인코더를 정렬하여 잔류 불일치를 제거합니다.
5. 최종화: 예측기를 네이티브 자기회귀 롤아웃(추론 조건을 모방) 을 사용하여 안정화된 인코더에서 훈련하여 정제된 표현에 대해 예측기를 완전히 수렴시킵니다.

훈련 목적은 예측된 다음 상태 임베딩과 실제 다음 상태에 대한 타겟 인코더의 임베딩 간의 $\ell_1$ 거리를 최소화하고, 다단계 안정성을 장려하기 위해 단기 롤아웃 손실을 결합합니다.

주요 기여

1. 안정적인 공동 훈련 커리큘럼: 표현 붕괴와 드리프트를 방지하면서 대규모 언어 모델 기반 인코더와 잠재 궤적 예측기를 성공적으로 공동 훈련하는 새로운 5 단계 일정.
2. 통합 텍스트 기반 표현: 보간, 정규화 또는 수동 설계 특징 없이 원시 임상 텍스트를 소비하며, 단일 세트의 잠재 임베딩에서 모든 하위 작업을 지원하는 프레임워크.
3. 예측기를 유지한 최초의 임상 JEPA: 이전 연구들이 예측기를 폐기한 것과 달리, CLIN-JEPA 는 안정적인 자기회귀 환자 상태 시뮬레이션을 가능하게 하기 위해 예측기를 유지합니다.

실험 결과

이 프레임워크는 84,497 건의 입원 기록을 가진 MIMIC-IV ICU데이터셋을 세 가지 축으로 평가했습니다.

1. 훈련 안정성 및 수렴

• 롤아웃 드리프트: 48 시간 자기회귀 범위에서 CLIN-JEPA 의 예측 오차는 수렴하여 1 시간 기준 대비 15.7% 감소했습니다. 반면, 기준선 및 제거 실험 (워밍업 또는 정렬 제거) 은 치명적으로 발산하여 드리프트 증가폭이 +3% 에서 +4951% 에 이르렀습니다.
• 표현 붕괴: 전체 CLIN-JEPA 커리큘럼만이 인코더 출력의 건강한 표준 편차 ($z_{std}$) 를 유지하여 워밍업 단계가 없는 변형에서 관찰된 붕괴를 피했습니다.

2. 잠재 기하학 진단

• 임상 구분력: 인코더는 악화되는 환자 (SOFA 점수 증가 $\ge$ 3) 가 안정적인 환자보다 잠재 공간에서 4.83 배 더 멀리 이동하는 기하학을 학습했습니다. 기준선 인코더는 훨씬 낮은 분리 비율 (2.62 배 및 1.03 배) 을 보였습니다.
• 시간적 동역학: CLIN-JEPA 의 구분력은 관찰 범위가 늘어남에 따라 증가하여 66 시간에서 정점을 찍었으며, 이는 고차원 동적 시스템인 중환자실 진행의 느리고 통합적인 시간 척도를 내부화했음을 반영합니다.

3. 하위 다중 작업 성능

작업별 미세 조정 없이 두 가지 벤치마크 (ICareFM 조기 사건 예측 및 8 가지 이진 위험 작업) 에서 평가되었습니다.

• ICareFM EEP: 평균 AUROC 0.851을 달성하여 가장 강력한 기준선 (LightGBM, 0.827) 및 인코더 제거 실험을 능가했습니다.
• 이진 위험 작업: 평균 AUROC 0.883을 달성하여 LSTM 기준선 (0.865) 을 능가했습니다.
• 예측기의 가치: 과거 정보만 입력하는 것보다 자기회귀 롤아웃 (미래 예측) 을 포함하는 것이 일관되게 성능을 향상시켰으며, 특히 이질적인 궤적 정보의 구성이 필요한 작업 (예: 신장 손상, 패혈증, 탈보상) 에서 두드러졌습니다.

의의

이 논문은 CLIN-JEPA 가 예측기를 유지함으로써 JEPA 사전 학습과 임상 추론 간의 격차를 성공적으로 해소한 최초의 프레임워크라고 주장합니다. 단일 임상 백본이 안정적인 자기회귀 롤아웃, 임상적으로 구분력 있는 잠재 기하학, 그리고 경쟁력 있는 다중 작업 성능을 동시에 제공할 수 있음을 보여줍니다. 이 연구는 고차원 동적 시스템인 중환자실과 같은 장기 시뮬레이션을 위한 표현을 기반으로 하기 위해서는 특정 커리큘럼에 의해 안내된 공유 목적 하에 인코더와 예측기를 공동 훈련하는 것이 필수적임을 확립합니다.