More than 5 years have passed since last update.

PRML第10章のベイズ混合ガウスモデルに対する変分推論をPythonで実装

Last updated at 2019-04-20Posted at 2019-03-21

この記事では、PRMLの第10章で述べられている、混合ガウスモデルに対する変分推論をpythonで実装します。
対応するjupyter notebookは筆者のgithubリポジトリにあります。連載全体の方針や、PRMLの他のアルゴリズムの実装については、連載のまとめページをご覧いただければと思います。

なかなか計算がハードな章ですが、ここでは結果だけを利用して、さくさく実装したいと思います。
なお、変分下界の計算については、PRMLの内容より少し詳しく扱っています。具体的には、教科書の式(10.70)-(10.77)の式をさらに変形し、実装に適した簡潔な形にしてあります。

1 理論のおさらい

1.1 設定

前回と同様、教師なし学習を考えます。

$N \in \mathbb{N}$ : データ点の個数
$D \in \mathbb{N}$ : データの次元
$x_0, x_1, \dots , x_{N-1} \in \mathbb{R}^D$ : データ点
データ点をまとめて、行列$X$と表します。ただし、$X_{n,i}$は$x_n$の$i$番目の要素とします。

1.2 モデル

$N$個のデータ点に対し、混合要素数$K \in \mathbb{N}$のベイジアン混合ガウスモデルは次で定義されます：

$$
\begin{align}
& p(\pi) = {\rm Dir}(\pi | \alpha_0)\
& p(Z | \pi) = \prod_{n=0}^{N-1} \prod_{k=0}^{K-1} \pi_{k}^{z_{n,k}} \
& p\left(X \middle| Z, \mu, \Lambda, \pi \right) =
\prod_{n=0}^{N-1} \prod_{k=0}^{K-1} \mathcal{N}\left(x_n \middle| \mu_k , \Lambda_{k}^{-1} \right)^{z_{n,k}} \
& p(\Lambda) = \prod_{k=0}^{K-1} \mathcal{W}\left(\Lambda_k \middle| W_0, \nu_0 \right)\
& p(\mu|\Lambda) = \prod_{k=0}^{K-1} \mathcal{N}\left(\mu_k \middle| m_0, (\beta_0 \Lambda_k)^{-1} \right)
\end{align}
$$

ただし、

$Z = (z_{n,k})_{ n \in \{ 0, 1, \dots, N-1 \} k \in \{ 0,1, \dots, K-1 \} }$は潜在変数を表す行列で、$\ z_{n,k}\in \{0,1 \}, \sum_{k=0}^{K-1}z_{n,k}=1$を満たします。
$\pi_k \geq 0, \sum_{k=0}^{K-1} \pi_k = 1$であり、$(\pi_k)_{k \in \left\{ 0,1,\dots, K-1 \right\}}$をまとめて$\pi$と表します。
$\mathcal{N}\left(x \middle| \mu_k , \Lambda_{k}^{-1} \right)$は、平均$\mu_k \in \mathbb{R}^D$、$D \times D$の精度行列$\Lambda_k$を持つガウス分布の確率密度関数です。まとめて、$(\mu_k)_{k \in \left\{ 0,1,\dots, K-1 \right\}}$を$\mu$、 $(\Lambda_k)_{k \in \left\{ 0,1,\dots, K-1 \right\}}$を$\Lambda$と表すことにします。
${\rm Dir}(\pi | \alpha_0)$はディリクレ分布の確率密度関数です。
$\mathcal{W}\left(\Lambda_k \middle| W_0, \nu_0 \right)$はウィシャート分布の確率密度関数です。

これらの確率分布を用いると、全ての変数の同時確率分布は次のように表されます：

$$
\begin{align}
p\left(X, Z, \pi, \mu, \Lambda \right)
= p\left(X \middle| Z, \mu, \Lambda, \pi \right)
p(Z | \pi)
p(\pi)
p(\mu|\Lambda)
p(\Lambda)
\end{align}
$$

1.3 変分分布

ここでの目標は、パラメーターの事後分布$p(Z, \pi,\mu,\Lambda| X)$を変分分布$q(Z, \pi, \mu, \Lambda)$を用いて近似することです。なお、陽には書かれていませんが、$q$はもちろん観測されたデータ$X$に依存します。

導出は本をご覧いただくとして、結果をまとめると

$$
\begin{align}
q(Z, \pi, \mu, \Lambda) = q(Z) q(\pi) \prod_{k=0}^{K-1} q(\mu_k, \Lambda_k),
\end{align}
$$

ただし、

$$
\begin{align}
& q(Z) = \prod_{n=0}^{N-1} \prod_{k=0}^{K-1} r_{n,k}^{z_{n,k}} \
& q(\pi) = {\rm Dir}(\pi|\alpha) \
& q(\mu_k, \Lambda_k) =
\mathcal{N}\left(\mu_k \middle| m_k, (\beta_k \Lambda_k)^{-1}\right)
\mathcal{W}\left( \Lambda_k \middle| W_k, \nu_k \right)
\end{align}
$$

であり、パラメーターの満たすべき条件は以下のようになります：

まず、$q(Z)$については

$$
\begin{align}
& r_{n,k} := \frac{\rho_{n,k}}{\sum_{j=0}^{K-1} \rho_{n,j}} \
& \rho_{n,k} := \frac{\tilde{\pi}_k \tilde{\Lambda}_{k}^{1/2}}{(2\pi)^{D/2}}
\exp\left[ - \frac{D}{2\beta_k} - \frac{\nu_k}{2} (x_n - m_k)^T W_k (x_n - m_k) \right] \
& \tilde{\pi}_k := \exp \left[ \psi(\alpha_k) - \psi(\hat{\alpha}) \right] \\
& \tilde{\Lambda}_k := \exp\left[ \sum_{i=0}^{D-1} \psi\left( \frac{\nu_k-i}{2} \right) + D\log 2 + \log(\det W_k) \right] \
& \psi(a) := \frac{d}{da} \log \Gamma(a) \
& \hat{\alpha} := \sum_{k=1}^{K} \alpha_k
\end{align}
$$
で、$q(\pi, \mu, \Lambda)$については
$$
\begin{align}
& \alpha_k := \alpha_0 + N_k \
& \beta_k := \beta_0 + N_k \
& \nu_k := \nu_0 + N_k \
& m_k := \frac{1}{\beta_k}(\beta_0 m_0 + N_k \bar{x}_k) \
& W_{k}^{-1} := W_{0}^{-1} + N_k S_k
{} + \frac{\beta_0 N_k}{\beta_0 + N_k}(\bar{x}_k - m_0)(\bar{x}_k - m_0)^T \
& N_k := \sum_{n=0}^{N-1} r_{n,k} \
& \bar{x}_k := \frac{1}{N_k} \sum_{n=0}^{N-1} r_{n,k} x_n \
& S_k := \frac{1}{N_k} \sum_{n=0}^{N-1} r_{n,k}(x_n - \bar{x}_k)(x_n - \bar{x}_k)^T
\end{align}
$$

なお、これらの式はパラメーターの定義を直ちに与えるものではありません。というのも、これらの式が表すのは、あくまで「$Z$についての分布のパラメーターは、$\pi, \mu, \Lambda$についての分布のパラメーターが決まれば決まり、逆もまた成り立つ」ことに過ぎないからです。換言すると、これらの式はパラメーターが満たすべき条件を表すself-consistent equation¹です。

このself-consistent equationの解は、次のような繰り返しによって求めることができます²。

パラメーター$\alpha_k, \beta_k, \nu_k, m_k, W_k$を初期化します³。
分担率$r_{n,k}$を$\alpha_k, \beta_k, \nu_k, m_k, W_k$から計算します。これをe-like-stepと呼ぶことにします。
$\alpha_k, \beta_k, \nu_k, m_k, W_k$を$r_{n,k}$から計算します。これをm-like-stepと呼ぶことにします。
2、3を収束するまで繰り返します。収束判定には、変分下界を用いることにします。

1.4 変分下界

収束判定やモデル選択(後述)に用いるために、変分下界$\mathcal{L}(q)$も計算しておきます。
変分下界の表式はPRMLの(10.70)-(10.77)に与えられていますが、実は計算を進めると更に簡単にすることができます。導出は省略しますが⁴、結果は次のようになります：

$$
\begin{align}
\mathcal{L}(q) = -\sum_{n=0}^{N-1}\sum_{k=0}^{K-1} r_{n, k} \log r_{n, k}
+ \log \frac{C(\alpha_0)}{C(\alpha)}
+ \frac{D}{2} \sum_{k=0}^{K-1} \log \frac{\beta_0}{\beta_k}
+ \sum_{k=0}^{K-1} \log \frac{B(W_0, \nu_0)}{B(W_k, \nu_k)}
- \frac{DN}{2} \log(2\pi),
\end{align}
$$
ただし、
$$
\begin{align}
C(\alpha) := \frac{\Gamma\left( \sum_{k=0}^{K-1} \alpha_k \right)}{\prod_{k=0}^{K-1} \Gamma(\alpha_k)}
\end{align}
$$
はディリクレ分布の規格化定数で、
$$
\begin{align}
B(W, \nu) := \left( \det W \right)^{-\nu/2} \left[ 2^{\nu D/2} \pi^{D(D-1)/4} \prod_{i=0}^{D-1} \Gamma\left( \frac{\nu - i}{2} \right) \right]^{-1}
\end{align}
$$
はウィシャート分布の規格化定数です。

1.5 モデル比較

さて、ここまでハイパーパラメーター$\alpha_0, \beta_0, \nu_0, m_0, W_0$はgivenとして話を進めてきました。実際の応用では、これらを選ぶ必要があります。

ここでは変分下界を最大化するハイパーパラメーターを選ぶことにします。これは、ハイパーパラメーターに対する事前分布$p(m)$が一様と仮定したときに、(10.36)の事後分布$q(m)$を最大化する$m$を取ることに対応します⁵。

1.6 予測分布

さて、モデル選択と学習が無事に済んだとして、最後に予測のための数式をおさらいしておきましょう。

新しい入力を$\hat{x}$、対応する潜在変数を$\hat{z}$としたとき、これらに対する予測分布は次の式で与えられます：

$$
\begin{align}
p(\hat{z} | \hat{x}, X) &\propto p(\hat{x}, \hat{z} | X) \simeq \prod_{k=0}^{K-1} \left[ \frac{\alpha_k}{\sum_{j=0}^{K-1} \alpha_j} \mathrm{St}\left(\hat{x} \middle| m_k, L_k, \nu_k + 1 - D \right) \right]^{\hat{z}_k} \
p(\hat{x} | X) &\simeq \frac{1}{\sum_{k=0}^{K-1} \alpha_k} \sum_{k=0}^{K-1} \alpha_k \mathrm{St}(\hat{x} | m_k, L_k, \nu_k + 1 - D) \
L_k &:= \frac{(\nu_k + 1 - D) \beta_k}{ 1 + \beta_k} W_k,
\end{align}
$$
ただし、$\mathrm{St}$はStudentのt分布を表します。導出はPRML 10.2.3節をご覧ください。

2. 数式からコードへ

では、ここまでの理論を元に、BayesianGaussianMixtureModelクラスや、いくつか関数を実装していきます。

2.1 概要

まず、BayesianGaussianMixtureModelの概観を示します。このクラスには、次のようなデータ属性とメソッドを持たせることにします。

データ属性

K : $K$, i.e., 潜在要素の個数
D : $D$, i.e., 入力の次元
alpha0 : $\alpha_0$, 正のfloat
beta0 : $\beta_0$, 正のfloat
nu0 : $\nu_0$, nu0 > $D - 1$を満たすfloat
m0 : $m_0$, ($D$,) array.
W0 : $W_0$, ($D$, $D$) array, 実対称正定値行列。
alpha : ($K$,) array, alpha[k] = $\alpha_k$
beta : ($K$,) array, beta[k] = $\beta_k$
nu : ($K$,) array, nu[k] = $\nu_k$
m : ($K$, $D$) array, m[k, i] = $m_k$の第$i$成分
W : ($K$, $D$, $D$) array, W[k, i, j] = $W_k$の(i,j)成分.
lower_bound : float, 変分下界

メソッド

_init_params : fittingを行う際に、パラメーターalpha, beta, nu, m, Wの初期化を行うメソッド。初期化の仕方は入力$X$に依存させることにします。
_e_like_step : e-like-stepの計算を行い、alpha, beta, nu, m, Wから分担率$r$を求めて返すメソッドです。
_m_like_step : m-like-stepを行うメソッドです。$r$からalpha, beta, nu, m, Wを計算します。
calc_lower_bound : 現在のパラメーターの値について、変分下界を計算して返すメソッドです。
fit : 入力データに対してfittingを行うメソッドです。具体的には、_e_like_stepと_m_like_stepを交互に実行し、変分下界の値の変化が十分小さくなれば停止します。
_predict_joint_proba : 与えられたデータ点$\hat{x}$(複数)に対し、同時分布$p(\hat{x}, \hat{z} | X)$を求めて返すメソッドです。
predict_proba : 与えられたデータ点$\hat{x}$(複数)に対し、それぞれの点が各潜在要素に属する確率を計算して返すメソッドです。
predict : 与えられたデータ点$\hat{x}$(複数)に対し、それぞれの点が属する潜在要素の予測値を返します。
calc_prob_density : 与えられたデータ点$\hat{x}$(複数)に対し、その予測確率密度を返します。

以下では、これらのメソッドの実装方法について述べていきます。

2.2 パラメーターの初期化

パラメーターalpha, beta, nu, m, Wの初期化方法はいろいろあるかと思いますが、ここでは次のような初期化を行うことにします：

alpha, beta, nu : alpha[k] = alpha0 + $N/K$のように初期化します。各潜在要素に「属する」データ点の個数はだいたい同じくらいでしょう(もしくは、そうだと考えて初期値にとっても問題ないでしょう)ということを暗に仮定しています⁶。
mu : 入力$X$から一様ランダムにK個の点を取り、muとして用いることにします。
W : 各$W_k$は同一の対角行列とし、その$(i, i)$成分は1/(入力$X$の第i成分の分散)と取ることにします⁷。

なお、「k-meansを先に走らせ、そのクラスター中心をmuの初期値として、クラスターの分散をWの初期値として用いる」などの手法も考えられますし、パラメーターではなく分担率$r$をランダムに初期化するなどの手法も取れるかと思います⁸。ただ、ここではあまり深く考えずに上記の手法を用いることにします。

これをコードに直すと、次のようになります：


    def _init_params(self, X, random_state=None):
        '''
        Method for initializing model parameterse based on the size and variance of the input data array. 
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        random_state : int
            int, specifying the seed for random initialization of m
        '''
        N, D = X.shape
        rnd = np.random.RandomState(seed=random_state)
        
        self.alpha = (self.alpha0 + N / self.K) * np.ones(self.K)
        self.beta = (self.beta0 + N / self.K) * np.ones(self.K)
        self.nu = (self.nu0 + N / self.K) * np.ones(self.K)
        self.m = X[rnd.randint(low=0, high=N, size=self.K)]
        self.W = np.tile(np.diag(1.0/np.var(X, axis=0)), (self.K, 1, 1))

2.3 E-like-step

_e_like_stepでは、次の式に従って$r_{n,k}$を計算します(1節の変分分布のところで見た式です。)。
$$
\begin{align}
r_{n,k} &:= \frac{\rho_{n,k}}{\sum_{j=0}^{K-1} \rho_{n,j}} \
\rho_{n,k} &:= \frac{\tilde{\pi}_k \tilde{\Lambda}_{k}^{1/2}}{(2\pi)^{D/2}}
\exp\left[ - \frac{D}{2\beta_k} - \frac{\nu_k}{2} (x_n - m_k)^T W_k (x_n - m_k) \right] \
\tilde{\pi}_k &:= \exp \left[ \psi(\alpha_k) - \psi(\hat{\alpha}) \right] \\
\tilde{\Lambda}_k &:= \exp\left[ \sum_{i=0}^{D-1} \psi\left( \frac{\nu_k-i}{2} \right) + D\log 2 + \log(\det W_k) \right] \
\psi(a) &:= \frac{d}{da} \log \Gamma(a) \ \ (\mbox{digamma function}) \
\hat{\alpha} &:= \sum_{k=1}^{K} \alpha_k
\end{align}
$$

計算にあたっては、次の3つの配列を用意しておきます。

tpi : (K,) array, tpi[k] = $\tilde{\pi}_k$
tlam : (K,) array, tlam[k] = $\tilde{\Lambda}_k$
rho : (N,K) array, rho[n,k] = $\rho_{n,k}$

tlamの計算のために、次のような配列arg_digammaを定めます：arg_digamma[k,i] = $\nu_k - i$

また、rhoを求めるにあたっては、まず($N, K, D$)配列diffをdiff[n,k,i] = $(x_n - m_k)_{i}$のように計算し、これを用いて$(N, K)$配列exponentをexponent[n, k] = $D/\beta_k + \nu_k (x_n-m_k)^T W_k (x_n-m_k)$と計算します。

これらの配列を用いて、rhoを求めることができます。なお、規格化のための$2\pi$のfactorは省略することにします。というのも、この部分は$r_{n,k}$の結果には影響しないためです。
また、rhoの定義の指数関数の影響から、数値精度の関係で$\sum_{j=0}^{K-1}\rho_{n,j} = 0$という結果が得られる可能性があります。この場合、$r_{n,k}$の規格化が正しくできなくなってしまいます。そこで、rhoをexponentから計算する際には、exponentの代わりに次のような配列exponent_subtractedを用います⁹：
To prevent such trouble, we subtract a constant (for each $n$) as exponent[n, k] - $\min_j$(exponent[n, j]).

コードにすると、次のように書くことができます。

    def _e_like_step(self, X):
        '''
        Method for calculating the array corresponding to responsibility.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
            
        Returns
        ----------
        r : 2D numpy array
            2D numpy array representing responsibility of each component for each sample in X, 
            where r[n, k] = $r_{n, k}$.
        
        '''
        N, _ = np.shape(X)
        
        tpi = np.exp( digamma(self.alpha) - digamma(self.alpha.sum()) )
        
        arg_digamma = np.reshape(self.nu, (self.K, 1)) - np.reshape(np.arange(0, self.D, 1), (1, self.D))
        tlam = np.exp( digamma(arg_digamma/2).sum(axis=1)  + self.D * np.log(2) + np.log(np.linalg.det(self.W)) )
        
        diff = np.reshape(X, (N, 1, self.D) ) - np.reshape(self.m, (1, self.K, self.D) )
        exponent = self.D / self.beta + self.nu * np.einsum("nkj,nkj->nk", np.einsum("nki,kij->nkj", diff, self.W), diff)
        
        exponent_subtracted = exponent - np.reshape(exponent.min(axis=1), (N, 1))
        rho = tpi*np.sqrt(tlam)*np.exp( -0.5 * exponent_subtracted )
        r = rho/np.reshape(rho.sum(axis=1), (N, 1))
        
        return r

2.4 M-like-step

_m_like_stepでは、今度は$r$からパラメーターを計算します。

$$
\begin{align}
\alpha_k &:= \alpha_0 + N_k \
\beta_k &:= \beta_0 + N_k \
\nu_k &:= \nu_0 + N_k \
m_k &:= \frac{1}{\beta_k}(\beta_0 m_0 + N_k \bar{x}_k) \
W_{k}^{-1} &:= W_{0}^{-1} + N_k S_k + \frac{\beta_0 N_k}{\beta_0 + N_k}(\bar{x}_k - m_0)(\bar{x}_k - m_0)^T \
N_k &:= \sum_{n=0}^{N-1} r_{n,k} \
\bar{x}_k &:= \frac{1}{N_k} \sum_{n=0}^{N-1} r_{n,k} x_n \
S_k &:= \frac{1}{N_k} \sum_{n=0}^{N-1} r_{n,k}(x_n - \bar{x}_k)(x_n - \bar{x}_k)^T
\end{align}
$$

計算にあたっては、次の配列を定義します。

n_samples_in_component : ($K$,)配列 n_samples_in_component[k] = $N_k$
barx : ($K, D$)配列,　barx[k, i] = $\bar{x}_{k, i}$
S : ($K, D, D$)配列, S[k,i,j] = $(S_k)_{i,j}$
diff : ($N,K,D$)配列, diff[n,k,i] = $x_{n,i} - \bar{x}_{k,i}$
diff2 : ($K, D$)配列, where diff2[k, i] = $\bar{x}_{k,i} - (m_0)_i$


    def _m_like_step(self, X, r):
        '''
        Method for calculating the model parameters based on the responsibility.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        r : 2D numpy array
            2-D numpy array representing responsibility of each component for each sample in X, 
            where r[n, k] = $r_{n, k}$.
        '''
        N, _ = np.shape(X)
        n_samples_in_component = r.sum(axis=0)
        barx = r.T @ X / np.reshape(n_samples_in_component, (self.K, 1))
        diff = np.reshape(X, (N, 1, self.D) ) - np.reshape(barx, (1, self.K, self.D) )
        S = np.einsum("nki,nkj->kij", np.einsum("nk,nki->nki", r, diff), diff) / np.reshape(n_samples_in_component, (self.K, 1, 1))
        
        self.alpha = self.alpha0 + n_samples_in_component
        self.beta = self.beta0 + n_samples_in_component
        self.nu = self.nu0 + n_samples_in_component
        self.m = (self.m0 * self.beta0 + barx * np.reshape(n_samples_in_component, (self.K, 1)))/np.reshape(self.beta, (self.K, 1))
        
        diff2 = barx - self.m0
        Winv = np.reshape(np.linalg.inv( self.W0 ), (1, self.D, self.D)) + \
            S * np.reshape(n_samples_in_component, (self.K, 1, 1)) + \
            np.reshape( self.beta0 * n_samples_in_component / (self.beta0 + n_samples_in_component), (self.K, 1, 1)) * np.einsum("ki,kj->kij",diff2,diff2) 
        self.W = np.linalg.inv(Winv)

2.5 変分下界の計算

次に、収束判定とhyperparameter選択のために用いる変分下界の計算を行う関数を実装します。

まず、コードを簡潔にするために、$\log C(\alpha)$と$\log B(W, \nu)$を計算する関数を別に書いておきます。$C$と$B$の定義は、次のものでした：
$$
\begin{align}
C(\alpha) &:= \frac{\Gamma\left( \sum_{k=0}^{K-1} \alpha_k \right)}{\prod_{k=0}^{K-1} \Gamma(\alpha_k)} \
B(W_k, \nu_k) &:= \left( \det W_k \right)^{-\nu_k/2} \left[ 2^{\nu_k D/2} \pi^{D(D-1)/4} \prod_{i=0}^{D-1} \Gamma\left( \frac{\nu_k - i}{2} \right) \right]^{-1}
\end{align}
$$

実装にあたっては、

$\log C(\alpha)$を計算する関数は配列を入力とし、
$\log B(W, \nu)$を計算する関数は、単一の$(W, \nu)$の組と、複数の$(W, \nu)$の組(e.g. $(W_1, \nu_{1}), (W_2, \nu_{2}), \dots$)いずれにも対応できるようにしておきました。

$\log B$については、これは$\log B(W_0, \nu_0)$を計算するためです。
$\log C$についても同様にfloatを入力として$\log C(\alpha_0)$を計算できるようにしたかったのですが、そのためには$K$の値が必要となり、入力を増やしたくなかったのでやめておきました。後の計算では、$\log C(\alpha_0)$の計算はlogC( self.alpha0 * np.ones(self.K))というふうに対応します¹⁰。


from scipy.special import gamma, digamma, gammaln

def logC(alpha):
    '''
    Function for calculating the natural log of C(alpha), the normalization constant for Dirichlet distributions.
    See equations (B.16) and (B.27) of PRML.
    
    Parameters
    ----------
    alpha : 1D numpy array
        1D numpy array representing the parameter alpha for Dirichlet distribution. 
            
    Returns
    ----------
    logC : flaot
        log of C(alpha), the normalization constant for Dirichlet distributions
    '''
    return gammaln(alpha.sum()) - gammaln(alpha).sum()
    
def logB(W, nu):
    '''
    Function for calculating the natural log of B(W, nu), the normalization constant for Wishart distributions.
    See equations (B.78) and (B.79) of PRML.
    
    Parameters
    ----------
    W : 2D or 3D numpy array
        numpy array representing the parameter W for Wishart distributions.
        When W is a 2D array, it must be a (D, D) symmetric positive definite matrix, where D is a integer.
        When W is a 3D array, it must be a (K, D, D) array, and W[k] are symmetric positive definite matrices, where K and D are integers. 
        In the latter case, it is understood that the function deals with several Wishart distributions at once.
    nu : float, or 1D numpy array
        Number(s) representing the degrees of freedom of Wishart distributions.
        When W is (K, D, D) (resp. 2D) numpy array, it must be (K,) numpy array (resp. float).
            
    Returns
    ----------
    logB : flaot or 1D numpy array
        log of B(alpha), the normalization constant for Wishart distributions.
    '''
    Wshape = W.shape
    if len(Wshape) == 2:
        D, _ = Wshape
        arg_gamma = nu - np.arange(0, D, 1)
        return -nu/2 * np.log(np.linalg.det(W)) - D/2 * nu * np.log(2) - D* (D - 1)/4 * np.log(np.pi) - gammaln(arg_gamma/2).sum()
    else:
        K, D, _ = Wshape
        arg_gamma = np.reshape(nu, (K, 1)) - np.reshape(np.arange(0, D, 1), (1, D))
        return -nu/2 * np.log(np.linalg.det(W)) - D/2 * nu * np.log(2) - D* (D - 1)/4 * np.log(np.pi) - gammaln(arg_gamma/2).sum(axis=1)

1節で述べたように、変分下界は次のように計算できます：
$$
\begin{align}
\mathcal{L}(q) = -\sum_{n=0}^{N-1}\sum_{k=0}^{K-1} r_{n, k} \log r_{n, k}
+ \log \frac{C(\alpha_0)}{C(\alpha)}
+ \frac{D}{2} \sum_{k=0}^{K-1} \log \frac{\beta_0}{\beta_k}
+ \sum_{k=0}^{K-1} \log \frac{B(W_0, \nu_0)}{B(W_k, \nu_k)}
- \frac{DN}{2} \log(2\pi),
\end{align}
$$

先ほどの2つの関数を用いると、変分下界を計算する関数は次のように実装できます：


    def _calc_lower_bound(self, r):
        '''
        Method for calculating the variational lower bound.
        
        Parameters
        ----------
        r : 2D numpy array
            2-D numpy array representing responsibility of each component for each sample in X, 
            where r[n, k] = $r_{n, k}$.
        Returns
        ----------
        lower_bound : float
            The variational lower bound, where the final constant term is omitted.
        '''
        return - (r * np.log(r)).sum() + \
            logC(self.alpha0*np.ones(self.K)) - logC(self.alpha) +\
            self.D/2 * (self.K * np.log(self.beta0) - np.log(self.beta).sum()) + \
            self.K * logB(self.W0, self.nu0) - logB(self.W, self.nu).sum()

2.6 fit

ここまで書いた4つのメソッドをまとめ、fittingのためのメソッドfitを定義します。
基本的には定義通りまとめるだけですが、

繰り返しの最大回数max_iter
収束判定の閾値tol(lower_boundの変化がtolより小さくなったら停止します。)

を引数として渡すことにします。

    def fit(self, X, max_iter=1e3, tol=1e-4, random_state=None, disp_message=False):
        '''
        Method for fitting the model.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        max_iter : int
            The maximum number of iteration
        tol : float
            The criterion for juding the convergence. 
            When the change of lower bound becomes smaller than tol, the iteration is stopped.
        random_state : int
            An integer specifying the random number seed for random initialization
        disp_message : Boolean
            Whether to show the message on the result.
        '''
        self._init_params(X, random_state=random_state)
        r = self._e_like_step(X)
        lower_bound = self._calc_lower_bound(r)
        
        for i in range(max_iter):
            self._m_like_step(X, r)
            r = self._e_like_step(X)
            
            lower_bound_prev = lower_bound
            lower_bound = self._calc_lower_bound(r)
            
            if abs(lower_bound - lower_bound_prev) < tol:
                break
        
        self.lower_bound = lower_bound
        
        if disp_message:
            print(f"n_iter : {i}")
            print(f"convergend : {i < max_iter}")
            print(f"lower bound : {lower_bound}")
            print(f"Change in the variational lower bound : {lower_bound - lower_bound_prev}")

2.7 予測

まず、Studentのt分布の確率密度関数を与える関数を定義しておきます。
Studentのt分布の確率密度関数は、次の式で与えられます(PRMLの式(B.68)-(B.72)参照)
$$
\begin{align}
\mathrm{St}\left(x \middle| \mu, L, \nu \right) &= \frac{\Gamma\left(\frac{\nu + D}{2} \right)}{\Gamma\left(\frac{\nu}{2}\right)}
\frac{(\det L)^{1/2}}{(\nu \pi)^{D/2}}
\left( 1 + \frac{\Delta^2}{\nu} \right)^{-\frac{\nu}{2} - \frac{D}{2}} \
\Delta^2 &:= (x - \mu)^T L (x- \mu)
\end{align}
$$

これをコードにします。


def multi_student_t(X, mu, L, nu):
    '''
    Function for calculating the probability density of multivariate Student's t distribution.
    See equations (B.68) and (B.72) of PRML.
    
    Parameters
    ----------
    X : 2D numpy array
        2D numpy array representing input data points on which the density are calculated.
        X[n, i] represents the i-th element of n-th point in X.
    mu : 1D numpy array
        1D numpy array representing the mean of the distribution. 
        Its length must coincide with X.shape[1]
    L : 2D numpy array
        2D numpy array representing "accuracy".
        It must be a (D, D) array, where D = X.shape[1]
    nu : float, must be positive
        a positive number representing the degrees of freedom of the distribution
            
    Returns
    ----------
    prob_density : 1D numpy array
        The probability density of the Student's t distribution with the given parameters mu, L, nu, evaluated on X, 
        where prob_density[n] = the density on n-th data of X.
    '''
    _, D = X.shape
    diff = X - mu
    delta2 = ((diff @ L ) * diff).sum(axis=1)
    coeff = gamma( (nu + D)/2 )/gamma(nu/2) * (np.linalg.det(L)**0.5) / ((nu*np.pi)**(D/2))
    return coeff/((1 + delta2/nu)**(0.5*nu + 0.5*D)) \\\

この関数を用いて、もろもろの予測を行うメソッドを実装していきましょう。

実装にあたっては、次の数式を用います。

$$
\begin{align}
p(\hat{z} | \hat{x}, X) &\propto p(\hat{x}, \hat{z} | X) \simeq \prod_{k=0}^{K-1} \left[ \frac{\alpha_k}{\sum_{j=0}^{K-1} \alpha_j} \mathrm{St}\left(\hat{x} \middle| m_k, L_k, \nu_k + 1 - D \right) \right]^{\hat{z}_k}\
p(\hat{x} | X) &\simeq \frac{1}{\sum_{k=0}^{K-1} \alpha_k} \sum_{k=0}^{K-1} \alpha_k \mathrm{St}(\hat{x} | m_k, L_k, \nu_k + 1 - D) \
L_k &:= \frac{(\nu_k + 1 - D) \beta_k}{ 1 + \beta_k} W_k
\end{align}
$$

ここでは具体的には、次の4種類の予測を返すメソッドを実装します。

_predict_joint_proba : 同時分布$p(\hat{x}, \hat{z} | X)$
predict_proba : 入力が各潜在要素に属する確率$p(\hat{z} | \hat{x}, X)$
predict : 入力が属する潜在要素の予測値
calc_prob_density : 確率密度

まず_predict_joint_probaで同時分布を計算し、その結果を規格化してpredict_proba の出力とします。また、predict_probaのargmaxをとってpredictの出力とします。

calc_prob_densityは$x$についての確率密度を返せばよいので、同時分布で潜在変数について和を取れば得ることができます。


    def _predict_joint_proba(self, X):
        '''
        Method for calculating and returning the joint probability. 
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        joint_proba : 2D numpy array
            A numpy array with shape (len(X), self.K), where joint_proba[n, k] = joint probability p(X[n], z_k=1 | training data)
        '''
        L = np.reshape( (self.nu + 1 - self.D)*self.beta/(1 + self.beta), (self.K, 1,1) ) * self.W
        tmp = np.zeros((len(X), self.K))
        for k in range(self.K):
            tmp[:,k] = multi_student_t(X, self.m[k], L[k], self.nu[k] + 1 - self.D)
        return tmp * np.reshape(self.alpha/(self.alpha.sum()), (1, self.K))
    
    def calc_prob_density(self, X):
        '''
        Method for calculating and returning the predictive density.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        prob_density : 1D numpy array
            A numpy array with shape (len(X), ), where proba[n] =  p(X[n] | training data)
        '''
        joint_proba = self._predict_joint_proba(X)
        return joint_proba.sum(axis=1)
    
    def predict_proba(self, X):
        '''
        Method for calculating and returning the probability of belonging to each component.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        proba : 2D numpy array
            A numpy array with shape (len(X), self.K), where proba[n, k] =  p(z_k=1 | X[n], training data)
        '''
        joint_proba = self._predict_joint_proba(X)
        return joint_proba / joint_proba.sum(axis=1).reshape(-1, 1)
    
    def predict(self, X):
        '''
        Method for predicting which component each input data belongs to.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        pred : 1D numpy array
            A numpy array with shape (len(X), ), where pred[n] =  argmax_{k} p(z_k=1 | X[n], training data)
        '''
        proba = self.predict_proba(X)
        return proba.argmax(axis=1)

2.8 `BayesianGaussianMixtureModel`全体のコード

まとめると、BayesianGaussianMixtureModelのコードは次のように書けます。

class BayesianGaussianMixtureModel:
    
    def __init__(self, K, D, alpha0, beta0, nu0, m0, W0):
        self.K = K
        self.D = D
        self.alpha0 = alpha0
        self.beta0 = beta0
        self.nu0 = nu0
        self.m0 = m0
        self.W0 = W0
        
        self.alpha = None
        self.beta = None
        self.nu = None
        self.m = None
        self.W = None
        self.lower_bound = None
            
    def _init_params(self, X, random_state=None):
        '''
        Method for initializing model parameterse based on the size and variance of the input data array. 
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        random_state : int
            int, specifying the seed for random initialization of m
        '''
        N, D = X.shape
        rnd = np.random.RandomState(seed=random_state)
        
        self.alpha = (self.alpha0 + N / self.K) * np.ones(self.K)
        self.beta = (self.beta0 + N / self.K) * np.ones(self.K)
        self.nu = (self.nu0 + N / self.K) * np.ones(self.K)
        self.m = X[rnd.randint(low=0, high=N, size=self.K)]
        self.W = np.tile(np.diag(1.0/np.var(X, axis=0)), (self.K, 1, 1))

    def _e_like_step(self, X):
        '''
        Method for calculating the array corresponding to responsibility.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
            
        Returns
        ----------
        r : 2D numpy array
            2D numpy array representing responsibility of each component for each sample in X, 
            where r[n, k] = $r_{n, k}$.
        
        '''
        N, _ = np.shape(X)
        
        tpi = np.exp( digamma(self.alpha) - digamma(self.alpha.sum()) )
        
        arg_digamma = np.reshape(self.nu, (self.K, 1)) - np.reshape(np.arange(0, self.D, 1), (1, self.D))
        tlam = np.exp( digamma(arg_digamma/2).sum(axis=1)  + self.D * np.log(2) + np.log(np.linalg.det(self.W)) )
        
        diff = np.reshape(X, (N, 1, self.D) ) - np.reshape(self.m, (1, self.K, self.D) )
        exponent = self.D / self.beta + self.nu * np.einsum("nkj,nkj->nk", np.einsum("nki,kij->nkj", diff, self.W), diff)
        
        exponent_subtracted = exponent - np.reshape(exponent.min(axis=1), (N, 1))
        rho = tpi*np.sqrt(tlam)*np.exp( -0.5 * exponent_subtracted )
        r = rho/np.reshape(rho.sum(axis=1), (N, 1))
        
        return r
    
    
    def _m_like_step(self, X, r):
        '''
        Method for calculating the model parameters based on the responsibility.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        r : 2D numpy array
            2-D numpy array representing responsibility of each component for each sample in X, 
            where r[n, k] = $r_{n, k}$.
        '''
        N, _ = np.shape(X)
        n_samples_in_component = r.sum(axis=0)
        barx = r.T @ X / np.reshape(n_samples_in_component, (self.K, 1))
        diff = np.reshape(X, (N, 1, self.D) ) - np.reshape(barx, (1, self.K, self.D) )
        S = np.einsum("nki,nkj->kij", np.einsum("nk,nki->nki", r, diff), diff) / np.reshape(n_samples_in_component, (self.K, 1, 1))
        
        self.alpha = self.alpha0 + n_samples_in_component
        self.beta = self.beta0 + n_samples_in_component
        self.nu = self.nu0 + n_samples_in_component
        self.m = (self.m0 * self.beta0 + barx * np.reshape(n_samples_in_component, (self.K, 1)))/np.reshape(self.beta, (self.K, 1))
        
        diff2 = barx - self.m0
        Winv = np.reshape(np.linalg.inv( self.W0 ), (1, self.D, self.D)) + \
            S * np.reshape(n_samples_in_component, (self.K, 1, 1)) + \
            np.reshape( self.beta0 * n_samples_in_component / (self.beta0 + n_samples_in_component), (self.K, 1, 1)) * np.einsum("ki,kj->kij",diff2,diff2) 
        self.W = np.linalg.inv(Winv)
        
    def _calc_lower_bound(self, r):
        '''
        Method for calculating the variational lower bound.
        
        Parameters
        ----------
        r : 2D numpy array
            2-D numpy array representing responsibility of each component for each sample in X, 
            where r[n, k] = $r_{n, k}$.
        Returns
        ----------
        lower_bound : float
            The variational lower bound, where the final constant term is omitted.
        '''
        return - (r * np.log(r)).sum() + \
            logC(self.alpha0*np.ones(self.K)) - logC(self.alpha) +\
            self.D/2 * (self.K * np.log(self.beta0) - np.log(self.beta).sum()) + \
            self.K * logB(self.W0, self.nu0) - logB(self.W, self.nu).sum()
            
        
    def fit(self, X, max_iter=1e3, tol=1e-4, random_state=None, disp_message=False):
        '''
        Method for fitting the model.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        max_iter : int
            The maximum number of iteration
        tol : float
            The criterion for juding the convergence. 
            When the change of lower bound becomes smaller than tol, the iteration is stopped.
        random_state : int
            An integer specifying the random number seed for random initialization
        disp_message : Boolean
            Whether to show the message on the result.
        '''
        self._init_params(X, random_state=random_state)
        r = self._e_like_step(X)
        lower_bound = self._calc_lower_bound(r)
        
        for i in range(max_iter):
            self._m_like_step(X, r)
            r = self._e_like_step(X)
            
            lower_bound_prev = lower_bound
            lower_bound = self._calc_lower_bound(r)
            
            if abs(lower_bound - lower_bound_prev) < tol:
                break
        
        self.lower_bound = lower_bound
        
        if disp_message:
            print(f"n_iter : {i}")
            print(f"convergend : {i < max_iter}")
            print(f"lower bound : {lower_bound}")
            print(f"Change in the variational lower bound : {lower_bound - lower_bound_prev}")
            
    
    def _predict_joint_proba(self, X):
        '''
        Method for calculating and returning the joint probability. 
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        joint_proba : 2D numpy array
            A numpy array with shape (len(X), self.K), where joint_proba[n, k] = joint probability p(X[n], z_k=1 | training data)
        '''
        L = np.reshape( (self.nu + 1 - self.D)*self.beta/(1 + self.beta), (self.K, 1,1) ) * self.W
        tmp = np.zeros((len(X), self.K))
        for k in range(self.K):
            tmp[:,k] = multi_student_t(X, self.m[k], L[k], self.nu[k] + 1 - self.D)
        return tmp * np.reshape(self.alpha/(self.alpha.sum()), (1, self.K))
    
    def calc_prob_density(self, X):
        '''
        Method for calculating and returning the predictive density.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        prob_density : 1D numpy array
            A numpy array with shape (len(X), ), where proba[n] =  p(X[n] | training data)
        '''
        joint_proba = self._predict_joint_proba(X)
        return joint_proba.sum(axis=1)
    
    def predict_proba(self, X):
        '''
        Method for calculating and returning the probability of belonging to each component.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        proba : 2D numpy array
            A numpy array with shape (len(X), self.K), where proba[n, k] =  p(z_k=1 | X[n], training data)
        '''
        joint_proba = self._predict_joint_proba(X)
        return joint_proba / joint_proba.sum(axis=1).reshape(-1, 1)
    
    def predict(self, X):
        '''
        Method for predicting which component each input data belongs to.
        
        Parameters
        ----------
        X : 2D numpy array
            2D numpy array representing input data, where X[n, i] represents the i-th element of n-th point in X.
        
        Returns
        ----------
        pred : 1D numpy array
            A numpy array with shape (len(X), ), where pred[n] =  argmax_{k} p(z_k=1 | X[n], training data)
        '''
        proba = self.predict_proba(X)
        return proba.argmax(axis=1)

3. 実験

それでは、ここまで書いてきたコードを実際に動かしてみましょう!

3.1 データ

使うデータは、make_blobsで作った、次の2次元のtoy dataです。
前回のEM algorithmの記事で用いたものと全く同じデータです。

3つのblobを生成しました。データ点の個数は100にしてあります。

以下の結果では、左のパネルでは混合要素についての分類結果をカラープロットで、右のパネルでは確率密度関数を等高線プロットで示します。

3.2 使ってみる

まず、単純に使ってみた結果をさくっと見てみましょう。

K=3, tol=1e-4としてあります。
また、hyperparameterの値はalpha0=1.0, beta0=1.0, nu0=2.0, m0=np.zeros(2), W0=np.eye(2)としました。

実行結果は次のようになりました。

n_iter : 10
convergend : True
lower bound : -389.1578289309176
Change in the variational lower bound : -3.9195783756440505e-05

期待通りにクラスタリングされていますね。ただし、いつもこううまく動くわけではなく、初期化の仕方を変えるとてんで違う結果が返ってくることもあります。

n_iter : 8
convergend : True
lower bound : -397.75619724106684
Change in the variational lower bound : 1.2367827594061964e-05

直感的には、あまりうまく学習できているようには見えませんね。
実際、変分下界の値の観点からも、前者の結果がより適切であることが分かります。

この結果から、実際のデータに混合ガウスモデルを適用する際には、複数の初期化を試してみて、最も良い結果を選択することが必要になることが見て取れます。

3.3 hyperparameter tuning

次に、hyperparameterのtuningを行います。
具体的には、grid searchを行い、変分下界を最大化するhyperparameterの値をとることにします。

gridは

alpha0 : [0.1, 1.0, 10.0, 100.0]
beta0 : [0.01, 0.1, 1.0]
nu0 : [1.1, 2.0, 11.0]
W0 : [0.01*np.eye(2), 0.1*np.eye(2), np.eye(2), 10.0*np.eye(2)]

として、各値の組に対して10回random_stateを変えて実行し、最も大きい変分下界を与えるものを選びました。結果は以下の通りです：

lower_bound : -367.9533872980828
alpha0 : 100.0
beta0 : 0.1
nu0 : 2.0
m0 : [0. 0.]
W0 : [[0.1 0. ]
 [0.  0.1]]

前の節の結果よりも変分下界値が改善しており、確率密度を見てもくっきり3つのクラスターに分かれていることが見て取れます。

3.4 大きい$K$についての結果

さて、ここまではK=3と「真のクラスター数」と等しくとってきました。しかし、実際のデータを扱う際には「真のクラスター数」を知ることはできませんし、そもそも「真のクラスター数」など存在しないかもしれません。

そのような場合に何が起きるのかを見てみましょう。

ここではK=10ととり、beta0=1.0, nu0=2.0, m0=np.zeros(2), W0=np.eye(2)と固定します。そして、alpha0の値を動かして結果がどう変わるか見てみましょう。というのも、PRMLの本文((10.69)式の下の段落)にもあるように、$\alpha_0$は、混合要素の実効的な個数を司っているからです。
なお、実験にあたってはパラメーター初期化のseedとしてrandom_stateを0から9までの10通り試し、最も大きな変分下界を与える結果を採用しています。

まず、alpha0=10.0としてみましょう。モデルの定義からも分かるように、こうすると、より多くの混合要素が効いてくることが期待されます。この設定では、余計な混合要素が多数現れると言っても良いでしょう。

lower bound : -452.0379619828742

予想通り、多数の混合要素が現れた複雑なクラスタリングが得られました。これは今のデータに対しては適切なクラスタリングとは言いがたいかと思います。実際、変分下界の値は3.2節の結果と比べて悪くなっています。

次に、alpha0=0.1としてみましょう。これは、いくつかの混合要素の実効的な個数を減らすことに対応します。

lower bound : -396.83871150380935

この結果は、さきほどと比べると適切なクラスタリングと呼べるかと思います。データ点は全て3つの主なクラスターのうちいずれかに分類されています。
また、変分下界の値も大きく改善しています。

4 まとめ

この記事では、ベイズ混合ガウスモデルに対する変分推論を実装しました。数式の導出がしんどい¹¹章ですが、一歩一歩実装すると全体像が見て取れるのではないかと思います。
特に、PRMLの変分下界の表式を更に整理したおかげで、複数の初期条件から適切な結果を選んだり、モデル選択を行ったりすることもできました。
また実験では、$K$を大きく、$\alpha_0$を小さくとることにより、混合要素の個数を自動で選べることを確かめられました。

日本語だと自己整合方程式? ↩
なお、筆者はこの繰り返しの収束性については理解していないため、ここでは論じません。 ↩
具体的な初期化方法については、2.2で述べます。 ↩
特にひねったことは必要ないのですが、長いです。。。 ↩
エビデンス近似と同様の手法ですね。 ↩
今回は時間がなかったのですが、この仮定がてんで成り立たなさそうな例で試して失敗するか見てみたいところです。 ↩
この取り方は、各「クラスター」の分散が、$X$全体の分散よりも大幅に小さい場合は上手くいかないかもしれません ↩
scikit-learnのBayesianGaussianMixtureでは、このどちらかを選べるようになっています。 ↩
この処理で得られる$r$の値が(理論的には)変わらないことは容易に確認できるかと思います。 ↩
美しくはないですし、$K$が大きいときには全くよろしくないのですが、コードと数式の対応を明確にするためにこの対応にしておきました。$K$が大きいときには、計算量を考えるとgammaln(self.K * self.alpha0) - self.K * gammaln(self.alpha0)と書いたほうが良いかと思います。 ↩
記事には書きませんでしたが、念のため導出を全てノートに展開したところ20ページくらいになりました。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up