Tae Hyun Kim (Lowell)

S-Learner

3분 읽기 #causal-inference#hte#meta-learner

정의

S-Learner(Single Learner)는 처치 지시변수(treatment indicator)를 하나의 특성(feature)으로 포함한 단일 모델로 반응함수(response function)를 추정한 뒤 CATE를 계산하는 Meta-learners다.

알고리즘:

  1. 단일 모델로 결합 반응함수를 추정한다: μ^(x,w)=E^[YX=x,W=w]\hat{\mu}(x, w) = \hat{E}[Y | X = x, W = w]

  2. CATE를 추정한다: τ^S(x)=μ^(x,1)μ^(x,0)\hat{\tau}_S(x) = \hat{\mu}(x, 1) - \hat{\mu}(x, 0)

직관적 이해

핵심 아이디어:

처치 WW를 또 하나의 특성으로 취급하고, 단일 모델로 전체 데이터를 함께 학습한다

Data:  (X, W, Y) for all observations

Model: μ̂(x, w) = f(x, w)  (single model)

CATE:  τ̂(x) = μ̂(x, 1) - μ̂(x, 0)

장점:

  • 가장 간단한 방법이다
  • 모든 데이터를 함께 사용한다(data sharing)
  • 처치군과 대조군 사이의 공통 패턴을 활용한다

단점:

  • 처치효과(treatment effect)가 작으면 무시될 수 있다(regularization이 WW를 떨어뜨린다)
  • μ0\mu_0μ1\mu_1의 구조가 매우 다르면 부적합하다

핵심 성질

데이터 공유(data sharing)

  • 전체 데이터 (n+m)(n + m)개로 하나의 모델을 학습한다
  • 대조군과 처치군의 공통 패턴을 학습할 수 있다

정규화 편향(regularization bias)

μ^(x,w)μ^(x)if treatment effect is small\hat{\mu}(x, w) \approx \hat{\mu}(x) \quad \text{if treatment effect is small}

  • regularization이 강할수록 WW의 영향을 무시하는 경향이 있다
  • CATE ≈ 0일 때 적합하며, 그렇지 않으면 편향(bias)이 발생한다

수렴 속도(convergence rate)

반응함수의 smoothness aμa_\mu에 의존한다: Rate=O((n+m)aμ)\text{Rate} = O((n+m)^{-a_\mu})

알고리즘 상세

def s_learner(X, W, Y, base_learner):
    # Step 1: Combine treatment as feature
    X_combined = np.column_stack([X, W])

    # Step 2: Fit single model
    model = base_learner.fit(X_combined, Y)

    # Step 3: Predict CATE
    def predict_cate(X_new):
        X_treat = np.column_stack([X_new, np.ones(len(X_new))])
        X_ctrl = np.column_stack([X_new, np.zeros(len(X_new))])
        return model.predict(X_treat) - model.predict(X_ctrl)

    return predict_cate

활용

적합한 경우

  • CATE가 대부분 0에 가까울 때: regularization이 올바르게 작동한다
  • 반응함수가 서로 유사할 때: μ0(x)μ1(x)+c\mu_0(x) \approx \mu_1(x) + c
  • 데이터가 제한적일 때: data sharing의 이점을 얻는다

부적합한 경우

  • 처치효과가 명확할 때: 효과가 무시될 수 있다
  • 반응함수가 매우 다를 때: 구조적 차이를 포착하기 어렵다
  • 이질적 효과(heterogeneous effect)가 중요할 때: 미묘한 차이를 놓친다

T-Learner와의 비교

AspectS-LearnerT-Learner
Models12
Data usageAll togetherSplit by treatment
SharingYesNo
Best whenCATE ≈ 0Different response structures
RiskIgnore small effectsNo data sharing

예시

시뮬레이션 설정:

  • μ0(x)=x\mu_0(x) = x
  • μ1(x)=x\mu_1(x) = x (즉 τ(x)=0\tau(x) = 0)

S-Learner 결과:

  • 정규화된 모델이 WW를 무시 → τ^(x)0\hat{\tau}(x) \approx 0
  • 올바른 추정이다

반대 시나리오:

  • μ0(x)=x\mu_0(x) = x, μ1(x)=x+2\mu_1(x) = x + 2 (즉 τ(x)=2\tau(x) = 2)
  • S-Learner가 regularization으로 WW의 영향을 축소할 수 있다 → bias

관련 개념

구현

Python (econml):

from econml.metalearners import SLearner
from sklearn.ensemble import RandomForestRegressor

s_learner = SLearner(overall_model=RandomForestRegressor())
s_learner.fit(Y, T, X=X)
cate = s_learner.effect(X_test)

R:

library(causalToolbox)
s_rf <- S_RF(feat = X, tr = W, yobs = Y)
cate <- EstimateCate(s_rf, X_test)

참고 문헌

  • kunzelMetalearnersEstimatingHeterogeneous2019 - S-learner 분석

연결 그래프