Tae Hyun Kim (Lowell)

R-Learner

정의

R-Learner(Residualized Learner)는 Robinson Transformation을 토대로, 결과와 처치(treatment)에서 각각 공변량(covariate)의 영향을 걷어낸 잔차(residual)를 써서 CATE를 추정하는 메타러너(meta-learner)다.

알고리즘:

Step 1: nuisance 함수 추정 (cross-fitting 사용) m^(x)=E^[YX=x],e^(x)=P^(W=1X=x)\hat{m}(x) = \hat{E}[Y|X=x], \quad \hat{e}(x) = \hat{P}(W=1|X=x)

Step 2: R-Loss 최소화 τ^=argminτL^n(τ)+Λn(τ)\hat{\tau} = \arg\min_\tau \hat{L}_n(\tau) + \Lambda_n(\tau)

여기서: L^n(τ)=1ni=1n[{Yim^(q(i))(Xi)}{Wie^(q(i))(Xi)}τ(Xi)]2\hat{L}_n(\tau) = \frac{1}{n} \sum_{i=1}^n \left[ \{Y_i - \hat{m}^{(-q(i))}(X_i)\} - \{W_i - \hat{e}^{(-q(i))}(X_i)\} \tau(X_i) \right]^2

  • m^(q(i))\hat{m}^{(-q(i))}: ii번째 관측치를 뺀 fold에서 추정한 m^\hat{m}
  • Λn(τ)\Lambda_n(\tau): 정규화(regularization) 항

직관적 이해

핵심 아이디어:

결과와 처치 양쪽에서 공변량의 영향을 모두 제거한 뒤, 순수한 처치효과(treatment effect)만 학습한다.

Step 1: Estimate nuisance functions
        m̂(x) = E[Y|X]    (outcome model)
        ê(x) = P(W=1|X)  (propensity model)

Step 2: Compute residuals (via cross-fitting)
        Ỹᵢ = Yᵢ - m̂(Xᵢ)        (outcome residual)
        W̃ᵢ = Wᵢ - ê(Xᵢ)        (treatment residual)

Step 3: Minimize R-loss
        τ̂ = argmin Σ[Ỹᵢ - W̃ᵢ·τ(Xᵢ)]² + regularization

왜 “R”인가?

  • Residual(잔차)을 쓴다.
  • Robinson transformation에 기반한다.
  • 또는 저자 이름의 이니셜이다.

핵심 성질

Quasi-Oracle 성질

핵심 정리 (Theorem 1): nuisance 성분을 o(n1/4)o(n^{-1/4}) rate로 추정하면, R-learner는 참 nuisance 함수를 아는 oracle과 같은 수렴 속도(convergence rate)에 도달한다.

τ^τ2=OP(lognn)+oP(1)nuisance error\|\hat{\tau} - \tau^*\|^2 = O_P\left(\frac{\log n}{n}\right) + o_P(1) \cdot \text{nuisance error}

의미:

  • nuisance 추정의 오차가 1차 항으로는 영향을 주지 않는다.
  • nuisance 추정이 느려도 괜찮다(n1/4n^{-1/4}만 충족하면 된다).

직교성(orthogonality)

Robinson Transformation의 직교성 조건은 다음과 같다. E[(Ym(X))(We(X))X]=0E[(Y - m^*(X)) \cdot (W - e^*(X)) | X] = 0

이 조건 덕분에 nuisance 오차에 대한 강건성(robustness)을 확보한다.

관심사의 분리

  1. 교란 통제: Step 1에서 m^,e^\hat{m}, \hat{e}를 추정한다.
  2. 처치효과 추정: Step 2에서 CATE에만 집중한다.

각 단계마다 서로 다른 ML 방법을 쓸 수 있다.

알고리즘 상세

def r_learner(X, W, Y, base_learner, n_folds=5):
    from sklearn.model_selection import KFold

    n = len(Y)
    m_hat = np.zeros(n)  # outcome residuals
    e_hat = np.zeros(n)  # treatment residuals

    # Step 1: Cross-fitted nuisance estimation
    kf = KFold(n_splits=n_folds, shuffle=True)

    for train_idx, val_idx in kf.split(X):
        # Fit outcome model
        outcome_model = base_learner.fit(X[train_idx], Y[train_idx])
        m_hat[val_idx] = outcome_model.predict(X[val_idx])

        # Fit propensity model
        propensity_model = base_learner.fit(X[train_idx], W[train_idx])
        e_hat[val_idx] = propensity_model.predict(X[val_idx])

    # Compute residuals
    Y_tilde = Y - m_hat  # outcome residual
    W_tilde = W - e_hat  # treatment residual

    # Step 2: Minimize R-loss
    # τ(x) = argmin Σ(Ỹᵢ - W̃ᵢ·τ(Xᵢ))²
    # This is equivalent to weighted least squares:
    # Ỹᵢ/W̃ᵢ ≈ τ(Xᵢ) with weight W̃ᵢ²

    # Pseudo-outcome for regression
    pseudo_outcome = Y_tilde / np.clip(W_tilde, 1e-6, None)
    weights = W_tilde ** 2

    # Fit CATE model (weighted regression)
    tau_model = base_learner.fit(X, pseudo_outcome, sample_weight=weights)

    return tau_model.predict

다른 메타러너와의 비교

AspectS-LearnerT-LearnerX-LearnerR-Learner
Models1243 (m, e, τ)
Data usageAll togetherSplitCross-groupAll + cross-fitting
TargetsResponseResponseImputed effectsResidualized outcome
Key featureSimpleSeparateImbalance handlingOrthogonality
Best whenCATE ≈ 0Different μ₀, μ₁Unbalanced groupsSimple CATE, complex nuisance

언제 쓰는가

잘 맞는 상황

  • CATE가 nuisance보다 단순할 때: 직교성이 nuisance의 복잡도가 미치는 영향을 최소화한다.
  • 교란(confounding)은 복잡하지만 처치효과는 간단할 때
  • cross-validation이 중요할 때: 각 단계마다 하이퍼파라미터를 따로 튜닝할 수 있다.

맞지 않는 상황

  • 성향점수(propensity score)가 극단적일 때: e(x)0e(x) \approx 0 또는 11이면 불안정하다.
  • 표본 크기가 매우 작을 때: cross-fitting 과정에서 데이터를 잃는다.
  • nuisance 추정이 어려울 때: n1/4n^{-1/4} rate를 못 맞추면 보장이 사라진다.

Rate 조건

Quasi-Oracle 달성 조건

m^m2e^e2=oP(n1/2)\|\hat{m} - m^*\|_2 \cdot \|\hat{e} - e^*\|_2 = o_P(n^{-1/2})

또는 이와 동치로: m^m2=oP(n1/4),e^e2=oP(n1/4)\|\hat{m} - m^*\|_2 = o_P(n^{-1/4}), \quad \|\hat{e} - e^*\|_2 = o_P(n^{-1/4})

수렴 속도

적절한 정칙성(regularity) 조건에서는 다음과 같다.

  • Penalized kernel regression: O(nα/(2α+d))O(n^{-\alpha/(2\alpha+d)}), 여기서 α\alpha는 smoothness다.
  • Linear CATE: parametric rate O(n1/2)O(n^{-1/2})가 가능하다.

시뮬레이션 결과 (원논문)

Setup A (복잡한 nuisance, 단순한 CATE):

  • R-learner가 가장 강한 성능을 보인다.
  • 직교성이 교란 효과를 제거한다.

Setup B (RCT, 일정한 성향점수):

  • R-learner ≈ T-learner.
  • 특별한 이점이 없다.

Setup C (쉬운 성향점수, 복잡한 baseline):

  • R-learner가 경쟁력이 있다.
  • X-learner와 비슷한 성능을 낸다.

Setup D (서로 무관한 처치군):

  • T-learner가 유리하다.
  • R-learner는 데이터 공유의 이점을 보지 못한다.

관련 개념

  • Meta-learners - 전체 framework
  • Robinson Transformation - 이론적 기반
  • R-Loss - 최적화 목적 함수
  • Quasi-Oracle Property - 핵심 이론적 보장
  • Cross-fitting - 과적합(overfitting) 편향(bias) 제거
  • S-Learner, T-Learner, X-Learner - 대안 방법들
  • DR-Learner - 이중 강건(doubly robust) 접근
  • CATE - 추정 대상
  • Propensity Score - nuisance 성분

구현

Python (econml):

from econml.dml import NonParamDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier

# R-learner is closely related to DML
r_learner = NonParamDML(
    model_y=RandomForestRegressor(),
    model_t=RandomForestClassifier(),
    model_final=RandomForestRegressor(),
    cv=5  # cross-fitting folds
)
r_learner.fit(Y, T, X=X)
cate = r_learner.effect(X_test)

R (rlearner package):

library(rlearner)

# Using Random Forest as base learner
r_rf <- rlasso(X, W, Y)  # or rboost, etc.
cate <- predict(r_rf, X_test)

참고 문헌

  • nieQuasiOracleEstimationHeterogeneous2020 - R-learner 원논문
  • chernozhukovDoubleDebiasedMachine2018 - 관련 DML 이론

연결 그래프