[논문 리뷰] SurvivalGAN: Generating Time-to-Event Data for Survival Analysis

Alexander Norcliffe, Bogdan Cebere외 3명|arXiv (Cornell University)|2023. 02. 24.

Machine Learning in Healthcare인용 수 5

한 줄 요약

SurvivalGAN은 생존 분석을 위한 고품질의 시간-사건 데이터를 합성하기 위해 설계된 생성적 적대 신경망이며, 특별히 시간-사건 분포의 불균형과 케이싱을 다루기 위해 전용의 시간-사건 회귀기, 불균형 샘플링, 조건부 GAN 학습을 통합하였다. 이는 표준 기준보다 더 높은 현실성 있는 생존 데이터를 생성하고, C-index 및 브리어 스코어 측면에서 후행 생존 모델의 성능을 향상시킨다.

ABSTRACT

Synthetic data is becoming an increasingly promising technology, and successful applications can improve privacy, fairness, and data democratization. While there are many methods for generating synthetic tabular data, the task remains non-trivial and unexplored for specific scenarios. One such scenario is survival data. Here, the key difficulty is censoring: for some instances, we are not aware of the time of event, or if one even occurred. Imbalances in censoring and time horizons cause generative models to experience three new failure modes specific to survival analysis: (1) generating too few at-risk members; (2) generating too many at-risk members; and (3) censoring too early. We formalize these failure modes and provide three new generative metrics to quantify them. Following this, we propose SurvivalGAN, a generative model that handles survival data firstly by addressing the imbalance in the censoring and event horizons, and secondly by using a dedicated mechanism for approximating time-to-event/censoring. We evaluate this method via extensive experiments on medical datasets. SurvivalGAN outperforms multiple baselines at generating survival data, and in particular addresses the failure modes as measured by the new metrics, in addition to improving downstream performance of survival models trained on the synthetic data.

연구 동기 및 목표

• 케이싱된 표준 생존 데이터에 대한 효과적인 합성 데이터 생성 방법의 부족을 해결한다.
• 생존 데이터 생성에서 유일한 세 가지 실패 유형을 정식화한다: 위험군 인원의 과소 또는 과잉 생성, 조기 케이싱.
• 이러한 실패 유형을 정량적으로 측정할 수 있는 새로운 평가 지표를 개발한다.
• 케이싱된 데이터를 통합하고 시간 축 및 케이싱 불균형을 처리할 수 있는 견고한 생성 프레임워크를 제안한다.
• SurvivalGAN이 생성한 합성 데이터가 실제 의료 데이터셋에서 후행 생존 모델의 성능을 향상시킴을 입증한다.

제안 방법

• 시간-사건 및 케이싱 상태의 연합 분포를 모델링하기 위해 시간-사건/케이싱 회귀기를 도입한다.
• 실제 데이터 분포를 반영하기 위해 시간 축과 케이싱 상태에 대한 불균형 샘플링을 구현한다.
• 시간-사건 및 케이싱 예측에 조건부된 표본 특성 생성을 위한 조건부 GAN 아키텍처를 사용한다.
• 생성된 생존 시간의 시간적 정확도를 향상시키기 위해 GAN과 함께 시간 회귀기를 종단 간 학습한다.
• 생존 함수를 모델링하여 학습 중 케이싱된 데이터를 통합하고, 이를 생성 과정을 안내하는 데 활용한다.
• 시간 회귀기, 불균형 샘플링, 조건부 GAN을 통합한 다단계 학습 파이프라인을 도입하여 공동 최적화를 수행한다.

(a) Over-optimism. The area under the synthetic curve is larger than the true curve, the model is over-optimistic.

실험 결과

연구 질문

• RQ1 케이싱 및 시간 축 불균형과 관련하여 생존 데이터 생성에서 특유의 핵심 실패 유형은 무엇인가?
• RQ2 기존의 생성 지표 외에 생존 데이터의 품질을 정량적으로 측정할 수 있는 방법은 무엇인가?
• RQ3 생성 모델이 케이징된 표준 생존 데이터를 효과적으로 처리하면서 실제 데이터의 통계적 성질을 유지할 수 있는가?
• RQ4 SurvivalGAN의 구성 요소들 — 시간 회귀기, 불균형 샘플링, 조건부 GAN — 가 데이터 품질 향상과 후행 모델 성능 향상에 기여하는 정도는 어느 정도인가?
• RQ5 SurvivalGAN가 생성한 합성 데이터는 기존의 생성 모델 대비 더 나은 성능을 보이는 생존 모델을 도출할 수 있는가?

주요 결과

• SurvivalGAN은 표준 생성 지표와 새로 제안된 생존 전용 지표를 기준으로 다섯 가지 기준 기반 생성 모델보다 현실적인 생존 데이터를 생성하는 데서 뛰어난 성능을 보였다.
• 시간 회귀기는 특히 샘플의 순서를 정확히 분류하는 데 기여하여 C-index 성능을 크게 향상시키며, 생존 예측의 단기적 성향을 감소시켰다.
• 시간 축과 케이싱 상태에 대한 불균형 샘플링은 생존 함수의 형태 유지에 필수적이며, 특히 PHEART 및 SEER와 같은 대규모 데이터셋에서 C-index 향상에 기여한다.
• 조건부 GAN 구성 요소는 정확한 校정(캘리브레이션)을 위해 필수적이며, 이 요소가 생략될 경우 모든 데이터셋에서 브리어 스코어가 심각하게 악화된다.
• 절단 분석 결과, 시간 회귀기, 불균형 샘플링, 조건부 GAN의 세 구성 요소 모두 최적의 성능을 내기 위해 필수적이며, 각각 다른 측면의 데이터 품질 향상에 독자적인 기여를 한다.
• SurvivalGAN가 생성한 데이터를 기반으로 한 후행 생존 모델은 기준 생성 모델의 합성 데이터를 사용한 모델보다 뛰어난 성능을 보였으며, 더 높은 C-index와 더 낮은 브리어 스코어를 기록했다.

(b) Over-pessimism. The area under the synthetic curve is lower than the true curve, the model is over-pessimistic.