Online Survival Analysis: A Bandit Approach under Cox PH Model

1. 문제 정의 (Problem Statement)

• 배경: 생존 분석 (시간 - 사건 데이터) 은 의학, 공학, 고객 이탈 분석 등 다양한 분야에서 널리 사용되지만, 기존 연구는 주로 오프라인 (Offline) 환경에서 수행되었습니다. 즉, 모든 데이터 수집이 완료된 후 분석을 수행하는 방식입니다.
• 도전 과제:
  1. 온라인 학습의 부재: 새로운 데이터가 순차적으로 들어오는 환경에서 치료법을 실시간으로 최적화하는 연구가 부족합니다.
  2. 지연된 피드백 (Delayed Feedback): 사건 (예: 사망, 재발) 이 발생할 때까지 시간이 걸리므로 즉각적인 보상 (Reward) 을 얻기 어렵습니다.
  3. 오른쪽 중도절단 (Right Censoring): 관측 기간 내에 사건이 발생하지 않아 데이터가 불완전하게 관측되는 문제가 발생합니다.
  4. 단계적 진입 (Staggered Entry): 연구 대상자들이 서로 다른 시점에 연구에 참여하여, 위험 집합 (Risk Set) 이 시간에 따라 복잡하게 변합니다. 이는 기존 밴디트 알고리즘의 이론적 가정 (마팅게일 성질 등) 을 무효화합니다.
• 목표: 위와 같은 제약 조건 (중도절단, 지연 피드백, 단계적 진입) 하에서 Cox PH 모델을 기반으로 한 컨텍스트 밴디트 (Contextual Bandit) 알고리즘을 개발하여, 누적 후회 (Regret) 를 최소화하는 최적의 치료 정책을 학습하는 것입니다.

2. 방법론 (Methodology)

저자는 Cox PH 모델을 온라인 밴디트 환경에 적용하기 위해 다음과 같은 프레임워크를 제안합니다.

A. 모델 설정 및 업데이트

• Cox PH 모델: 위험 함수를 $\lambda(s | X) = \lambda_0(s) \exp(X^\top \beta)$로 정의합니다. 여기서 $X$는 공변량과 행동 (치료법) 의 결합 특징 벡터입니다.
• 온라인 부분우도 (Partial Likelihood) 업데이트:
  • 기존 오프라인 방식은 매번 데이터를 다시 피팅하여 계산 비용이 $O(N^2)$로 비효율적입니다.
  • 제안된 방법은 증분 업데이트 (Incremental Update) 방식을 사용합니다. 새로운 데이터가 들어오거나 중도절단 상태가 변경될 때, 이전의 부분우도 함수를 재사용하여 $P_1$ (새로운 사건 기여) 과 $P_2$ (위험 집합 변화 보정) 항을 추가하는 방식으로 효율적으로 업데이트합니다. 이 방식은 계산 복잡도를 $O(N)$ 수준으로 줄여줍니다.
• 단계적 진입 처리: 캘린더 시간 ($\tau$) 과 생존 시간 ($s$) 의 두 가지 시간 척도를 고려하여, 각 시점에서의 위험 집합 $R(\tau, s)$을 정확히 정의하고 이를 부분우도 함수에 반영합니다.

B. 탐험 - 활용 (Exploration-Exploitation) 전략

세 가지 고전적인 밴디트 알고리즘을 생존 분석 환경에 적응시켰습니다.

1. $\epsilon$-Greedy (EG): 확률적으로 무작위 행동을 선택하거나 현재 추정된 최적 행동을 선택합니다.
2. UCB (Upper Confidence Bound): Cox 모델의 계수 추정치 $\hat{\beta}$에 대한 **신뢰 타원 (Confidence Ellipsoid)**을 구성하여 불확실성을 고려한 상한선을 계산합니다.
3. Thompson Sampling (TS): $\beta$의 사후 분포를 근사 (Laplace 근사 또는 HMC 사용) 하여 샘플링하고, 이를 기반으로 행동을 선택합니다.

C. 이론적 보장 (Theoretical Guarantees)

• 신뢰 타원 (Confidence Ellipsoid): 단계적 진입으로 인해 기존 마팅게일 집중 부등식이 성립하지 않는 문제를 해결하기 위해, 조정된 스코어 프로세스 (Adjusted Score Process) $Q(\tau)$를 정의하여 이를 마팅게일로 변환했습니다. 이를 통해 $\hat{\beta}_t$와 참값 $\beta$ 사이의 오차에 대한 균일한 상한을 유도했습니다.
• 후회 (Regret) 한계: 제안된 알고리즘 (EG, UCB, TS) 에 대해 서브선형 (Sublinear) 후회 한계 $\tilde{O}(d\sqrt{T})$를 증명했습니다. 여기서 $d$는 특징 차원, $T$는 라운드 수입니다. 이는 중도절단과 지연 피드백이 있더라도 최적 정책에 빠르게 수렴함을 의미합니다.

3. 주요 기여 (Key Contributions)

1. 통합된 온라인 생존 분석 프레임워크: 오른쪽 중도절단, 단계적 진입, 지연된 피드백을 동시에 처리하는 최초의 밴디트 프레임워크를 제안했습니다.
2. 유연한 알고리즘 적용: EG, UCB, TS 세 가지 표준 밴디트 전략을 Cox PH 모델에 성공적으로 적용하고, 생존 데이터의 특수성 (위험 집합의 동적 변화) 을 반영한 탐험 전략을 개발했습니다.
3. 이론적 혁신: 단계적 진입으로 인해 기존 마팅게일 도구가 무효화되는 문제를 해결하기 위해 새로운 조정된 스코어 프로세스와 신뢰 타원을 구축하여, 서브선형 후회 한계를 수학적으로 증명했습니다.
4. 효율적인 계산: 매번 모델을 처음부터 학습하는 대신, 부분우도 함수를 증분적으로 업데이트하는 알고리즘을 제안하여 대규모 데이터셋에서의 실시간 학습을 가능하게 했습니다.

4. 실험 결과 (Results)

• SEER 암 데이터 (실제 데이터 기반 실험):
  • SEER (Surveillance, Epidemiology, and End Results) 데이터베이스의 유방암 환자 데이터 (약 29 만 명) 를 사용하여 실험했습니다.
  • 치료법 (유방보존수술 vs 유방절제술) 을 선택하는 문제로 설정하고, 알고리즘이 시간이 지남에 따라 최적의 생존 확률에 빠르게 수렴하는 것을 확인했습니다.
  • Cox PH 모델이 실제 데이터 생성 과정과 완벽히 일치하지 않더라도 (모델 오명세), 제안된 알고리즘은 안정적인 성능을 보이며 생존 확률을 향상시켰습니다.
• 시뮬레이션:
  • 다양한 중도절단 비율과 모델 오명세 상황에서 EG, UCB, TS 알고리즘의 성능을 평가했습니다.
  • 모든 알고리즘이 모수 ($\beta$) 를 빠르게 정확하게 추정하고, 평균 생존 확률을 최적화하는 것을 확인했습니다.
  • 단순한 베이스라인 (관측된 데이터만 사용하는 편향된 방법) 과 비교하여 제안된 방법이 훨씬 낮은 MSE 와 더 높은 생존 확률을 보였습니다.

5. 의의 및 결론 (Significance)

• 임상 및 의사결정 지원: 임상 시험이나 정밀 의학에서 치료 효과를 기다리는 긴 기간 동안에도 실시간으로 치료 전략을 최적화할 수 있게 하여, 환자 생존율을 높이고 비용을 절감할 수 있는 잠재력을 가집니다.
• 이론적 확장: 생존 분석과 강화 학습 (밴디트) 의 교차 영역에서 중요한 이론적 난제 (단계적 진입과 중도절단 하의 집중 부등식) 를 해결함으로써, 향후 관련 연구의 기초를 마련했습니다.
• 실용성: 계산 효율성을 고려하여 대규모 데이터에서도 적용 가능한 알고리즘을 제시함으로써, 실제 산업 현장 (고객 이탈 예측, 기계 고장 예측 등) 에 적용 가능한 실용적인 솔루션을 제공합니다.

요약하자면, 이 논문은 Cox 비례위험 모델을 기반으로 한 온라인 생존 분석을 가능하게 하는 이론적, 알고리즘적, 실증적 토대를 구축하여, 지연된 피드백과 불완전한 데이터 환경에서도 최적의 의사결정을 내릴 수 있음을 입증했습니다.