[논문 리뷰] Optimal Transport for Treatment Effect Estimation

Hao Wang, Zhichao Chen외 8명|arXiv (Cornell University)|2023. 10. 27.

Advanced Causal Inference Techniques인용 수 5

한 줄 요약

이 논문은 최소 제곱 오차(MSE)와 관측되지 않은 혼동요인 영향(UCE)을 해결하기 위해 유연한 질량 보존 정규화자와 프록시 사실적 결과 정규화자를 도입한, 최적 운반 기반의 새로운 조건부 평균 치료 효과(CATE) 추정 방법인 ESCFR를 제안한다. 이는 확률적 최적 운반을 통해 전체 데이터 공간에서 치료 그룹을 정렬함으로써 관찰 데이터에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Estimating conditional average treatment effect from observational data is highly challenging due to the existence of treatment selection bias. Prevalent methods mitigate this issue by aligning distributions of different treatment groups in the latent space. However, there are two critical problems that these methods fail to address: (1) mini-batch sampling effects (MSE), which causes misalignment in non-ideal mini-batches with outcome imbalance and outliers; (2) unobserved confounder effects (UCE), which results in inaccurate discrepancy calculation due to the neglect of unobserved confounders. To tackle these problems, we propose a principled approach named Entire Space CounterFactual Regression (ESCFR), which is a new take on optimal transport in the context of causality. Specifically, based on the framework of stochastic optimal transport, we propose a relaxed mass-preserving regularizer to address the MSE issue and design a proximal factual outcome regularizer to handle the UCE issue. Extensive experiments demonstrate that our proposed ESCFR can successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.

연구 동기 및 목표

• 관측 데이터에서 치료 선택 편향을 해결하기 위해, 치료군과 대조군 간의 공변량 이동으로 인한 문제를 다루는 것.
• 기존 방법에서 미니배치 샘플링 오차(MSE)로 인해 분포 차이 추정이 왜곡되는 문제를 해결하기 위해, 이상치나 불균형한 배치에 민감하지 않은 방법을 개발하는 것.
• 기본적인 무혼동성 가정 하에서 분포 차이 계산이 편향되는 관측되지 않은 혼동요인 영향(UCE)을 완화하기 위한 방법 개선.
• 지역적 미니배치가 아닌 전체 데이터 공간에서 작동하는, 확률적 최적 운반에 기반한 체계적인 인과 추론 프레임워크를 개발하는 것.
• 관측되지 않은 혼동요인의 영향을 고려한 프록시 정규화자를 통합하여 반사적 결과 추정을 향상시키는 것.

제안 방법

• 관측 데이터에서 CATE를 추정하기 위해 확률적 최적 운반에 기반한 전체 공간 반사적 회귀(ESCFR) 프레임워크를 제안한다.
• 이상치나 불균형한 미니배치에 의한 영향을 줄이기 위해, 유연한 질량 보존 정규화자를 도입하여 분포 정렬의 안정성을 높인다.
• 관측되지 않은 혼동요인의 결과에 대한 단조성 영향을 모델링함으로써, 그들의 영향을 고려한 프록시 사실적 결과 정규화자를 설계한다.
• 전체 데이터 공간에서의 분포 차이를 계산하기 위해 일반화된 싱크호른 차이를 사용하여, 잠재적으로 편향된 미니배치 추정에 의존하지 않는다.
• 대표성 매핑과 반사적 결과를 동시에 학습하기 위해 이중 최적화 프레임워크를 활용하여, 분포 정렬과 결과 충실도를 모두 확보한다.
• 관측되지 않은 혼동요인과 결과 간의 단조성 가정을 적용하여, 가능한 운반 전략의 집합을 줄이고 추정 정확도를 향상시킨다.

실험 결과

연구 질문

• RQ1 어떻게 미니배치 샘플링 오차가 분포 차이 추정에 영향을 미치는지를 완화함으로써 관측 데이터에서 CATE 추정 성능을 향상시킬 수 있는가?
• RQ2 기존 최적 운반 기반 CATE 추정기에서 관측되지 않은 혼동요인이 얼마나 편향을 유발하는가? 이를 어떻게 수정할 수 있는가?
• RQ3 지역적 미니배치가 아닌 전체 데이터 공간에서 작동하는 체계적인 최적 운반 프레임워크를 설계할 수 있는가?
• RQ4 사실적 결과에 프록시 정규화자를 통합함으로써, 관측되지 않은 혼동이 존재하는 상황에서 반사적 예측 정확도는 어떻게 향상되는가?
• RQ5 관측되지 않은 혼동요인과 결과 간의 단조성 가정을 활용하여 운반 전략 선택의 모호성을 줄일 수 있는가?

주요 결과

• ESCFR는 다양한 벤치마크 데이터셋에서 최신 기술 수준의 방법들을 뛰어넘는 성능을 보이며, 분포 이동에 대한 뛰어난 내구성을 입증한다.
• 유연한 질량 보존 정규화자는 이상치와 결과 불균형이 있는 미니배치에 미치는 영향을 효과적으로 줄여 훈련 과정의 안정성과 정렬 정확도를 향상시킨다.
• 프록시 사실적 결과 정규화자는 특히 단조성 가정이 성립할 경우 관측되지 않은 혼동요인에 의한 편향을 줄여 더 정확한 반사적 예측을 가능하게 한다.
• 실험 결과, 특히 높은 수준의 관측되지 않은 혼동이 존재하는 환경에서 ESCFR는 기존 기준 방법들보다 더 낮은 평균 제곱 오차(MSE)를 기록한다.
• 비선형적 결과 관계와 숨겨진 혼동요인이 존재하는 다양한 데이터셋에서도 일관된 성능 향상을 보였다.
• 제거 실험 결과, 질량 보존 정규화자와 프록시 정규화자가 성능 향상에 기여하며, 이들의 설계 논리가 검증됨을 확인했다.