Gradient-Based Causal Tree Ensembles: A Backbone Architecture for Heterogeneous Treatment Effects

Motivation

• Known: representation learning 기반 neural network 방법들이 HTE 추정에서 주목받아 왔으나, Curth et al. (2021) 등의 최근 연구는 tabular 데이터에서 Causal Forests 같은 tree-based 방법이 특정 data-generating process에서 neural network보다 우수할 수 있음을 보였다. 또한 tabular 데이터 지도학습 전반에서 tree-based 모델이 uninformative feature가 존재할 때 neural network보다 강건한 경향이 알려져 있다.
• Gap: 기존 differentiable tree 모델(Deep Neural Decision Trees, NODE, GRANDE 등)은 supervised learning에는 적용되었지만 causal/HTE 추정에 특화된 아키텍처가 부재했고, 특히 multi-way split을 지원하는 dense representation의 일반화와 routing 연산의 높은 계산 비용 문제가 미해결 상태였다.
• Why: 헬스케어와 정책 결정 분야에서 무작위 실험이 불가능한 상황에서 observational data로부터 정확하고 강건한 HTE 추정이 필수적이며, tree 기반 구조의 장점과 neural network의 end-to-end 최적화 이점을 결합한 모듈형 백본은 다양한 기존 HTE 아키텍처에 손쉽게 통합될 수 있어 실용적 가치가 크다.
• Approach: GRACE는 multi-way, oblique, soft split을 가진 다수의 트리를 완전히 파라미터화하여 PyTorch 등 최신 ML 프레임워크에서 backpropagation으로 end-to-end 학습 가능하게 하며, 기존 FCNN layer를 직접 대체할 수 있도록 설계되었다.

Achievement

Figure 4. Performance comparison using semi-synthetic datasets.

1. GRACE 아키텍처 제안: multi-way, oblique, soft split을 결합한 미분 가능 causal tree ensemble로, FCNN layer와 seamless하게 교체 가능한 모듈형 구조를 구현.
2. 광범위한 벤치마크에서 성능 우위: binary 및 non-binary(multi-valued, bundle, continuous, multi-dimensional) treatment 설정을 포함한 다양한 synthetic 및 semi-synthetic 데이터셋에서 neural network 및 tree-based baseline 대비 지속적이고 상당한 성능 향상을 달성.
3. FCNN과의 이론적 연결 및 ablation 분석: GRACE layer를 FCNN layer의 확장으로 해석하는 표기법 비교를 제시하고, 각 아키텍처 요소(multi-way split, oblique split, soft split 등)가 성능 향상에 기여하는 정도를 정량적으로 분리.

How

• N개의 관측 샘플 (xi, t_i^f, y_i^f)을 Neyman-Rubin potential outcomes framework 하에서 정식화하고, Conditional Treatment Response Function(CTRF)과 CATE(Conditional Average Treatment Effect)를 추정 목표로 설정
• Unconfoundedness와 Overlap 가정 하에 GRACE layer를 여러 층으로 쌓아 binary/non-binary treatment 설정에 대응하는 전용 아키텍처 구성 (Figure 1)
• 각 GRACE layer는 random feature subset, tree weighting, multi-way/oblique/soft split routing, multi-dimensional leaf score 산출 단계로 구성되며 GRANDE 등 기존 differentiable tree 모델의 dense representation을 multi-way split 지원 형태로 일반화
• routing 연산의 계산 비용을 줄이는 cost-reduction 전략을 도입해 효율성 개선
• RMSE(PEHE) 및 multi-valued/continuous/multi-dimensional treatment에 대응하는 일반화된 RMSE metric을 사용해 synthetic·semi-synthetic 벤치마크에서 성능 평가 및 ablation study 수행

Originality

• differentiable tree 모델을 causal inference/HTE 추정에 특화하여 재설계한 최초 시도 중 하나로, GRANDE 등 기존 supervised learning용 differentiable tree와 차별화
• dense tree representation을 multi-way split으로 일반화하여 깊이 증가 없이 표현력을 확장한 점이 참신
• routing 연산의 계산 비용을 줄이는 새로운 cost-reduction 메커니즘을 제안해 효율성과 표현력을 동시에 확보
• FCNN layer의 일반화로서 GRACE를 이론적으로 재해석하고 구성 요소별 기여도를 분리하는 ablation 프레임워크 설계

Limitation & Further Study

• 초록과 본문 발췌만으로는 실제 실험 데이터셋의 규모, 실제 real-world 데이터 적용 결과, computational overhead의 정량적 비교가 충분히 제시되지 않아 실용성 검증이 제한적일 수 있음
• synthetic 및 semi-synthetic 벤치마크 위주의 평가로, 실제 관측 데이터에서의 unconfoundedness 가정 위반 시 강건성에 대한 논의가 부족
• 후속 연구로는 대규모 real-world 헬스케어/정책 데이터셋에서의 검증, treatment 공간이 매우 큰 경우의 확장성 분석, 그리고 이론적 수렴성/일관성(consistency) 보장에 대한 통계적 분석이 필요

Evaluation

총평: GRACE는 differentiable tree ensemble을 HTE 추정에 특화하여 확장한 견고한 아키텍처로, 모듈성과 실험적 우수성을 겸비한 유망한 백본으로 평가되나 real-world 검증과 이론적 보장에 대한 추가 연구가 필요하다.