ベイズ線形回帰を実装してみた

今回は，Pythonを使ってベイズ線形回帰を実装してみた．

はじめに
モデルの定義
推論の手順
- 事後分布の計算
- 予測分布の計算
やってみた
参考

はじめに

ベイズ推論についてはそれなりに勉強したつもりでいたが，そういえばちゃんと手を動かして実装したことはあまりなかった．

ちょうど，次の須山先生の本を読んでいる時にベイズ線形回帰が扱われていたので，Pythonの勉強も兼ねて実装してみる．

ベイズ深層学習 (機械学習プロフェッショナルシリーズ)

作者: 須山敦志
出版社/メーカー: 講談社
発売日: 2019/08/08
メディア: 単行本
この商品を含むブログを見る

モデルの定義

入力 $\mathbf{X} = \{\mathbf{x}_n \in \mathbb{R}^D; n=1,2,\cdots,N \}$ から連続値 $\mathbf{Y} = \{\mathbf{y}_n \in \mathbb{R}^1; n=1,2,\cdots,N \}$ を予測する問題を考える．このとき，ベイズ線形回帰モデルの同時分布は次のようになる．

$p(\mathbf{Y}, \mathbf{w}|\mathbf{X}) = p(\mathbf{\mathbf{w}})p(\mathbf{Y}|\mathbf{X}, \mathbf{w}) = p(\mathbf{\mathbf{w}})\prod_{n=1}^{N} p(y_n|\mathbf{x}_n, \mathbf{w}) \tag{1}$

ここで，ラベル $y_n$ は固定の分散 $\sigma_{y}^{2}$ をもつガウス分布に従って出力されるとする． $p(y_n|\mathbf{x}_n, \mathbf{w}) = \mathcal{N}(y_n|\mathbf{w}^{\mathrm{T}} \phi(\mathbf{x}_n), \sigma_{y}^{2}) \tag{2}$

ただし， $\phi : \mathbb{R}^{D} \rightarrow \mathbb{R}^{H}$ は特徴量関数であり，今回は $M$ 次の多項式関数を用いる．このように，特徴量関数を設計することで，入力と出力の間の非線形な関係を表現可能となる．

「え，線形回帰じゃないの？」と思う人もいるかもしれないが，このモデルは基底関数のパラメータによる線形結合に予測を行っているため，紛れもなく“線形”回帰である． “線形”と呼ばれているのは予測に用いる関数が直線であるという意味ではないということに注意しよう．

ここでは各特徴量に対する重みパラメータ $\mathbf{w} \in \mathbb{R}^{H}$ のみを学習するとする．この $\mathbf{w}$ の事前分布として，次のような平均 $\mathbf{m}_0$ ，共分散 $\boldsymbol{\Sigma}_0$ のガウス分布を与える．

$p(\mathbf{w}) = \mathcal{N}(\mathbf{w}| \mathbf{m}_0, \boldsymbol{\Sigma}_0) \tag{3}$

なお，参考元の書籍では事前分布の平均を $\mathbf{0}$ として式展開しているが，今回は陽に変数 $\mathbf{m}_0$ を与えた．そのため，以降の式展開が一部異なることに注意．

推論の手順

ベイズ推論は，一般的に 事後分布の計算 → 予測分布の計算 という手順で行われる（ただし必須というわけではない）．

今回は事前分布として共役な分布を仮定しているため，これらの分布は解析的に求めることができる．

事後分布の計算

重みパラメータ $\mathbf{w}$ の事後分布はベイズの定理を用いて算出できる．

$\displaystyle{p(\mathbf{w}|\mathbf{Y}, \mathbf{X}) = \frac{p(\mathbf{Y}|\mathbf{X}, \mathbf{w})p(\mathbf{w})}{p(\mathbf{Y}|\mathbf{X})}} \tag{4}$

この手の計算は慣れが必要だが，指数部に着目して式を整理していけば比較的簡単に計算できる．結果として，事後分布は次のようなガウス分布となる．

$\displaystyle{p(\mathbf{w}|\mathbf{Y}, \mathbf{X}) = \mathcal{N}(\mathbf{w}|\hat{\mathbf{\mu}}, \hat{\boldsymbol{\Sigma}} ) } \tag{5}$

ただし，

$\begin{align} \hat{\boldsymbol{\Sigma}}^{-1}\ &= \ \sigma_{y}^{-2} \sum_{n=1}^{N} \phi(\mathbf{x}_n)\phi(\mathbf{x}_n)^{\mathrm{T}} + \boldsymbol{\Sigma}_0^{-1} \tag{6} \\ \hat{\mathbf{\mu}}\ &= \ \hat{\boldsymbol{\Sigma}} \left( \sigma_{y}^{-2} \sum_{n=1}^{N} y_n \phi(\mathbf{x}_n) + \boldsymbol{\Sigma}_0^{-1} \mathbf{m}_0 \right) \tag{7} \end{align}$

予測分布の計算

学習後，新たに得られたテスト入力 $\mathbf{x}_{*}$ に対して，予測値 $y_{*}$ の分布 $p(y_{*}|\mathbf{x}_{*})$ を計算する．

$\begin{align} p(y_{*}|\mathbf{x}_{*}) &= \int p(y_{*}|\mathbf{x}_{*}, \mathbf{w})p(\mathbf{w}|\mathbf{Y}, \mathbf{X}) \mathrm{d}\mathbf{w} \\ &= \int \mathcal{N}(y_{*}|\mathbf{w}^{\mathrm{T}} \phi(\mathbf{x}_{*}), \sigma_{y}^{2}) \mathcal{N}(\mathbf{w}|\hat{\mathbf{\mu}}, \hat{\boldsymbol{\Sigma}}^{-1}) \mathrm{d}\mathbf{w} \\ &= \mathcal{N}(y_{*}|\mu_{*}(\mathbf{x}_{*}), \sigma^2_{*}(\mathbf{x}_{*})) \tag{8} \end{align}$

ただし，

$\begin{align} \mu_{*}(\mathbf{x}_{*}) &= \hat{\mathbf{\mu}}^{\mathrm{T}} \phi(\mathbf{x}_{*}) \tag{9} \\ \sigma^2_{*}(\mathbf{x}_{*}) &= \sigma^2_y + \phi(\mathbf{x}_{*})^{\mathrm{T}} \hat{\boldsymbol{\Sigma}}\phi(\mathbf{x}_{*}) \tag{10} \end{align}$

やってみた

ベイズ線形回帰クラスの実装

まず，ベイズ線形回帰用のクラスを次のように実装した．

import numpy as np
import matplotlib.pyplot as plt

class BayesianLinearRegression:
    """Bayesian linear regression class"""
    def __init__(self):
        """Self variables
        
        self.mu_0 (Dx1 column vector): Mean vector
        self.Sigma_0 (DxD matrix): Covariance matrix
        self.sigma2_y: Noise term
        """
        self.mu_0 = None
        self.Sigma_0 = None
        self.sigma2_y = None
        
    def set_prior(self, mu_0, Sigma_0, sigma2_y):
        """Set parameters for prior distribution
        """
        self.mu_0 = mu_0 if mu_0.shape[0] == 1 else mu_0.T
        self.Sigma_0 = Sigma_0
        self.sigma2_y = sigma2_y
        
    def learn(self, x_train, y_train, M):
        """Learn model and update posterior distribution
        """
        phi_ = self.base_polynomial(x_train, M).T
        
        Sigma_ = np.linalg.inv((self.sigma2_y**(-1))*np.dot(phi_, phi_.T) \
                               + np.linalg.inv(self.Sigma_0))
        mu_ = np.dot(Sigma_, self.sigma2_y**(-1)*np.sum(y_train*phi_.T, axis=0) \
                     + np.dot(np.linalg.inv(self.Sigma_0), self.mu_0))
        
        self.mu_0 = mu_
        self.Sigma_0 = Sigma_
        
        return mu_, Sigma_
        
    def get_sample_from_posterior(self, x_lin):
        """Get sample from posterior distiburion
        """
        w_ = np.random.multivariate_normal(mean=self.mu_0, cov=self.Sigma_0)
        
        phi_ = self.base_polynomial(x_lin, M).T
        sample_mean_ = np.einsum('h,hn->n', w_, phi_)
        
        return sample_mean_
        
    def predict(self, x_test):
        """Calculate predictive distribution
        """
        phi_ = self.base_polynomial(x_test, M).T
        
        y_est_ = np.dot(self.mu_0.T, phi_)
        sigma2_y_est_ = self.sigma2_y + np.diag(np.dot(phi_.T.dot(self.Sigma_0), phi_))
        
        return y_est_, sigma2_y_est_
    
    def base_polynomial(self, x, M):
        """Calculate polynomial basis function
        """
        D_ = x.shape[1]
        phi_ = np.zeros(shape=(N_, int((M+1)*D_)), dtype='float')
        for m in range(M+1):
            phi_[:,m::(M+1)] = x**m
        return phi_

Pythonの勉強も兼ねて，極力forループは使わないように工夫したつもりだ．なお，上記クラス内では書籍での表記に従い，ベクトルは縦ベクトルとして定義していることに注意．

学習データの作成

学習用データとして，正弦波をもとに作られた $N = 10$ 個のデータ点を用いる．入力データは1次元（ $D = 1$ ）とし，多項式の次数は $M = 3$ に設定した．

N_train = 10 # number of samples
x_train = 2*np.pi*np.random.rand(N_train).reshape(-1,1)
y_train =  np.sin(x_train) + np.random.normal(0,0.1,N_train).reshape(-1,1)

D = x_train.shape[1] # number of dimensions
M = 3 # degree of polynimial
H = (M+1)*D # number of dimensions after transformation

事前分布の設定

事前分布のパラメータは，ひとまず $\mathbf{\mu}_0 = \mathbf{0}$ ， $\mathbf{\Sigma}_0 = \mathbf{I}$ とおいた．ラベル $y_n$ の分散（今回は定数として扱っている）は， $\sigma^2_y = 1$ とした．

さらに，実際にデータを学習（事後分布を計算）する前に，事前分布からサンプリングを行ってみた．

# Parameters for prior distribution
mu_0 = np.zeros(H)
Sigma_0 = np.identity(H)

BayesLR = BayesianLinearRegression()
BayesLR.set_prior(mu_0, Sigma_0, 1) # Set parameters

x_lin = np.linspace(-1, 1, 30).reshape(-1,1)

# Get 5 samples from prior
for i in range(5):
    y_sample = BayesLR.get_sample_from_posterior(x_lin)
    plt.plot(x_lin, y_sample)

サンプルされたいくつかの $\mathbf{w}$ の値に基づき，多項式関数をプロットした結果が次の図である．

このように，ベイズ統計における回帰モデルでは，データを観測する前に候補となる関数を事前分布からサンプリングすることで，モデルで想定されている関数の具体例を示すことができる．

事後分布の計算と予測分布の算出

さて，では実際に学習データを用いて事後分布の計算および予測分布の算出を行ってみよう．

# Calculate posterior distribution
post_mu, post_Sigma = BayesLR.learn(x_train, y_train, M)

# Calculate predictive distribution
x_lin = np.linspace(-1, 7, 30).reshape(-1,1)
pred_mean, pred_var = BayesLR.predict(x_lin)

upper = pred_mean + np.sqrt(pred_var)
lower = pred_mean - np.sqrt(pred_var)

# Get five samples from posterior
for i in range(5):
    y_sample = BayesLR.get_sample_from_posterior(x_lin)
    if i==0:
        plt.plot(x_lin, y_sample, color='green', linewidth = 1.2, label='sample from posterior')
    else:
        plt.plot(x_lin, y_sample, color='green', linewidth = 1.2)

plt.plot(x_lin, pred_mean, color='orangered', linewidth = 2.5, linestyle='dashed', label='predictive mean')
plt.fill_between(x_lin[:,0], upper, lower, color='lightgrey', label='predictive variance')
plt.scatter(x_train, y_train, color='black', marker='o')

plt.legend(ncol=2)
plt.ylim(-3,3)

次の図中において，黒丸プロットが学習データ，緑の実線が事後分布からサンプルした多項式関数，赤い破線が予測分布の平均，灰色の部分が予測分布から求めた予測の標準偏差を表している．

図から分かるように，学習データにうまくフィッティングした多項式関数が得られていることが分かる．

逐次学習させてみる

さて，一つ前の例では学習データを全て一気に学習に用いたが，サンプル点一つずつ学習させるとどうなるだろうか．

このような学習を逐次学習やオンライン学習という（ベイズ学習の枠組みではベイズ更新とも呼ばれる）．

手順としては，得られた学習サンプルを用いて事後分布を算出した後，それを次のステップにおける事前分布に設定し，また同様に事後分布を算出する，という流れを反復して行っていく．つまり，事前分布を（事後分布により）逐次的に更新していくわけである．

# Parameters for prior distribution
mu_0 = np.zeros(H)
Sigma_0 = np.identity(H)

x_lin = np.linspace(-1, 7, 30).reshape(-1,1)

BayesLR = BayesianLinearRegression()
BayesLR.set_prior(mu_0, Sigma_0, 1)

plt.figure(figsize=(10,20))

# Learn and predict one sample at a time
for i in range(N_train):
    # Calculate posterior
    BayesLR.learn(x_train[i].reshape(-1,1), y_train[i].reshape(-1,1), M)
    # Calculate predictive
    pred_mean, pred_var = BayesLR.predict(x_lin)
    upper = pred_mean + np.sqrt(pred_var)
    lower = pred_mean - np.sqrt(pred_var)
    
    plt.subplot(5,2,i+1)
    plt.plot(x_lin, pred_mean, color='orangered', linewidth = 2.5, linestyle='dashed', label='predictive mean')
    plt.fill_between(x_lin[:,0], upper, lower, color='lightgrey', label='predictive variance')
    plt.scatter(x_train[0:i+1], y_train[0:i+1], color='black', marker='o')
    plt.ylim(-2.5,2.5)
    plt.title('N = {}'.format(i+1))

次に結果を列挙する．予測分布の平均の変化は勿論だが，予測分布の標準偏差の変化にも注意してみる．

サンプルを新しく学習する度に，予測分布の平均が目標となる関数（今回であれば正弦波）に近づいていることがわかる．

一方，データが少ない時には予測の標準偏差は大きく，予測における不確実性が大きい．しかし，新たな観測値を得ると，観測サンプル付近の不確実性が小さくなっていることがわかる．観測したデータ数が多い場合，モデルが予測に対して自信を持っていることを示している．

このように，観測データ数に応じた予測の不確実性を評価できることは，ベイズ推論の大きなメリットだと感じた．

また，よく見ると $N=10$ の場合の結果は，一つ前の全データを一括して学習した場合の結果と非常に似ている．今回のように共役事前分布を使った学習では，データを逐次的に与えた場合と一度にすべて与えた場合とで，最終的に得られる事後分布が一致するらしい（ただし，データの生成過程に順序の依存性を仮定しない場合のみ）．