5.2 線形回帰のベイズ推論 - 第5章ベイズ推論プログラミング

はじめに

Pythonでスラスラわかる ベイズ推論「超」入門(赤石 雅典 (著), 須山 敦志 (監修))の5.2節のPyMCコードをNumPyroで書き直しました。
アイリス・データセットのversicolorのがく片の長さsepal_lengthとがく片の幅sepal_widthの1次関数近似のベイズ推論を行います。
単純な線形回帰であっても一目ではベイズの定理と確率モデルとの関係性が分かりにくいので、数式で考えることの恩恵が大きくなります。

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

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

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

5.2 線形回帰のベイズ推論

5.2.1 問題設定

versicolorのがく片の長さ(sepal_length)と幅(sepal_width)の関係を1次関数近似します。

5.2.2 データ準備

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

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

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

# sepal_lengthとsepal_widthの列を抽出
X = jnp.array(df_versicolor['sepal_length'], dtype = float)
Y = jnp.array(df_versicolor['sepal_width'], dtype = float)

抽出したデータの散布図をプロットします。
少しばらつきは大きいですが1次関数近似してもよさそうです。

plt.title('２つの項目間の関係')
plt.scatter(X, Y, label='ベイズ推論で利用', c='b', marker='o')
plt.legend()
plt.xticks([4, 5, 6, 7, 8])
plt.yticks([1, 2, 3, 4, 5])
plt.xlabel('sepal_length')
plt.ylabel('sepal_width')

5.2.3 確率モデル定義1

中間変数を設定しないモデル作成の項です。
少しモデル定義の関数が短くなりますが、可読性が下がるデメリットの方が大きいです。
個人的に実用性を感じないので省略します。

5.2.4 確率モデル2

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

まず、散布図より $N$ 組の sepal_length $x_{i}$ とがく片の幅 sepal_width $y_{i}$ は1次関数で近似できると仮定します。

$$y_{i} \approx \omega_{0} + \omega_{1} x_{i}$$

確率モデルを作成するためのテクニックとして、右辺を $\mu_{i}$ とおくことにします。

$$\mu_{i} \equiv \omega_{0} + \omega_{1} 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}, \omega_{1} \}$ と正規分布の標準偏差 $\sigma$ を求める問題になりました。
これらのパラメータの推論精度を高める情報はありません。
広めの事前分布を与えておきましょう。
1次関数のパラメータ $\{ \omega_{0}, \omega_{1} \}$ はそれぞれ平均0, 標準偏差10の正規分布に従うと仮定します。
標準偏差 $\sigma$ は標準偏差10の半正規分布に従うと仮定します。
さらに、これらのパラメータは独立であると仮定します。

$$\begin{aligned} \omega_{0} \sim& N(0,10^2)\\ \omega_{1} \sim& N(0,10^2)\\ \sigma \sim& HN(0,10^2)\\ P(\omega_{0}, \omega_{1}, \sigma) =& P(\omega_{0}) \times P(\omega_{1}) \times P(\sigma) \end{aligned}$$

ここまでの結果を条件付確率 $P(H|D)=P(D,H)/P(D)$ に当てはめます。
データは説明変数と目的変数の組 $D = (x, y)$ ですが、 $y$ に関する確率分布を仮定しているので、良い感じに分離します。
説明変数 $x$ と今回のパラメータは互いに独立であると仮定しておきます。

$$\begin{aligned} P(H = (\omega_{0}, \omega_{1}, \sigma)|D = (x, y)) / P(D = (x,y)) =& P(D = (x, y),H = (\omega_{0}, \omega_{1}, \sigma)) / \left( P(y|x) \times P(x) \right)\\ =& P(y,(x, \omega_{0}, \omega_{1}, \sigma)) / \left( P(y|x) \times P(x) \right)\\ =& P(y | (x, \omega_{0}, \omega_{1}, \sigma)) \times P(x, \omega_{0}, \omega_{1}, \sigma) / \left( P(y|x) \times P(x) \right)\\ =& P(y | (x, \omega_{0}, \omega_{1}, \sigma)) \times P(x) \times P(\omega_{0}) \times P(\omega_{1}) \times P(\sigma) / \left( P(y|x) \times P(x) \right)\\ =& P(y | (x, \omega_{0}, \omega_{1}, \sigma)) \times P(\omega_{0}) \times P(\omega_{1}) \times P(\sigma) / P(y|x)\\ \propto& P(y | (x, \omega_{0}, \omega_{1}, \sigma)) \times P(\omega_{0}) \times P(\omega_{1}) \times P(\sigma) \end{aligned}$$

説明変数がある場合、条件付確率から計算すると説明変数の確率がきれいに約分されます。
PyMCやNumPyroのプログラミングは、右辺の最後の式を後ろから記述します。

def model_linear_regression(X, Y = None):
    '''
        5.2節のversicolorのがく片長さと幅の確率分布モデル
    '''
    # 1次関数の切片は標準偏差10の半正規分布に従うと仮定します
    ω0 = numpyro.sample("ω0", dist.Normal(loc = 0, scale = 10))
    # 1次関数の傾きは標準偏差10の半正規分布に従うと仮定します
    ω1 = numpyro.sample("ω1", dist.Normal(loc = 0, scale = 10))
    # 1次関数の右辺を $\mu_{i}$ とおきます
    μ = numpyro.deterministic("μ", ω0 + ω1 * 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
}
try_render_model(model_linear_regression, render_name = "線形回帰", **model_args)

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

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

model_args = {
    "X": X,
    "Y": Y
}
idata = run_mcmc(
    model_linear_regression,
    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次関数の切片と傾きの事後分布をプロットします。

ax = az.plot_posterior(idata, var_names = ["ω0", "ω1"])
plt.suptitle("1次関数近似の係数の分布")
plt.tight_layout()
plt.show()

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

サンプリングしたそれぞれの切片と傾きを使って近似直線をプロットします。
今回は4,000個のサンプリングデータがあるので、4,000本の直線をプロットします。

このセルの計算はNumPyのテクニックがつまった計算です。
何度も読み返すことになるでしょう。

# データ量を節約するために2点間を結ぶ線分をプロットします。
# 説明変数 x の値域より少し広い範囲を選択します。
x_values = np.array([X.min()-0.1, X.max()+0.1])
print(x_values, x_values.shape) # (2点,)

# サンプリング結果から切片と傾きを取り出しshapeを加工する
ω0_posterior = idata['posterior']['ω0'].values.reshape(-1, 1)
ω1_posterior = idata['posterior']['ω1'].values.reshape(-1, 1)

# shapeの確認
print(ω0_posterior.shape, ω1_posterior.shape)# (サンプリングサイズ,1列)

# 4000パターンそれぞれで、２点の１次関数値の計算
y_preds = x_values * ω1_posterior + ω0_posterior
print(y_preds.shape)# (サンプリングサイズ,2点!?)

for y_pred in y_preds:
    # 各行のデータから近似直線をぷろｗする。
    plt.plot(x_values, y_pred, lw=1, alpha=0.01, c='g')
plt.scatter(X, Y)
plt.title('ベイズ推論による回帰直線')
plt.xlabel('sepal_length')
plt.ylabel('sepal_width');

5.2.7 少ない観測値でのベイズ推論

省略。

終わりに

1次関数近似のベイズ推論を行いました。
機械学習の最初に学ぶ線形回帰もこんなに複雑になるのですね。