ShrinkageTrees: An R Package for Bayesian Tree Ensembles for Survival Analysis and Causal Inference

기술 요약: "ShrinkageTrees: 생존 분석 및 인과 추론을 위한 베이지안 트리 앙상블 R 패키지"

문제 제기
고차원 설정($p \gg n$)에서의 생존 분석은 세 가지 주요 과제에 직면합니다: (1) 유연한 비모수 모델을 요구하는 공변량과 사건 발생 시간 간의 복잡하고 알려지지 않은 관계; (2) 우측 검열(right-censored) 및 구간 검열(interval-censored) 데이터와 같은 불분명한 관측치의 존재; (3) 검열 및 교란 요인이 존재하는 상황에서 인과적 처치 효과를 추정하는 어려움, 특히 처치 효과가 하위 그룹마다 다를 때의 문제. 기존의 베이지안 가법 회귀 나무(BART) 구현체들은 이 세 가지를 동시에 해결하지 못합니다. 구체적으로, 현재 구간 검열된 결과에 대해 베이지안 트리 앙상블을 제공하는 R 패키지는 없으며, 예후와 처치 효과를 분리하는 $\tau$-learner 아키텍처는 시점(time-to-event) 엔드포인트에 대해 구현된 바 없고, 표준 BART 규제화는 오버피팅 위험이 있는 고차원 유전체 설정에서 종종 불충분합니다.

방법론
본 논문은 가속 실패 시간(Accelerated Failure Time, AFT) 프레임워크 내에서 베이지안 트리 앙상블을 구현하는 ShrinkageTrees R 패키지를 소개합니다. 핵심 모델은 표준 BART 공식을 다음과 같이 확장합니다:
$$\log T_i = \sum_{j=1}^m g(x_i; T_j, H_j) + \epsilon_i, \quad \epsilon_i \sim N(0, \sigma^2)$$
여기서 $T_i$는 사건 발생 시간, $x_i$는 공변량이며, 합계는 $m$개의 얕은 결정 나무들의 앙상블을 나타냅니다. 이 패키지는 깁스 샘플러(Gibbs sampler) 내의 데이터 증강(data augmentation)을 통해 우측 및 구간 검열 데이터를 처리합니다.

방법론은 구현되어 있으며 결합 가능한 두 가지 유형의 규제화를 구분합니다:

1. 트리 구조 규제화: 관련 있는 변수의 부분집합에 대해서만 분할을 유도하기 위해, 패키지는 디리클레 분할 사전 분포(Dirichlet splitting prior, DART)를 구현합니다. 이는 분할 확률에 디리클레 사전 분포를 배치하여 구조적 희소성(structural sparsity)을 유도합니다.
2. 단계 높이 규제화: 정보가 없는 리프(leaf) 값을 0으로 수축시키면서 강한 신호는 보존하기 위해, 패키지는 리프 값($H_j$)에 대한 글로벌-로컬 수축 사전 분포(global–local shrinkage priors)를 구현합니다. 논문은 단계 높이에에 대한 호스슈(horseshoe) 사전 분포를 배치하는 호스슈 포레스트(Horseshoe Forest)를 강조합니다. 이는 글로벌 수축 파라미터($\tau$)와 로컬 스케일 파라미터($\lambda_\ell$)를 사용하는 정규 분포의 척도 혼합(scale mixture of normals) 표현을 활용합니다. 호스슈 사전 분포는 적응형 수축을 허용하여 $p \gg n$ 환경을 효과적으로 다룹니다.

인과 추론을 위해, 이 패키지는 생존 결과에 적응된 $\tau$-learner(베이지안 인과 포레스트)를 구현합니다. 이는 결과를 예후 함수 $\mu(x, \hat{e}(x))$와 처치 효과 함수 $\tau(x)$로 분해하며, 각각 별도의 포레스트에 의해 모델링됩니다. 이를 통해 예후 성분과 인과 성분에 대해 별도의 규제화 및 복잡도 제어가 가능합니다.

주요 기여

• 생존 분석을 위한 최초의 호스슈 포레스트 구현: 이 패키지는 고차원 생존 분석을 가능하게 하기 위해 단계 높이에 호스슈 사전 분포를 배치하는 호스슈 포레스트의 첫 번째 구현체를 제공합니다.
• 구간 검열 지원: 구간 검열된 생존 결과에 대해 베이지안 트리 앙상블을 제공하는 최초의 R 패키지입니다.
• 생존 분석을 위한 인과 추론: 시점 엔드포인트에 대한 $\tau$-learner를 구현하여, 생존 맥락에서의 이질적 처치 효과(CATE) 및 평균 처치 효과(ATE) 추정을 가능하게 합니다.
• 유연한 규제화: 사용자가 구조적 규제화(디리클레 분할)와 단계 높이 규제화(호스슈, 하프-코시 또는 표준 정규 분포)를 결합할 수 있는 모듈형 프레임워크를 제공합니다.
• 소프트웨어 구현: Rcpp를 통한 효율적인 C++ 백엔드를 특징으로 하며, 멀티 체인 MCMC를 지원하고, 전체 S3 워크플로우(적합, 예측, coda를 통한 수렴 진단 및 시각화)를 제공합니다.

결과
논문은 반합성 난소암 데이터셋과 시뮬레이션 연구를 통해 방법론을 검증합니다:

• 난소암 데이터셋: 고차원 설정($n=357, p=1004$)에서 표준 BART는 거의 완벽한 훈련 C-index(0.989)를 달성하여 오버피팅을 나타냈습니다. 반면, 호스슈 포레스트는 더 낮은 훈련 C-index(0.727)를 기록했으나 더 잘 보정된 불확실성을 제공했습니다. 인과 분석은 카보플라틴 대 시스플라틴에 대해 작고 유의하지 않은 평균 처치 효과를 추정하였으며, 좁은 신뢰 구간을 보여주었습니다.
• 시뮬레이션 연구: 알려진 실제값과 구간 검열된 결과를 가진 시뮬레이션에서 표준 BART, DART, 호스슈 포레스트를 차원 $p \in \{50, 500, 5000\}$에 걸쳐 비교했습니다. 호스슈 포레스트는 모든 차원에서 일관되게 가장 낮은 평균 제곱 오차(RMSE), 가장 좁은 신뢰 구간, 그리고 거의 명목적인 95% 커버리지 비율을 달-성했습니다. 표준 BART는 지나치게 넓은 구간(과잉 커버리지)을 생성했고, DART는 차원이 증가함에 따라 커버리지가 부족해졌습니다.
• 계산 성능: 벤치마크 결과, 호스슈 포레스트는 (가역 점프 MCMC 및 로컬 스케일 샘플링으로 인해) 공액 사전 분포(conjugate-prior) 모델보다 계산 비용이 많이 들지만, 표본 크기에 따라 선형적으로 확장되며 기존 BART 구현체의 속도와 유사한 수준임을 보여줍니다.

의의 및 주장
본 논문은 ShrinkageTrees가 생존 분석, 인과 추론, 고차원 규제화를 통합함으로써 기존 베이지안 트리 앙상블 문헌의 중요한 공백을 메운다고 주장합니다. 이 패키지는 표준 BART가 적응하기 어려운 고차원 환경으로 베이지안 트리 앙상블을 확장하는 도구로서 자리매김합니다. 저자들은 이 패키지가 생존 엔드포인트를 위한 $\tau$-learner를 제공하는 최초의 패키지이며, 생존 분석을 위해 호스슈 포레스트를 구현한 최초의 패키지임을 강조합니다. 본 연구는 Jacobs 등(2025)이 개발한 방법론의 실용적인 구현으로 제시되었으며, 베이지안 트리 앙상블에 새로 입문하는 실무자뿐만 아니라 BART, dborts, bcf와 같은 기존 패키지에서 제공되지 않는 기능이 필요한 숙련된 사용자 모두를 대상으로 합니다. 저자들은 수치형 공변량이 필요하다는 점(순서 없는 범주형 변수에 대한 더미 인코딩 필요)과 AFT 프레임워크 내의 정규 분포 오차 가정과 같은 제한 사항을 겸허히 언급하였습니다.