因果推論に入門( 処置効果が非線形なデータで ) / Meta-Learnersの実装・比較

因果分析に興味があるということでこちらの書籍を勉強中。

book.mynavi.jp

こちらの書籍で紹介されている因果推論の手法を整理していきます。
処置効果が線形な場合は一度整理したので、今回は処置効果が非線形場合に、適用可能な手法の整理となります。

目次

Meta-Learnersとは

以前の記事で整理した手法である回帰分析、IPTW法、DR法は因果関係にある変数間の処置効果が線形の場合を想定していました。
それに対してMeta-Learnersとは非線形な因果関係および処置効果を扱える因果推論手法であり、いわゆるCATE(条件付き処置効果)を推定する手法となります。
ただし樹木モデル( ランダムフォレスト、勾配ブースティング木など )やニューラルネットワークといった分析モデルを指すのでは無く、処置効果を求めるために"分析モデルを用いた解き方・使い方"のようなものとなります。
ちなみに今回は分析モデルにはランダムフォレストを用います。

データの説明

以下の構造をもつデータで因果推論を行います。
職場での上司の研修受講が部下の面談満足度に影響するという場面を想定します。
また双方とも上司の部下育成の熱心さと仕事のパフォーマンスにも影響を受けるとします。
加えて、研修受講の効果は上司の部下育成の熱心さが高いと効果も高くなるとします。
( そこそこ妥当、そこそこ違和感というような場面設定ですが気にせず )

データの因果関係

とはいっても、これだけではイメージしきれない部分もあると思うので、
データの生成コードと可視化結果を載せます。
非線形な処置効果を表しているのは、部下の面談の満足度のコードになります。
熱心さの場合分けをif文で処理し、研修受講したかどうかにかかる係数を指定しています。
やっていることは、熱心さが-1 から 1の値を取る中、
-1以上0未満では0.5、
0以上0.5未満では0.7、
0.5以上1以下では1
の処置効果となるようにしています。
ですのでおおよそ平均的には0.675の処置効果となります。

import random
import numpy as np

# 乱数のシードを固定
random.seed(9999)
np.random.seed(9999)

# Scipy 平均0、分散1に正規化(標準化)関数
import scipy.stats

# シグモイド関数をimport
from scipy.special import expit

# 標準正規分布の生成用
from numpy.random import *

# グラフの描画用
import matplotlib.pyplot as plt
%matplotlib inline
import seaborn as sns

# その他
import pandas as pd

# データ数
num_data = 500

# enthusiam: 上司の部下育成の熱心さ。-1から1の一様乱数。
x1 = np.random.uniform(low=-1, high=1, size=num_data)

# performance: 上司の仕事のパフォーマンス。-1から1の一様乱数。
x2 = np.random.uniform(low=-1, high=1, size=num_data)
# x2 = np.random.choice([-1, -0.5, 0, 0.5, 1], num_data, p=[0.1, 0.2, 0.1, 0.3, 0.3])

# train_attend: 上司が「上司向け:部下とのキャリア面談のポイント研修」に参加したかどうか
e_z = randn(num_data)  # ノイズの生成
# シグモイド関数により受講するかどうかを熱心さに影響を受ける確率で表現(あくまで確率)
z_prob = expit(5*x1 + 5*x2 + 5*e_z)
Z = np.array([])
for i in range(num_data):
    Z_i = np.random.choice(2, size=1, p=[1-z_prob[i], z_prob[i]])[0] # 
    Z = np.append(Z, Z_i)

# stsf_level: 部下の面談の満足度
## 研修の効果は熱心さで効果が変わってくるとする(介入効果の非線形性)
t = np.zeros(num_data)
for i in range(num_data):
    if x1[i] < 0:
        t[i] = 0.5
    elif x1[i] >= 0 and x1[i] < 0.5:
        t[i] = 0.7
    elif x1[i] >= 0.5:
        t[i] = 1.0
e_y = randn(num_data)
Y = 2.0 + t*Z + 0.3*x1 + 0.3*x2 + 0.2*e_y

# テーブルを作成
df_B = pd.DataFrame({'enthusiam': x1,
                   'performance': x2,
                   'train_attend': Z,
                   'effect': t,
                   'stsf_level': Y
                   })

df_B.head()  # 先頭を表示
[output]
# データの先頭行
    enthusiam  performance  train_attend  effect  stsf_level
0   0.646778   -0.201372        1.0         1.0    2.968777
1   -0.562795  0.574169         1.0         0.5    3.018853
2   -0.937078  0.979220         1.0         0.5    2.252320
3   0.091852   -0.512450        0.0         0.7    2.098059
4   -0.691089  -0.040364        0.0         0.5    1.863024

データの可視化

データをseabornのpairplotで簡単に可視化します。
満足度(sfsf_level)に対して、上司の熱心さ(enthusiam)、上司のパフォーマンス(performance)、研修を受講したかどうか(train_attend)は正の相関が見られます。
また処置効果(effect、受講->満足度)は熱心さに影響を受けて階段状になっており、非線形な影響を受けていることがわかります。ただし勿論ですが、この部分の可視化は現実では確認することはできません。

sns.pairplot(df_B)

各変数間の可視化

また上司が研修を受講したかどうかでの部下の面談の満足度のヒストグラムも確認します。 研修の受講自体は良い影響が出ていそうとわかります。
しかし各分布の平均値をみると、満足度の差が1.0近くもあります(全体平均では0.675程度のはず)。
部下育成に熱心な上司は、そもそも研修を受講しやすいというセレクションバイアス(比較している集団の潜在的な傾向の違いによるバイアス)がかかっています(となるようなデータの生成をしています)。
また当然ですがこれらの可視化は、単変量や2変量を確認するものなので、他の変量による非線形な影響にも気付けません。

# 研修を受けたかどうかで購入量のヒストグラムを比較
df_B[df_B["train_attend"]==1]["stsf_level"].hist(color="r", alpha=0.3, label="train_attend=1")
df_B[df_B["train_attend"]==0]["stsf_level"].hist(color="b", alpha=0.3, label="train_attend=0")
plt.show()

上司が研修を受けたかどうか別の部下の面談の満足度のヒストグラム

[output]
enthusiam       0.130303
performance     0.205564
train_attend    1.000000
effect          0.715789
stsf_level      2.838375
dtype: float64
======
enthusiam      -0.223002
performance    -0.167072
train_attend    0.000000
effect          0.597233
stsf_level      1.883290
dtype: float64

傾向スコアによるIPTWでわかること、わからないこと

それではまず傾向スコアによるIPTW(逆確率重み付け法)です。 以前の記事でも整理したので割愛しますが、介入有り・無しそれぞれの期待値を傾向スコアの調整・推定し、その差分を処置効果とするものです。

以下で、実行します。

from sklearn.linear_model import LogisticRegression

# 傾向スコアを求めるために、原因変数を目的変数としてモデリング

# 説明変数
X = df_B[["enthusiam", "performance"]]

# 目的変数(原因変数)
Z = df_B["train_attend"]

# 学習(回帰の実施)
model = LogisticRegression()
model.fit(X,Z)

# 平均処置効果 ATEを計算
Y = df_B["stsf_level"]
ATE_i = Y / Z_pred[:, 1] * Z - Y / Z_pred[:, 0] * (1-Z)
ATE = 1 / len(Y) * ATE_i.sum()
print("推定したATE", ATE)
[output]
推定したATE 0.687

推定したATE(平均処置効果)は0.687と結構良い推定結果になっており、セレクションバイアスも考慮できていそうです。
しかしIPTWでは全体の平均的な処置効果を推定できるだけであり、 非線形な影響を確認できません。

Meta-Learnersの実装

そこでここからがMeta-Learnersの実装になります。
今回は4つの手法を比較していきます。
再喝ですが、分析モデルにはランダムフォレストを用いていきます。

T-Learner

T-Learnerは、原因変数をモデル構築には使用せず、集団ごとの分割の基準のみに使います。
Tは2=Twoの略らしく、2つのモデルを構築します。
手順としては以下です。
1. まず介入有り無しでデータを分ける
2. それぞれのデータで、結果変数と共変量(原因変数(=介入)以外)でモデル構築
3. それぞれのモデルで、介入有り無し分割前の共変量(原因変数(=介入)以外)で結果変数の推定
 介入有りの推定値:mu_1
 介入無しの推定値:mu_2
4. 各サンプルにおける各モデルの推定値の差の平均がATE
サンプルごとの推定処置効果(各サンプルにおける各モデルの推定値の差) = mu_1(i) - mu_2(i)
ATE(平均) = (サンプルごとの推定処置効果).mean()

また以下のようにATT(処置群における平均処置効果)とATU(対象群における平均処置効果)も推定できます。
ATT = ( 介入有り結果変数 - 介入無しモデル(介入有り共変量) ).mean()
ATU = ( 介入有りモデル(介入無し共変量) - 介入無し結果変数).mean()

それでは以下のコードで、T-Learnerを実装し、推定処置効果の可視化まで行います。
model0とmodel1という様に、2つのモデルを構築しています。

# ライブラリをインポート
from sklearn.ensemble import RandomForestRegressor

# 介入の有りと無しでデータを分割
df_0 = df_B[df_B["train_attend"]==0]
df_1 = df_B[df_B["train_attend"]==1]

# それぞれでモデル構築(目的変数: 結果変数、説明変数: 共変量)
model0 = RandomForestRegressor(max_depth=4)
M_0 = model0.fit( df_0[["enthusiam", "performance"]], df_0["stsf_level"]) # 介入を受けてないモデル

model1 = RandomForestRegressor(max_depth=4)
M_1 = model1.fit( df_1[["enthusiam", "performance"]], df_1["stsf_level"]) # 介入を受けたモデル

# 各モデルで分割前の全データで結果変数を推定
mu_0 = M_0.predict(df_B[["enthusiam", "performance"]])
mu_1 = M_1.predict(df_B[["enthusiam", "performance"]])

# サンプルごとの処置効果の推定とATEの算出
ATE_i = mu_1 - mu_0
ATE = ATE_i.mean()
print(f"ATE: {ATE}")

# 推定した処置効果を描画
t_estimated = ATE_i
plt.scatter(df_B["enthusiam"], df_B["effect"], c="k", alpha=1, label="true_effect") # 正解を描画
plt.scatter(df_B["enthusiam"], t_estimated, c="r", alpha=0.3, label="estimated_effect") # 推定値を描画
plt.legend()
plt.show()
[output]
ATE: 0.6874744694414802

推定したATEは0.687とIPTWと同様の性能を示しました。
ただMeta-Learnersはサンプルごとの処置効果も推定でき、以下が可視化結果になります。
綺麗に正解に沿っている訳ではありませんが、非線形性を表現できており、熱心さが非線形に効いていることは充分に判断できます。

T-Learnerによる熱心さに対する推定処置効果(研修受講->満足度)

またもう一つの共変量である上司のパフォーマンスに対する処置効果の可視化です。
こちらは比較的一様な分布になっていることから、同じ共変量であっても熱心さの大小が処置効果に大きく影響しており、パフォーマンスはさほど(もしくは全く)影響していないことが判断できます。

T-Learnerによるパフォーマンスに対する推定処置効果(研修受講->満足度)

S-Learner

T-LearnerだけでもMeta-Learnersのメリットを感じられたと思います。
ここからは他の手法についても眺めるように見ていきます。

まずS-Learnerです。SはSingleの頭文字の様で、T-Learnerに対する感じですね。
モデルが1つということで、T-Learnerと異なり、原因変数をモデル構築に使用します。
手順としては以下です。
1. 結果変数と共変量+原因変数の全データでモデル構築
2. 介入が全て有り or 無しとして置き換えたデータセットを作成
  X_0 : 介入を全て無しに置き換えた原因変数+共変量
  X_1 : 介入を全て有りに置き換えた原因変数+共変量
3. サンプルごとの推定処置効果 = モデル(X_1) - モデル(X_0)
4. ATE = (サンプルごとの推定処置効果).mean()

以下で可視化まで実装します。

# ライブラリをインポート
from sklearn.ensemble import RandomForestRegressor

# モデルを構築
model = RandomForestRegressor(max_depth=4)
X = df_B[["enthusiam", "performance", "train_attend"]]
model.fit(X, df_B["stsf_level"])

# 目的変数を各処置の場合を1つのモデルから求めて、ATEを計算
X_0 = X.copy()
X_0["train_attend"] = 0 

X_1 = X.copy()
X_1["train_attend"] = 1

ATE_i = model.predict(X_1) - model.predict(X_0)
ATE = ATE_i.mean()
print(f"ATE: {ATE}")

# 推定した処置効果を描画
t_estimated = ATE_i
plt.scatter(df_B["enthusiam"], df_B["effect"], c="k", alpha=0.3, label="true_effect") # 正解を描画
plt.scatter(df_B["enthusiam"], t_estimated, c="r", alpha=0.3, label="estimated_effect") # 推定値を描画
[output]
ATE: 0.706241672477917

推定したATEは0.706と性能は若干低い結果でした。ただ良し悪しに再現性は無さそうな程度感かなとも思います。
以下に可視化結果です。T-Learnerとの差はあまりわかりませんが、S-Learnerも非線形性を表現できており、熱心さが非線形に効いていることは充分に判断できるかと思います。

S-Learnerによる熱心さに対する推定処置効果(研修受講->満足度)

X-Learner

X-Learnerは、T-Learnerの結果を傾向スコアを用いて補正する手法となります。
気持ちとしては、自分のデータを使わずに構築されたモデルでの推定結果に重きを置くことで、トータルとしての推定精度を上げたいというものの様です(そこにさらに傾向スコアでの補正をしているイメージかなと認識してます)。
手順としては以下です。
1. まず介入有り無しでデータを分ける
2. それぞれのデータセットで、結果変数と共変量(原因変数(=介入)以外)でモデル構築
3. それぞれのデータセットとモデル推定値の各サンプルにおいて、ITTとITUを求める
 ITU = 介入有りモデル(介入無し共変量) - 介入無し結果変数
 ITT = 介入有り結果変数 - 介入無しモデル(介入有り共変量)
  (ATU = ITU.mean(),、ATT = ITT.mean())
4. 各データセットの共変量を説明変数、上記ITTおよひITUを目的変数として、モデル構築
   ITUモデル = model.fit(介入無し共変量, ITU)
   ITTモデル = model.fit(介入有り共変量, ITT)
5. 上記の各モデルを全データの共変量を使って、ITT(hat)とITE(hat)を推定
   ITU(hat) = ITUモデル(全データの共変量)
   ITT(hat) = ITTモデル(全データの共変量)
6. 原因変数と共変量で傾向スコア( g(x) : 介入有りの確率)算出
7. 各効果の推定値を傾向スコアで調整
 (サンプルごとの推定した処置効果) = g(x) * ITU(hat) + ( 1 - g(x) ) * ITT(hat)
 ATE = (サンプルごとの推定した処置効果).mean()

以下で可視化まで実装します。
大まかには前半はT-Learnerを行い、そこから傾向スコアの算出、処置効果の補正という流れです。

# ライブラリをインポート
from sklearn.ensemble import RandomForestRegressor

# 介入の有りと無しでデータを分割
df_0 = df_B[df_B["train_attend"]==0]
df_1 = df_B[df_B["train_attend"]==1]

# それぞれでモデル構築(目的変数: 結果変数、説明変数: 共変量)
model0 = RandomForestRegressor(max_depth=4)
M_0 = model0.fit( df_0[["enthusiam", "performance"]], df_0["stsf_level"]) # 介入を受けてないモデル

model1 = RandomForestRegressor(max_depth=4)
M_1 = model1.fit( df_1[["enthusiam", "performance"]], df_1["stsf_level"]) # 介入を受けたモデル

# ITUとITTを求める
tau_0 = M_1.predict(df_0[["enthusiam", "performance"]]) - df_0["stsf_level"]
tau_1 = df_1["stsf_level"] - M_0.predict(df_1[["enthusiam", "performance"]])

# ITUとITTの推定モデルを構築
model_ITU = RandomForestRegressor(max_depth=4)
M_ITU = model_ITU.fit(df_0[["enthusiam", "performance"]], tau_0)

model_ITT = RandomForestRegressor(max_depth=4)
M_ITT = model_ITT.fit(df_1[["enthusiam", "performance"]], tau_1)

# 傾向スコアを求める
from sklearn.linear_model import LogisticRegression

# 説明変数: 共変量
X = df_B[["enthusiam", "performance"]]

# 目的変数: 結果変数
Z = df_B["train_attend"]

# モデルを構築
model = LogisticRegression()
model.fit(X, Z)

# 傾向スコア算出
Z_proba = model.predict_proba(X)

# それぞれのモデルで全データの処置効果を推定し、傾向スコアで調整
ITU_pred = M_ITU.predict(df_B[["enthusiam", "performance"]])
ITT_pred = M_ITT.predict(df_B[["enthusiam", "performance"]])
ATE_i = Z_proba[:, 1]*ITU_pred + Z_proba[:, 0]*ITT_pred
ATE = ATE_i.mean()
print(f"ATE: {ATE}")

# 推定した処置効果を描画
t_estimated = ATE_i
plt.scatter(df_B["enthusiam"], df_B["effect"], c="k", alpha=0.3, label="estimated_TE") # 正解を描画
plt.scatter(df_B["enthusiam"], t_estimated, c="r", alpha=0.3, label="estimated_TE") # 推定値を描画
[output]
ATE: 0.6757366813565363

推定したATEは0.676と良好な結果でした。
以下に可視化結果です。他の手法と比べて振れ幅も狭くなり、より良好な推定をしている様に見えます。

X-Learnerによる熱心さに対する推定処置効果(研修受講->満足度)
傾向スコアでの補正前の推定値とも比べました。傾向スコアによる補正前について、推定ITUは熱心さが大きい範囲で性能が良く・分散が小さい、推定ITTは熱心さが小さい範囲で性能が良く・分散が小さくなっています。そして傾向スコアでの補正で良いとこどりができていることが確認できます。
(ただなぜこうなるかは考察しきれませんでした。処置効果が大きい範囲は反実仮想の値を推定する際の精度によっては分散が大きくなるイメージはつくんですが、、、)
傾向スコアによる補正前後_X-Learnerによる熱心さに対する推定処置効果(研修受講->満足度)

DR-Learner

DR-Learnerは、傾向スコアを反実仮想の推定に利用し、より性能の良い機械学習モデルを構築する手法の様です(正直腹落ちしておらず、結果変数の調整に傾向スコアと反実仮想を用いているイメージです)。

X-Learnerとの違いは、傾向スコアの使い方が異なります。

X-Learnerは、モデルはよくある流れで構築して、推定したITTとITUに傾向スコアを用いて効果を推定しています。
1. T-Learner的に算出したITTおよびITUを目的変数に、共変量を説明変数としてモデル構築
2. そのモデルに再度共変量を入力して、(サンプルごとの推定した処置効果)を推定
3. 推定した処置効果を介入有り無しで逆側の傾向スコアで加重平均して、介入効果を求める

DR-Learnerは、モデル構築に用いる結果変数の調整に用い、その調整された結果変数をもとにしたITTとITUを各サンプルのITEとして、モデル構築および推定を行っています。
1. T-Learner的にITTとITUを算出する際に、結果変数側を傾向スコアと推定結果変数で調整して算出する(精度良く推定しながら、推定誤差も調整することで、個別の誤差を少しでも緩和)
2. 算出したITTとITUを目的変数に、共変量を説明変数にしてモデル構築
3. そのモデルに再度共変量を入力して、推定処置効果を求める

以下が手順となります。
1. T-learner的に機械学習モデルを構築する
2. 結果変数を傾向スコアと1の機械学習モデルによる推定した結果変数で調整する
3. 潜在的結果変数も1の機械学習モデルから求め、介入有り無しでそれぞれのITEを求める
4. 3で求められたITEと共変量でモデル構築
5. 共変量で介入効果を推定する

# ライブラリをインポート
from sklearn.ensemble import RandomForestRegressor

# T-Learnerで処置のケースごとにモデル構築
# 介入の有りと無しでデータを分割
df_0 = df_B[df_B["train_attend"]==0]
df_1 = df_B[df_B["train_attend"]==1]

# それぞれでモデル構築(目的変数: 結果変数、説明変数: 共変量)
model0 = RandomForestRegressor(max_depth=4)
M_0 = model0.fit( df_0[["enthusiam", "performance"]], df_0["stsf_level"]) # 介入を受けてないモデル

model1 = RandomForestRegressor(max_depth=4)
M_1 = model1.fit( df_1[["enthusiam", "performance"]], df_1["stsf_level"]) # 介入を受けたモデル

# 傾向スコアを求めるモデル構築
from sklearn.linear_model import LogisticRegression

X = df_B[["enthusiam", "performance"]]
Z = df_B["train_attend"]
model_ps = LogisticRegression()
model_ps.fit(X, Z)

g_x = LogisticRegression().fit(X, Z)
g_x_val = g_x.predict_proba(X)

# DR法に基づく推定
# 処置群について
Y_1 = (
    M_1.predict(df_1[["enthusiam", "performance"]])
    + (df_1["stsf_level"] - M_1.predict(df_1[["enthusiam", "performance"]]))
    /g_x.predict_proba(df_1[["enthusiam", "performance"]])[:, 1]
    )
df_1["ITE"] = Y_1 - M_0.predict(df_1[["enthusiam", "performance"]])

# 非処置群について
Y_0 = (
    M_0.predict(df_0[["enthusiam", "performance"]]) 
    + (df_0["stsf_level"] - M_0.predict(df_0[["enthusiam", "performance"]]))
    /g_x.predict_proba(df_0[["enthusiam", "performance"]])[:, 0]
)
df_0["ITE"] = M_1.predict(df_0[["enthusiam", "performance"]]) - Y_0

 # データを結合
df_DR = pd.concat([df_0, df_1])
display(df_DR.head())

# ITEに対してモデル構築し、サンプルごとの処置効果を推定する
model_DR = RandomForestRegressor(max_depth=4)
M_DR = model_DR.fit(df_DR[["enthusiam", "performance"]], df_DR["ITE"])
ATE_i = M_DR.predict(df_DR[["enthusiam", "performance"]])
ATE = ATE_i.mean()
print(f"ATE: {ATE}")

# 推定した処置効果を描画
t_estimated = ATE_i
plt.scatter(df_DR["enthusiam"], df_DR["effect"], c="k", alpha=1, label="true_effect") # 正解を描画
plt.scatter(df_DR["enthusiam"], t_estimated, c="r", alpha=0.3, label="estimated_effect") # 推定値を描画
plt.legend()
plt.show()
[output]
ATE: 0.6705706151150326

推定したATEは0.671と良好な結果でした。 以下に可視化結果です。大きく外れてしまうサンプルもある様な結果でした。
y軸を拡大した可視化も見てみると、他手法と比べて各熱心さの範囲でx軸に並行な推定結果となっており、処置効果の傾向にはより近づこうとしている印象です。

DR-Learnerによる熱心さに対する推定処置効果(研修受講-&gt;満足度)
y軸拡大_DR-Learnerによる熱心さに対する推定処置効果(研修受講->満足度)

モデルのハイパラ調整(X-Learner、木の最大深さと葉ノードに属する最小サンプル数)

Meta-Learnerは使用する分析モデルは何でも良いので、実際はモデル比較やハイパーパラメータ調整を行うのだろうかと思います。
ということでハイパーパラーメータ調整を行い、より精度の高い機械学習モデルを用いると結果がどう変わるのか確認しました。
- 用いる手法: X-Learner
- 調整するハイパーパラメータ: 木の最大深さ、葉ノードに属する最小サンプル数
とします。
ハイパーパラメータ調整を行う箇所は、T-Learnerでの2つのモデル構築とITUとITTの推定モデルの構築の部分です。
またハイパーパラメータの最適化は、評価指標をMAE(平均絶対誤差)とした交差検証による平均値比較でのグリッドサーチを行います。

以下に実装コードです。
交差検証とグリッドサーチは多用するので関数を用意しておきます。

from sklearn.model_selection import KFold
from sklearn.metrics import mean_absolute_error

# 交差検証の評価指標の結果を返す。評価指標はMAE(絶対平均誤差)で行う。
def cv_mse(X, y, model, split_num = 5, shuffle = True):
  cv = KFold(n_splits=split_num, random_state=0, shuffle=shuffle)
  mae_list = []
  for train_index, test_index in cv.split(X):
      X_train, X_test = X.iloc[train_index, :], X.iloc[test_index, :]
      y_train, y_test = y.iloc[train_index], y.iloc[test_index]

      model.fit(X_train, y_train)
      y_pred = model.predict(X_test)

      mae = np.round(mean_absolute_error(y_test, y_pred), 3)
      mae_list.append(mae)

  return mae_list

# 交差検証とグリッドサーチをして、最も良いパラメータを返す
def rf_cv_grid_search(X, y, depth=range(1, 9), leaf=range(1, 11)): 
  best_val_score = np.inf  # best_scoreを初期化
  best_params = {}   # best_paramを初期化
  for i in depth:
    for j in leaf:
      model = RandomForestRegressor(
          max_depth=i, 
          min_samples_leaf=j, 
          random_state=0)
      score_array = np.array(cv_mse(X, y, model))
      mean_score = score_array.mean()
      
      if mean_score < best_val_score:
        best_val_score = mean_score
        best_params = {'max_depth': i, 'min_samples_leaf': j}

  print('best_val_score: {:0.3f}'.format(best_val_score))
  print('best_param: ', best_params)
  
  return best_params

まず今回使ってきたmax_depth=4のランダムフォレストでの交差検証による精度を確認します。

# ライブラリをインポート
from sklearn.ensemble import RandomForestRegressor

# 介入の有りと無しでデータを分割
df_0 = df_B[df_B["train_attend"]==0]
df_1 = df_B[df_B["train_attend"]==1]

print("【T-Learnerの精度評価_MSE】")
print("==介入無し==")
score_TL_0 = cv_mse(df_0[["enthusiam", "performance"]], df_0["stsf_level"],
       RandomForestRegressor(max_depth=4))
print(f"mse: {score_TL_0}")
print(f"mean_mse: {round(np.array(score_TL_0).mean(), 3)}")

print("==介入有り==")
score_TL_1 = cv_mse(df_1[["enthusiam", "performance"]], df_1["stsf_level"],
       RandomForestRegressor(max_depth=4))
print(f"mse: {score_TL_1}")
print(f"mean_mse: {round(np.array(score_TL_1).mean(), 3)}")

# それぞれでモデル構築(目的変数: 結果変数、説明変数: 共変量)
model0 = RandomForestRegressor(max_depth=4)
M_0 = model0.fit( df_0[["enthusiam", "performance"]], df_0["stsf_level"]) # 介入を受けてないモデル

model1 = RandomForestRegressor(max_depth=4)
M_1 = model1.fit( df_1[["enthusiam", "performance"]], df_1["stsf_level"]) # 介入を受けたモデル

# ITUとITTを求める
tau_0 = M_1.predict(df_0[["enthusiam", "performance"]]) - df_0["stsf_level"]
tau_1 = df_1["stsf_level"] - M_0.predict(df_1[["enthusiam", "performance"]])

# 推論時には異なる集団の推論も実施するためのモデルであるが評価ができないので、同じ集団内で評価する
print("【ITUとITTのモデル構築の精度評価_MSE】")

print("==ITU==")
score_ITU = cv_mse(df_0[["enthusiam", "performance"]], tau_0,
       RandomForestRegressor(max_depth=4))
print(f"mse: {score_ITU}")
print(f"mean_mse: {round(np.array(score_ITU).mean(), 3)}")

print("==ITT==")
score_ITT = cv_mse(df_1[["enthusiam", "performance"]], tau_1,
       RandomForestRegressor(max_depth=4))
print(f"mse: {score_ITT}")
print(f"mean_mse: {round(np.array(score_ITT).mean(), 3)}")
[output]
【T-Learnerの精度評価_MSE】
==介入無し==
mse: [0.151, 0.174, 0.144, 0.167, 0.171]
mean_mse: 0.161
==介入有り==
mse: [0.164, 0.192, 0.164, 0.17, 0.143]
mean_mse: 0.167
【ITUとITTのモデル構築の精度評価_MSE】
==ITU==
mse: [0.163, 0.168, 0.13, 0.168, 0.173]
mean_mse: 0.16
==ITT==
mse: [0.173, 0.189, 0.181, 0.163, 0.151]
mean_mse: 0.171

以下が、グリッドサーチを用いたモデルで順次処理していく実装です。

# ライブラリをインポート
from sklearn.ensemble import RandomForestRegressor

# 介入の有りと無しでデータを分割
df_0 = df_B[df_B["train_attend"]==0]
df_1 = df_B[df_B["train_attend"]==1]

# 介入無しのモデルのハイパラのグリッドサーチ
print("==介入無しモデルハイパラ調整の結果==")
best_params_0 = rf_cv_grid_search(df_0[["enthusiam", "performance"]], df_0["stsf_level"])

# 最も良かった精度のパラメータで介入無しのモデル構築
model0 = RandomForestRegressor(**best_params_0)
M_0 = model0.fit(df_0[["enthusiam", "performance"]], df_0["stsf_level"]) 

# 介入有りのモデルのハイパラのグリッドサーチ
print("==介入有りモデルハイパラ調整の結果==")
best_params_1 = rf_cv_grid_search(df_1[["enthusiam", "performance"]], df_1["stsf_level"])

# 最も良かった精度のパラメータで介入有りのモデル構築
model1 = RandomForestRegressor(**best_params_1)
M_1 = model1.fit(df_1[["enthusiam", "performance"]], df_1["stsf_level"]) 

# ITUとITTを求める
tau_0 = M_1.predict(df_0[["enthusiam", "performance"]]) - df_0["stsf_level"]
tau_1 = df_1["stsf_level"] - M_0.predict(df_1[["enthusiam", "performance"]])

# ITUのモデルのハイパラのグリッドサーチ
print("==ITUモデルのハイパラ調整の結果==")
best_params_ITU = rf_cv_grid_search(df_0[["enthusiam", "performance"]], tau_0)

# ITUの推定モデルを構築
model_ITU = RandomForestRegressor(**best_params_ITU)
M_ITU = model_ITU.fit(df_0[["enthusiam", "performance"]], tau_0)

# ITTのモデルのハイパラのグリッドサーチ
print("==ITTモデルのハイパラ調整の結果==")
best_params_ITT = rf_cv_grid_search(df_1[["enthusiam", "performance"]], tau_1)

# ITTの推定モデルを構築
model_ITT = RandomForestRegressor(**best_params_ITT)
M_ITT = model_ITT.fit(df_1[["enthusiam", "performance"]], tau_1)

# 傾向スコアを求める
from sklearn.linear_model import LogisticRegression

# 説明変数: 共変量
X = df_B[["enthusiam", "performance"]]

# 目的変数: 結果変数
Z = df_B["train_attend"]

# モデルを構築
model = LogisticRegression()
model.fit(X, Z)

# 傾向スコア算出
Z_proba = model.predict_proba(X)

# それぞれのモデルで全データの処置効果を推定し、傾向スコアで調整
ITU_pred = M_ITU.predict(df_B[["enthusiam", "performance"]])
ITT_pred = M_ITT.predict(df_B[["enthusiam", "performance"]])
ATE_i = Z_proba[:, 1]*ITU_pred + Z_proba[:, 0]*ITT_pred
ATE = ATE_i.mean()
print("==ATE==")
print(f"ATE: {ATE}")

# 推定した処置効果を描画
t_estimated = ATE_i
plt.scatter(df_B["enthusiam"], df_B["effect"], c="k", alpha=0.3, label="true_ITE") # 正解を描画
plt.scatter(df_B["enthusiam"], ITU_pred, c="b", alpha=0.3, label="estimated_ITU")
plt.scatter(df_B["enthusiam"], ITT_pred, c="g", alpha=0.3, label="estimated_ITT")
plt.scatter(df_B["enthusiam"], t_estimated, c="r", alpha=0.3, label="estimated_effect") # 推定値を描画
plt.legend()
plt.show()
[output]
==介入無しモデルハイパラ調整の結果==
best_val_score: 0.159
best_param:  {'max_depth': 4, 'min_samples_leaf': 4}
==介入有りモデルハイパラ調整の結果==
best_val_score: 0.164
best_param:  {'max_depth': 4, 'min_samples_leaf': 7}
==ITUモデルのハイパラ調整の結果==
best_val_score: 0.159
best_param:  {'max_depth': 3, 'min_samples_leaf': 3}
==ITTモデルのハイパラ調整の結果==
best_val_score: 0.165
best_param:  {'max_depth': 2, 'min_samples_leaf': 1}
==ATE==
ATE: 0.6706889541529172

X-Learner(ハイパラ調整)による熱心さに対する推定処置効果(研修受講->満足度)

精度に関しては、若干良くなりました。可視化結果についても、正解に近いところでハイパーパラメータ調整前よりも推定範囲が狭くなっていて良好になっているように見受けられます。ただ階段状に推定できていません(=熱心さが0.0~0.5の範囲が特に)。効果の大きさが切り替わって遷移している部分は難しいとかあるのでしょうか?
境界付近(熱心さが0または0.5あたり)は推定が難しくなるのはそれはそうだと思うので、両端に境界があると推定がうまくいかないイメージももつことはできますが。

おわりに

主にMeta-Learnersと呼ばれる、非線形な因果関係および処置効果を扱える因果推論手法を整理・実装しました。Meta-Learnersを用いることで、共変量の処置効果への非線形な影響を捉えられることが確認できました。また各手法で少しづつ挙動の違いがあり、今回のケースではX-Learnerが良い手法だなという印象でした。 しかしX-Learnerでの推定ITTおよびITUの挙動の違い、DR-Learnerでの大きく推定が外れるサンプルがあること、Meta-Learnersでも階段状にはうまく推定できていないところはうまく説明できそうになかったです。このあたりはさらに集計で深掘りをしたり、そもそも論文を読んでみたりすることで色々わかるのかもしれません。
今回は入門ということで、とりあえず実装できたことに満足です。