🤖 machine learning

SurvPFN: Towards Foundation Models for Survival Predictions

이 논문은 수백만 개의 합성 태스크를 통해 사전 학습되어 파운데이션 모델이 생존 예측을 위한 검측 데이터(censored data)를 처리할 수 있도록 하는 사전 데이터 적합 네트워크인 SurvPFN을 소개하며, 이는 데이터셋별 적합 과정이나 특화된 아키텍처 없이도 실제 데이터셋에서 경쟁력 있는 성능을 달성한다.

원저자: Samuel Böhm (Institute of Epidemiology and Prevention, Medical Center - University of Freiburg, Faculty of Medicine, University of Freiburg, Freiburg, Germany), Lennart Purucker (Department of Compute

게시일 2026-06-04

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Samuel Böhm (Institute of Epidemiology and Prevention, Medical Center - University of Freiburg, Faculty of Medicine, University of Freiburg, Freiburg, Germany), Lennart Purucker (Department of Computer Science, University of Freiburg, Freiburg, Germany, PriorLabs, Freiburg, Germany), Frank Hutter (Department of Computer Science, University of Freiburg, Freiburg, Germany, PriorLabs, Freiburg, Germany), Pascal Schlosser (Institute of Epidemiology and Prevention, Medical Center - University of Freiburg, Faculty of Medicine, University of Freiburg, Freiburg, Germany, Department of Epidemiology, Johns Hopkins Bloomberg School of Public Health, Baltimore, Maryland, US, CIBSS - Centre for Integrative Biological Signalling Studies, University of Freiburg, Freiburg, Germany)

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 전구의 수명이 정확히 언제 다할지 예측하려고 노력하고 있다고 상상해 보십시오. 완벽한 세상이라면, 모든 전구가 수명을 다할 때까지 기다렸다가 정확한 시간을 기록할 수 있을 것입니다. 하지만 현실 세계에서는 영원히 기다릴 수 없습니다. 실험을 중단할 때까지도 여전히 작동 중인 전구들이 있고, 어떤 전구들은 수명이 다하기 전에 버려지기도 합니다. 통계학에서는 이를 **검측(Censoring)**이라고 부릅니다. 당신은 그 전구가 적어도 실험을 중단할 때까지는 작동했다는 사실은 알지만, 정확히 어느 순간에 수명이 다했을지는 모릅니다.

오랫동안, '파운데이션 모델(Foundation Models)'이라 불리는 강력하고 새로운 AI 도구들(특히 "TabPFN")은 "이 이메일이 스팸인가?" 또는 "이 집은 얼마에 팔릴 것인가?"와 같은 것을 예측하는 데 매우 뛰어난 성능을 보여왔습니다. 하지만 이 모델들은 생존 예측(예: "이 환자가 언제 재발할 것인가?")에는 어려움을 겪었습니다. 왜냐하면 이들은 저 "여전히 작동 중인" 전구들을 처리하는 방법을 몰랐기 때문입니다. 만약 이들에게 검측된 데이터를 입력한다면, 모델은 혼란에 빠져 편향된 추측을 하게 될 것입니다.

여기 SurvPFN이 있습니다. 저자들은 이 "전구 문제"를 이해할 수 있도록 설계된 새로운 버전의 AI를 구축했습니다.

이들이 어떻게 이 일을 해냈는지, 일상적인 비유를 통해 설명하겠습니다.

1. "훈련 체육관" (합성 데이터)

격투기 선수를 훈련시킬 때 실제 경기에서만 싸우게 할 수는 없습니다. 선수는 체육관에서 연습이 필요합니다. 저자들은 수백만 개의 가짜 생존 시나리오를 사용하여 거대한 합성 체육관을 만들었습니다.

설정: 그들은 수학적 레시피(구조적 인과 모델, Structural Causal Model)를 사용하여 서로 다른 수명을 가진 가짜 "전구들"을 생성했습니다.
반전: 그들은 의도적으로 이 가짜 전구들 중 일부를 "검측"했습니다. 즉, 실제 상황처럼 타이머를 일찍 멈춘 것입니다.
교훈: 그들은 AI에게 "멈춘" 전구는 고장 난 전구가 아니라, 단지 적어도 그 시간까지는 살아남았다는 것을 가르쳤습니다. 또한, 특정 사망 시간을 추측하는 대신 이러한 하한선(lower bounds)에 대해 정직하게 반응하도록 보상하는 특별한 채점 방식("검측 손실", censored loss)을 사용했습니다.

2. "두 갈래" 뇌 (모델 아키텍처)

그들이 만든 AI는 기존 모델의 작고 효율적인 버전인 NanoTabPFN을 기반으로 합니다. 생존 분석을 위해 작동하게 만들고자, 그들은 특별한 기능을 추가했습니다.

이벤트 지표 (The Event Indicator): 이것은 AI가 모든 데이터 포인트에 대해 확인하는 작은 깃발과 같습니다.
- 깃발 = 0: "이 전구는 시간 X에 수명이 다했다." (정확한 시간을 알고 있음).
- 깃발 = 1: "이 전구는 시간 X에도 여전히 작동 중이었다." (X 이후로도 생존했다는 것만 알고 있음).
AI는 이 두 종류의 깃발을 처리하기 위해 뇌 안에 두 개의 별도 "경로"를 가지고 있습니다. 이를 통해 AI는 "생존한" 전구가 "사망한" 전구와 (비록 같은 시간에 발생했더라도) 다르다는 것을 학습할 수 있습니다.

3. "만능" 접근 방식 (One-Size-Fits-All)

보통 생존을 예측하려면 의사나 데이터 과학자는 특정 데이터셋(예: 특정 라디오 방송국에 맞추기 위해 라디오를 튜닝하는 것처럼)에 맞춰 특정 모델을 튜닝하는 데 몇 주를 소비해야 합니다.

SurvPFN의 초능력: 이 모델은 저 거대한 합성 체육관에서 단 한 번 훈련되었습니다. 훈련이 끝난 후, 모델은 새로운 실제 데이터셋(예: 병원 기록이나 기계 신뢰성 로그)에 바로 투입되어 즉시 예측을 시작할 수 있습니다. 특정 그룹을 위해 다시 훈련하거나 "미세 조정(fine-tuning)"할 필요가 없습니다. 이는 마치 새로운 나라에 갈 때마다 사전 없이도 모든 언어를 구사하는 만능 통역사와 같습니다.

4. 결과: 얼마나 뛰어난가?

저자들은 22개의 실제 데이터셋(암 연구부터 실업 데이터까지)을 사용하여 SurvPFN을 테스트하고, 이를 "골드 스탠다드(표준)" 방법들과 비교했습니다.

Cox 비례 위험 모델 (Cox Proportional Hazards): 고전적이고 오래된 통계적 방법.
랜덤 생존 포레스트 (Random Survival Forests): 인기 있는 머신러닝 방법.
DeepSurv: 현대적인 딥러닝 접근 방식.
BinSurv: 생존 문제를 일련의 "예/아니오" 질문으로 변환하려는 최신 방법.

결론:
SurvPFN은 이 모든 전문화되고 정교하게 튜닝된 전문가들과 대등한 수준의 성능을 보여주었습니다.

새로운 데이터셋을 위해 다시 훈련할 필요가 없었습니다.
특별한 "특징 공학(feature engineering, 데이터를 수동으로 조작하는 것)"이 필요하지 않았습니다.
복잡하고 맞춤 제작된 아키텍처조차 필요하지 않았습니다.

한계점 (Limitations)

이 논문은 이 AI가 현재 가진 한계에 대해서도 솔직하게 밝히고 있습니다.

크기: 이 모델은 작은 데이터셋(최대 1,000행 및 10개 특징)에서 가장 잘 작동합니다. 이는 스포츠카와 같습니다. 빠르고 민첩한 경주에는 뛰어나지만, 아직 거대한 트럭 함대를 운용하기 위해 만들어진 것은 아닙니다.
"와이불(Weibull)" 편향: AI는 특정 수학적 형태의 생존 곡선(와이불 분포)을 기반으로 훈련되었습니다. 따라서 많은 실제 시나리오를 잘 처리하지만, 실제 데이터가 체육관에서 본 형태와 전혀 다를 경우 어려움을 겪을 수 있습니다.
정적 관점: 이 모델은 생존에 영향을 미치는 요인(예: 연령이나 직업 유형)이 시간이 지나도 변하지 않는다고 가정합니다. 즉, "경쟁 위험(competing risks)"(예: 질병이 진행되기 전에 교통사고로 사망하는 경우)을 처리하지 못합니다.

핵심 요약

이 논문은 생존 분석이 "하한선(lower bounds)"에 대한 특별한 규칙이 있는 분포 회귀(distributional regression)의 한 종류라는 점을 입증합니다. "이벤트 지표" 깃발을 통해 "검측" 규칙을 존중하도록 일반 목적의 AI를 가르침으로써, 매번 새로운 모델을 처음부터 구축할 필요 없이 강력한 사전 훈련된 파운데이션 모델을 시간-사건(time-to-event) 예측에 사용할 수 있게 되었습니다.

기술 요약: SurvPFN: 생존 예측을 위한 파운데이션 모델을 향하여

문제 정의

TabPFN과 같은 테이블형 파운데이션 모델(Tabular Foundation Models, TFMs)은 합성 데이터로 학습된 사전 데이터 적합 네트워크(prior-data fitted networks, PFNs)를 활용하여 표준 분류 및 회귀 작업에서 빠른 성공을 거두었습니다. 그러나 생존 예측 작업은 이러한 패러다임의 영향을 거의 받지 못했습니다. 주요 장벽은 **중도 절단(censoring)**입니다. 생존 분석에서는 관심 있는 사건(예: 사망, 고장)이 연구 기간 내에 모든 대상자에게 관찰되지 않는 경우가 많습니다(우측 중도 절단). 표준 TFM 아키텍처와 회귀 목적 함수는 중도 절단된 데이터를 자연스럽게 처리할 수 없습니다. 중도 절단된 시간을 회귀 타겟으로 취급하거나 이를 폐기하는 것은 체계적인 편향과 부정확한 예측을 초래합니다. 기존의 생존 방법론(예: Cox 비례 위험 모델, Random Survival Forests, DeepSurv)은 라이클리후드 함수(likelihood functions)에서 중도 절단을 명시적으로 모델링하지만, 일반적으로 데이터셋마다 개별적인 피팅(fitting)이 필요하며 "파운데이션 모델"로서의 제로샷 적응(zero-shot adaptation) 능력은 부족합니다.

방법론

저자들은 생존 예측을 위해 특별히 설계된 사전 데이터 적합 네트워크인 SurvPFN을 제안합니다. 이 접근 방식은 생존을 중도 절단된 손실 함수를 사용하는 연속 시간 분포 회귀 문제로 취급합니다.

1. 합성 사전 생성 (Synthetic Prior Generation)

데이터셋마다 별도의 피팅 없이 모델을 학습시키기 위해, 저자들은 **구조적 인과 모델(Structural Causal Models, SCMs)**을 사용하여 수백만 개의 합성 생존 태스크를 생성합니다.

사건 발생 시간: 생존 분석의 표준 선택지인 **와이불 분포(Weibull distribution)**로부터 생성됩니다. 모델은 비례 위험(PH) 및 비비례 위험(non-PH) 체제를 모두 지원합니다. 비비례 위험 시나리오에서는 두 번째 결과 노드가 대상자별로 형상 매개변수(shape parameter)를 수정하여 생존 곡선이 교차하도록 유도합니다.
중도 절단 메커니즘: 비정보적 중도 절단 메커니즘이 적용됩니다. 여기에는 다음이 포함됩니다:
- 행정적 중도 절단(Administrative censoring): 사건 시간 분포의 분위수를 기준으로 설정된 데이터셋 수준의 추적 관찰 기간($FU$).
- 순차적 등록(Staggered enrollment): 관찰 시작 시간의 가변성.
- 연구 중간 탈락(Mid-study dropout): 관찰이 행정적 중도 절단 시간 이전에 균등하게 추출된 시간에 종료될 수 있음.
데이터 구조: 합성 데이터셋은 2~10개의 공변량(범주형 특성 포함)과 최대 1,000개의 행을 포함합니다.

2. 모델 아키텍처 및 학습

SurvPFN은 경량화되고 사용자 친화적인 TFM 백본인 NanoTabPFN(4.4M 파라미터, 6개 레이어)을 기반으로 구축되었습니다.

입력 표현: 사건 지표(사건 발생 0, 중도 절단 1)가 명시적인 특성 열으로 추가됩니다. 이를 통해 모델은 관찰된 사건 시간과 하한값(중도 절단된 시간)을 구분할 수 있습니다.
타겟 인코딩: 두 갈래의 타겟 인코더가 동일한 아키텍처의 별도 인코더를 통해 관찰된 사건과 중도 절단된 행을 처리하며, 이를 통해 모델은 실제 사건 시간과 하한값에 대한 구별된 표현을 학습할 수 있습니다.
학습 목적 함수: 모델은 **중도 절단된 음의 로그 가능도(negative log-likelihood, NLL)**를 최소화하도록 학습됩니다:
- 관찰된 사건의 경우: 정보적 중도 절단을 보정하기 위해 역 확률 가중치(IPCW)를 적용합니다.
- 중도 절단된 행의 경우: 모델이 중도 절단 시간 이전에 확률 질량을 배치하는 것에 대해 페널티를 부여합니다.
- 랭킹 항(Ranking Term): 예측된 생존 시간과 SCM에 의해 유도된 기저 위험 순서를 정렬하기 위해 미분 가능한 쌍별 랭킹 항을 선택적으로 추가합니다.
추론: 모델은 이산 밀도 $\hat{p}_j$ 를 출력하며, 이로부터 생존 함수 $\hat{S}(t)$ 와 예측 중앙값을 직접 도출합니다.

3. 평가 프로토콜

모델은 임상, 경제 및 신뢰성 영역을 아우르는 22개의 큐레이션된 실제 생존 데이터셋인 SurvSet에서 평가되었습니다.

제약 조건: 소규모 백본의 학습 체제에 맞추기 위해 데이터셋을 $\le 1000$ 행 및 $\le 10$ 특성으로 필터링했습니다.
베이스라인: 다음 모델들과 비교되었습니다:
- Cox 비례 위험 모델 (CoxPH): 지배적인 고전적 베이스라인.
- Random Survival Forests (RSF): 표준 비매개변수 앙상블.
- DeepSurv: 가장 널리 사용되는 딥러닝 접근 방식.
- BinSurv: 시간을 이진 분류 작업으로 이산화하는 최신 TFM 기반 접근 방식.
지표: 성능은 Uno's C-index(변별력), 통합 브라이어 점수(Integrated Brier Score, IBS)(정확도), 통합 교정 지수(Integrated Calibration Index, ICI)(교정도)를 통해 평가되었습니다.

주요 결과

경쟁력 있는 성능: SurvPFN은 22개 데이터셋 전체에서 모든 베이스라인과 매우 경쟁적인 성능을 보였습니다. 평균적으로 RSF, DeepSurv, BinSurv와 대등하거나 이를 약간 상회하며, CoxPH와도 유사한 성능을 보였습니다.
제로샷 능력: 결정적으로, SurvPFN은 데이터셋별 피팅, 생존 특화 아키텍처, 또는 특성 공학 없이 이러한 성능을 달성했습니다. 이는 오직 사전 학습된 프리어(prior)에만 의존합니다.
절제 연구(Ablation Studies):
- 사건 지표: 사건 지표를 특성에서 제거(1로 강제)하면 변별력(C-index)은 유지되지만 교정도(ICI)가 크게 저하됩니다. 이는 상대적인 위험 순서는 시간만으로 추론할 수 있지만, 절대적인 교정에는 어떤 데이터 포인트가 중도 절단되었는지에 대한 명시적 지식이 필요함을 확인시켜 줍니다.
- 손실 함수 변형: (합성 사전 학습 중에만 사용 가능한) 오라클 타겟을 사용한 변형 모델은 실제 데이터에서도 네이티브 중도 절단 손실 함수와 유사한 성능을 보였으며, 이는 네이티브 손실 함수가 전이(transfer)에 충분함을 시사합니다.

의의 및 주장

본 논문은 생존 분석이 PFN에 의해 해결 가능한 연속 시간 분포 회귀 문제로 성공적으로 프레임화될 수 있음을 주장합니다. 핵심 기여는 다음과 같습니다:

격차 해소: 파운데이션 모델이 기존의 표준 회귀 PFN에는 부적합하다고 여겨졌던 중도 절단 데이터를 처리할 수 있음을 입증했습니다.
단순성: 사건 지표 특성과 중도 절단 손실 함수만 추가하면 표준 회귀 PFN 아키텍처로도 전문화된 생존 모델의 성능을 회복할 수 있음을 보여주었습니다.
효율성: 새로운 데이터셋에 대한 하이퍼파라미터 튜닝이나 재학습이 필요 없는, "플러그 앤 플레이(plug-and-play)" 솔루션을 제공합니다.

저자들은 현재의 평가가 소규모 데이터셋( $\le 1000$ 행) 및 정적 공변수로 제한되어 있다는 점을 언급하며 한계점에 대해 겸허한 태도를 유지하였으며, 현재의 합성 프리어가 와이불 분포에 국한되어 있고 경쟁 위험(competing risks)이나 시변 공변량(time-varying covariates)을 아직 다루지 못하고 있음을 인정하며, 이를 향후 확장성 및 표현력을 위한 방향으로 제시하였습니다.