Scalable Random Survival Forests via Risk Set Sampling

기술 요약: 리스크 세트 샘플링을 통한 확장 가능한 랜덤 생존 포레스트(Scalable Random Survival Forests)

문제 정의
랜덤 생존 포레스트(Random Survival Forests, RSF)는 검열된 시간-사건 데이터(censored time-to-event data)를 위해 설계된 강력한 머신러닝 확장 기법이다. 그러나 대규모 또는 고차원 데이터셋에 대한 RSF의 적용은 계산적으로 매우 부담스러울 수 있다. 주요 병목 현상은 트리 구축 과정에서 로그-랭크 분할 규칙(log-rank splitting rule)을 계산하기 위해 반복적으로 "리스크 세트(risk sets)"를 평가해야 한다는 점이다. 모든 이벤트 시간마다 아직 위험에 처해 있는 모든 개체를 평가해야 하므로, 계산 비용은 샘플 크기와 고유한 이벤트 시간의 수에 따라 급격히 증가한다. 리스크 세트 샘플링은 Cox 기반 신경망(예: Cox-Time)에서 성공적으로 통합되었으며 고전 역학(nested case-control)의 표준 설계로 사용되어 왔으나, RSF에는 아직 적응되지 않았다.

방법론
저자들은 리스크 세트 샘플링을 트리 구축 과정에 직접 통합하는 RSF 알고리즘의 수정을 제안한다.

• 핵심 메커니즘: 각 이벤트 시간 $t_k$에서 전체 리스크 세트($R_k$)를 사용하는 대신, 리스크 세트에서 정해진 비율의 개체를 샘플링한 $\tilde{R}_k$를 사용한다.
• 통계적 조정: 로그-랭크 통계량은 노드 내의 기대 이벤트 수와 분산에 의존한다. 저자들은 전체 리스크 세트 대신 샘플링된 리스크 세트를 사용하여 이러한 계산을 수정한다. 구체적으로, 딸 노드(daughter node) $\alpha$에 대한 기대 이벤트 수 $E(D_{\alpha,k})$와 분산 $Var(D_{\alpha,k})$는 샘플링된 위험 인원의 비율($\tilde{R}_{\alpha,k} / \tilde{R}_k$)을 사용하여 재계산된다.
• 구현: 이 접근 방식은 고성능 랜덤 포레스트 구현체인 ranger R 패키지에 구현되었다. 저자들은 10%에서 100% 사이의 샘플링 비율을 테스트하였으며, 여기서 100%는 표준적인 전체 리스크 세트를 의미한다.
• 실험 설계:
  • 시뮬레이션: 지수 분포(Exponential) 및 와이불 혼합 분포(Weibull mixture distributions)를 사용하여 샘플 크기 $n=500, 1000, 2000$ 및 공변량 $p=15$인 이질적 생존 데이터를 생성하였다.
  • 실제 데이터: 두 가지 데이터셋인 METABRIC(1904개의 유방암 관측치)과 독일 유방암 연구 그룹(GBSG) 데이터셋(686개 관측치)을 활용하였다.
  • 지표: 성능은 5-겹 교차 검증(5-fold cross-validation) 및 Out-of-Bag(OOB) 추정치를 통한 일치도 지수(C-index)와 계산 실행 시간(runtime)을 통해 평가되었다.

주요 결과

• 예측 성능: 모든 시뮬레이션 및 실제 데이터셋에서, 샘플링된 리스크 세트로 훈련된 RSF 모델(특히 25%~40% 사용 시)은 전체 리스크 세트로 훈련된 모델과 거의 구별할 수 없을 정도로 유사한 C-index 값을 달-성했다.
• 계산 효율성: 리스크 세트를 25~40% 범위로 줄이면 훈련 시간이 50% 이상 단축되었다. 실행 시간은 리스크 세트 비율에 따라 단조 증가하였으며, 가장 큰 이득은 더 큰 데이터셋에서 관찰되었다.
• 안정성 및 편향: 연구 결과, 리스크 세트 샘플링은 체계적인 편향을 도입하기보다는 주로 추정치의 분산에 영향을 미치는 것으로 나타났다. C-index의 신뢰 구간은 모든 샘플링 비율에 걸쳐 상당 부분 겹쳤으며, 이는 예측 변별력 측면에서 통계적으로 유의미한 차이가 없음을 나타낸다.
• 규제 효과(Regularization Effect): 여러 시나리오에서 중간 정도의 샘플링 비율(25~40%)로 훈련된 모델이 전체 리스크 모델보다 약간 더 높은 C-index 값을 보였는데, 이는 노이즈를 줄이는 유익한 규제 효과가 존재할 가능성을 시사한다.

의의 및 주장
본 논문은 리스크 세트 샘플링이 예측 성능이나 모델 안정성을 저해하지 않으면서도 RSF의 확장성을 실질적으로 개선하는 간단하고 효과적인 수정 방법이라고 주장한다. 저자들의 주장은 다음과 같다:

1. 학문 간 가교: 이 접근 방식은 고전적 역학 샘플링 설계(nested case-control)와 현대의 생존 분석용 머신러닝 사이의 방법론적 가교를 제공한다.
2. 실용적 유용성: 전체 리스크 계산은 RSF 모델링에서 우수한 예측 성능을 달성하기 위해 실질적인 이점을 거의 제공하지 못하면서도 훨씬 높은 계산 비용을 발생시킨다.
3. 접근성: 이 방법론을 널리 사용되는 ranger 패키지에 구현함으로써, 저자들은 대규모 및 고차원 생존 데이터에 확장 가능한 RSF를 실제로 적용할 수 있도록 돕는다.

저자들은 이 방법이 견고하지만, 향후 연구에서 더 복적인 프레임워크(경쟁 위험, 다상 모델)에 대한 적용을 탐색하고 다른 RSF 하이퍼파라미터와의 상호작용을 조사할 수 있다고 결론짓는다.