使用标准化流进行变分推断.

在变分推断中,优化目标通常是变分下界ELBO

\[\mathbb{E}_{z \text{~} q(z|x)} [\log p(x | z)] - KL[q(z|x)||p(z)]\]

而变分下界与原目标$\log p(x)$之间存在的gap为$\Bbb{E}_{z \text{~} q(z|x)}[\log \frac{q(z|x)}{p(z|x)}]=KL(q(z|x)||p(z|x))$。让上式为$0$的条件是$q(z|x)=p(z|x)$,即找到一个与真实数据的后验分布$p(z|x)$相同的分布$q(z|x)$。

然而后验分布$q(z|x)$通常假设为较为简单的分布形式(如正态分布),不能拟合足够复杂的数据后验分布。本文作者设计了Normalizing flow,用一系列可逆映射将原始分布转换成新的分布。通过优化这一系列分布,即可实现将简单分布转换为复杂的真实后验分布。

具体地,给定可逆映射$f:\Bbb{R}^d\to \Bbb{R}^d$,用它将随机变量$z~q(z)$变换为新的变量$z’=f(z)$,根据概率密度的变量替换公式,新变量$z’$的分布为:

\[q(z') = q(z)|\det \frac{d f^{-1}}{d z'}|= q(z)|\det \frac{d f}{d z}|^{-1}\]

为了构建足够复杂的分布,可以用多个类似的可逆映射$z_k=f_k(z_{k-1})$进行嵌套:

\[z_0 \leftrightarrow z_1 \leftrightarrow z_2 \leftrightarrow \cdots \leftrightarrow z_{K-1} \leftrightarrow z_{K}\]

根据链式法则,变量$z_K$的分布为:

\[q(z_K) = q(z_0)|\det \prod_{k=1}^{K} \frac{d f_k}{d z_{k-1}}|^{-1}\]

其对数似然函数为:

\[\log q(z_K) = \log q(z_0) - \sum_{k=1}^{K}|\det \frac{d f_k}{d z_{k-1}}|\]

此时不需要显式地计算分布$q(z_K)$的概率密度函数,而是通过初始分布$q(z_0)$的概率密度以及映射过程产生的Jacobian行列式计算即可。

一般函数的Jacobian行列式计算复杂。因此需要设计Jacobian行列式可以被高效运算的可逆映射函数。作者设计了两种映射函数:

① Planar Flow

平面流定义如下:

\[f(z) = z+uh(w^Tz+b)\]

其中$u,w,b$是可学习参数,$h(\cdot)$是element-wise非线性函数。该映射的Jacobian行列式为:

\[|\det \frac{d f}{d z}| = |\det (I+u(h'w)^T)| = |1+u^Twh'|\]

② Radial Flow

径向流定义如下:

\[f(z) = z+\frac{b}{a+||z-z_0||^2}(z-z_0)\]

planar flowradial flow分别是在平面空间和球形空间中对原分布进行变换。下图表示嵌套$K$次变换对分布的影响: