SurvDiff: 생존 분석을 위한 확산 기반 합성 데이터 생성 모델

초록

SurvDiff는 혼합형 공변량, 사건 시간, 그리고 오른쪽 검열 정보를 동시에 생성하도록 설계된 최초의 엔드‑투‑엔드 확산 모델이다. 생존 데이터 특유의 검열 메커니즘을 손실 함수에 직접 반영하고, 조기 사건에 가중치를 높이는 희소성‑인식 가중치 스킴을 도입해 학습 안정성을 확보한다. 다중 의료 데이터셋에서 기존 GAN·표준 확산 기반 베이스라인을 능가하며, 합성 데이터로 훈련한 생존 모델이 실제 데이터에 대해 높은 C‑index와 낮은 Brier score를 기록한다.

상세 분석

SurvDiff는 기존 표준 확산 모델이 연속·이산 변수 혼합과 검열 정보를 동시에 다루지 못한다는 한계를 정확히 짚어낸다. 논문은 먼저 생존 데이터의 구조를 (연속 공변량, 이산 공변량, 사건 지표 E, 관측 시간 T) 네 요소로 정의하고, 이 네 변수를 하나의 벡터로 결합해 확산 과정에 투입한다. 전방 노이즈 단계에서는 Gaussian 노이즈와 마스크‑확산을 병행해 연속 변수는 연속적인 가우시안 노이즈, 이산 변수는 마스크를 통해 카테고리별 확률 분포를 유지한다. 역방향 단계에서는 스코어 네트워크 µθ(x,u)를 학습하는데, 여기서 핵심은 “생존‑특화 손실”이다. 손실 함수는 두 부분으로 구성된다. 첫째, 일반적인 MSE/크로스엔트로피 기반 스코어 매칭 손실에 더해 사건 시간 순서를 보존하도록 설계된 순위 손실(Ranking Loss)을 추가한다. 이는 사건이 발생한 샘플들 사이의 시간 차이가 실제 데이터와 동일한 순서를 유지하도록 강제한다. 둘째, 검열 메커니즘을 반영하기 위해 검열된 샘플에 대한 로그우도 항을 별도로 계산하고, 검열되지 않은 샘플에 비해 가중치를 낮춘다. 특히, “희소성‑인식 가중치 스킴”은 초기(조기) 사건이 데이터에서 더 많이 관측되는 특성을 이용해, 시간 t가 작을수록 손실 가중치를 크게 설정하고, 늦은 시간일수록 가중치를 감소시켜 학습이 불안정해지는 것을 방지한다.

모델 아키텍처는 최신 표준인 TabDiff의 변형을 채택해, 각 변수 유형별 임베딩 레이어와 교차‑ attention 메커니즘을 결합한다. 이를 통해 변수 간 복잡한 상관관계와 비선형 상호작용을 효과적으로 학습한다. 또한, 검열 지표 E를 이진 마스크로 처리해, 역확산 과정에서 검열 여부를 조건부로 샘플링하도록 설계했다.

실험에서는 3개의 실제 의료 데이터셋(예: SEER 암 데이터, MIMIC‑III ICU 데이터, 그리고 다기관 임상시험 데이터)을 사용해 4가지 평가 지표를 제시한다. (1) 공변량 분포 유사성(Jensen‑Shannon, Wasserstein 거리), (2) 사건‑시간 발산(Event‑Time Divergence, ETD), (3) 전체 형태 지표(Shape metric), (4) 생존 모델 성능(TSTR 시나리오에서 C‑index와 Brier score). 모든 지표에서 SurvDiff는 기존 SurvivalGAN, Ashhad 프레임워크, 그리고 TabDiff·CTGAN·TVAE 등 일반‑목적 베이스라인을 크게 앞선다. 특히, 검열 비율이 40% 이상인 데이터셋에서도 사건‑시간 분포와 검열 패턴을 정확히 재현해, 임상 연구에서 합성 데이터 활용 가능성을 크게 높인다.

한계점으로는 현재 모델이 오른쪽 검열만을 다루며, 좌측·간격 검열 등 복합 검열 상황에는 추가적인 확장이 필요하다는 점을 언급한다. 또한, 고차원 이미지·시계열 결합 데이터에 대한 적용은 아직 검증되지 않았다. 향후 연구에서는 다중 검열 유형을 포괄하는 손실 설계와, 멀티모달 의료 데이터에 대한 확산 기반 통합 모델링을 목표로 제시한다.