CONVERSE 변분 대조 학습 기반 위험군 구분 및 생존 예측 모델

초록

CONVERSE는 변분 오토인코더와 다중 대조 손실을 결합해 환자 데이터를 저차원 잠재공간에 매핑하고, 클러스터링과 자기‑페이스 학습을 통해 위험군을 해석 가능하게 구분한다. 각 클러스터에 특화된 생존 헤드를 두어 시간‑to‑event 예측 정확도를 높이며, 네 개의 공개 생존 데이터셋에서 기존 딥 서바이벌 모델 대비 동등하거나 우수한 성능을 보였다.

상세 분석

본 논문은 임상 의사결정에 필수적인 생존 분석을 고성능과 해석 가능성 사이의 트레이드오프 없이 달성하고자 하는 목표 아래 설계되었다. 핵심 아이디어는 변분 오토인코더(VAE)를 이용해 환자 특성을 압축한 잠재 표현 z 를 얻고, 이를 다중 뷰(단일 또는 시암형)로 확장한 뒤 클러스터링과 대조 학습을 동시에 적용한다는 점이다. VAE는 재구성 손실 L_REC과 표준 정규분포와의 KL 발산 L_KLD를 최소화함으로써 과적합을 방지하고, 잠재 공간을 정규화한다. 시암형 인코더를 선택하면 두 개의 독립 파라미터를 가진 인코더가 동일 입력을 서로 다른 뷰로 변환해, 뷰 간 일관성을 강화하는 IVIW 손실과 뷰‑클러스터 간 분포 일치를 유도하는 IVCW 손실을 추가로 적용한다.

클러스터링 단계에서는 K‑means, 계층적, GMM, 스펙트럴 등 다양한 알고리즘을 하이퍼파라미터 탐색으로 선택할 수 있다. 클러스터 중심 M과 할당 c_i를 이용해 L_CLUS 손실로 잠재 벡터와 중심 사이 거리를 최소화한다. 이후 대조 손실 L_CL은 세 가지 구성요소를 가중합한다. 첫 번째 IV‑CG는 검열된 환자를 앵커로 삼고, 같은 클러스터에 속한 검열되지 않은 환자를 양성, 다른 클러스터를 음성으로 삼아 InfoNCE 형태로 학습한다. 이는 동일 위험군 내 표현 일치를 촉진한다. 두 번째 IVIW는 시암형 뷰의 동일 환자 두 표현을 양성 쌍으로, 다른 환자를 음성 쌍으로 삼아 뷰 간 정합성을 강화한다. 세 번째 IVCW는 각 클러스터에 대한 소프트 할당(q) 분포를 뷰 간에 일치시키는 대조 손실로, 클러스터링 신뢰도를 반영한다. 가중치 α는 학습 과정에서 자동 조정 가능하도록 설계되었다.

생존 예측은 이산 시간 프레임을 채택하고, 잠재 표현과 원본 특성을 결합한 h_i를 입력으로 하는 피드포워드 네트워크를 사용한다. 공유 헤드(shared‑head)와 클러스터‑특화 헤드(cluster‑specific‑heads) 두 가지 구성을 제공하는데, 후자는 각 클러스터에 별도 네트워크를 할당해 서브팝ulation 별 위험 패턴을 보다 정밀히 모델링한다. 손실 함수는 음의 로그우도(NLL)와 순위 손실(L_RANK)의 가중합으로, 캘리브레이션과 판별력을 동시에 최적화한다.

학습 절차는 (1) 사전 학습 단계에서 VAE와 생존 헤드를 공동 최적화해 의미 있는 초기 표현을 확보하고, (2) 클러스터 초기화 단계에서 선택된 알고리즘으로 중심과 할당을 고정한 뒤, (3) 자기‑페이스 학습(SPL) 기반의 단계적 재학습을 수행한다. SPL은 인스턴스별 손실(L_i)이 평균 ± 표준편차 기반 임계값 λ_e 이하인 ‘쉬운’ 샘플부터 점진적으로 포함시켜, 초기 클러스터 경계가 불안정한 상황에서도 안정적인 최적화를 가능하게 한다.

실험에서는 METABRIC, SUPPORT, SEER 등 네 개의 공개 데이터셋에 대해 Concordance Index(C‑index), Integrated Brier Score(IBS) 등 표준 지표를 사용해 DeepSurv, Cox‑PH, DCM, VadeSC, DVCSurv 등 최신 딥 서바이벌 모델과 비교하였다. CONVERSE는 특히 클러스터‑특화 헤드와 다중 대조 손실을 결합했을 때 C‑index에서 1~3% 향상을 보였으며, 위험군 시각화와 클러스터 별 생존 곡선이 임상적으로 의미 있는 차이를 나타냈다. 또한, 자기‑페이스 학습을 적용하지 않은 경우 대비 학습 안정성이 크게 개선되고, 클러스터링 품질(Adjusted Rand Index)도 상승하였다.

본 모델의 강점은 (1) 변분 프레임워크를 통한 정규화된 잠재 공간 제공, (2) 다중 대조 손실로 클러스터 일관성과 뷰 간 정합성을 동시에 강화, (3) 자기‑페이스 학습으로 초기 클러스터링 불안정을 완화, (4) 클러스터‑특화 생존 헤드로 서브팝ulation 별 위험 패턴을 정밀히 포착한다는 점이다. 한계점으로는 클러스터 수 K를 사전에 지정해야 하는 점, 대규모 데이터에서 대조 쌍 샘플링 비용이 증가할 수 있다는 점, 그리고 이산 시간 프레임 선택이 결과에 민감할 수 있다는 점을 들 수 있다. 향후 연구에서는 베이지안 비모수 클러스터링을 도입해 K를 자동 추정하고, 메모리 효율적인 대조 샘플링 전략을 개발하며, 연속 시간 서바이벌 모델과의 통합을 모색할 여지가 있다.