带结构推理网络的变分信息传递（Variational Message Passing with Structured Inference Networks ICLR 2018）文章阅读

引言

深度模型与概率图形模型相结合的推理方法可以提供既灵活又易于解释的模型。文章提出了一种变分推理算法，用于在混合模型中的变分推理。首先，文章提出了一种结合了图模型和变分自编码（VAE）的结构化推理网络。其次，文章建立了一些条件，在这些条件下，这种推理网络可以实现类似于 VAE 的快速摊销推理（amortized inference）。最后，文章推导了变分消息传递算法来执行高效的自然梯度推理，同时保留了有效的摊销推理。通过同时支持对于深度结构化模型的结构化、摊销、自然梯度推理，本文的方法简化和扩展了目前的推理方法。

模型和挑战

考虑使用局部隐变量${\mathbf{x}}_n$建模数据向量${\mathbf{y}}_n$。文章通过给定${\mathbf{x}}_n$和神经网络参数${\theta }_{\mathrm{NN}}$建模${\mathbf{y}}_n$，并使用概率图模型（PGM）表示数据 x : = { x 1 , x 2 , … , x N } \mathbf{x}:=\{\mathbf{x}_1,\mathbf{x}_2,\ldots,\mathbf{x}_N\} x:={x1​,x2​,…,xN​}和隐变量 y : = { y 1 , y 2 , … , y N } \mathbf{y}:=\{\mathbf{y}_1,\mathbf{y}_2,\ldots,\mathbf{y}_N\} y:={y1​,y2​,…,yN​}之间的关系。具体来说，联合分布如何
$$\begin{array}{c}p(\mathbf{y},\mathbf{x},\theta ):=\underset{\text{DNN}}{\underbrace{[\prod _{n=1}^{N}p({\mathbf{y}}_n|{\mathbf{x}}_n,{\theta }_{\mathrm{NN}})]}}\underset{\text{DNN}}{\underbrace{[p(\mathbf{x}|{\theta }_{\mathrm{PGM}})]}}\underset{\text{Hyperprior}}{\underbrace{[p({\theta }_{\mathrm{PGM}})]}}\end{array}$$
其中${\theta }_{\mathrm{NN}}$和${\theta }_{\mathrm{PGM}}$分别是深度神经网络（DNN）和概率图模型（PGM）的参数，并且 θ : = { θ N N , θ P G M } \mathbf{\theta}:=\{\mathbf{\theta}_{\mathrm{NN}}, \mathbf{\theta}_{\mathrm{PGM}}\} θ:={θNN​,θPGM​}。

这种结合了DNN和PGM的模型被称为结构化变分自编码（SVAE），该模型使用结构化先验$p(\mathbf{x}|{\theta }_{\mathrm{PGM}})$来从数据中提取有用的结构信息。SVAE与VAE最大的区别为，VAE的先验分布为一个简单的多变量高斯分布$p(\mathbf{x})=\mathcal{N}(\mathbf{x}|0,\mathbf{I})$，而没有结构化分布。给出下列混合先验示例。

示例：混合模型先验

假设我们想要将数据${\mathbf{y}}_n$分组到$K$个聚类团。对于这样一个任务，标准的高斯先验不是一个合适的先验。我们应该使用${\mathbf{x}}_n$的混合模型先验，如下
$$\begin{array}{c}p(\mathbf{x}|{\theta }_{\mathrm{PGM}})=\prod _{n=1}^{N}p({\mathbf{x}}_n|{\theta }_{\mathrm{PGM}})=\prod _{n=1}^N\left[\sum _{k=1}^{K}p({\mathbf{x}}_n|z_n=k){\pi }_k\right]\end{array}$$
其中， z k ∈ { 1 , 2 , … , K } z_k \in\{1,2,\ldots,K\} zk​∈{1,2,…,K}为第$n$个数据的混合下标，${\pi }_k$为混合比例，并且${\sum }_{k=1}^K{\pi }_k=1$。每一个混合分量可以被进一步建模，例如，采用高斯分布，$p({\mathbf{x}}_n|z_n=k):=\mathcal{N}({\mathbf{x}}_n|{\mu }_n,{\mathbf{\Sigma }}_n)$，如此，先验分布为高斯混合模型，PGM的超参数为 θ P G M : = { μ n , Σ n , π k } k = 1 K \mathbf{\theta}_{\mathrm{PGM}}:= \{ \mathbf{\mu}_n, \mathbf{\Sigma}_n,\pi_k \}_{k=1}^K θPGM​:={μn​,Σn​,πk​}k=1K​。图1给出了SVAE的先验模型。这类型的结构化先验对于发现数据中的聚类团是有效的，比VAE的推理更加简单。

图1 （a）为结合了深度模型和PGM的生成式模型，（b）为结构化推理模型

文章的主要目的是近似后验分布$p(\mathbf{x},\theta |\mathbf{y})$。具体来说，类似于VAE，我们想要通过推理网络近似$\mathbf{x}$后验分布。在VAE，这通过DNN表示的参数化函数实现，给出
$$\begin{array}{c}p(\mathbf{x}|\mathbf{y},{\theta }_{\mathrm{NN}})\approx \prod _{n=1}^{N}q({\mathbf{x}}_n|f_{\varphi }({\mathbf{y}}_n))\end{array}$$
其中，$p(\mathbf{x}|\mathbf{y},{\theta }_{\mathrm{NN}})$表示$\mathbf{x}$的后验分布，$q({\mathbf{x}}_n|f_{\varphi }({\mathbf{y}}_n))$为编码器的分布，一般来说，$q$是指数族分布，它的参数使用参数为$\varphi$的DNN$f_{\varphi }$给出。对所有使用相同的函数$f_{\varphi }(\cdot )$，从而减少了变分参数的数量，并能在之间共享统计优势。

结构化推理网络

本文首先设计了一种结合了PGM和VAE的推理网络。然后，提出了包含了两种类型因子的结构化推理网络（SIN），
$$\begin{array}{c}q(\mathbf{x}|\mathbf{y},\varphi ):=\frac{1}{\mathcal{Z}(\varphi )}\underset{\mathrm{DNN}\mathrm{Factor}}{\underbrace{[\prod _{n=1}^{N}q({\mathbf{x}}_n|f_{{\varphi }_{\mathrm{NN}}}({\mathbf{y}}_n))]}}\underset{\mathrm{PGM}\mathrm{Factor}}{\underbrace{[q(\mathbf{x}|{\varphi }_{\mathrm{PGM}})]}}\end{array}$$
这里的DNN Factor类似于公式(3)，PGM Factor为一个与PGM先验$p(\mathbf{x}|{\mathbf{\theta }}_{\mathrm{PGM}})$具有相似图结构的指数族分布。DNN因子的作用是实现灵活性，PGM因子的作用是模型的PGM结构结合到推理网络。两个因子都有自己的参数。${\varphi }_{\mathrm{NN}}$为DNN的参数，${\varphi }_{\mathrm{PGM}}$为PGM的参数。参数集合记为 ϕ : = { ϕ N N , ϕ P G M } \phi:=\{ \phi_{\rm NN}, \phi_{\rm PGM} \} ϕ:={ϕNN​,ϕPGM​}。

为了快速的摊销推理，这些DNN因子和PGM因子需要满足下面两个条件。第一个条件是归一化因子$\log \mathcal{Z}(\varphi )$容易计算和可微的。第二个条件是可以从SIN中采样，即，${\mathbf{x}}^{\ast }(\varphi )\sim q(\mathbf{x}|\mathbf{y}.\varphi )$。还有一个非必要的条件，希望可以使用重参数化技巧计算${\mathbf{x}}^{\ast }(\varphi )$的梯度。

当上面两个条件成立时，下界的随机梯度可以采用与VAE相似的方法计算。下面，我们假设$\theta$是一个确定变量（下一节松弛这个条件）。这种情况下，变分下界可以计算为
L S I N ( θ , ϕ ) : = E q [ log ⁡ p ( y , x ∣ θ ) q ( x ∣ y , ϕ ) ] = E q [ log ⁡ ∏ n { p ( y n ∣ x n , θ N N ) } p ( x ∣ θ P G M ) Z ( ϕ ) ∏ n q ( x n ∣ f ϕ N N ( y n ) ) q ( x ∣ ϕ P G M ) ] = ∑ n = 1 N E q [ log ⁡ p ( y n ∣ x n , θ N N ) q ( x n ∣ f ϕ N N ( y n ) ) ] + E q [ log ⁡ p ( x ∣ θ P G M ) ] − E q [ log ⁡ q ( x ∣ ϕ P G M ) ] + log ⁡ Z ( ϕ ) \begin{equation} \begin{split} \mathcal{L}& _{\rm SIN}(\theta, \phi) := \mathbb{E}_q \left[ \log \frac{p(\mathbf{y}, \mathbf{x} | \mathbf{\theta})}{q(\mathbf{x}| \mathbf{y}, \mathbf{\phi })} \right] = \mathbb{E}_q \left[ \log \frac{{ \prod_{n} \{ p(\mathbf{y}_n |\mathbf{x}_n, \mathbf{\theta}_{\mathrm{NN}}) \} } { p(\mathbf{x}| \mathbf{\theta}_{\mathrm{PGM}}) } \mathcal{Z}(\mathbf{\phi })} { { \prod_{n} q(\mathbf{x}_n | f_{\mathbf{\phi}_{\rm NN}} (\mathbf{y}_n)) } { q(\mathbf{x} | {\mathbf{\phi}_{\rm PGM}}) } } \right] \\ & = \sum_{n=1}^N \mathbb{E}_q \left[ \log \frac{ p(\mathbf{y}_n |\mathbf{x}_n, \mathbf{\theta}_{\mathrm{NN}})}{q(\mathbf{x}_n | f_{\mathbf{\phi}_{\rm NN}} (\mathbf{y}_n))} \right] + \mathbb{E}_q \left[ \log p(\mathbf{x}| \mathbf{\theta}_{\mathrm{PGM}}) \right] - \mathbb{E}_q \left[ \log q(\mathbf{x} | {\mathbf{\phi}_{\rm PGM}}) \right] + \log \mathcal{Z}(\mathbf{\phi }) \end{split} \end{equation} L​SIN​(θ,ϕ):=Eq​[logq(x∣y,ϕ)p(y,x∣θ)​]=Eq​[log∏n​q(xn​∣fϕNN​​(yn​))q(x∣ϕPGM​)∏n​{p(yn​∣xn​,θNN​)}p(x∣θPGM​)Z(ϕ)​]=n=1∑N​Eq​[logq(xn​∣fϕNN​​(yn​))p(yn​∣xn​,θNN​)​]+Eq​[logp(x∣θPGM​)]−Eq​[logq(x∣ϕPGM​)]+logZ(ϕ)​​​
上式的第一部分与标准VAE的下界是一致的；第二部分为PGM先验先验下界；剩下两部分为PGM结构。如果可以计算最后三项的梯度并从SIN中采样${\mathbf{x}}^{\ast }(\varphi )$，可以进行类似VAE的摊销推理。

GMM先验的SIN

公式（2）中有额外的隐变量$z_n$。为了在SIN中使用GMM模型，我们选择的参数为 ϕ P G M : = { μ ˉ k , Σ ˉ k , π ˉ k } k = 1 K \phi_{\rm PGM} := \{ \mathbf{\bar{\mu}}_k, \mathbf{\bar{\Sigma}}_k, \bar{\pi}_k \}_{k=1}^K ϕPGM​:={μˉ​k​,Σˉk​,πˉk​}k=1K​，同时保持DNN部分为高斯分布
$$\begin{array}{c}q(\mathbf{x}|\mathbf{y},\varphi ):=\frac{1}{\mathcal{Z}(\varphi )}\prod _{n=1}^N\underset{\mathrm{DNN}\mathrm{Factor}}{\underbrace{[\mathcal{N}({\mathbf{x}}_n|{\mathbf{m}}_n,{\mathbf{V}}_n)]}}\underset{\mathrm{PGM}\mathrm{Factor}}{\underbrace{[\sum _{k=1}^K\mathcal{N}({\mathbf{x}}_n|{\bar{\mu }}_k,{\bar{\mathbf{\Sigma }}}_k){\bar{\pi }}_k]}}\end{array}$$
​幸运的是，因为高斯分布和多模分布是共轭的，我们可以边缘化${\mathbf{x}}_n$获得闭式表达式$\log \mathcal{Z}(\varphi ):={\sum }_n\log {\sum }_k\mathcal{N}({\mathbf{m}}_n|{\bar{\mu }}_n,{\mathbf{V}}_n+{\bar{\mathbf{\Sigma }}}_n){\bar{\pi }}_k$。我们可以从SIN中采用通过边缘化$q(z_n=k|\mathbf{y},\varphi )\propto \mathcal{N}({\mathbf{m}}_n|{\bar{\mu }}_n,{\mathbf{V}}_n+{\bar{\mathbf{\Sigma }}}_n){\bar{\pi }}_k$。给定$z_n$，我们可以从分布$q({\mathbf{x}}_n|z_n=k,\mathbf{y},\varphi )=\mathcal{N}({\mathbf{x}}_n|{\tilde{\mu }}_n,{\tilde{\mathbf{\Sigma }}}_n)$中采样${\mathbf{x}}_n$，其中，${\tilde{\mathbf{\Sigma }}}_n^{-1}={\mathbf{V}}_n^{-1}+{\bar{\mathbf{\Sigma }}}_n^{-1}$以及${\tilde{\mu }}_n={\tilde{\mathbf{\Sigma }}}_n({\mathbf{V}}_n^{-1}{\mathbf{m}}_n+{\bar{\mathbf{\Sigma }}}_n^{-1}{\mu }_n)$。
GMM分布的变分下界为
$$\begin{array}{c}\begin{array}{rl}{\mathcal{L}}_{\mathrm{SIN}}(\theta ,\varphi )& ={\mathbb{E}}_q\left[\log \frac{{\prod }_{n}p({\mathbf{y}}_n|{\mathbf{x}}_n,{\theta }_{\mathrm{NN}}){\sum }_k\mathcal{N}({\mathbf{x}}_n|{\mu }_n,{\mathbf{\Sigma }}_n){\pi }_k\mathcal{Z}(\varphi )}{{\prod }_n\mathcal{N}({\mathbf{x}}_n|{\mathbf{m}}_n,{\mathbf{V}}_n){\sum }_k\mathcal{N}({\mathbf{x}}_n|{\bar{\mu }}_k,{\bar{\mathbf{\Sigma }}}_k){\bar{\pi }}_k}\right]\\ =& \sum _{n=1}^N{\mathbb{E}}_q\left[\log p({\mathbf{y}}_n|{\mathbf{x}}_n,{\theta }_{\mathrm{NN}})\right]-\sum _{n=1}^N{\mathbb{E}}_q\left[\log \mathcal{N}({\mathbf{x}}_n|{\mathbf{m}}_n,{\mathbf{V}}_n)\right]+\\ & \sum _{n=1}^N{\mathbb{E}}_q\left[\log \sum _k\mathcal{N}({\mathbf{x}}_n|{\mu }_n,{\mathbf{\Sigma }}_n){\pi }_k\right]-{\mathbb{E}}_q\left[\log \sum _k\mathcal{N}({\mathbf{x}}_n|{\bar{\mu }}_k,{\bar{\mathbf{\Sigma }}}_k){\bar{\pi }}_k\right]+\log \mathcal{Z}(\varphi )\end{array}\end{array}$$

自然梯度变分推理的变分信息传递

上面我们假设${\theta }_{\mathrm{PGM}}$是确定的。在本节中，我们松弛这个条件，并假设${\theta }_{\mathrm{PGM}}$服从指数族先验模型$p({\theta }_{\mathrm{PGM}}|{\eta }_{\mathrm{PGM}})$，该指数族先验模型的自然参数为${\eta }_{\mathrm{PGM}}$。我们假设平均场近似$q(\mathbf{x},\theta |\mathbf{y}):=q(\mathbf{x}|\mathbf{y},\varphi )q({\theta }_{\mathrm{PGM}}|{\lambda }_{\mathrm{PM}})$，其中，第一个公式为SIN，第二部分为指数族分布。

我们首先推导变分下界。
$$\begin{array}{c}\mathcal{L}({\lambda }_{\mathrm{PGM}},{\theta }_{NN},\theta )={\mathbb{E}}_{q({\theta }_{\mathrm{PGM}}|{\lambda }_{\mathrm{PM}})}[{\mathcal{L}}_{\mathrm{SIN}}(\theta ,\varphi )]-{\mathbb{E}}_{\mathrm{KL}}[q({\theta }_{\mathrm{PGM}}|{\lambda }_{\mathrm{PM}})||p({\theta }_{\mathrm{PGM}}|{\lambda }_{\mathrm{PM}})]\end{array}$$
推理算法流程图如下。

由于算法支持结构化、摊销和自然梯度（SAN）更新，因此称之为 SAN 算法。更新方便地将 PGM 和 DNN 计算分开。步骤 3-6 操作 PGM 部分，我们可以使用现有的 PGM 实现。第 7 步操作 DNN 部分，我们可以重复使用 VAE 实现。算法不仅概括了以前的工作，还通过重用现有软件简化了实现过程。