DR-Learner

정의

DR-Learner는 의사결과(pseudo-outcome)를 공변량(covariate)에 회귀시켜 CATE를 추정하는 2단계 이중 강건(doubly robust) 추정량(estimator)이다.

Stage 1: nuisance 추정

성향점수(propensity score): $\hat{\pi}(x) = P(A = 1 | X = x)$
결과 회귀(outcome regression): $\hat{\mu}_a(x) = E[Y | X = x, A = a]$ for $a \in \{0, 1\}$

Stage 2: 의사결과 회귀 $\hat{\tau}_{DR}(x) = \hat{E}_n[\hat{\varphi}(Z) | X = x]$

여기서 의사결과는 다음과 같다. $\hat{\varphi}(Z) = \underbrace{\hat{\mu}_1(X) - \hat{\mu}_0(X)}_{\text{plug-in}} + \underbrace{\frac{A - \hat{\pi}(X)}{\hat{\pi}(X)(1-\hat{\pi}(X))}(Y - A\hat{\mu}_1(X) - (1-A)\hat{\mu}_0(X))}_{\text{augmentation/correction}}$

직관적 이해

핵심 아이디어:

이중 강건 의사결과를 계산한다(ATE의 효율적 영향함수(EIF)).
이 의사결과를 $X$ 에 회귀시켜 smoothing한다.
CATE의 구조(smoothness, sparsity)를 별도로 활용한다.

Stage 1:  Estimate π̂(x), μ̂₁(x), μ̂₀(x) using any ML method
              ↓
Stage 2:  Compute pseudo-outcome φ̂(Z) for each observation
              ↓
Stage 3:  Regress φ̂ on X to get τ̂(x)

왜 “이중 강건”인가?

$\hat{\pi}$ 또는 $\hat{\mu}$ 중 하나만 정확해도 편향(bias)이 사라진다.
둘 다 틀려도 두 오차의 곱만 남는다: $O(||\hat{\pi} - \pi_0|| \cdot ||\hat{\mu} - \mu_0||)$ .

핵심 성질

이중 강건성

편향 항은 성향점수와 결과 회귀 오차의 곱에만 의존한다. $\text{Bias} = O(||\hat{\pi} - \pi_0|| \cdot ||\hat{\mu} - \mu_0||)$

rate 적응성

CATE의 smoothness $\gamma$ 에 적응한다.
개별 nuisance function의 smoothness $\alpha, \beta$ 와 분리된다.
plug-in 추정량보다 빠른 rate를 달성할 수 있다.

oracle 효율성

다음 조건이 성립하면 oracle rate를 달성한다. $\sqrt{\alpha\beta} \geq \frac{d/2}{\sqrt{1 + \frac{d}{\gamma}}\sqrt{1 + \frac{d}{2s}}}$

여기서 각 기호는 다음을 뜻한다.

$\alpha$ : 성향점수 smoothness
$\beta$ : 결과 회귀 smoothness
$\gamma$ : CATE smoothness
$s$ : harmonic mean smoothness
$d$ : 공변량 차원

알고리즘

# DR-Learner Algorithm
def dr_learner(X, A, Y, n_folds=5):
    # Stage 1: Cross-fitted nuisance estimation
    pi_hat = cross_fit_estimate(X, A, model='classifier')
    mu1_hat = cross_fit_estimate(X[A==1], Y[A==1], model='regressor')
    mu0_hat = cross_fit_estimate(X[A==0], Y[A==0], model='regressor')

    # Stage 2: Compute pseudo-outcomes
    phi_hat = (mu1_hat - mu0_hat) + \
              (A - pi_hat) / (pi_hat * (1 - pi_hat)) * \
              (Y - A * mu1_hat - (1 - A) * mu0_hat)

    # Stage 3: Regress pseudo-outcome on X
    tau_hat = regress(X, phi_hat, model='smoother')

    return tau_hat

다른 learner와의 비교

Method	Key Idea	Pros	Cons
T-Learner	Separate models per treatment	Simple	No sharing across groups
S-Learner	Single model with A as feature	Shares info	May miss heterogeneity
X-Learner	Two-stage imputation	Good for imbalance	Complex
R-Learner	Residualize then regress	Orthogonality	Requires product rate
DR-Learner	DR pseudo-outcome regression	Double robustness, rate adaptation	Stability condition needed

이론적 보장

주요 오차 한계(Theorem 2): $\hat{\tau}_{DR}(x) - \tilde{\tau}(x) = \hat{E}_n[\hat{b}(X) | X = x] + o_P(\sqrt{R_n^*(x)})$

여기서 각 항은 다음을 뜻한다.

$\tilde{\tau}(x)$ : oracle 추정량(참 의사결과를 사용)
$\hat{b}(x)$ : nuisance 추정에서 오는 편향
$R_n^*(x)$ : oracle 분산(variance)

안정성 조건이 필요하다: 2단계 회귀 추정량은 입력 섭동(input perturbation)에 안정적이어야 한다.

비교: DR-Learner vs R-Learner

Aspect	DR-Learner	R-Learner
Pseudo-outcome	$\hat{\mu}_1 - \hat{\mu}_0 + \text{correction}$	$(Y - \hat{\mu})(A - \hat{\pi})/\text{var}$
Rate condition	Product rate	Product rate
Oracle condition	$\sqrt{\alpha\beta} \geq \ldots$	Weaker for lp-R-Learner
Implementation	Simpler	More complex (lp version)

활용

의료: 임상시험에서의 이질적 처치효과(treatment effect)
정책: 하위집단별 정책 효과
마케팅: 개인화된 처치 반응
사회과학: 인과효과의 이질성

구현

Python (econml):

from econml.dr import DRLearner
dr = DRLearner(model_propensity=LogisticRegression(),
               model_regression=RandomForestRegressor(),
               model_final=RandomForestRegressor())
dr.fit(Y, T, X=X, W=W)
cate = dr.effect(X_test)

R (grf):

library(grf)
# grf의 causal_forest가 유사한 doubly robust 속성 가짐
cf <- causal_forest(X, Y, W)
tau_hat <- predict(cf)$predictions

참고 문헌

kennedyOptimalDoublyRobust2023 - DR-Learner 이론 및 oracle efficiency
chernozhukovDoubleDebiasedMachine2018 - DML framework
nieQuasiOracleEstimationHeterogeneous2020 - R-Learner

정의