Tabular Foundation Models Can Do Survival Analysis

기술 요약: 표 형식 파운데이션 모델의 생존 분석 수행 능력

문제 정의
생존 분석(survival analysis) 또는 시간-사건 분석(time-to-event analysis)은 단순히 사건의 발생 여부를 넘어, 관심 있는 사건이 '언제' 발생하는지를 모델링하는 것을 목표로 한다. 이 영역의 결정적인 과제는 **우측 검열(right-censoring)**이다. 이는 관찰 기간이 사건이 발생하기 전에 종료되어 실제 사건 발생 시간을 알 수 없는 상황을 의미한다. 표 형식 파운데이션 모델(Tabular Foundation Models, TFMs)은 분류 및 회귀 작업에서 XGBoost와 같은 방법론을 능가하며 소규모 데이터 환경에서 놀라운 성공을 거두어 왔으나, 이를 생존 분석에 적응시키는 것은 간단하지 않다. 표준 회귀 아키텍처는 검열을 본래적으로 처리할 수 없으며, 기존의 생존 특화 모델들은 대개 단일 데이터셋을 위해 처음부터 학습해야 하므로 저데이터 환경에서의 활용도가 제한적이다. 또한, TFM은 일반적으로 정확도나 평균 제곱 오차(MSE)로 평가되지만, 생존 모델링에는 시간 경과에 따른 위험 역학을 특징짓기 위한 잘 보정된 확률이 필요하다.

방법론
저자들은 별도의 생존 특화 사전 학습, 미세 조정 또는 아키텍처 변경 없이, 인컨텍스트 학습(in-context learning)을 통해 기성(off-the-shelf) TFM을 활용할 수 있도록 생존 분석을 일련의 이진 분류 문제로 재구성하는 프레임을 제안한다.

1. 시간 이산화(Temporal Discretization): 연속적인 시간 지평을 $K$개의 이산 구간으로 분할한다.
2. 누적 실패 공식화(Cumulative Failure Formulation): 특정 구간 내에서 사건이 발생할 확률인 이산 시간 위험률(discrete-time hazard rate)을 모델링하는 전통적인 방식과 달리, 이 방법은 **누적 실패 확률(cumulative failure probability)**을 직접 모델링한다. 각 시점 $t_k$에 대해, 레이블 $Y_{i,k}$는 $1(T_i \le t_k)$로 정의되며, 이는 시점 $t_k$까지 사건이 발생했는지 여부를 나타낸다.
3. 검열 처리: 검열된 관측치는 검열 시점 이후의 레이블이 누락된 사례로 취급된다. 훈련 목적 함수는 관찰된 레이블(즉, $t_k < C_i$인 경우)에 대해서만 제한적으로 적용되는 표준 이진 교차 엔트로피 손실(binary cross-entropy loss)을 사용한다.
4. 인컨텍스트 학습: 재구성된 데이터셋(특징, 시점 및 이진 레이블)은 TFM(예: MITRA, TabPFN) 또는 표준 분류기에 입력된다. 모델은 공변량(covariates)이 주어졌을 때 특정 시점 $t_k$까지의 실패 확률을 학습한다.
5. 동적 확장(Dynamic Extension): 시변 공변량(time-varying covariates)을 갖는 동적 설정의 경우, 모델은 $t_k$까지의 관측 이력을 바탕으로, 대상자가 $t_k$까지 생존했다는 조건하에 미래 지평 $\Delta$에 대한 누적 실패 확률을 예측한다.

이론적 보장
논문은 **조건부 독립 검열(conditionally independent censoring)**이라는 표준 가설 하에, 제안된 이진 교차 엔트로피 손실을 최소화하는 것이 훈련 세트의 크기가 무한히 커짐에 따라 실제 기저 생존 확률을 회복한다는 이론적 증명(정리 4.2)을 제공한다. 이는 분류 공식화가 생존 목적 함수와 통계적으로 일치함을 입증한다.

주요 기여

1. 분류 공식화: 저자들은 누적 실패 확률을 직접 모델링하는 특정 분류 공식화를 도입하였다. 경험적으로 이 방식은 (구간 수가 증가함에 따라 곱셈적 오차 누적 문제를 겪는) 이산 위험률 기반 방식이나 다중 클래스(multi-class) 방식보다 우수한 성능을 보인다.
2. 일관성 증명: 표준 검열 가정 하에서 이진 분류 손실을 최소화하는 것이 일관된 생존 확률 추정치를 산출함을 증명하였다.
3. 경험적 평가: 이 접근법은 SurvSet 저장소의 48개 실제 데이터셋(정적 43개, 동적 5개)을 통해 평가되었다. 연구는 기성 TFM을 이 공식을 사용하여 사용했을 때를 고전적 베이스라인(CoxPH, Random Survival Forests) 및 딥러닝 모델(DeepSurv, DeepHit, Joint Models)과 비교하였다.

결과

• 정적 설정: 43개 데이터셋 전체에서 누적 실패(CF) 공식을 사용한 TFM(특히 MITRA-CF 및 XGBoost-CF)이 여러 지표(C-Index, Integrated AUC, Integrated Brier Score)에서 가장 우수한 종합 성능을 달いました. CF 공식은 구간 수($K$)가 증가함에 따라 성능이 저하되는 위험 기반(DH) 접근 방식에 비해 더 높은 견고성을 보였다.
• 동적 설정: 5개의 동적 데이터셋에서 MITRA-CF는 랜드마크(landmarked) CoxPH/RSF 및 Dynamic DeepHit을 포함한 모든 베이스라인을 능가하였다. 이는 TFM에 내재된 시간적 사전 지식(temporal priors)이 종단적 데이터의 복잡성을 모델링하는 데 특히 효과적임을 시사한다.
• 지표 상관관계: 연구 결과 이진 교차 엔트로피 손실과 통합 브라이어 점수(Integrated Brier Score) 사이에 강한 상관관계($r=0.89$)가 발견되었으며, 이는 분류 목적 함수를 최적화하는 것이 생존 확률의 보정(calibration)을 직접적으로 개선한다는 이론적 주장을 뒷받침한다.

의의 및 주장
본 논문은 적절한 분류 공식화를 제공함으로써 기성 표 형식 파운데이션 모델로부터 생존 분석 능력을 이끌어낼 수 있다고 주장한다. 이 방식은 생존 모델을 처음부터 훈련해야 하는 전문화된 모델에 비해 계산 효율적인 대안을 제공하며, 특히 파운데이션 모델이 강점을 보이는 소규모 데이터 환경에서 유용하다.

저자들은 합성 생존 태스크에 맞춰 트랜스포머를 특별히 사전 학습시키는 다른 연구들(Survival In-Context 및 SurvivalPFN)과 본 연구를 차별화한다. 대조적으로, 본 연구는 범용 TFM이 생존 특화 사전 학습이나 아키텍처 수정 없이도, 올바른 분류 공식이 제공될 때 생존 태스크에서 최첨단(state-of-the-art) 성능을 달성할 수 있음을 보여준다. 이 방법은 시간 빈(bin)의 수에 유연하며 정적 및 동적 설정 모두에 적용 가능하지만, 조건부 독립 검열 가설을 충족하는 시나리오로 제한된다.