PyMC3を使ってベイズ推定、テストデータ(out-of-sample)に対する検証を実施する

はじめに

今度、マーケティング的な分析もすることになりそうです。私は今までベイズ推定をさわり程度に勉強してきましたが、実務でも使う機会が出てきそうです。時間的に余裕がある今のうちに勉強して手を動かし、自分用の備忘を兼ねて内容をここに書きます。
Pythonでベイズ推論を行うライブラリとしてPyMC3を使います。この記事では、PyMC3を使って、モデルの推定とテストデータに対する検証を行うまでを記載します。

参考

主に以下のPyMC3チュートリアルを参考にしています。
[1] Prior and Posterior Predictive Checks
[2] [Getting started with PyMC3]
(https://docs.pymc.io/notebooks/getting_started.html)

準備

AnacondaにPyMC3をインストールします。私の場合、Windows10 64bitに入れます。
ここで私はハマってしまいました。これはこちらの記事に書いていますので、参考にしてみてください。
PyMC3はバージョン3.11.2を使っています。

線形回帰モデル

説明変数が$x$一つだけの線形モデルに対して、PyMC3でのモデルの表現方法、パラメータ推定方法を扱います。
目的変数を$y$とすると、モデルは以下のように数式で表現されます。$\epsilon$
$y = a + bx + \epsilon$
ここで$\epsilon$は誤差で、標準偏差$\sigma$の正規分布に従うとします。これは、以下のように表現できます
$\mu = a + bx$
$y \sim N(\mu, \sigma^2)$

線形回帰モデルで使うデータと事前分布

パラメータの真の値を$a^* =0.5, b^* = 3, \sigma^* = 2$ とします。
これに従って、サンプルサイズ100のデータを乱数を使って生成します。

N = 100

true_a, true_b, x = 0.5, 3.0, np.random.normal(loc=2, scale=6, size=N)
true_mu = true_a + true_b * x
true_sd = 2.0

# numpyの乱数シード固定
np.random.seed(seed=RANDOM_SEED)
y = np.random.normal(loc=true_mu, scale=true_sd, size=N)

f"{x.mean():.2f}, {x.std():.2f}, {y.mean():.2f}, {y.std():.2f}"

結果は'0.84, 5.76, 3.17, 17.07'となり、説明変数と目的変数のどちらもばらつきが1より大きくなっています。[1]によると、ベイズ推定ではこのようなおおきなばらつきはよくないらしいです。そこで、説明変数と目的変数のどちらも平均0、標準偏差1となるように標準化します。

def standardize(series):
    """Standardize a pandas series"""
    return (series - series.mean()) / series.std()

x_scaled = standardize(x)
y_scaled = standardize(y)

f"{x_scaled.mean():.2f}, {x_scaled.std():.2f}, {y_scaled.mean():.2f}, {y_scaled.std():.2f}"

結果は'-0.00, 1.00, -0.00, 1.00'となり、標準化されました。
この前処理後のデータの分布をプロットすると以下のようになります。

ベイズ推定では事前の知識などに基づいて事前分布を設定します。
ここでは仮に$a$と$b$を平均0、標準偏差10とおいてみます。また、誤差$\sigma$も未知ですので、事前に分布を設定しておきます。誤差$\sigma$は正の値しかとらないことを反映した分布にする必要があり、今回は指数分布とします。妥当性は後述のprior predictive checkで確認して、適切でなさそうなら再設定することにします。

PyMC3での設定

上記の設定は、PyMC3では以下のように実装して表現します。標準化したx_scaled, y_scaledを使います。

with pm.Model() as model_1:
    #a,bの事前分布
    a = pm.Normal("a", 0.0, 10.0)
    b = pm.Normal("b", 0.0, 10.0)

    mu = a + b * x_scaled

    #sigmaの事前分布
    sd = pm.Exponential("sd", 1.0)

    y_obs = pm.Normal("obs", mu=mu, sigma=sd, observed=y_scaled)

    #Prior Predictive Check
    prior_checks = pm.sample_prior_predictive(samples=50, random_seed=RANDOM_SEED)

1行目でPyMC3のModelオブジェクトをmodel_1として新しく作っています。
正規分布、指数分布をPyMC3の関数pm.Normal、pm.Exponentialを使って定義しています。

Prior Predictive Check

上のコード最後にあるprior_checks = pm.sample_prior_predictive(samples=50, random_seed=RANDOM_SEED)で、パラメータ$a$, $b$, $\sigma$の事前分布に従ってデータを生成しています。以下のコードで生成したデータをプロットします。

_, ax = plt.subplots()

xx = np.linspace(-2, 2, 50)

for a, b in zip(prior_checks["a"], prior_checks["b"]):
    yy = a + b * xx
    ax.plot(xx, yy, c="k", alpha=0.4)

ax.set_xlabel("x (stdz)")
ax.set_ylabel("Mean y (stdz)")
ax.set_title("Prior predictive checks -- Flat priors")

上の x_scaledとy_scaledの散布図と比較すると、y軸のとる範囲が1桁大きくなっています。これは、$a$と$b$の事前分布に設定した標準偏差の値10が適切ではないからです。よって、事前分布を以下のように再設定します。

with pm.Model() as model_1:
    a = pm.Normal("a", 0.0, 0.5)
    b = pm.Normal("b", 0.0, 1.0)

    mu = a + b * x_scaled
    sd = pm.Exponential("sd", 1.0)

    y_obs = pm.Normal("obs", mu=mu, sigma=sd, observed=y_scaled)
    prior_checks = pm.sample_prior_predictive(samples=50, random_seed=RANDOM_SEED)

このprior predictive checkの結果は以下のようになり、妥当な範囲に収まっています。

事後分布を求める

以下のようにPyMCのsample()を使ってサンプリングにより事後分布を推定します。

with model_1:
    trace_1 = pm.sample(1000,#The number of samples to draw
                        tune=2000, #Number of iterations to tune 
                        random_seed=RANDOM_SEED, return_inferencedata=True)

私のPC(AMD RYZEN 7)で、およそ27秒ほどでサンプリングが実施できました。

次に、plot_traceを使って、サンプリングを可視化します。

with model_1:
    az.plot_trace(trace_1);

左に得られた事後分布、右にサンプルの軌跡が描画されます。

右の図の軌跡が一定の分布に収束していないなら、それは一定の分布がサンプルから得られていないことを意味し、サンプルのやり方を検討する、統計モデル自体を再検討することが必要となります。
この例では、それなりに一定の分布に収束していると考えられます。ですからサンプリングは妥当として、次に進むことにします。

Posterior Predictive Check

次にサンプルされたパラメータの事後分布に従ってデータを発生させ、実際のデータと比較するPosterior Predictive Checkを実施します。訓練データを適切に表現できているかを確認するチェックと理解しています。
sample_posterior_predictiveを使ってPosterior Predictive Checkを実施し、

with model_1:
    ppc = pm.sample_posterior_predictive(
        trace_1, var_names=["a", "b", "sd", "obs"], random_seed=RANDOM_SEED
    )

az.plot_ppc(az.from_pymc3(posterior_predictive=ppc, model=model_1))

※[1]ではver_namesにsdが入っていいなかったのですが、気持ちが悪かったので自分は入れてみました。結果はsdを入れない場合とほとんど変わりませんでした。間違えていたら、指摘いただけると助かります。

結果は以下のようになりました。

横軸がyの値、縦軸が頻度（のはず）です。黒い線が実際のデータ、青い線がサンプルに基づいて生成されたものです。データとあっていて、推定されたモデルはそれなりに妥当だと考えられます。

次に、実際のデータと、モデルの信頼区間を比較するプロットを作成します。

mu_pp = (ppc["a"] + ppc["b"] * x_scaled[:, None]).T

_, ax = plt.subplots()

ax.plot(x_scaled, y_scaled, "o", ms=4, alpha=0.4, label="Data")
ax.plot(x_scaled, mu_pp.mean(0), label="Mean outcome", alpha=0.6)
az.plot_hdi(
    x_scaled,
    mu_pp,
    ax=ax,
    fill_kwargs={"alpha": 0.8, "label": "Mean outcome 94% HDI"},
)
az.plot_hdi(
    x_scaled,
    ppc["obs"],
    ax=ax,
    fill_kwargs={"alpha": 0.8, "color": "#a1dab4", "label": "Outcome 94% HDI"},
)

ax.set_xlabel("x (stdz)")
ax.set_ylabel("y (stdz)")
ax.set_title("Posterior predictive checks")
ax.legend(ncol=2, fontsize=10)

※[1]にはaz.plot_hpdとあったのですが、PyMC3 バージョン3.11.2だとaz.plot_hdiとする必要があります。

結果は下図のようになり、こちらからも得られたモデルの妥当性を確認できます。

ちょっと調べて、ドキュメントには記載はなかったのですが、GitHubのコードをみると、plot_hdiの信頼区間のデフォルト値は0.94(94%)らしいです。（なんで0.95じゃないんだろう？私の感覚がおかしいのかな？？）

前処理の標準化を実施せずにベイズ推定を実行したらどうなる？

気になるので、標準化する前のxとyを使って、上と同様のサンプリングを実施してみました。生成するモデルオブジェクトの名前をmodel_0としてやってみました。

with pm.Model() as model_0:
    a = pm.Normal("a", 0.0, 2.0)
    b = pm.Normal("b", 0.0, 2.0)

    mu = a + b * x
    sd = pm.Exponential("sd", 1.0)

    y_obs = pm.Normal("obs", mu=mu, sigma=sd, observed=y)
    prior_checks = pm.sample_prior_predictive(samples=50, random_seed=RANDOM_SEED)

結果、計算にかかった時間は25秒ほどで、前処理をしたときと比較してほとんど変わりませんでした。このくらいの値のばらつきならあまり影響はないようです。ただし、遅くなる時もあるようですので、前処理→サンプリングによるベイズ推定の流れはおさえておきたいです。

テストデータ(out-of-sample）に対する検証

推定したモデルがテストデータ（モデル推定に一切使われていないデータ、out-of-sampleデータとも言います）に対しても有効化の検証を行います。ここでは、[1]に従ってロジスティック回帰モデルでやります。
（少し検索した範囲で、ほかの記事でも、out-of-sampleに対する例はロジスティック回帰モデルが使われているのはなぜだろう？）

$p = \frac{1}{1+ \exp{-(intercept + slope * x)}}$
$y$は確率$p$の二項分布に従い、1と0の2値をとるとします。

この設定で、乱数を使って400サンプルの推定に使うデータを作ります。

N = 400
true_intercept = 0.2
true_slope = 1.0

np.random.seed(seed=RANDOM_SEED)
x = np.random.normal(size=N)

true_p = logistic(true_intercept + true_slope * x)

np.random.seed(seed=RANDOM_SEED)
y = np.random.binomial(1, true_p)

これをベイズ推定でロジスティック回帰モデルのパラメータの分布を求めます。

with pm.Model() as model_2:
    betas = pm.Normal("betas", mu=0.0, sigma=np.array([0.5, 1.0]), shape=2)

    # set predictors as shared variable to change them for PPCs:
    x_data = pm.Data("x_data", x) # pred -> x_data
    p = pm.Deterministic("p", pm.math.invlogit(betas[0] + betas[1] * x_data))
    
    y_pbs = pm.Bernoulli("y_pbs", p=p, observed=y)
    
    trace_2 = pm.sample(1000, tune=2000, random_seed=RANDOM_SEED, return_inferencedata=True)
    
az.summary(trace_2, var_names=["betas"], round_to=2)

次に、50件のout-of-sampleデータを作ります。

x_out_of_sample = np.random.normal(size=50)

y_out_of_sample = np.random.binomial(
    1, logistic(true_intercept + true_slope * x_out_of_sample)
)

with model_2:
    # update values of predictors:
    pm.set_data({"x_data": x_out_of_sample})
    # use the updated values and predict outcomes and probabilities:
    posterior_predictive = pm.sample_posterior_predictive(
        trace_2, var_names=["p"], random_seed=RANDOM_SEED
    )
    model_preds = posterior_predictive["p"]

ここで、pm.set_dataでout-of-sampleデータの説明変数をセットしています。次に、trace_2に格納されているサンプリングの結果をこのout-of-sampleデータに適用して、モデルの予測結果であるmodel_predsを求めています。
model_predsは4000(MCMCのサンプル数）×50(out-of-sampleのサンプル数）の配列となっています。

この予測結果と、実際のout-of-sampleデータを比較するプロットを作成します。

_, ax = plt.subplots(figsize=(12, 6))

# uncertainty about the estimates:
ax.plot(
    [x_out_of_sample, x_out_of_sample],
    az.hdi(model_preds).T,
    lw=6,
    color="#00204C",
    alpha=0.8,
)
# expected probability of success:
ax.plot(
    x_out_of_sample,
    model_preds.mean(0),
    "o",
    ms=5,
    color="#FFE945",
    alpha=0.8,
    label="Expected prob.",
)

# actual outcomes:
ax.scatter(
    x=x_out_of_sample,
    y=y_out_of_sample,
    marker="x",
    color="#A69C75",
    alpha=0.8,
    label="Observed outcomes",
)

# true probabilities:
x = np.linspace(x_out_of_sample.min() - 0.1, x_out_of_sample.max() + 0.1)
ax.plot(
    x,
    logistic(true_intercept + true_slope * x),
    lw=2,
    ls="--",
    color="#565C6C",
    alpha=0.8,
    label="True prob.",
)

ax.set_xlabel("Predictor")
ax.set_ylabel("Prob. of succes")
ax.set_title("Out-of-sample Predictions")
ax.legend(fontsize=10, frameon=True, framealpha=0.5);

真の確率(True prob.)と、サンプリングから求められた予測された確率の平均(Expected prob.)はそれなりに一致しており、妥当なモデルが得られていると考えられます。

まとめ

PyMC3を使って、モデルオブジェクトを生成、prior predictive check、MCMCのサンプリングによるモデル推定、posterior predictive check、そしてテストデータ（out-of-sampleデータ）に対する検証までを行いました。
コードの書き方はscikit-learnとは違いますので、慣れる必要はあるなと思いました。ただし、少なくとも線形回帰やロジスティック回帰ならばそれほど難しいところはなく、MCMCサンプリングの収束にも問題ありませんでした。
仮説を設定して、データ数が少ないうちからパラメータの信頼区間を求めたり仮説検証するときなどに、使っていければと思います。

なお、今回のコードを記述したjupyter notebookはGitHubにあげていますので、興味ある方はご参照ください。