CleanSurvival: Automated data preprocessing for time-to-event models using reinforcement learning

기술 요약: CleanSurvival

문제 정의
데이터 전처리는 머신러닝 파이프라인에서 중요하지만 종종 간과되는 구성 요소로, 모델 성능에 지대한 영향을 미칩니다. 자동 머신러닝 (AutoML) 프레임워크가 분류 및 회귀 작업을 위한 자동화된 전처리를 통합하기 시작했지만, 생존 분석(시간-사건 모델링)을 위한 맞춤형 자동화 솔루션은 명확히 부족합니다. 생존 분석은 사건 시간이 부분적으로만 관찰되는 중도절단 데이터로 인해 고유한 도전을 제시합니다. 기존 AutoML 도구는 중도절단 데이터를 네이티브로 처리하지 못하거나, 수동 전처리에 의존하거나, 시간 이산화를 통해 생존 작업을 분류 문제로 변환하여 정보 손실을 초래합니다. 또한, 현재 자동화된 전처리 도구는 Cox 비례위험 모델이나 Random Survival Forests 와 같은 생존 분석에 필요한 특정 보상 구조와 모델 유형을 지원하지 않는 경우가 많습니다.

방법론
본 논문은 강화 학습(RL) 접근법, 구체적으로 Q-learning을 기반으로 구축된 생존 분석을 위한 자동화된 데이터 전처리 프레임워크인 CleanSurvival을 소개합니다.

• 아키텍처: 이 프레임워크는 전처리 단계 선택을 순차적 의사결정 문제로 간주합니다. 에이전트는 데이터셋의 특성에 의해 정의된 상태 공간을 탐색하여 보상 함수를 최대화하는 행동(전처리 기법) 시퀀스를 선택합니다.
• 전처리 행동: 이 프레임워크는 다음을 위한 다양한 기법을 지원합니다:
  • 결측치 대체: 완전 사례 분석 (CCA), 평균/중앙값 대체, 체인 방정식 다중 대체 (MICE), k-최근접 이웃 (KNN).
  • 이상치 탐지: Mahalanobis 거리를 기반으로 한 타원형 포락선 (Elliptic Envelope) 및 Martingale 잔차.
  • 특성 선택/추출: 단변량 Cox 비례위험 선택, LASSO, 재귀적 특성 제거 (RFE), 정보 이득.
• 생존 모델: 이 시스템은 Cox 비례위험 모델, Random Survival Forests(RSF), DeepHit(딥러닝 접근법), 사용자 제공 모델을 포함한 다양한 시간-사건 모델과 호환됩니다.
• 보상 구조: Q-learning 에이전트는 홀드아웃 세트에서 평가된 일관성 지수(C-index) 를 기반으로 한 보상 신호에 의해 안내됩니다. 이는 선택된 전처리 파이프라인이 하류 생존 모델의 예측 성능을 위해 특별히 최적화되도록 보장합니다.
• 작동 모드: 이 도구는 네 가지 모드를 제공합니다:
  1. 주 알고리즘: Q-learning 최적화.
  2. 무작위 클리닝: 기준 비교를 위한 무작위 탐색.
  3. 사용자 정의 파이프라인: 구성 파일을 통한 사용자 정의 고정 파이프라인.
  4. 준비 없음: 기준 성능을 위한 원시 데이터 입력.

주요 기여

1. 생존 분석을 위한 최초의 RL 기반 전처리: CleanSurvival 은 원래 회귀/분류를 위해 설계된 Learn2Clean 프레임워크를 확장하여 중도절단 시간-사건 데이터의 특정 제약을 처리합니다.
2. 모델별 최적화: 정적 전처리와 달리, 이 프레임워크는 특정 모델 클래스에 최적화된 파이프라인을 동적으로 선택합니다(예: Cox 에 최적화된 파이프라인은 RSF 에 최적화된 파이프라인과 다를 수 있음).
3. 오픈소스 구현: 저자들은 scikit-survival과 같은 표준 라이브러리와 통합된 GitHub 에서 제공되는 Python 패키지를 제공합니다.
4. 중도절단 처리: 이 프레임워크는 평가 중 중도절단 관측치를 고려하기 위해 통합 Brier 점수 (IBS) 와 같은 지표에 중도절단 확률 역가중치 (IPCW) 를 명시적으로 포함합니다.

실험 결과
이 프레임워크는 Rotterdam(유방암), Flchain(혈청 자유 경쇄), GBSG(독일 유방암 연구 그룹) 의 세 가지 실제 세계 데이터셋에서 평가되었습니다. 견고성을 테스트하기 위해 인공 결측 (MCAR, MAR, MNAR) 이 도입되었습니다.

• 성능: CleanSurvival 은 단순 기준선 (완전 사례 분석 및 평균 대체) 일관적으로 능가했으며, C-index 측면에서 종종 무방향 무작위 탐색 및 Optuna(베이지안 최적화) 와 일치하거나 초과했습니다.
  • Flchain데이터셋에서 CleanSurvival 은 0.67 의 C-index 를 달성한 반면, CCA 는 0.39, 평균 대체는 0.62 였습니다.
  • 50% 결측이 있는 Rotterdam데이터셋에서 CleanSurvival 은 0.820.83 사이의 C-index 를 달성하여 CCA(0.770.79) 및 평균 대체 (0.78~0.79) 를 크게 능가했습니다.
• 특성 선택: Q-learning 정책은 특성 선택 연산자를 빈번히 선택했으며, LASSO 는 최적화된 파이프라인의 약 68% 에서, RFE 는 약 42% 에서 나타났습니다.
• 보정: 통합 Brier 점수 (IBS) 는 보정이 일반적으로 결측 메커니즘 전반에 걸쳐 안정적이었음을 나타냈으며, 40% MCAR 조건은 더 높은 분산을 보인 것을 제외합니다.
• 실행 시간: 실행 시간 행동은 조건에 의존적이었습니다. 프레임워크가 고성능 파이프라인을 성공적으로 식별했지만, 최상 관측 C-index 의 95% 에 도달하는 시간은 데이터셋과 결측 메커니즘에 따라 다양했습니다.

의의 및 주장
본 논문은 CleanSurvival이 자동화된 전처리를 생존 분석으로 일반화함으로써 AutoML 생태계의 중요한 격차를 해소한다고 주장합니다. 그 의의는 다음과 같습니다:

• 견고성: 수동 또는 정적 접근법과 비교하여 다양한 수준의 결측 및 노이즈에서 자동화된 RL 기반 전처리가 더 견고한 생존 모델을 산출할 수 있음을 입증합니다.
• 접근성: 중도절단 데이터에 대한 복잡한 전처리 파이프라인을 수동으로 튜닝할 도메인 전문 지식이 부족할 수 있는 실무자들에게 생존 분석 연구를 "더 쉽고 빠르게 수행"할 수 있게 합니다.
• 통합: 데이터 클리닝과 하류 생존 모델링 간의 원활한 연결을 제공하여, 클리닝 결정이 일반적인 데이터 품질 지표가 아닌 생존 예측에 미치는 영향에 의해 직접 안내되도록 보장합니다.

저자들은 겸손한 어조를 유지하며, 결과가 유망하지만 고차원적 다중 모달 임상 데이터셋으로의 일반화는 추가 검증이 필요하며, Q-learning 루프의 계산 오버헤드(특히 KNN 및 딥러닝 모델과 함께 사용 시) 는 여전히 제한 사항임을 인정합니다. 향후 작업은 프레임워크를 추가 전처리 작업으로 확장하고 확장성을 개선하는 것을 제안합니다.