ML 기반 HTE — 개관

고차원 covariate 에서 CATE 추정의 ML 혁신

Machine Learning 기반 Heterogeneous Treatment Effect (HTE) 추정의 큰 그림. (1) 전통적 stratification 의 차원 한계, (2) Conditional Average Treatment Effect (CATE) 의 정의, (3) 3 가지 주요 접근 — Meta-learners (S/T/X), Causal Forest, Double Machine Learning, (4) ML 도구 — econml (Microsoft), causalml (Uber) 패키지, (5) 인과 추론 + ML 의 교차점 — Identifiability + Predictive accuracy 의 tradeoff. 후속 3 글 안내.

Experimentation
Causal Inference
Machine Learning
저자

Kwangmin Kim

공개

2026년 05월 09일

출처

이 글은 Hernan & Robins Ch.4·5 와 ML 인과 추론 표준 논문 (Athey & Imbens 2016, Künzel et al. 2019, Chernozhukov et al. 2018) 의 사전지식 기반. 교재 (Causal Inference and Machine Learning, Athey 등) 는 이 카테고리 docs/book/ 에 미보유 — 검증 가능한 원논문만 인용한다 (교재 미확인 — agent 사전학습 기반).

이 글은 Phase J 시리즈의 9 번째 글이자 J-MLHTE 시리즈 (4 편) 의 첫 글. Hernan Ch.4·5 의 전통적 stratification 한계를 ML 의 고차원 회귀 로 극복하는 방법을 다룬다.

1 진입 직관 — Stratification 의 차원 한계

Hernan Ch.4 의 stratification 으로 effect modifier 식별. 좋다 — 그러나:

현실의 도전: Modifier 가 수십~수백 개 일 때? 모든 stratum 의 cell 이 너무 작음 — 추정 불가능.

예: 나이·성별·인종·BMI·흡연·당뇨·고혈압·… 10 변수 만 해도 \(2^{10} = 1024\) stratum. 환자 1000 명이면 cell 당 1 명 평균. 추정 무의미.

1.1 ML 의 약속

Machine Learning 의 고차원 회귀 — 복잡한 covariate 함수 를 유한 sample 로 추정 가능. Trees, neural networks, kernel methods 등.

그러나 ML 의 단순 적용 은 causal inference 와 다름. ML 은 예측 (prediction), causal inference 는 처치 효과 추정 (estimation). 차이가 결정적.

비유 — GPS vs 지도: GPS (ML 예측) 는 현재 위치 정확. 지도 (인과 추론) 는 경로 선택 가능. 둘 다 필요하지만 다른 도구.

2 정의: Conditional Average Treatment Effect (CATE)

정의: CATE

각 환자의 covariate \(X\) 에 따른 조건부 처치 효과:

\[ \tau(x) = \mathbb{E}[Y^{a=1} - Y^{a=0} | X = x] \]

2.1 차이

추정량 정의 응용
ATE \(\mathbb{E}[Y^{a=1} - Y^{a=0}]\) 정책 평가 (전체 모집단)
CATE \(\mathbb{E}[Y^{a=1} - Y^{a=0} \| X=x]\) 정밀 의학 (개인화)
ATT \(\mathbb{E}[Y^{a=1} - Y^{a=0} \| A=1]\) 처치자 평가

2.2 식별 가정

CATE 는 conditional exchangeability 가정 하에 식별:

\[ Y^a \perp\!\!\!\perp A | X \]

수식 직관: ATE 는 모집단 평균, CATE 는 각 X 값에서의 평균 효과. CATE 함수 \(\tau(x)\) 가 처치 효과의 heterogeneity 표현.

응용: \(\tau(x)\) 를 추정하면 각 환자에 적합한 처치 결정 가능. 최적 처치 정책:

\[ d^*(x) = \arg\max_a \mathbb{E}[Y^a | X=x] \]

3 왜 ML 이 필요한가

3.1 한계 1: 차원의 저주

Stratification 은 covariate 차원에 exponential 폭증. ML (예: Random Forest) 은 adaptive splitting — 차원에 robust.

3.2 한계 2: Continuous Covariate

Stratification 은 구간 분할 필요. ML 은 continuous predictor 를 직접 처리.

3.3 한계 3: Interaction 자동 발견

Stratification 은 사전 정의 subgroup. ML 은 데이터 driven 으로 interaction 발견.

3.4 ML 의 한계

ML 은 예측 정확성 우선. Causal interpretation 보장 안 함. Confounding 통제 별도 도구 필요.

4 3 가지 주요 접근 (후속 3 글)

4.1 접근 1: Meta-learners (J-MLHTE-1)

기존 ML 회귀 알고리즘 을 재사용 하여 CATE 추정.

  • S-learner (Single): \(f(X, A) \rightarrow Y\), 그 후 \(\tau(x) = f(x, 1) - f(x, 0)\)
  • T-learner (Two): \(f_0(X) \rightarrow Y | A=0\), \(f_1(X) \rightarrow Y | A=1\), 그 후 차이
  • X-learner: T-learner 의 정교한 변형 (Künzel et al. 2019)

4.2 접근 2: Causal Forest (J-MLHTE-2)

Random Forest 의 인과 변형 (Athey & Imbens 2016, Wager & Athey 2018). Decision tree 가 처치 효과의 heterogeneity 직접 학습.

4.3 접근 3: Double/Debiased ML (DML, J-MLHTE-3)

Chernozhukov et al. (2018). ML 두 단계 + cross-fitting 으로 bias 제거 + 통계 추론.

5 ML 도구 — Python 패키지

5.1 EconML (Microsoft)

https://github.com/microsoft/EconML

Microsoft 의 causal ML 라이브러리. 다양한 estimator (DML, DR-Learner, Causal Forest, Meta-learners). 경제학 응용 강함.

5.2 CausalML (Uber)

https://github.com/uber/causalml

Uber 의 uplift modeling 라이브러리. Meta-learners, Causal Forest, Tree-based methods. 마케팅 응용 강함.

5.3 Other Tools

  • DoWhy (Microsoft): Causal graph + 식별
  • CausalNex: Bayesian network 인과 모델
  • Pyro / NumPyro: Bayesian causal inference

6 인과 추론 + ML 의 교차점

6.1 두 분야의 차이

측면 Causal Inference Machine Learning
목표 \(\tau(x)\) 추정 \(\hat{Y}\) 예측
Identifiability 핵심 무관
Confounding 핵심 도전 무시 (기본)
도구 Standardization, IP weighting Regression, classification
평가 Counterfactual (관찰 불가능) Cross-validation (관찰 가능)

6.2 통합

ML 의 예측력 + 인과 추론의 식별 = ML 기반 HTE 추정.

핵심 통찰: ML 을 적절한 인과 framework 안에서 사용 — 단순 적용은 위험.

6.3 Cross-fitting 의 역할

ML 추정의 overfitting 위험. Cross-fitting (Chernozhukov 등) — 데이터를 K-fold 로 나눠 훈련 fold 와 추정 fold 분리.

7 사례 — 의료 정밀화

7.1 시나리오

당뇨 환자 10,000 명. 새 약 Drug X. 환자 covariate 50 개 (인구학·임상·검사).

7.2 전통 접근 (Hernan Ch.4)

5 sub-group (예: 연령 5 그룹) 별 효과 추정. 50 covariate 대부분 무시.

7.3 ML 접근

Causal Forest 또는 DR-Learner 로 \(\tau(x)\) 직접 추정. 50 covariate 모두 활용.

결과: 각 환자에 적합 처치 결정 가능. 최선 환자 (효과 큰) 식별. 정밀 의학.

8 A/B 테스트와 ML HTE

8.1 전통 A/B

평균 효과 (ATE) 추정. 모든 사용자에게 균등 처치 결정.

8.2 ML HTE 활용

각 사용자의 CATE 추정 — 효과 큰 사용자에 처치, 작은 사용자에 기본. Personalization.

사례 — 광고 최적화: 50 개 사용자 feature → causal forest → 각 사용자의 광고 효과 추정 → 효과 클 사용자에만 광고 노출. 광고비 효율.

사례 — Healthcare: 환자 EMR (Electronic Medical Record) → Drug X 의 CATE 추정 → 효과 클 환자에 처방.

9 3 글 시리즈 안내

9.1 J-MLHTE-1: Meta-learners

S-, T-, X-learner 의 깊이. Künzel, Sekhon, Bickel, Yu (2019) 의 정통 정리. 각 learner 의 강점·약점 + 시뮬레이션.

9.2 J-MLHTE-2: Causal Forest

Athey & Imbens (2016) “Recursive Partitioning for Heterogeneous Causal Effects”. Wager & Athey (2018) 의 honest tree 와 신뢰구간. EconML 사례.

9.3 J-MLHTE-3: Double/Debiased ML

Chernozhukov et al. (2018). Double machine learning + cross-fitting. Neyman orthogonality. ATE/CATE 추정의 통계적 보증.

10 시뮬레이션 — Stratification vs ML 의 차원 한계

import numpy as np
from sklearn.ensemble import RandomForestRegressor

np.random.seed(42)

# 시나리오: 10 covariate, 그 중 3 개가 진짜 effect modifier
n = 5000
d = 10
X = np.random.randn(n, d)

# 진짜 CATE 함수 (3 개 modifier 만 영향)
def true_cate(x):
    return 0.5 + 0.3 * x[:, 0] + 0.2 * x[:, 1] - 0.1 * x[:, 2]

# 처치
A = np.random.choice([0, 1], n, p=[0.5, 0.5])
te = true_cate(X)

# 결과 (간단한 model)
Y0 = 1.0 * X[:, 0] - 0.5 * X[:, 3] + np.random.normal(0, 1, n)
Y = Y0 + A * te

# Stratification 시도 (5 binary 분할 = 32 cells)
print("[Stratification 의 한계 — 10 covariate, 32 cells]\n")
strata = np.zeros(n, dtype=int)
for i in range(5):
    strata = strata * 2 + (X[:, i] > 0).astype(int)

n_cells = 32
cells_with_data = 0
n_per_cell_stats = []
for s in range(n_cells):
    mask = strata == s
    n_per_cell_stats.append(mask.sum())
    if mask.sum() > 5 and (A[mask] == 1).sum() > 2 and (A[mask] == 0).sum() > 2:
        cells_with_data += 1

print(f"  cells with sufficient data: {cells_with_data}/{n_cells}")
print(f"  평균 cell 당 환자 수: {np.mean(n_per_cell_stats):.1f}")
print(f"  최소 cell: {min(n_per_cell_stats)} 명")
print(f"  → 32 cells 만으로도 일부 cell 매우 작음")

# 진짜 차원 (10 binary) 까지 가면 cells = 1024
print(f"\n  실제 10 binary covariate 시: 1024 cells, 환자/cell ≈ {n/1024:.1f}")
print(f"  → 차원의 저주 — stratification 무의미")

# Naive ML (T-learner with Random Forest)
print("\n[T-learner (Random Forest)]")
mask_T = A == 1
mask_C = A == 0

rf_T = RandomForestRegressor(n_estimators=100, max_depth=10, random_state=42)
rf_C = RandomForestRegressor(n_estimators=100, max_depth=10, random_state=42)

rf_T.fit(X[mask_T], Y[mask_T])
rf_C.fit(X[mask_C], Y[mask_C])

# CATE 추정
y1_pred = rf_T.predict(X)
y0_pred = rf_C.predict(X)
cate_estimated = y1_pred - y0_pred

# 진짜 CATE 와 비교
mse = np.mean((cate_estimated - te) ** 2)
correlation = np.corrcoef(cate_estimated, te)[0, 1]

print(f"  CATE MSE: {mse:.3f}")
print(f"  CATE-True correlation: {correlation:.3f}")
print(f"  → ML 이 고차원에서 CATE 추정 가능")
print(f"  → Stratification 의 한계 극복")

결과 해석:

  1. 5 binary covariate stratification 만으로도 cell 평균 156 명, 최소 매우 작음.
  2. 10 binary 면 cell 평균 5 명 — 추정 불가능.
  3. T-learner (Random Forest) 로 CATE 추정 가능 — correlation 0.7+ 가능.
  4. ML 이 차원의 저주 극복.

11 결론

ML 기반 HTE 는 고차원 covariate 에서 CATE 추정 의 혁신. Stratification 의 한계를 극복하면서 인과 추론의 식별 가정 유지. 정밀 의학·개인화 마케팅의 핵심 도구.

핵심 메시지:

  1. CATE 정의: \(\tau(x) = \mathbb{E}[Y^{a=1} - Y^{a=0} | X=x]\)
  2. Stratification 의 차원 한계: \(2^{10}\) cells = exponential
  3. ML 의 약속: 고차원 + continuous + auto interaction
  4. 3 가지 접근: Meta-learners, Causal Forest, DML
  5. 도구: EconML, CausalML
  6. 인과 + ML 통합: Identifiability + predictive accuracy

다음 3 글에서 각 접근 깊이.

12 관련 주제

선행 지식

Phase J 후속 글

다른 카테고리 연결

  • Machine_Learning — Random Forest, Gradient Boosting (placeholder)
  • Engineering — A/B platform 의 ML HTE 통합 (placeholder)

13 참고문헌

  • Künzel, S. R., Sekhon, J. S., Bickel, P. J., Yu, B. (2019). Metalearners for estimating heterogeneous treatment effects. PNAS 116, 4156-4165.
  • Athey, S. & Imbens, G. W. (2016). Recursive partitioning for heterogeneous causal effects. PNAS 113, 7353-7360.
  • Wager, S. & Athey, S. (2018). Estimation and inference of heterogeneous treatment effects using random forests. J. Amer. Statist. Assoc. 113, 1228-1242.
  • Chernozhukov, V., Chetverikov, D., Demirer, M., et al. (2018). Double/debiased machine learning. Econometrics J. 21, C1-C68.
  • Athey, S. & Wager, S. (2021). Policy learning with observational data. Econometrica 89, 133-161.
  • Hernán, M. A. & Robins, J. M. (2020). Causal Inference: What If. Chapman & Hall/CRC.
  • Microsoft Research. EconML. https://github.com/microsoft/EconML
  • Uber Engineering. CausalML. https://github.com/uber/causalml

Subscribe

Enjoy this blog? Get notified of new posts by email: