[논문 리뷰] Deep Cox Mixtures for Survival Regression

Chirag Nagpal, Steve Yadlowsky외 2명|arXiv (Cornell University)|2021. 01. 16.

Statistical Methods and Inference참고 문헌 9인용 수 5

한 줄 요약

이 논문은 환자를 잠재적 하위군으로 모델링하고, 각 하위군이 고유한 비례 위험 구조를 가지며, 딥 뉴럴 네트워크와 비모수적 기준 위험 함수를 통해 추정되는 유연한 생존 회귀 모델인 Deep Cox Mixtures를 제안한다. 이 방법은 특히 소수 민족에 대해 보다 뛰어난 校정성(calibration) 향상을 이룩하며, 기존의 고전적 및 딥 러닝 기반 모델들을 능가한다.

ABSTRACT

Survival analysis is a challenging variation of regression modeling because of the presence of censoring, where the outcome measurement is only partially known, due to, for example, loss to follow up. Such problems come up frequently in medical applications, making survival analysis a key endeavor in biostatistics and machine learning for healthcare, with Cox regression models being amongst the most commonly employed models. We describe a new approach for survival analysis regression models, based on learning mixtures of Cox regressions to model individual survival distributions. We propose an approximation to the Expectation Maximization algorithm for this model that does hard assignments to mixture groups to make optimization efficient. In each group assignment, we fit the hazard ratios within each group using deep neural networks, and the baseline hazard for each mixture component non-parametrically. We perform experiments on multiple real world datasets, and look at the mortality rates of patients across ethnicity and gender. We emphasize the importance of calibration in healthcare settings and demonstrate that our approach outperforms classical and modern survival analysis baselines, both in terms of discriminative performance and calibration, with large gains in performance on the minority demographics.

연구 동기 및 목표

• 비례 위험 가정이 성립하지 않을 경우 의료 환경에서 생존 모델의 낮은 校정성 문제를 해결한다.
• 인종 및 성별 소수 민족과 같은 대표성이 부족한 인구 집단에서의 예측 성능—특히 校정성—을 향상시킨다.
• 잠재적 환자 하위군을 학습함으로써 비례 위험 가정을 초월하는 민감한 생존 모델을 개발한다.
• 향상된 校정성이 분류 성능(콘cordance 지수로 측정)에 손해를 주지 않음을 입증한다.
• 실제 임상 데이터셋에 적용 가능한 실용적이고 오픈소스인 딥 러닝 프레임워크를 제공한다.

제안 방법

• 각 구성 요소가 고유한 위험 함수를 가지는 잠재적 환자 하위군을 나타내는 코크스 회귀 성분의 혼합 모델을 사용해 생존 결과를 모델링한다.
• 환자 특징 기반으로 그룹별 위험 비율을 추정하기 위해 딥 뉴럴 네트워크를 사용하여 비선형 관계를 유연하게 모델링한다.
• 비모수적 스퍼링 기반 스무딩을 통해 각 그룹 내에서 기준 생존 함수를 추정하여 모수적 분포 가정을 피한다.
• 최적화 효율성을 높이기 위해 혼합 성분에 대한 하드 할당을 사용하는 근사적인 몬테카를로 기대값최대화(EM) 알고리즘을 구현한다.
• E단계는 캡킨-마이어 추정치를 사용해 근사하고, M단계는 그룹별 파라미터에 대해 확률적 경사 하강법을 통해 수행한다.
• 특히 인종과 성별에 따른 하위군 수준의 성능에 중점을 두어 불확실성과 校정성 메트릭을 모델 평가에 통합한다.

실험 결과

연구 질문

• RQ1 딥 뉴럴 네트워크 기반 위험 비율을 갖는 코크스 모델 혼합 모델이 표준 코크스 및 딥 러닝 기반 모델 대비 생존 예측의 校정성을 향상시킬 수 있는가?
• RQ2 제안된 모델이 블랙 및 여성 환자와 같은 대표성이 부족한 인구 집단에서 생존 예측의 성능 격차를 줄일 수 있는가?
• RQ3 잠재적 하위군을 통해 비례 위험 가정을 완화함으로써 분류 성능 및 校정성 향상에 어느 정도 기여하는가?
• RQ4 실제 데이터셋인 SUPPORT와 같은 데이터에서 추정된 하위군별 생존 곡선이 경험적 캡킨-마이어 추정치와 어떻게 비교되는가?
• RQ5 모델의 해석 가능성과 하위군 탐지 기능이 실질적인 임상 의사결정 지원에 어떤 영향을 미치는가?

주요 결과

• Deep Cox Mixtures는 모든 인구 집단에서 校정성을 크게 향상시키며, 특히 블랙 및 여성 환자와 같은 소수 민족에서 가장 뚜렷한 성능 향상을 보였다.
• 최첨단 기준 모델의 콘코르던스 지수(C-index)를 유지하거나 초월하여, 분류 성능와 校정성 간의 상충관계가 없음을 입증했다.
• 하위군별 생존 곡선이 서로 교차함으로써 비비례 위험 구조가 아님을 나타내었으며, 이는 모델이 잠재적 클러스터링을 통해 성공적으로 이를 포착함을 시사한다.
• 하위군에 대한 캡킨-마이어 추정치가 모델 예측 생존 곡선과 밀접하게 일치하여, 모델이 의미 있는 환자 하위집단을 회복할 수 있음을 검증했다.
• 기존 도구인 2013 ACC/AHA 풀드 코hort 방정식에서 알려진 바와 같이, 블랙 인구 집단에서 심장혈관 위험을 과대평가하는 문제를 모델이 줄였다.
• GitHub에 공개된 오픈소스 구현을 통해 재현성과 향후 의료 및 연구 응용 분야에서의 확장 가능성이 보장된다.