1変数ガウス分布の変分推論 - 機械学習基礎理論独習

はじめに

本記事では、以下の式を使います。

$\begin{eqnarray} \ln q(z_i)=\langle\ln p({\mathcal D},{\bf Z})\rangle_{q({\bf Z}_{\backslash i})}+{\rm const.}\tag{1}\\ \end{eqnarray}$

式 $(1)$ については、こちらで解説しています。
また、1次元ガウス分布の平均と分散の事後分布の記事を読んでおくと、本記事の理解の助けになるかもしれません。

1変数ガウス分布

1変数 $x$ についてのガウス分布を用いて、分解による変分近似の例を示します。

ガウス分布から独立に発生したと仮定する観測値 $x$ のデータ集合 ${\mathcal D}=\{x_1,\ldots,x_N\}$ が与えられたとします。
この時、もともとのガウス分布の平均 $\mu$ と精度 $\tau$ の事後分布を求めてみます。
精度 $\tau$ とは、分散 $\sigma^2$ 逆数で、 $\tau=(\sigma^2)^{-1}$ と表せます。

尤度関数は、以下の式です。

$\begin{eqnarray} p({\mathcal D}|\mu,\tau)&=&\prod_{n=1}^Np(x_n|\mu,\tau)\\ &=&\prod_{n=1}^N\left(\frac{\tau}{2\pi}\right)^{1/2}\exp\left(-\frac{\tau}{2}(x_n-\mu)^2\right)\\ &=&\left(\frac{\tau}{2\pi}\right)^{N/2}\exp\left(-\frac{\tau}{2}\sum_{n=1}^N(x_n-\mu)^2\right)\tag{2} \end{eqnarray}$

$\mu,\tau$ に関する共役事前分布を導入します。

$\begin{eqnarray} p(\mu|\tau)&=&{\mathcal N}(\mu|\mu_0,(\lambda_0\tau)^{-1})\\ &=&\left(\frac{\lambda_0\tau}{2\pi}\right)^{1/2}\exp\left(-\frac{\lambda_0\tau}{2}(\mu-\mu_0)^2\right)\tag{3}\\ \end{eqnarray}$

$\begin{eqnarray} p(\tau)&=&{\rm Gam}(\tau|a_0,b_0)\\ &=&\frac{1}{\Gamma(a_0)}b_0^{a_0}\tau^{a_0-1}\exp(-b_0\tau)\tag{4} \end{eqnarray}$

同時分布は、以下の式です。

$\begin{eqnarray} p({\mathcal D}|\mu,\tau)=p({\mathcal D}|\mu,\tau)p(\mu|\tau)p(\tau)\tag{5} \end{eqnarray}$

グラフィカルモデルは、以下の図1です。

図1
f:id:olj611:20211005085823p:plain:w250

変分推論の適用

以上の式から事後分布 $p(\mu,\tau|{\mathcal D})$ を求めると、ガウス-ガンマ分布となるのですが、
本記事は、変分推論の記事なので、以下のように事後分布を分解した変分近似を考えることにします。

$\begin{eqnarray} p(\mu,\tau|{\mathcal D})\simeq q(\mu,\tau)=q_\mu(\mu)q_\tau(\tau)\tag{6} \end{eqnarray}$

以下の計算で使う $\ln p({\mathcal D}|\mu,\tau),\ln p(\mu|\tau),\ln p(\tau)$ を先に計算しておきます。

$\begin{eqnarray} \ln p({\mathcal D}|\mu,\tau)&=&\ln\left(\left(\frac{\tau}{2\pi}\right)^{N/2}\exp\left(-\frac{\tau}{2}\sum_{n=1}^N(x_n-\mu)^2\right)\right)\\ &=&\frac{N}{2}\ln\tau-\frac{N}{2}\ln(2\pi)-\frac{\tau}{2}\sum_{n=1}^N(x_n-\mu)^2\tag{7} \end{eqnarray}$

$\begin{eqnarray} \ln p(\mu|\tau)&=&\ln\left(\left(\frac{\lambda_0\tau}{2\pi}\right)^{1/2}\exp\left(-\frac{\lambda_0\tau}{2}(\mu-\mu_0)^2\right)\right)\\ &=&\frac{1}{2}\ln\lambda_0\tau-\frac{1}{2}\ln(2\pi)-\frac{\lambda_0\tau}{2}(\mu-\mu_0)^2\tag{8} \end{eqnarray}$

$\begin{eqnarray} \ln p(\tau)&=&\ln\left(\frac{1}{\Gamma(a_0)}b_0^{a_0}\tau^{a_0-1}\exp(-b_0\tau)\right)\\ &=&-\ln\Gamma(a_0)+a_0\ln b_0+(a_0-1)\ln\tau-b_0\tau\tag{9} \end{eqnarray}$

$\ln q_\mu^\star(\mu)$ は、式 $(1)$ より、以下の式で表されます。

$\begin{eqnarray} \ln q_\mu^\star(\mu)&=&\langle\ln p({\mathcal D},\mu,\tau)\rangle_{q_\tau(\tau)}+{\rm const}\\ &=&\langle\ln p({\mathcal D}|\mu,\tau)p(\mu|\tau)p(\tau)\rangle_{q_\tau(\tau)}+{\rm const}\\ &=&\langle\ln p({\mathcal D}|\mu,\tau)+ \ln p(\mu|\tau) + \ln p(\tau)\rangle_{q_\tau(\tau)}+{\rm const}\\ &=&\langle\ln p({\mathcal D}|\mu,\tau)+ \ln p(\mu|\tau)\rangle_{q_\tau(\tau)}+{\rm const}\tag{10}\\ &=&\left\langle \frac{N}{2}\ln\tau-\frac{N}{2}\ln(2\pi)-\frac{\tau}{2}\sum_{n=1}^N(x_n-\mu)^2+\frac{1}{2}\ln\lambda_0\tau-\frac{1}{2}\ln(2\pi)-\frac{\lambda_0\tau}{2}(x_n-\mu)^2\right\rangle_{q_\tau(\tau)}+{\rm const}\\ &=&\left\langle -\frac{\tau}{2}\sum_{n=1}^N(x_n-\mu)^2- \frac{\lambda_0\tau}{2}(\mu-\mu_0)^2\right\rangle_{q_\tau(\tau)}+{\rm const}\tag{11}\\ &=&-\frac{\langle\tau\rangle_{q_\tau(\tau)}}{2}\left(\lambda_0(\mu-\mu_0)^2+\sum_{n=1}^N(x_n-\mu)^2\right)+{\rm const}\tag{12} \end{eqnarray}$

式 $(10),(11)$ で $\mu$ に関係のない項は ${\rm const}$ にまとめています。
式 $(12)$ は $q_\mu^\star(\mu)$ はガウス分布なので、 $q_\mu(\mu)$ は以下の式で表せます。

$\begin{eqnarray} q_\mu(\mu)={\mathcal N}(\mu)|\mu_N,\lambda_N^{-1})\tag{13} \end{eqnarray}$

$\begin{eqnarray} \mu_N=\frac{\lambda_0\mu_0+N\overline{x}}{\lambda_0+N}\tag{14} \end{eqnarray}$

$\begin{eqnarray} \lambda_N=(\lambda_0+N)\langle\tau\rangle_{q_\tau(\tau)}\tag{15} \end{eqnarray}$

式 $(14)$ で、 $\overline{x}=\dfrac{1}{N}\displaystyle\sum_{n=1}^Nx_n$ とおきました。
式 $(14),(15)$ の導出については、PRML演習問題 10.7(標準)を参照してください。
※式 $(15)$ の因数 $\langle\tau\rangle_{q_\tau(\tau)}$ ですが、この時点では、 $q_\tau(\tau)$ がどのような分布か明らかでない為、計算できません。

$\ln q_\tau^\star(\tau)$ は、式 $(1)$ より、以下の式で表されます。

$\begin{eqnarray} \ln q_\tau^\star(\tau)&=&\langle\ln p({\mathcal D},\mu,\tau)\rangle_{q_\mu(\mu)}+{\rm const}\\ &=&\langle\ln p({\mathcal D}|\mu,\tau)p(\mu|\tau)p(\tau)\rangle_{q_\mu(\mu)}+{\rm const}\\ &=&\langle\ln p({\mathcal D}|\mu,\tau)+ \ln p(\mu|\tau) + \ln p(\tau)\rangle_{q_\mu(\mu)}+{\rm const}\\ &=&\Biggl\langle \frac{N}{2}\ln\tau-\frac{N}{2}\ln(2\pi)-\frac{\tau}{2}\sum_{n=1}^N(x_n-\mu)^2+\frac{1}{2}\ln\lambda_0\tau-\frac{1}{2}\ln(2\pi)-\frac{\lambda_0\tau}{2}(\mu-\mu_0)^2\\ &&-\ln\Gamma(a_0)+a_0\ln b_0+(a_0-1)\ln\tau-b_0\tau\Biggl\rangle_{q_\mu(\mu)}+{\rm const}\\ &=&(a_0-1)\ln\tau-b_0\tau+\frac{N+1}{2}\ln\tau+\left\langle-\frac{\tau}{2}\sum_{n=1}^N(x_n-\mu)^2-\frac{\lambda_0\tau}{2}(\mu-\mu_0)^2\right\rangle_{q_\mu(\mu)}+{\rm const}\tag{16}\\ &=&(a_0-1)\ln\tau-b_0\tau+\frac{N+1}{2}\ln\tau-\frac{\tau}{2}\left\langle\sum_{n=1}^N(x_n-\mu)^2+\lambda_0(\mu-\mu_0)^2\right\rangle_{q_\mu(\mu)}+{\rm const}\tag{17} \end{eqnarray}$

式 $(16)$ で $\tau$ に関係のない項は ${\rm const}$ にまとめています。
式 $(17)$ はガンマ分布なので、 $q_\tau(\tau)$ は以下の式で表せます。

$\begin{eqnarray} q_\tau(\tau)={\rm Gam}(\tau|a_N,b_N)\tag{18}\\ \end{eqnarray}$

$\begin{eqnarray} a_N=a_0+\frac{N+1}{2}\tag{19} \end{eqnarray}$

$\begin{eqnarray} b_N=b_0+\frac{1}{2}\left\langle\sum_{n=1}^N(x_n-\mu)^2+\lambda_0(\mu-\mu_0)^2\right\rangle_{q_\mu(\mu)}\tag{20} \end{eqnarray}$

式 $(19),(20)$ の導出については、PRML演習問題 10.7(標準)を参照してください。

$q_\mu(\mu)$ と $q_\tau(\tau)$ の分布が判明したので、それぞれのパラメータ $\mu_N,\lambda_N,a_N,b_N$ を求めればよいことになります。
式 $(14),(19)$ より、 $\mu_N,a_N$ は固定なので、 $\lambda_N,b_N$ を求めればよいことになります。

式 $(20)$ を変形します。

$\begin{eqnarray} b_N&=&b_0+\frac{1}{2}\left\langle\sum_{n=1}^N(x_n-\mu)^2+\lambda_0(\mu-\mu_0)^2\right\rangle_{q_\mu(\mu)}\\ &=&b_0+\frac{1}{2}\left\langle\sum_{n=1}^N(x_n^2-2x_n\mu+\mu^2)+\lambda_0(\mu^2-2\mu_0\mu+\mu_0^2)\right\rangle_{q_\mu(\mu)}\\ &=&b_0+\frac{1}{2}\left\langle N\overline{x^2}-2N\overline{x}\mu+N\mu^2+\lambda_0\mu^2-2\lambda_0\mu_0\mu+\lambda_0\mu_0^2\right\rangle_{q_\mu(\mu)}\\ &=&b_0+\frac{1}{2}(N\overline{x^2}+\lambda_0\mu_0^2)+\frac{1}{2}(N+\lambda_0)\langle\mu^2\rangle_{q_\mu(\mu)}-(N\overline{x}+\lambda_0\mu_0)\langle\mu\rangle_{q_\mu(\mu)}\tag{21} \end{eqnarray}$

式 $(21)$ で、 $\overline{x^2}=\dfrac{1}{N}\displaystyle\sum_{n=1}^Nx_n^2$ とおきました。

式 $(15)$ より、 $\lambda_N$ は $q_\tau(\tau)$ に依存しており、式 $(21)$ より、 $b_N$ は $q_\mu(\mu)$ に依存しています。
相互に依存しており、解析的に解くのが難しい為、以下の繰り返し法で解きます。

なお、 $q_\mu(\mu)$ がガウス分布であり、 $q_\tau(\tau)$ がガンマ分布であるので、
式 $(15)$ の $\langle\tau\rangle_{q_\tau(\tau)}$ 、式 $(21)$ の $\langle\mu^2\rangle_{q_\mu(\mu)},\langle\mu\rangle_{q_\mu(\mu)}$ は以下のようになります。

$\begin{eqnarray} \langle\tau\rangle_{q_\tau(\tau)}=\frac{a_N}{b_N}\tag{22} \end{eqnarray}$

$\begin{eqnarray} \langle\mu^2\rangle_{q_\mu(\mu)}=\mu_N^2+\frac{1}{\lambda_N}\tag{23} \end{eqnarray}$

$\begin{eqnarray} \langle\mu\rangle_{q_\mu(\mu)}=\mu_N\tag{24} \end{eqnarray}$

式 $(22)$ は、ガンマ分布 ${\rm Gam}(\tau|a,b)$ の平均 $\langle\tau\rangle$ が $\langle\tau\rangle=\dfrac{a}{b}$ であることを用いました。
式 $(23)$ は、一般に成り立つ ${\rm var}[x]=\langle x^2\rangle-\langle x\rangle^2$ を用いました。

アルゴリズム

1.初期化

$\mu_0,\lambda_0,a_0,b_0$ を初期化します。
$\mu_N,a_N,b_N$ を以下のように初期化します。

$\begin{eqnarray} &&\mu_N\leftarrow\frac{\lambda_0\mu_0+N\overline{x}}{\lambda_0+N}\tag{25}\\ &&a_N\leftarrow a_0+\frac{N+1}{2}\tag{26}\\ &&b_N\leftarrow b_0\tag{27} \end{eqnarray}$

2. $\lambda_N$ の更新

$\lambda_N$ を更新します。

$\begin{eqnarray} \lambda_N\leftarrow(\lambda_0+N)\langle\tau\rangle_{q_\tau(\tau)}\tag{28} \end{eqnarray}$

ただし、

$\begin{eqnarray} \langle\tau\rangle_{q_\tau(\tau)}=\frac{a_N}{b_N}\tag{29} \end{eqnarray}$

とします。

3. $b_N$ の更新

$b_N$ を更新します。

$\begin{eqnarray} b_N\leftarrow b_0+\frac{1}{2}(N\overline{x^2}+\lambda_0\mu_0^2)+\frac{1}{2}(N+\lambda_0)\langle\mu^2\rangle_{q_\mu(\mu)}-(N\overline{x}+\lambda_0\mu_0)\langle\mu\rangle_{q_\mu(\mu)}\tag{30} \end{eqnarray}$

ただし、

$\begin{eqnarray} \langle\mu^2\rangle_{q_\mu(\mu)}=\mu_N^2+\frac{1}{\lambda_N}\tag{31} \end{eqnarray}$

$\begin{eqnarray} \langle\mu\rangle_{q_\mu(\mu)}=\mu_N\tag{32} \end{eqnarray}$

$\begin{eqnarray} \overline{x^2}=\frac{1}{N}\displaystyle\sum_{n=1}^Nx_n^2\tag{33} \end{eqnarray}$

$\begin{eqnarray} \overline{x}=\dfrac{1}{N}\displaystyle\sum_{n=1}^Nx_n\tag{34} \end{eqnarray}$

とします。

4.終了条件
対数尤度を再計算し、前回との差分があらかじめ設定していた収束条件を満たしていなければ2に戻り、満たしていれば終了します。

※ $\lambda_N$ を初期化し、2と3の処理を入れ替えても同じです。
※対数尤度の代わりに、繰り返し回数を決めて、それを終了条件とするのもよいと思います。
※対数尤度の代わりに、変分下限(変分下界)を計算してもよいと思います。
※対数尤度計算時の $\mu,\tau$ は、 $q_\mu(\mu),q_\tau(\tau)$ の平均やモードなどを使えばよいと思います。