はじめに

こちらの記事でSurrogate indexについて初めて知りました。

developers.cyberagent.co.jp

Netflixでも活用されていて、一定の成果を上げているようです。

arxiv.org

業務で扱う課題を解決してくれる可能性があったため、理解しておきたいと思いました。

論文等を読みつつRで挙動などを確認していきます。

はじめに
概要
定義
関係性
因果効果の推定
論文での適用事例
Rで確認

概要

元の論文はこちらです。

www.nber.org

よりわかりやすい説明は冒頭のブログ記事を参照してください。

課題感

施策の因果効果を検討するときに、A/Bテスト等を行う場合があります。

しかし、この時に検証する因果効果は短期的なものであるときが多く、長期的な因果効果を測るのは難しい場合が多く存在します。

施策の長期的な影響を測るためには、時間がかかることや単純な中間指標を用いたような推定では不正確な場合があるためです。

そこで短期的な実験の結果を利用して、長期的な因果効果を算出することを考えています。

方針

基本的な方針は、短期に観測される中間指標を多く組み合わせた結果変数「代替指標(Surrogate index)」を作成し評価をしていきます。

この時に、それぞれの変数に統計的な代用性があるかには注目せずに、複数の中間変数があることが重要になります。

具体的には、複数の中間指標 $S_i$ を説明変数とした線形回帰を用いて長期の結果の予測値を形成し、その予測値に対して治療効果を推定することになります。

多くの変数を組み合わせることで、処置 $W$ から長期結果 $Y$ は生まれるまでの全ての因果経路を網羅する可能性が高まる。

これにより、長期的な治療効果をより迅速かつ正確に推定できる（標準誤差が小さくなる）可能性があります。

この時、もし処理 $W$ の効果が短期的にはプラスであっても長期的にはマイナスであったり、先物喰いがある場合には短期の代替指標 $S_i$ の長期結果 $Y$ への影響が小さく、長期の代替指標 $S_i$ の長期結果 $Y$ への影響が大きく評価されることになる。

仮定

はじめに2つの異なるデータセットを持っていると仮定する。

実験データセット( $E$ )：処置 $W$ と中間指標 $S$ に関するデータで、 $W$ がランダムに割り当てられている
観察データセット( $O$ )：Sと長期の結果 $Y$ に関するデータ、 $W$ も含まれる場合もあるが、 $W$ はランダムに割り当ていない

これらのデータに対してルービンの因果モデルに従った潜在的な結果変数が二つ存在する。

$\displaystyle{ \begin{split} Y_i=\left\{\begin{array}{ll}Y_i(0) & \text { if } W_i=0, \\ Y_i(1) & \text { if } W_i=1,\end{array} \quad\right. \quad \text{and } S_i = \left\{ \begin{array}{ll}S_i(0) & \text { if } W_i=0, \\ S_i(1) & \text { if } W_i=1,\end{array} \right . \end{split} }$

今回は実験データセットにおける因果効果を推定したい。

$\displaystyle{ \tau = \mathbb{E}[ Y_i(1) - Y_i(0) | P_i = E ] }$

しかし、実験開始から長時間が経過しないと得られないため、代理変数を用いることになる。

Surrogate indexを用いて因果関係を推定するには 3 つの仮定が必要となる

仮定1 強く無視できる割り当て条件

実験グループにおいて、処置に関する背景情報を調整することにより、処置の平均的な因果関係を推定できることを意味する。

$\displaystyle{ (i)\quad W_i \perp\left(Y_i(0), Y_i(1), S_i(0), S_i(1)\right) \mid X_i, P_i=\mathrm{E} }$

ここで $e(X_i)$ は、処置を受ける条件付き確率：傾向スコア。

仮定2 代理性

代理変数を用いて、処置効果を推定するためには、代理変数と結果変数の間の因果経路を補足している必要がある。

これは実験グループにおいて、効果処置 $W$ と $Y$ が代理変数 $S$ や共変量により独立になることに対応する。

$\displaystyle{ W_i \perp Y_i \mid S_i, X_i, P_i = \mathrm E. }$

操作変数法における除外制約に相当する。

仮定3 互換性

以上の仮定に加えて実験サンプル( $E$ )と観察サンプル( $O$ )の関係を定義する必要があります。

$\tau$ の関係を推定するために、実験サンプルと観察サンプルにおける結果変数の条件付き確率分布が一致することが求められます。

$\displaystyle{ Y_i\left|S_i, X_i, P_i=\mathrm{O} \sim Y_i\right| S_i, X_i, P_i=\mathrm{E} }$

この仮定は次のように言い換えることができます。

$\displaystyle{ P_i \perp Y_i | S_i, X_i }$

定義

Surrogate indexを用いて処理効果ATEを推定するには、いくつかの値の定義がされています。

Surrogate index：代理指数

代理指数 $h$ )は、観察サンプルの中間変数と共変量を考慮した場合の結果変数の条件付き期待値です。

$\displaystyle{ h_0(s, x) = \mathbb{E} [ Y_i | S_i = s, X_i = x, P_i = O ] }$

線形モデルでは、中間指標 $S$ 、共変量 $X$ に対する主な結果の回帰からの予測値となる。

代理スコア：Surrogate Score

代理スコアは、代理指標と共変量の値が与えられた場合の処置を受ける条件付き確率

傾向スコアと同様に、代理スコアは多次元の統計的な情報をスカラーに変換し、他のサンプルとの比較を簡易的にできる。

$\displaystyle{ r(s, x) = pr(W_i = 1 | S_i = s, X_i = s, P_i = E) }$

Sampling Score

観察サンプルと実験サンプルを比較できるようにする重み。

$\displaystyle{ t(s, x) = \mathrm{pr}\left(P_ = E| S_i = s, X_i = x\right) = \frac{ \mathrm{pr(S_i = s, X_i = x | P_i = E})q }{ \mathrm{ pr(S_i = s, X_i = x| P_i = E) q + pr(S_i = s, X_i = x | P_i = \mathrm{O})(1 - q) } } .}$

潜在的条件付き期待値

治療前の変数とサロゲートが与えられた場合の潜在的な結果の条件付き期待も考えておく。

$\displaystyle{ \mu _E(s, x, w) = \mathbb{E} \left[Y_i | S_i = s、X_i = x、W_i = w、P_i = E \right ] }$

関係性

以上の定義と仮定から以下のような関係も導かれる

(i) 仮定2 が成り立つと仮定

$\displaystyle{ \mu_E (s, x, w) = h_E(s, x), \quad \text{all} s \in \mathbb{S}, \quad x \in \mathbb{X}, \text{and } w \in \mathbb{W} }$

(ii) 仮定3 が成り立つと仮定

$\displaystyle{ h_E(s, x) = h_O(s, x) \text{ all } s \in S、\text{ and } x \in X }$

(iii) 仮定2と3 が成り立つと仮定

$\displaystyle{ \mu_E(s, x, w) = h_O(s, x) \text{ all } s \in S、x \in X、\text{ and } w \in W。 }$

因果効果の推定

最終的な平均処置効果の推定方法を3種類提案している。

複数の推定方法があることで、実験やデータの取得状況で推定方法を変えることができるため。

一つ目の表現は、実験サンプルにおいてSurrogate indexで推定されたものとコントロールとの間の違いを傾向スコアで調整している。

この表現では代替指数の推定が必要だが、代替スコアを推定する必要がない。

$\displaystyle{ \tau^{\mathrm{E}}=\mathbb{E} \left [h_{\mathrm{O}}\left(S_i, X_i\right) \cdot \frac{W_i}{e\left(X_i\right)}-h_{\mathrm{O}}\left(S_i, X_i\right) \cdot \frac{1-W_i}{1-e \left(X_i\right)} \mid P_i= \mathrm{E} \right ] }$

ここで $h_{\mathrm{O}}$ は観測サンプルで得られた代理指標の確率分布/統計モデルである。所謂IPWによる推定に対応している。

第二の表現では観察サンプルの結果 2 つの加重平均の差の期待値として書かれる。重みは、実験サンプルとサンプリングスコアで推定された代理スコアの関数です。

この方法では代替スコアの推定が必要であるが、代替指数を推定する必要がない。

$\displaystyle{ \tau^{\mathrm{O}}=\mathbb{E} \left [Y_i \frac{r \left(S_i, X_i\right) t \left(S_i, X_i\right) \left(1-q\right)}{e \left(X_i \right) (1 - t \left(S_i, X_i\right) )q }- Y_i \frac{(1-r \left(S_i, X_i\right))t \left(S_i, X_i\right)(1-q)}{(1-e \left(X_i\right)) (1-t \left(S_i, X_i\right))q} \mid P_i= \mathrm{O} \right ] }$

第三の表現では両方の推定が必要です。

$\displaystyle{ \tau^{\mathrm{O}, \mathrm{E}} = \mathbb{E} [ \psi (P_i, Y_i, S_i, W_i, X_i) ]}$

$\psi$ は次のように定義される。

$\displaystyle{ \begin{split} \begin{aligned} \psi(p, y, s, w, x) =&\frac{1_{p=\mathrm{E}}}{q}\left(\frac{h_{\mathrm{O}}(s, x) w}{e(x)}-\frac{h_{\mathrm{O}}(s, x)(1-w)}{1-e(x)}\right) \\ &+ \frac{1_{p=\mathrm{O}}}{1-q}\left(\frac{t(s, x)}{1-t(s, x)} \frac{1-q}{q}\right) \frac{\left(y-h_{\mathrm{O}}(s, x)\right)(r(s, x)-e(x))}{e(x)(1-e(x))} . \end{aligned} \end{split} }$

ここで $1_{ p }$ は、実験サンプルと観測サンプルのどちらに該当するかを示すインデックスである。

仮定 1 ～ 4 が成立する場合、平均治療効果は、データの次の 3 つの推定可能な関数と等しくなります。

$\displaystyle{ \tau \equiv \mathbb{E} [ Y_i(1) - Y_i(0)|P_i = \mathrm{E} ] = \tau ^ \mathrm{E} = \tau ^ \mathrm{O} =\tau ^ {\mathrm{O, E}}. }$

論文での適用事例

論文では、扶養家族手当の受給による雇用率の改善に関するデータに対して適用されています。

参加者のうち22％が雇用され、ランダム割り当て前の四半期の平均収入は453ドルでした。

処置グループには4405人、対照グループには1040人の参加者が含まれていました。

処置の割り当ての後9年間、失業保険データベースからの四半期ごとの雇用率と収益を測定し、研究参加者を追跡しています。

プログラムの雇用率と収益に対する処置効果は最初は大きかったものの、時間とともに低下していました。

そのため短期の雇用への影響を代理指標として使用することで、9年間の雇用全体への影響をより迅速に推定できる可能性があるかどうかを検証しています。

なお、結果変数の推定には線形回帰モデルを利用している、

$\displaystyle{ \bar{Y}_{iT} = \beta_0 +\sum_{t=1}^S \beta_t Y_{it} + \epsilon_i }$

そして、因果効果は次のように算出される

$\displaystyle{ \hat \tau_{T,S} = \frac{1}{N_1} \sum_{i=1}^N \hat Y_{i,T,S} W_i- \frac{1}{N_0} \sum_{i=1}^N \hat Y_{i,T,S} (1-W_i) }$

結果として推定された9年間の全体の雇用への影響は+6.2%であった。

Surrogate indexを利用した推定では、ナイーブな推定よりも早い段階でバイアスの小さい推定が行うことができている。ナイーブな推定：処置実施後tステップ後の雇用率の単純な差

信頼区間の推定も可能。

Rで確認

Rでサンプルデータを作成してSurrogate indexを試してみました。

残念ながらパッケージはありませんでしたが、Susan Athey自身がRでの実行を解説しているページがありましたので、こちらのページを参考にしています。

d2cml-ai.github.io

データの生成

以下のようなデータを生成しました。

生成したサンプルは観測データ、実験データともに2000としています。

$\displaystyle{ X_1 \sim Normaal(0, 1) \\ X_2 \sim Normaal(0, 1) \\ X_3 \sim Poisson(3) \\ X_4 \sim Gamma(1, 1) \\ X_5 \sim Gamma(1, 1) \\ B = 7.3X_1 +4.1X_2 + 7.2X_3 - 8.8X_4 +2.1X_5 - 2.5X_2X_3 + 1000 \\ Y_{t} \sim Gamma(B+E_t, 100) \\ }$

$E_i$ は設定しているそのタイミングの処置効果です。

そして観測期間は10ステップとして、9ステップ目まではそれまでの結果変数の平均値を中間指標、10ステップ目までの平均値 $\bar Y_t$ を結果変数としている。

推定

推定には、単純な観測時刻の期待値からの算出を行うナイーブな場合、 Atheyのコードに習った回帰モデルを用いたSurrogate Indexを用いる場合、 Surrogate Indexを用いたIPW推定型を用いる場合の3パターンを試した。

なかなかのクソコードです。

y_reg = list()
surrogate_index = list()
surrogate_score = list()
sampling_score = list()
propensity_score = list()

naive_index = list()
emp_weight_p = list()
emp_weight_se = list()

for (i in 1:10) {
  if (i >1){
    example_df[,paste("Y", i, sep="")] = example_df[,paste("Y", i, sep="")] + example_df[,paste("Y", i-1, sep="")]
    observed_df[,paste("Y", i, sep="")] = observed_df[,paste("Y", i, sep="")] + observed_df[,paste("Y", i-1, sep="")]
  }
}

all_df = bind_rows(example_df, observed_df) 

## Store regressions
for (i in 1:9) {
  y_reg[[i]] = lm(formula = S_eq(i), data = observed_df) #%>% step(direction = "both")
  emp_weight_p[[i]] = as.matrix(y_reg[[i]]$coefficients)
  emp_weight_se[[i]] = confint(y_reg[[i]])
  
  example_df[,'y_surrogate'] = predict(y_reg[[i]], newdata = example_df)
  surrogate_index[[i]] = lm("y_surrogate ~ W", data = example_df)
  surrogate_score[[i]] = glm(formula = SS_eq(i), data = example_df, family = binomial(link = "logit")) #%>% step(direction = "both")
  sampling_score[[i]] = glm(formula = SP_eq(i), data = all_df, family = binomial(link = "logit"))
  example_df[, paste('surrogate_index',i, sep="")] = predict(y_reg[[i]], newdata = example_df)
  observed_df[, paste('surrogate_score',i, sep="")] = predict(surrogate_score[[i]], newdata = observed_df, type="response")
  observed_df[, paste('sampling_score',i, sep="")] = predict(sampling_score[[i]], newdata = observed_df, type="response")
  
  if(i ==1){
    propensity_score_model = glm(W ~ X1+X2+X3+X4+X5, data = example_df, family = binomial(link = "logit")) #%>% step(direction = "both")
    observed_df[,'obs_propensity_score'] = predict(propensity_score_model, newdata = observed_df, type="response")
    example_df[,'exp_propensity_score'] = predict(propensity_score_model, newdata = example_df, type="response")
  }
  
}