[논문 리뷰] Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation

Yifan Zhang, Hanlin Zhang외 3명|arXiv (Cornell University)|2022. 02. 02.

Advanced Causal Inference Techniques인용 수 12

한 줄 요약

이 논문은 트랜스포머 기반 백본을 사용하여 치료법과 공변량 간의 상호작용을 자기주의 주의(self-attention)를 통해 모델링함으로써 이질적 치료 효과 추정의 유연성, 강건성, 해석 가능성 향상을 도모하는 TransTEE를 제안한다. 이는 이산, 연속, 구조적 치료법 전반에 걸쳐 효과적으로 작용한다. TransTEE는 다양한 벤치마크에서 강력한 기준 모델들을 능가하며, 사전 훈련된 언어 모델의 해석 가능한 공변량 조정 및 감시 기능을 제공한다.

ABSTRACT

Previous works on Treatment Effect Estimation (TEE) are not in widespread use because they are predominantly theoretical, where strong parametric assumptions are made but untractable for practical application. Recent work uses multilayer perceptron (MLP) for modeling casual relationships, however, MLPs lag far behind recent advances in ML methodology, which limits their applicability and generalizability. To extend beyond the single domain formulation and towards more realistic learning scenarios, we explore model design spaces beyond MLPs, i.e., transformer backbones, which provide flexibility where attention layers govern interactions among treatments and covariates to exploit structural similarities of potential outcomes for confounding control. Through careful model design, Transformers as Treatment Effect Estimators (TransTEE) is proposed. We show empirically that TransTEE can: (1) serve as a general purpose treatment effect estimator that significantly outperforms competitive baselines in a variety of challenging TEE problems (e.g., discrete, continuous, structured, or dosage-associated treatments) and is applicable to both when covariates are tabular and when they consist of structural data (e.g., texts, graphs); (2) yield multiple advantages: compatibility with propensity score modeling, parameter efficiency, robustness to continuous treatment value distribution shifts, explainable in covariate adjustment, and real-world utility in auditing pre-trained language models

연구 동기 및 목표

• 기존의 MLP 기반 치료 효과 추정기의 한계를 해결하기 위해, 일반화 능력이 떨어지고, 파rameter 효율성이 떨어지며, 분포 이탈에 취약한 점을 개선한다.
• 이질적 치료 효과 추정에서 복잡한 치료-공변량 상호작용을 모델링하기 위해 MLP 대신 트랜스포머 백본을 유연한 대안으로 탐색한다.
• 임의의 치료 유형과 호환되는 통합 프레임워크를 개발하여, 성향 스코어 모델링과의 호환성, 파rameter 효율성, 연속적 치료 분포 이탈에 대한 강건성을 확보한다.
• 주의 기반 특성 중요도를 통해 공변량 조정을 해석 가능하게 하고, 실제 환경에서의 인과 추론을 지원한다.
• 사전 훈련된 언어 모델의 공정성 및 사회적 영향을 점검하기 위한 실용성을 입증한다.

제안 방법

• TransTEE는 모든 치료법과 공변량을 동시에 임bedding하는 단일 트랜스포머 백본을 사용하여, 다중 작업 아키텍처를 피하고 엔드 투 엔드 학습을 가능하게 한다.
• 자기주의 주의 메커니즘은 치료법과 공변량 간의 동적 상호작용을 모델링하여, 적응형 공변량 선택 및 혼란 요인 제어를 가능하게 한다.
• 주의 가중치를 활용해 혼란 요인과 관련 공변량을 식별하고 우선순위를 정함으로써, 조건부 무시 가능성을 위한 분리 원칙(disjunctive cause criterion)과 일치시킨다.
• 성향 스코어 모델링과의 호환성을 고려하여 관찰 연구에서의 강건성과 해석 가능성 향상을 높인다.
• 표현형 및 구조적 입력(예: 텍스트, 그래프)을 모두 지원하여 광범위한 적용 가능성을 확보한다.
• 실제 응용의 실용성을 평가하기 위해, 합성 벤치마크를 넘어서는 서브모델링 과제를 도입한다.

실험 결과

연구 질문

• RQ1 트랜스포머 백본은 이산, 연속, 구조적 치료법을 포함한 다양한 치료 유형에서 이질적 치료 효과를 효과적으로 모델링할 수 있는가?
• RQ2 TransTEE는 최신의 MLP 기반 기준 모델들과 비교해 성능, 강건성, 파rameter 효율성 측면에서 어떻게 다른가?
• RQ3 TransTEE의 주의 메커니즘은 인과 추론을 위한 해석 가능하고 의미 있는 공변량 조정을 어느 정도 제공하는가?
• RQ4 TransTEE는 사전 훈련된 언어 모델의 공정성 및 사회적 영향을 점검하는 데 효과적으로 활용될 수 있는가?
• RQ5 TransTEE는 연속적 치료 값의 분포 이탈에 얼마나 강건한가?

주요 결과

• TransTEE는 IHDP, News, TCGA 포함 네 가지 치료 유형에 걸쳐 여섯 개의 벤치마크에서 경쟁 기준 모델들을 뛰어넘으며, 일반화 능력과 성능 향상이 뚜렷하다.
• 반-합성 및 실제 세계 데이터셋에서 평균 치료 효과(Average Treatment Effect, ATE) 추정에서 최신 기술 수준의 성능을 달성했으며, 서브모델링 과제에서 최대 ATE가 0.9976에 도달했다.
• 연속적 치료 값의 분포 이탈에 대해 TransTEE는 강건성을 유지하며, MLP 기반 모델이 성능 저하를 보이는 상황에서도 안정적인 성능을 보였다.
• TransTEE의 주의 가중치는 혼란 요인과 관련 공변량을 효과적으로 식별하고 우선순위를 정함으로써, 인과 식별 원칙과 일치하는 해석 가능한 특성 중요도 점수를 제공한다.
• TransTEE는 사전 훈련된 언어 모델의 공정성 점검에 효과적으로 활용되었으며, 극단적인 ATE 값 0.9976과 0.0을 통해 편향 분석을 통한 사회적 편향 노출을 가능하게 하여 공정성 평가에서의 실용성을 입증했다.
• 이 프레임워크는 파rameter 효율적이며 성향 스코어 모델링과 호환되어 실제 인과 추론 환경에서의 실용적 구현을 강화한다.