Tabular Foundation Models for Clinical Survival Analysis via Survival-Aware Adaptation

기술 요약: 생존 인지형 적응을 통한 임상 생존 분석용 표 형식 파운데이션 모델

문제 정의
사망 또는 질병 진행과 같은 사건 발생 시간(time-to-event)을 예측하는 것은 임상적 의사결정의 핵심적인 과업이며, 일반적으로 생존 분석을 통해 다뤄진다. Cox 비례 위험 모델(Cox proportional hazards)과 같은 고전적 통계 방법이나 DeepSurv, DeepHit과 같은 딥러닝 모델이 널리 사용되고 있지만, 이들은 대개 작업 특화적 훈련, 광범한 하이퍼파라도미터 튜닝, 그리고 방대한 양의 레이블된 데이터를 필요로 한다. 이러한 자원은 실제 임상 환경에서 흔히 부족한 요소들이다. 또한, 이러한 모델들은 관심 있는 사건이 모든 환자에게서 관찰되지 않는 우측 검열(right-censored) 데이터를 처리할 수 있어야 한다.

TabPFN, TabDPT, TabICL과 같은 표 형식 파운데이션 모델(Tabular Foundation Models, FMs)의 최근 발전은 합성 표 형식 태스크로부터 범용적인 표현을 학습함으로써, 타겟 데이터셋에 대한 별도의 훈련 없이도 인컨텍스트 학습(in-context learning)을 통해 강력한 성능을 발휘할 수 있는 유망한 대안을 제시한다. 그러나 이들의 임상 생존 분석에 대한 적용 가능성은 아직 충분히 탐구되지 않았다. 기존의 FM들은 주로 이산 분류(discrete classification)나 회귀(regression)를 위해 설계되었으며, 검열된 시점-사건(time-to-event) 예측을 네이티브하게 지원하지 않는다. 이를 적응시키려는 기존의 시도들은 생존 문제를 이산화된 시간 구간에 대한 일련의 이진 분류 작업으로 재구성하는 방식을 취해왔는데, 이는 계산 오버헤드, 클래스 불균형, 그리고 시간적 모델링의 불일치 문제를 야기한다.

방법론
저자들은 임상 생존 분석에 표 형식 파운데이션 모델을 적용하기 위한 경량 적응 접근법을 제안한다. 핵심 전략은 파운데이션 모델의 사전 학습된 백본(backbone)을 고정된 상태로 유지하면서, 경량화된 생존 인지형 예측 헤드(survival-aware prediction head)만을 훈련하는 것이다. 본 연구는 두 가지 패러다임을 평가한다:

1. 제로샷 인컨텍스트 학습(Zero-Shot In-Context Learning): 사전 학습된 모델을 이산화된 생존 태스크에 직접 적용한다. 각 시간 빈(time bin)에 대해, 모델은 훈련 코호트를 인컨텍스트 서포트로 사용하여 사건이 발생할 확률을 예측한다. 생존 함수는 이러한 확률들로부터 재구성되며, 단조성(monotonicity)을 강제하기 위한 후처리 과정을 거친다.
2. 생존 인지형 적응(Survival-Aware Adaptation, 제안 방식): 사전 학습된 표현에 다중 작업 로지스틱 회귀(Multi-Task Logistic Regression, MTLR) 헤드를 부착한다. 이산 분류와 달리, MTLR은 생존 시퀀스를 유효한 생존 시퀀스에 대한 결합 확률로 모델링한다. 이 접근 방식은 다음과 같다:
   • 생존에서 사건 발생으로의 전이를 나타내는 이진 시퀀스로 생존 결과를 인코딩한다.
   • 유효한 시퀀스에 대해 로그-선형 모델을 정의하며, 이는 검열 시간 이후까지의 생존과 일치하는 시퀀스들을 주변화(marginalizing)함으로써 검열을 자연스럽게 처리한다.
   • 시간적 확장이나 사후 교정 없이 구조적으로 단조성을 강제한다.

저자들은 세 가지 대표적인 아키텍처(TabPFN, TabDPT, TabICL)를 다양한 공공 생존 벤치마크(SUPPORT, METABRIC, GBSG, WHAS, FLCHAIN, SEER, VETERANS)와 두 개의 대규모 ICU 코호트(MIMIC-IV 및 eICU)를 통해 평가한다.

주요 기여

• 적응 프레임워크: 사전 학습된 표현 위에 생존 인지형 헤드(특히 MTLR)를 훈련함으로써, 표 형식 파운데이션 모델을 검열된 시점-사건 예측으로 확장하는 매개변수 효율적인 방법을 소개한다.
• 포괄적 평가: 고전적 통계 모델, 트리 기반 방법, 딥 뉴럴 생존 모델을 포함한 강력한 베이스라인들과 소-중규모 벤치마크 및 대규모 전자 건강 기록(EHR) 데이터셋에 걸쳐 엄격한 비교를 수행한다.
• 실증적 검증: 사전 학습된 표 형식 표현을 생존 인지형 목적 함수와 결합하는 것이 강력한 베이스라인들과 비교하여 경쟁력 있거나 우수한 성능을 보임을 입증하며, 특히 크고 이질적인 데이터셋에서 그러하다.

결과
실험 결과, 생존 인지형 적응이 모든 데이터셋에서 가장 일관된 전략임이 나타났다:

• 성능: MIMIC-IV 데이터셋에서 TabDPT-FT-MTLR 모델은 C-index 0.856을 달성하였으며, 이는 최상의 비-FM 베이스라인(DeepSurv, 0.844) 대비 +1.4%의 상대적 개선을, 최상의 제로샷 모델(0.802) 대비 +6.7%의 개선을 나타낸다. 마찬가지로 eICU에서 TabICL-FT-MTLR은 C-index 0.797을 기록하며 DeepSurv(0.784)와 최상의 제로샷 모델(0.749)을 능가했다.
• 제로샷 vs. 적응: 제로샷 추론은 특히 규모가 작거나 정제된 데이터셋에서 경쟁력이 있으나, 생존 인지형 적응에 의해 지속적으로 상회된다. 저자들은 이를 시간 빈을 독립적으로 취급하고 클래스 불균형 문제를 겪는 제로샷 이산화 재구성의 한계로 돌린다.
• 표현 품질: 사전 학습된 표 형식 백본은 동일한 MTLR 목적 함수와 결합했을 때 표준 MLP 인코더보다 일관되게 우수한 성능을 보였으며, 이는 파운데이션 모델이 생존 태스크로 효과적으로 전이될 수 있는 유용한 귀납적 편향(inductive bias)을 인코딩하고 있음을 시사한다.
• 위험 층화(Risk Stratification): 임상 유용성 분석에서, 생존 인지형 모델은 제로샷 모델에 비해 저위험, 중위험, 고위험 환자 그룹 간의 더 명확하고 안정적인 분리를 생성했다. 적응된 모델은 위험 그룹 간의 조기 발산과 곡선 교차 현상이 적게 나타났으며, 이는 시간 경과에 따른 더 신뢰할 수 있는 위험 순위 산출을 의미한다.

의의 및 주장
본 논문은 표 형식 파운데이션 모델이 경량 생존 인지형 적응과 결합될 때 전통적인 생존 모델링 방식에 대한 실용적이고 효과적인 대안이 될 수 있다고 주장한다. 본 연구의 의의는 다음과 같다:

• 효율성: 이 접근법은 백본의 광범한 태스크 특화 훈련을 제거하여, 고성능 생존 예측을 위한 배포를 단순화하고 데이터 요구량을 줄인다.
• 강건성: 이 방법은 소규모 공공 벤치마크부터 대규모의 이질적인 ICU 코호트에 이르기까지 다양한 임상 도메인과 데이터 체계에서 강력한 일반화 능력을 보여준다.
• 임상적 유용성: 더 정확한 위험 층화와 고위험 환자의 조기 식별을 가능하게 함으로써, 이 접근법은 중환자실 환경에서 보다 시의적절하고 실행 가능한 의사결정을 지원한다.

저자들은 제로샷 추론이 자원이 제한된 환경에서 강력한 베이스라인 역할을 수행하지만, 데이터셋의 크기와 복잡성이 증가함에 따라 생존 인지형 적응이 더욱 유익해지며, 이는 사전 학습된 파운데이션 모델을 수정하지 않고도 검열과 시간적 의존성을 처리하는 원칙적인 방법을 제공한다고 결론짓는다.