🦆

5.3 階層ベイズモデル - 第5章ベイズ推論プログラミング

に公開

はじめに

Pythonでスラスラわかる ベイズ推論「超」入門(赤石 雅典 (著), 須山 敦志 (監修))の5.3節のPyMCコードをNumPyroで書き直しました。
アイリス・データセットの3種類の花のデータを3個ずつ抽出しました。
合計9個のデータで3種類の花のがく片の長さsepal_lengthとがく片の幅sepal_widthの1次関数近似のベイズ推論を行います。
関連するデータが少しずつあるという業務でありがちな状況でベイズ推論が輝きます。

フォルダ構造とユーティリティ関数、ライブラリimport

リンク集の記事にフォルダ構造とユーティリティ関数、ライブラリimportを掲載しました。
準備としてそちらのページをご覧ください。

  1. フォルダ構造とユーティリティ関数
  2. ライブラリimport

5.3 階層ベイズモデル

5.3.1 問題設定

3種類の花から3つずつデータを抽出して、がく片の長さと幅の1次関数近似を行います。

5.3.2 データ準備

問題設定の通りにデータを抽出します。

# データセットを読み込む
df = sns.load_dataset("iris")

# setosa を抽出する
df_setosa = df.query('species == "setosa"')
df_versicolor = df.query('species == "versicolor"')
df_virginica = df.query('species == "virginica"')

# 乱数により3個のインデックスを生成
import random
random.seed(42)
indexes =range(len(df_setosa))
sample_indexes=random.sample(indexes, 3)

# df0, df1, df2のデータ数をそれぞれ3行にする
df_setosa_sel = df_setosa.iloc[sample_indexes]
df_versicolor_sel = df_versicolor.iloc[sample_indexes]
df_virginica_sel = df_virginica.iloc[sample_indexes]

# 全部連結して一つにする
df_sel = pd.concat([df_setosa_sel, df_versicolor_sel, df_virginica_sel]).reset_index(drop=True)

# 加工結果の確認
display(df_sel)

抽出したデータの散布図をプロットします。
versicolorとvirginicaは1次関数近似できそうですが、setosaはデータが集中しすぎて1次関数近似が難しそうです。

sns.scatterplot(
    x='sepal_length', y='sepal_width', hue='species', style='species',
    data=df_sel, s=100)
plt.xticks([4, 5, 6, 7, 8])
plt.yticks([1, 2, 3, 4, 5])
plt.title('抽出した計9個の観測値の散布図');

X = jnp.array(df_sel['sepal_length'].values, dtype = float)
Y = jnp.array(df_sel['sepal_width'].values, dtype = float)
species = df_sel['species']
cl = jnp.array(pd.Categorical(species).codes, dtype = int)

# 結果確認
print(X)
print(Y)
print(species.values)
print(cl)

5.3.3 確率モデル定義

確率モデルをプログラミングするために、前章と同様に数式を使って状況を整理します。

まず、散布図より N 組の sepal_length x_{i} とがく片の幅 sepal_width y_{i} は1次関数で近似できると仮定します。
このとき1次関数の切片と傾きは花の種類の序数 s 毎に異なる値を取ると仮定します。

y_{i} \approx \omega_{0, s} + \omega_{1, s} x_{i}

ただでさえデータが少ないのにさらに分割するのは不安ですが先に進めます。
確率モデルを作成するためのテクニックとして、右辺を \mu_{i} とおくことにします。

\mu_{i} \equiv \omega_{0, s} + \omega_{1, s} x_{i}\\ \Rightarrow y_{i} \approx \mu_{i}

y_{i}x_{i} の値に応じた定数 \mu_{i} に近い値を取ると読むことができます。
5.1節の正規分布の仮定とほぼ同じですね。
y_{i}x_{i} の値に応じた平均値 \mu_{i} をパラメータに持つ正規分布に従うと仮定します。

y_{i} \sim N(\mu_{i}, \sigma^2)

1次関数近似のベイズ推論は、1次関数のパラメータ \{ \omega_{0, s}, \omega_{1, s} \} と正規分布の標準偏差 \sigma を求める問題になりました。
階層ベイズはさらにもう一段階の確率分布の仮定をおきます。
1次関数のパラメータ \{ \omega_{0, s}, \omega_{1, s} \} はそれぞれが花の種類ごとの平均と標準偏差を持つ正規分布に従うと仮定します。

\begin{aligned} \omega_{0, s} \sim& N(\mu_{\omega_{0}, s}, \sigma_{\omega_{0}, s}^2)\\ \mu_{\omega_{0}, s} \sim& N(0,10^2)\\ \sigma_{\omega_{0}, s} \sim& HN(10^2)\\ \omega_{1, s} \sim& N(\mu_{\omega_{1}, s}, \sigma_{\omega_{1}, s}^2)\\ \mu_{\omega_{1}, s} \sim& N(0,10^2)\\ \sigma_{\omega_{1}, s} \sim& HN(10^2) \end{aligned}

標準偏差 \sigma は標準偏差10の半正規分布に従うと仮定します。

\begin{aligned} \sigma \sim& HN(0,10^2) \end{aligned}

ここまで複雑だと数式がごちゃごちゃするので、一気にプログラミングを行います。
ここまでの仮定を、後ろの方から記述します。

def model_hierarchical_bayes(X, Y = None, cl = None, n_groups = None):
    '''
        5.3節の3種類の花の1次関数近似の階層ベイズモデル
    '''
    # 1次関数のパラメータ $\{ \omega_{0, s}, \omega_{1, s} \}$ はそれぞれが花の種類ごとの平均と標準偏差を持つ正規分布に従うと仮定します
    μ_ω0 = numpyro.sample("μ_ω0", dist.Normal(loc = 0, scale = 10))
    σ_ω0 = numpyro.sample("σ_ω0", dist.HalfNormal(scale = 10))
    # ベクトル化(階層ベイズの花の種類の数を確率変数に割り当てるためのNumPyroのお作法)
    with numpyro.plate("group", n_groups):
        ω0 = numpyro.sample("ω0", dist.Normal(loc = μ_ω0, scale = σ_ω0))
    # 傾きも切片と同様にモデルを定義する
    μ_ω1 = numpyro.sample("μ_ω1", dist.Normal(loc = 0, scale = 10))
    σ_ω1 = numpyro.sample("σ_ω1", dist.HalfNormal(scale = 10))
    with numpyro.plate("group", n_groups):
        ω1 = numpyro.sample("ω1", dist.Normal(loc = μ_ω1, scale = σ_ω1))
    # 1次関数の右辺を $\mu_{i}$ とおきます
    μ = numpyro.deterministic("μ", ω0[cl] + ω1[cl] * X)
    # 正規分布の標準偏差は標準偏差10の半正規分布に従うと仮定します
    σ = numpyro.sample("σ", dist.HalfNormal(scale = 10))
    # 目的変数 y は説明変数 x の値に応じた平均値 μ をパラメータとする正規分布に従うと仮定します。
    # ベクトル化(学習用データを確率変数に割り当てるためのNumPyroのお作法)
    N = len(X)
    with numpyro.plate("N", N):
        numpyro.sample("Y", dist.Normal(loc = μ, scale = σ), obs = Y)

作成したモデルをプロットします。
1次関数の確率分布にさらにもう一段階上の確率分布を仮定するのが階層ベイズモデルです。

model_args = {
    "X": X,
    "Y": Y,
    "cl": cl,
    "n_groups": len(species.unique())
}
try_render_model(model_hierarchical_bayes, render_name = "階層ベイズ", **model_args)

5.3.4 サンプリングと結果分析

データを用意してモデルを作成したら後はユーティリティ関数に渡すだけです。

model_args = {
    "X": X,
    "Y": Y,
    "cl": cl,
    "n_groups": len(species.unique())
}
idata = run_mcmc(
    model_hierarchical_bayes,
    num_chains = 4,
    num_warmup = 1000,
    num_samples = 1000,
    thinning = 1,
    seed = 42,
    target_accept_prob = 0.8,
    log_likelihood = False,
    **model_args
)

結果分析のコードは書籍とほぼ同じです。

まずはサンプリングが上手くいったか確認します。

az.plot_trace(idata, compact = False, var_names = ["ω0", "ω1", "σ"])
plt.tight_layout()

知りたかった1次関数の切片と傾きの事後分布をプロットします。

summary = az.summary(idata, var_names = ["ω0", "ω1"])
display(summary)

5.3.5 散布図と回帰直線の重ね描き

花の種類ごとの1次関数近似の直線と答えの散布図をプロットします。
たった9個のデータから計算したとは思えない精度です。

# alphaとbetaの平均値の導出
means = summary['mean']
ω0_0 = means['ω0[0]']
ω0_1 = means['ω0[1]']
ω0_2 = means['ω0[2]']
ω1_0 = means['ω1[0]']
ω1_1 = means['ω1[1]']
ω1_2 = means['ω1[2]']

# 回帰直線用座標値の計算
x_range = np.array([df['sepal_length'].min()-0.1, df['sepal_length'].max()+0.1+0.1])
y0_range = ω1_0 * x_range + ω0_0
y1_range = ω1_1 * x_range + ω0_1
y2_range = ω1_2 * x_range + ω0_2

# 散布図表示
sns.scatterplot(
    x='sepal_length', y='sepal_width', hue='species', style='species',
    data=df_sel, s=50)
plt.plot(x_range, y0_range, label='setosa')
plt.plot(x_range, y1_range, label='versicolor')
plt.plot(x_range, y2_range, label='virginica')
plt.legend();

# 散布図表示
sns.scatterplot(
    x='sepal_length', y='sepal_width', hue='species', style='species',
    data=df, s=50)
plt.plot(x_range, y0_range, label='setosa')
plt.plot(x_range, y1_range, label='versicolor')
plt.plot(x_range, y2_range, label='virginica')
plt.legend();

終わりに

階層ベイズモデルのベイズ推論を行いました。
関連するデータが少しずつあるという業務でありがちな状況でベイズ推論が輝くことが分かりました。

Discussion