반지도학습 기반 딥 Cox 모델로 암 예후 예측 정확도 향상

초록

본 연구는 Mean Teacher 프레임워크를 활용한 깊은 반지도학습(DSSL) 방식인 Cox‑MT를 제안한다. Cox‑MT는 라벨이 있는 생존 데이터와 라벨이 없는 유전체·이미지 데이터를 동시에 학습함으로써 기존 ANN 기반 Cox‑nnet보다 RNA‑seq 및 병리 슬라이드 이미지 두 모달리티에서 모두 높은 C‑index와 낮은 Brier 점수를 달성한다. 특히 라벨이 없는 샘플을 추가할수록 성능이 점진적으로 개선되며, 유전체와 이미지 정보를 결합한 다중모달 모델이 단일모달 모델을 능가한다.

상세 분석

본 논문은 전통적인 Cox 비례위험 모델을 딥러닝 기반으로 확장하면서, 라벨이 제한된 상황에서도 모델 성능을 극대화할 수 있는 새로운 학습 전략을 제시한다. 핵심 아이디어는 Mean Teacher(Mean Teacher, MT) 프레임워크를 생존 분석에 적용하는 것으로, 학생 네트워크와 교사 네트워크 사이에 EMA(Exponential Moving Average) 기반 파라미터 동기화를 수행한다. 학생 네트워크는 라벨이 있는 데이터에 대해 전통적인 부분가능도(partial likelihood) 손실 L_s를 최소화하고, 라벨이 없는 데이터에 대해서는 교사 네트워크의 예측과의 일관성(consistency) 손실 L_u를 가중치 w와 함께 결합한다(L = L_s + w·L_u). 이 구조는 라벨이 없는 데이터가 교사 네트워크를 통해 “가짜 라벨”을 제공하도록 유도해, 모델이 데이터 분포 전반을 학습하도록 만든다.

실험에서는 TCGA의 네 종류 암(BRCA, LUAD, LUSC, UCEC) RNA‑seq 데이터를 이용해 4,000개의 변동성이 큰 유전자를 선택하고, 다층 퍼셉트론(MLP) 구조를 적용하였다. Cox‑MT는 두 개의 은닉층(첫 번째 1,0001,500노드, 두 번째 1002,000노드)을 사용했으며, Cox‑nnet은 64노드 하나의 은닉층만을 사용했다. 20번의 랜덤 분할 교차검증 결과, 모든 암 유형에서 Cox‑MT의 평균 C‑index가 Cox‑nnet보다 0.090.18만큼 상승했고, Integrated Brier Score(IBS)도 0.0380.082만큼 감소하였다. 이는 라벨이 있는 검열(censored) 샘플만으로도 기존 모델보다 더 풍부한 위험 정보를 추출함을 의미한다.

라벨이 없는 샘플을 추가한 실험에서는 GEO 데이터셋(GSE96058)에서 3,409개의 유방암 RNA‑seq 샘플을 무라벨 데이터로 활용하였다. 무라벨 샘플 수를 1,000, 2,000, 3,000, 3,409개까지 점진적으로 늘릴 때, C‑index는 0.81→0.90으로, IBS는 0.087→0.061으로 개선되었다. 이는 MT 구조가 무라벨 데이터의 잠재적 구조를 효과적으로 정규화하고, 학생 네트워크가 보다 일반화된 표현을 학습하게 함을 보여준다.

다중모달 실험에서는 RNA‑seq과 Whole Slide Image(WSI) 특징을 각각 DINOv2와 자체 CNN으로 추출한 뒤, Transformer 기반의 Cross‑Attention 모듈을 통해 64차원 임베딩을 결합하였다. 다중모달 Cox‑MT는 C‑index 0.83, IBS 0.079를 기록했으며, 동일 구조의 Cox‑nnet(0.80, 0.091) 및 단일모달 Cox‑MT(RNA‑seq 0.81, 0.087; WSI 0.66, 0.151)보다 현저히 우수했다. 이는 서로 다른 데이터 모달리티 간의 상호보완적 정보를 효과적으로 통합할 수 있음을 시사한다.

하이퍼파라미터 민감도 분석에서는 EMA 상수 α(0.90.999), 일관성 가중치 w(0.13), 드롭아웃 비율(≤0.4), 가우시안 노이즈 σ(00.1), 학습률(0.0010.005) 등이 모델 성능에 큰 영향을 주지 않으며, 비교적 넓은 범위에서 안정적인 학습이 가능함을 확인하였다. 이미지 데이터에 대한 데이터 증강 실험에서는 색상 jitter(γ≈0.4)가 회전(β≈10°~30°)보다 더 큰 성능 향상을 제공했다.

전반적으로 이 논문은 반지도학습을 생존 분석에 성공적으로 적용함으로써, 라벨이 부족한 의료 데이터 환경에서도 딥러닝 기반 Cox 모델의 예측력을 크게 끌어올릴 수 있음을 입증한다. 특히, 검열 데이터와 무라벨 데이터를 동시에 활용하는 손실 설계, 그리고 다중모달 통합을 위한 Transformer 기반 어텐션 메커니즘은 향후 다양한 임상·비임상 분야에 확장 가능한 프레임워크로 활용될 가능성이 높다.