前回は状態空間モデルにおける変分下限を導出した。この関数をパラメータ $\phi$ について微分可能にするために必要な認識モデル $q$ への制約を考えることが今回の目標である。

いよいよ深層学習っぽい話になると期待している。

確率勾配変分ベイズ法の適用

さて、状態空間モデルの変分下限を再掲する。

$\displaystyle { \mathcal{L}(\mathbf{x}, \mathbf{u}; \theta,\phi) }$

$\displaystyle { \quad = \sum_{t=1}^{T} \mathbb{E}_{q(z_t|\mathbf{x},\mathbf{u},\phi)}[\log p(x_t|z_t,\theta) ] }$

$\displaystyle { \qquad - \mathrm{KL}\big( q_1(z_1|\mathbf{x},\mathbf{u},\phi) \;\|\; p(z_1) \big) }$

$\displaystyle { \qquad - \sum_{t=2}^{T} \mathbb{E}_{q(z_{t-1}|\mathbf{x},\mathbf{u},\phi)} \big[ \mathrm{KL} \big( q_t(z_t|z_{t-1},\mathbf{x},\mathbf{u},\phi) \;\|\; p(z_t|z_{t-1},u_{t-1},\theta) \big) \big] }$

まず、第二項以降のKL情報量を解析的に求めるために分解された各 $q_t$ が正規分布であると仮定する。

$\displaystyle { z_t|z_{t-1}, \mathbf{x}, \mathbf{u}, \phi \sim \mathcal{N}(\mu_t, \Sigma_t) }$

ここで $\mu_t = \mu(z_{t-1}, \mathbf{x}, \mathbf{u}, \phi)$ および $\Sigma_t = \Sigma(z_{t-1}, \mathbf{x}, \mathbf{u}, \phi)$ はパラメータ $\phi$ を持つニューラルネットワークであるとする。このように仮定すると、 VAE のときと同様に $z_t$ を次のように $\phi$ と無関係な確率変数の関数で表すことができる。

$\displaystyle { z_t = \mu(z_{t-1}, \mathbf{x},\mathbf{u},\phi) + \epsilon_t \cdot \sqrt{\Sigma(z_{t-1}, \mathbf{x}, \mathbf{u}, \phi)} ,\quad \epsilon_t \sim \mathcal{N}(0,I) }$

これらを代入すると、変分下限の各項は次のように求めることができる。

第一項

$\displaystyle { \sum_{t=1}^{T} \mathbb{E}_{q(z_t|\mathbf{x},\mathbf{u},\phi)}[\log p(x_t|z_t,\theta) ] = \sum_{t=1}^{T} \mathbb{E}_{p(\epsilon_t)}[\log p(x_t|\mu_t + \epsilon_t \sqrt{\Sigma_t},\theta) ] }$

第二項

$\displaystyle { \mathrm{KL}\big( q_1(z_1|\mathbf{x},\mathbf{u},\phi) \;\|\; p(z_1) \big) = \frac{1}{2} (-\log |\Sigma_1| + \mathrm{tr}\;\Sigma_1 + \| \mu_1 \|^2) + \text{const.} }$

第三項

$\displaystyle { \sum_{t=2}^{T} \mathbb{E}_{q(z_{t-1}|\mathbf{x},\mathbf{u},\phi)} \big[ \mathrm{KL} \big( q_t(z_t|z_{t-1},\mathbf{x},\mathbf{u},\phi) \;\|\; p(z_t|z_{t-1},u_{t-1},\theta) \big) \big] }$

$\displaystyle { \quad = \frac{1}{2} \sum_{t=2}^{T} \mathbb{E}_{q(z_{t-1}|\mathbf{x},\mathbf{u},\phi)} \big[ \log S_{\beta} - \log\Sigma_t + \mathrm{tr}\;S_{\beta}^{-1} \Sigma_t + (G_{\alpha} - \mu_t)^{\top} S_{\beta}^{-1} (G_{\alpha} - \mu_t) \big] + \text{const.} }$

$\displaystyle { \quad = \frac{1}{2} \sum_{t=2}^{T} \mathbb{E}_{p(\epsilon_{t-1})} \big[ f(\mu_{t-1} + \epsilon_{t-1} \sqrt{\Sigma}_{t-1},\mathbf{x},\mathbf{u},\phi)\big] + \text{const.} }$

ただし、 $f$ は次のようにおいた。

$f = \log S_{\beta} - \log\Sigma_t + \mathrm{tr}\;S_{\beta}^{-1} \Sigma_t + (G_{\alpha} - \mu_t)^{\top} S_{\beta}^{-1} (G_{\alpha} - \mu_t)$

以上の結果をまとめると、変分下限は次のような $\phi, \theta$ で微分可能な形にすることができる。

$\displaystyle { \mathcal{L}(\mathbf{x}, \mathbf{u}; \theta,\phi) = \sum_{t=1}^{T} \mathbb{E}_{p(\epsilon_t)}[\log p(x_t|\mu_t + \epsilon_t \sqrt{\Sigma_t},\theta) ] }$

$\displaystyle { \qquad -\frac{1}{2} (-\log |\Sigma_1| + \mathrm{tr}\;\Sigma_1 + \| \mu_1 \|^2) - \frac{1}{2} \sum_{t=2}^{T} \mathbb{E}_{p(\epsilon_{t-1})} \big[ f(\mu_{t-1} + \epsilon_{t-1} \sqrt{\Sigma_{t-1}},\mathbf{x},\mathbf{u},\phi)\big] + \text{const.} }$