[논문 리뷰] Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes

Ahmed M. Alaa, Mihaela van der Schaar|arXiv (Cornell University)|2017. 04. 10.

Gaussian Processes and Bayesian Inference참고 문헌 32인용 수 90

한 줄 요약

베이지안 다중 작업 가우시안 프로세스 프레임워크를 사용해 관찰 데이터에서 개인화된 치료 효과를 추정하고, 불확실성 정량화와 선택 편향 완화를 위한 위험 기반 경험적 베이즈를 제공합니다; IHDP 및 UNOS 데이터셋에서 baselines 대비 PEHE 성능이 향상됨을 보였습니다.

ABSTRACT

Predicated on the increasing abundance of electronic health records, we investi- gate the problem of inferring individualized treatment effects using observational data. Stemming from the potential outcomes model, we propose a novel multi- task learning framework in which factual and counterfactual outcomes are mod- eled as the outputs of a function in a vector-valued reproducing kernel Hilbert space (vvRKHS). We develop a nonparametric Bayesian method for learning the treatment effects using a multi-task Gaussian process (GP) with a linear coregion- alization kernel as a prior over the vvRKHS. The Bayesian approach allows us to compute individualized measures of confidence in our estimates via pointwise credible intervals, which are crucial for realizing the full potential of precision medicine. The impact of selection bias is alleviated via a risk-based empirical Bayes method for adapting the multi-task GP prior, which jointly minimizes the empirical error in factual outcomes and the uncertainty in (unobserved) counter- factual outcomes. We conduct experiments on observational datasets for an inter- ventional social program applied to premature infants, and a left ventricular assist device applied to cardiac patients wait-listed for a heart transplant. In both experi- ments, we show that our method significantly outperforms the state-of-the-art.

연구 동기 및 목표

• 관찰 데이터에서 개인화된 치료 효과(ITE) 추정의 필요성과 EHR의 선택 편향 문제를 다룸.
• 잠재 결과를 벡터 값을 갖는 함수로 모델링하고 ITE 추정을 vvRKHS의 다중 작업 학습 문제로 해석합니다.
• 처치와 대조 결과 간 공유 및 작업별 구조를 포착하기 위해 선형 코어그래니제이션 커널을 갖춘 CMGPs를 개발합니다.
• 핵심 하이퍼파라미터를 사실적 오차와 반사실적 불확실성의 균형으로 조정하기 위해 위험 기반 경험적 베이즈를 도입합니다.
• ITE에 대한 베이지안 신뢰 구간을 제공하고, semi-synthetic IHDP 및 UNOS 데이터셋에서 다수의 베이스라인과 비교 평가합니다.

제안 방법

• ITE를 벡터 값을 갖는 함수 f = [f_0, f_1]로 연결하기 위해 Y(i)(w) = f_w(X_i) + ε_i,w를 가우시안 노이즈로 형식화합니다.
• f를 벡터 값 커널 K_theta를 갖는 가우시안 프로세스로 모델링하고 선형 코어그래니제이션 모델을 사용하여 f_0와 f_1이 서로 다른 공분산(K_theta = A_0 k_0 + A_1 k_1)을 가질 수 있도록 합니다.
• ITE T(x) = e^T f(x)에서 e = [-1, 1]^T를 사용하고, T가 커널 평가의 스팬에 속함을 보이는 표현자 정리를 도출하며, 이를 베이지안 GP의 후방 평균과 연결합니다.
• 사실적 오차와 후방 반사실적 불확실성을 최소화하여 커널 하이퍼파라미터를 선택하기 위한 베이지안 PEHE 위험 및 위험 기반 경험적 베이즈를 도입합니다.
• 하이퍼파라미터의 ADAM 기반 최적화, 후방 ITE 계산 및 T(x)에 대한 신뢰 구간 추정 등을 포함한 CMGPs의 알고리즘을 설명합니다.
• 선택적으로: 베이지안 정규화를 이용한 leave-one-out 교차 검증의 사용으로 하이퍼파라미터 선택의 안정성을 보장합니다.

실험 결과

연구 질문

• RQ1 관찰 데이터에서 잠재 결과에 대한 다중 작업 GP 사전을 사용하여 개인화된 치료 효과를 정확히 추정할 수 있는가?
• RQ2 /counterfactual의 불확실성을 어떻게 측정하고 전파하여 명시적 경향성 모델링 없이도 신뢰 구간을 구성할 수 있는가?
• RQ3 핵심 하이퍼파라미터에 대한 위험 기반 경험적 베이스가 표준 GP 혹은 단일 작업 접근법에 비해 선택 편향에 대한 로버스트니스를 개선하는가?
• RQ4 제안된 CMGP 접근법이 IHDP의 반합성 및 실제 UNOS 데이터셋에서 PEHE 측면의 최첨단 베이스라인과 비교하여 어떤 성능을 보이는가?
• RQ5 임상 대기자 시나리오에서 이질적 치료 효과에 CMGP를 적용할 때 정밀 의료에 어떤 실용적 시사점이 생기는가?

주요 결과

벤치마크	IHDP 인-샘플 PEHE	IHDP 외부 샘플 PEHE	UNOS 인-샘플 PEHE	UNOS 외부 샘플 PEHE
CMGP	0.9 ± 0.07	1.0 ± 0.08	1.7 ± 0.10	1.8 ± 0.13
GP	2.1 ± 0.11	2.3 ± 0.14	4.1 ± 0.15	4.5 ± 0.20
BART	2.0 ± 0.13	2.2 ± 0.17	3.5 ± 0.17	3.9 ± 0.23
CF	2.3 ± 0.21	2.4 ± 0.23	3.8 ± 0.25	4.3 ± 0.31
VTRF	2.5 ± 0.26	2.9 ± 0.51	4.5 ± 0.35	4.9 ± 0.41
CFRF	2.7 ± 0.31	3.3 ± 0.72	4.7 ± 0.21	5.2 ± 0.32
BLR	5.9 ± 0.31	6.1 ± 0.41	5.7 ± 0.21	6.2 ± 0.30
BNN	2.1 ± 0.11	2.2 ± 0.13	3.2 ± 0.10	3.3 ± 0.12
CFRW	1.0 ± 0.07	1.2 ± 0.08	2.7 ± 0.07	2.9 ± 0.11
kNN	3.2 ± 0.12	4.2 ± 0.22	5.2 ± 0.11	5.4 ± 0.12
MS	2.8 ± 0.18	2.9 ± 0.20	4.6 ± 0.12	4.8 ± 0.16
TML	4.9 ± 0.23	4.9 ± 0.23	6.2 ± 0.31	6.2 ± 0.31

• CMGP는 IHDP 및 UNOS에서 인-샘플 및 외부 샘플 PEHE 모두에서 광범위한 베이스라인(GP, BART, CF, CFRF, Balancing regressors, kNN, MS, TML)을 능가합니다.
• 위험 기반 경험적 베이지스 하이퍼파라미터 선택은 사실적 결과를 적합시키면서 반사실적 불확실성을 규제하여 선택 편향을 감소시킵니다.
• 선형 코어그래니제이션 커널은 f0와 f1 간의 이질성 및 관련 특징을 서로 다르게 포착하게 하며 W–X 상호작용을 무시하는 방법보다 우수합니다.
• ITE에 대한 신뢰 구간이 생성되어 치료 결정에 대한 개별화된 신뢰도 평가를 가능하게 합니다.
• UNOS LVAD 대기자 맥락에서 실용적 임상 활용성을 보여주며 개인 위험을 고려한 배정 의사 결정의 잠재적 개선을 시사합니다.