Discrete Survival Knowledge Distillation for Competing Risks Analysis

Motivation

• Known: discrete-time neural network 기반 survival model(DeepHit 등)이 competing risks에 적용되어 왔고, knowledge distillation은 분류 문제에서 teacher-student 정보 전이를 통해 성능을 향상시키는 방법으로 널리 활용되어 왔다.
• Gap: 기존 knowledge distillation 방법은 uncensored outcome을 가정하여 censored 관측치에 직접 적용할 수 없으며, teacher와 student가 동일한 outcome 정의(granularity, event 정의)를 공유한다고 가정하기 때문에 competing risks 환경에서 outcome 이질성이 존재할 때 지식 전이가 어렵다는 한계가 있다.
• Why: 신장이식과 같이 rare event, heavy censoring, 소규모 표본, 정책·시기적 분포 변화(distribution shift), 개인정보 보호 제약이 동시에 존재하는 임상 문제에서 외부 정보를 안전하고 강건하게 활용해 예측 정확도를 높이는 것은 실질적 의사결정 지원에 중요하다.
• Approach: discrete-time competing risks 신경망 모델의 log-likelihood를 cause-specific hazard 기반으로 정의하고, 각 시점·원인별 조건부 확률(다항 분포)에 대해 teacher 예측과 student 예측 간 시간에 따른 KL divergence를 정규화항으로 추가하여 student를 학습시킨다.

Achievement

Figure 1. Ctd for each risk and predictive deviance on the simu-

시뮬레이션 및 신장이식 실제 데이터 적용을 통해 discrimination(Ctd)과 calibration(predictive deviance) 측면에서 기존 방법 대비 향상된 성능을 보였으며, model misspecification과 outcome-definition 이질성에도 강건함을 입증하였다.

How

Figure 2. Effect of distillation temperature T on the predictive per-

• discrete-time cause-specific hazard λj(τk; Zi)를 softmax 형태의 다항 로짓 함수로 모델링하고 내부 코호트의 log-likelihood ℓ(r)을 정의
• teacher 모델(Cox model, discrete failure-time model, deep learning model 등 다양한 형태 허용)로부터 cause-specific hazard 예측값 λ̃j(τk; Zi)을 추출
• 각 시점 τk에서 student와 teacher의 조건부 다항 분포 간 cause-specific, time-dependent KL divergence를 계산해 손실 함수에 distillation 항으로 추가
• distillation temperature T를 도입하여 teacher guidance의 강도를 조절하고, 여러 teacher가 존재할 때 적합도에 따라 가중치를 적응적으로 부여해 호환되는 teacher는 강조하고 비호환 teacher는 down-weighting
• teacher와 student 간 covariate 부분 중복 및 outcome granularity(예: overall failure vs cause-specific) 차이를 허용하는 구조 설계
• raw data나 모델 파라미터 공유 없이 teacher의 예측 확률만을 이용해 privacy-conscious 지식 전이 수행

Originality

• 기존 knowledge distillation이 uncensored classification 문제에 국한된 것과 달리, censored competing risks 생존분석에 특화된 discrete-time, cause-specific, time-dependent KL divergence 기반 distillation 목적함수를 최초로 제안
• teacher와 student 간 outcome 정의(granularity, event 정의)가 다를 수 있는 상황(예: 집계된 overall failure teacher vs cause-specific student)까지 포괄하는 프레임워크 설계
• 다양한 유형의 teacher(Cox model, discrete-time model, deep learning model)를 raw data 공유 없이 예측 확률만으로 통합할 수 있는 적응적 가중치 메커니즘 도입
• 기존 transfer learning 방법이 선형 회귀 계수 정렬에 의존하는 것과 달리, 파라미터 일대일 대응을 요구하지 않는 방식으로 deep neural network 기반 teacher-student 통합을 가능하게 함

Limitation & Further Study

• teacher 모델의 적합도 평가 및 가중치 산정 방식이 실제로 얼마나 다양한 teacher 이질성 수준에서 안정적으로 작동하는지에 대한 이론적 보장이 제한적일 수 있음
• distillation temperature T 및 teacher 가중치 등 hyperparameter 선택이 실제 응용에서 최적화하기 까다로울 수 있으며 이에 대한 민감도 분석이 추가로 필요
• 실제 적용 사례가 신장이식이라는 특정 도메인에 집중되어 있어 다른 competing risks 문제로의 일반화 가능성에 대한 추가 검증이 요구됨
• 매우 많은 수의 이질적 teacher가 존재하는 대규모 federated 환경에서의 확장성 및 강건성에 대한 추가 연구가 필요

Evaluation

총평: censored competing risks 환경에 특화된 knowledge distillation을 체계적으로 정식화한 실용적이고 참신한 연구로, 신장이식이라는 실질적 임상 문제에 대한 적용을 통해 방법론적 기여와 실무적 가치를 모두 보여준다.