在贝叶斯机器学习中,后验分布(posterior distribution)往往难以计算,因此通常需要变分推断(variational inference)。通过这种方法,我们在训练过程中最大化数据的对数似然的证据下界(evidence lower bound)。变分自动编码器(variational autoencoder)是使用变分推理的一个重要例子。在本文中,我们将一步一步推导标准变分自动编码器的变分下界损失函数。我们基于高斯隐先验(gaussian latent prior)和高斯近似后验(gaussian approximate posterior)来进行推导,在这种假设下,变分最小下界表达式中的KL项具有解析形式。

1. 贝叶斯定理(Bayes Theorem)

贝叶斯定理是一种在观测到新的证据(evidence)时更新信念(belief)的方式。$z$ 是某种“假设”,给定新观察到的数据 $x$,假设 $z$ 的概率为 $p(z|x)$,由以下公式给出:

$$ p(z | x)=\frac{p(x | z) p(z)}{p(x)} \tag{1} $$

其中 $p(x)$ 是数据 $x$ 的概率,$p(x|z)$ 是给定假设 $z$ 时数据 $x$ 的概率,$p(z)$ 是假设 $z$ 的概率。在上面的公式中,各项所对应的术语如下表所示:

SymbolName
$z$Latent variable
$x$Evidence of Data
$p(x)$Evidence probability
$p(z)$Prior probability
p(z | x)Posterior probability
p(x | z)Likelihood probability

虽然贝叶斯定理直观上不好理解,但其推导过程十分简单。对于事件 $X$ 和事件 $Y$,其联合概率为:

$$ p(X \cap Y)=p(X | Y) p(Y)\tag{2} $$

由于 $\text{‘AND’}$ 可以交换,因此:

$$ p(X \cap Y)=p(Y \cap X)=p(Y | X) p(X) \tag{3} $$

$$ p(X | Y) p(Y)=p(Y | X) p(X) \tag{4} $$

两边同时除以 $p(Y)$,即得到贝叶斯定理:

$$ p(X | Y)=\frac{p(Y | X) p(X)}{p(Y)} \tag{5} $$

2. KL散度(Kullback-Lerbler Divergence)

当我们比较两个分布时,我们需要度量两个分布之间的相似性。KL散度被广泛用于度量分布之间的相似性,它是两个分布之间信息差(information difference)的期望(expectation)。

事件 $x$ 关于概率分布 $p$ 的信息熵为:

$$ I_{p}(x)=-\log p(x) \tag{6} $$

事件 $x$ 关于概率分布 $q$ 的信息熵为:

$$ I_{q}(x)=-\log q(x) \tag{7} $$

$p(x)$ 与 $q(x)$ 之间的信息差为:

$$ \Delta I=I_{p}-I_{q}=-\log p(x)+\log q(x)=\log \left(\frac{q(x)}{p(x)}\right)\tag{8} $$

KL散度是上述信息差的期望,写作:

$$ D_{K L}(q(x) \| p(x)) :=E_{\sim q}[\Delta I]=\int(\Delta I) q(x) d x=\int q(x) \log \left(\frac{q(x)}{p(x)}\right) d x \tag{9} $$

相似地,

$$ D_{K L}(p(x) \| q(x)) :=E_{\sim p}[\Delta I]=\int(\Delta I) p(x) d x=\int p(x) \log \left(\frac{p(x)}{q(x)}\right) d x \tag{10} $$

注意,KL散度是不对称的,即

$$ D_{K L}(q(x) \| p(x)) \neq D_{K L}(p(x) \| q(x)) \tag{11} $$

KL散度总是非负的,即

$$ D_{K L}(q(x) \| p(x))=-\int q(x) \log \left(\frac{p(x)}{q(x)}\right) d x \geq 0 \tag{12} $$

为了证明这个结论,我们使用不等式(由下图可知显然成立)

$$ \log t \leq t-1 \tag{13} $$

1.png

因此,

$$ \begin{aligned} -D_{K L}(q(x) \| p(x))=\int q(x) \log \left(\frac{p(x)}{q(x)}\right) d x &\leq \int q(x)\left(\frac{p(x)}{q(x)}-1\right) d x\\ &=\int q(x) \frac{p(x)}{q(x)} d x-\int q(x) d x\\ &=\int p(x) d x-\int q(x) d x\\ &=1-1=0 \end{aligned} \tag{14} $$

$$ -D_{K L}(q(x) \| p(x)) \leq 0 \tag{15} $$

等价于

$$ D_{K L}(q(x) \| p(x)) \geq 0 \tag{16} $$

得证。

3. VAE目标函数(VAE Objective)

VAE由一个encoder和一个decoder组成。encoder部分产生一个近似后验分布(approximate posterior distribution) $q(z|x)$,并用一个神经网络的权重集合 $\theta$ 来进行参数化,因此可以表示为 $q_{\theta}(z | x)$;decoder部分产生一个似然分布(likelihood distribution) $p(x|z)$,并用另一个神经网络的权重集合 $\phi$ 进行参数化 ,因此可以表示为 $p_{\phi}(x | z)$。encoder的输出是隐分布(latent distribution)的参数,它经过采样后产生decoder的输入。一个典型的 VAE 结构如下图所示:

2.png

近似后验分布 $q_{\theta}\left(z | x_{i}\right)$ 与 真实后验分布 $p\left(z | x_{i}\right)$ 之间的KL散度为:

$$ D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p\left(z | x_{i}\right)\right)=-\int q_{\theta}\left(z | x_{i}\right) \log \left(\frac{p\left(z | x_{i}\right)}{q_{\theta}\left(z | x_{i}\right)}\right) d z \geq 0 \tag{17} $$

使用贝叶斯定理,可得

$$ D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p\left(z | x_{i}\right)\right)=-\int q_{\theta}\left(z | x_{i}\right) \log \left(\frac{p_{\phi}\left(x_{i} | z\right) p(z)}{q_{\theta}\left(z | x_{i}\right) p\left(x_{i}\right)}\right) d z \geq 0 \tag{18} $$

使用对数性质,可得

$$ D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p\left(z | x_{i}\right)\right)=-\int q_{\theta}\left(z | x_{i}\right)\left[\log \left(\frac{p_{\phi}\left(x_{i} | z\right) p(z)}{q_{\theta}\left(z | x_{i}\right)}\right)-\log p\left(x_{i}\right)\right] d z \geq 0 \tag{19} $$

将积分拆开,可得

$$ -\int q_{\theta}\left(z | x_{i}\right) \log \left(\frac{p_{\phi}\left(x_{i} | z\right) p(z)}{q_{\theta}\left(z | x_{i}\right)}\right) d z+\int q_{\theta}\left(z | x_{i}\right) \log p\left(x_{i}\right) d z \geq 0 \tag{20} $$

在上式中,我们注意到 $\log \left(p\left(x_{i}\right)\right)$ 是常量,可以从积分中提出来,得到

$$ -\int q_{\theta}\left(z | x_{i}\right) \log \left(\frac{p_{\phi}\left(x_{i} | z\right) p(z)}{q_{\theta}\left(z | x_{i}\right)}\right) d z+\log p\left(x_{i}\right) \int q_{\theta}\left(z | x_{i}\right) d z \geq 0 \tag{21} $$

由于 $q_{\theta}\left(z | x_{i}\right)$ 是一个概率分布,其积分为1,因此

$$ -\int q_{\theta}\left(z | x_{i}\right) \log \left(\frac{p_{\phi}\left(x_{i} | z\right) p(z)}{q_{\theta}\left(z | x_{i}\right)}\right) d z+\log p\left(x_{i}\right) \geq 0 \tag{22} $$

移项得

$$ \log p\left(x_{i}\right) \geq \int q_{\theta}\left(z | x_{i}\right) \log \left(\frac{p_{\phi}\left(x_{i} | z\right) p(z)}{q_{\theta}\left(z | x_{i}\right)}\right) d z \tag{23} $$

使用对数性质,可得

$$ \log p\left(x_{i}\right) \geq \int q_{\theta}\left(z | x_{i}\right)\left[\log p_{\phi}\left(x_{i} | z\right)+\log p(z)-\log q_{\theta}\left(z | x_{i}\right)\right] d z \tag{24} $$

将不等式右边写成期望的形式,得到

$$ \log p\left(x_{i}\right) \geq E_{\sim q_{\theta}\left(z | x_{i}\right)}\left[\log p_{\phi}\left(x_{i} | z\right)+\log p(z)-\log q_{\theta}\left(z | x_{i}\right)\right] \tag{25} $$

$$ \log p\left(x_{i}\right) \geq E_{\sim q_{\theta}\left(z | x_{i}\right)}\left[\log p\left(x_{i}, z\right)-\log q_{\theta}\left(z | x_{i}\right)\right] \tag{26} $$

根据等式(23),我们有

$$ \log p\left(x_{i}\right) \geq \int q_{\theta}\left(z | x_{i}\right) \log \left(\frac{p(z)}{q_{\theta}\left(z | x_{i}\right)}\right) d z+\int q_{\theta}\left(z | x_{i}\right) \log p_{\phi}\left(x_{i} | z\right) d z \tag{27} $$

$$ \log p\left(x_{i}\right) \geq-D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p(z)\right)+E_{\sim q_{\theta}\left(z | x_{i}\right)}\left[\log p_{\phi}\left(x_{i} | z\right)\right] \tag{28} $$

上述不等式的右侧即为证据下界(Evidence Lower Bound, ELBO),也被称为变分下界(variarional lower bound)。ELBO界定了数据的似然,而最大化数据的似然正是我们的目标。因此,最大化ELBO就是最大化数据的对数概率,这就是变分推断的核心思想。ELBO包含2项,第一个KL散度项起到正则化的作用,它约束了近似后验分布的形式;第二项被称为重建项(reconstruction term),因为它度量的是decoder重建的数据的似然值。

特别地,我们拥有选择隐变量 $z$ 形式的自由。如果我们选择高斯分布作为隐先验 $p(z)$ 和近似后验 $q_{\theta}\left(z | x_{i}\right)$ 的形式,就可以得到损失函数的解析形式。

4. VAE损失的解析形式:高斯隐变量(Closed form VAE Loss: Gaussian Latents)

假设我们选择

$$ p(z) \rightarrow \frac{1}{\sqrt{2 \pi \sigma_{p}^{2}}} \exp \left(-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}\right) \tag{29} $$

$$ q_{\theta}\left(z | x_{i}\right) \rightarrow \frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right) \tag{30} $$

则ELBO的KL散度项(正则项)变为:

$$ -D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p(z)\right)=\int \frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right) \log \left(\frac{\frac{1}{\sqrt{2 \pi \sigma_{p}^{2}}} \exp \left(-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}\right)}{\frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right)}\right) d z \tag{31} $$

使用对数性质对右侧进行简化:

$$ \begin{aligned} &\int \frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right) \log \left(\frac{\frac{1}{\sqrt{2 \pi \sigma_{p}^{2}}} \exp \left(-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}\right)}{\frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right)}\right) d z\\ = &\int \frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right)\left\{-\frac{1}{2} \log (2 \pi)-\log \left(\sigma_{p}\right)-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}+\frac{1}{2} \log (2 \pi)+\log \left(\sigma_{q}\right)+\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right\}dz\\ = &\frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \int \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right)\left\{-\log \left(\sigma_{p}\right)-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}+\log \left(\sigma_{q}\right)+\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right\} d z\\ = &\frac{1}{\sqrt{2 \pi \sigma_{q}^{2}}} \int \exp \left(-\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right)\left\{\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}+\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right\} d z \end{aligned} \tag{32} $$

写成期望的形式,可得

$$ \begin{aligned} -D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p(z)\right)&=E_{q}\left\{\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}+\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right\}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)+E_{q}\left\{-\frac{\left(x-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}+\frac{\left(x-\mu_{q}\right)^{2}}{2 \sigma_{q}^{2}}\right\}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{1}{2 \sigma_{p}^{2}} E_{q}\left\{\left(x-\mu_{p}\right)^{2}\right\}+\frac{1}{2 \sigma_{q}^{2}} E_{q}\left\{\left(x-\mu_{q}\right)^{2}\right\} \end{aligned} \tag{33} $$

由于

$$ \sigma_{q}^{2}=E_{q}\left\{\left(x-\mu_{q}\right)^{2}\right\} \tag{34} $$

因此

$$ \begin{aligned} -D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p(z)\right)&=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{1}{2 \sigma_{p}^{2}} E_{q}\left\{\left(x-\mu_{p}\right)^{2}\right\}+\frac{\sigma_{q}^{2}}{2 \sigma_{q}^{2}}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{1}{2 \sigma_{p}^{2}} E_{q}\left\{\left(x-\mu_{p}\right)^{2}\right\}+\frac{1}{2}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{1}{2 \sigma_{p}^{2}} E_{q}\left\{\left(x-\mu_{q}+\mu_{q}-\mu_{p}\right)^{2}\right\}+\frac{1}{2}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{1}{2 \sigma_{p}^{2}} E_{q}\left\{\left(x-\mu_{q}\right)^{2}+2\left(x-\mu_{q}\right)\left(\mu_{q}-\mu_{p}\right)+\left(\mu_{q}-\mu_{p}\right)^{2}\right\}+\frac{1}{2}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{1}{2 \sigma_{p}^{2}}\left[E_{q}\left\{\left(x-\mu_{q}\right)^{2}\right\}+2 E_{q}\left\{\left(x-\mu_{q}\right)\left(\mu_{q}-\mu_{p}\right)\right\}+E_{q}\left\{\left(\mu_{q}-\mu_{p}\right)^{2}\right\}\right]+\frac{1}{2}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{1}{2 \sigma_{p}^{2}}\left[\sigma_{q}^{2}+2 * 0 *\left(\mu_{q}-\mu_{p}\right)+\left(\mu_{q}-\mu_{p}\right)^{2}\right]+\frac{1}{2}\\ &=\log \left(\frac{\sigma_{q}}{\sigma_{p}}\right)-\frac{\sigma_{q}^{2}+\left(\mu_{q}-\mu_{p}\right)^{2}}{2 \sigma_{p}^{2}}+\frac{1}{2} \end{aligned} \tag{35} $$

如果取 $\sigma_{p}=1,\mu_{p}=0$,则得到

$$ \begin{aligned} -D_{K L}\left(q_{\theta}\left(z | x_{i}\right) \| p(z)\right) &=\log \left(\sigma_{q}\right)-\frac{\sigma_{q}^{2}+\mu_{q}^{2}}{2}+\frac{1}{2}\\ &=\frac{1}{2} \log \left(\sigma_{q}^{2}\right)-\frac{\sigma_{q}^{2}+\mu_{q}^{2}}{2}+\frac{1}{2}\\ &=\frac{1}{2}\left[1+\log \left(\sigma_{q}^{2}\right)-\sigma_{q}^{2}-\mu_{q}^{2}\right] \end{aligned} \tag{36} $$

回忆ELBO,即公式(28):

$$ \log p\left(x_{i}\right) \geq-D_{K L}\left(q_{\theta}\left(z | x_{i}\right)| | p(z)\right)+E_{\sim q_{\theta}\left(z | x_{i}\right)}\left[\log p_{\phi}\left(x_{i} | z\right)\right] $$

现在我们可以将右侧的最大化目标写为

$$ \frac{1}{2}\left[1+\log \left(\sigma_{j}^{2}\right)-\sigma_{j}^{2}-\mu_{j}^{2}\right]+E_{\sim q_{\theta}\left(z | x_{i}\right)}\left[\log p_{\phi}\left(x_{i} | z\right)\right] \tag{37} $$

其中 $\sigma_{j}^{2}$ 和 $\mu_{j}$ 是近似分布 $q$ 的参数,$j$ 是隐向量 $z$ 的索引。对于一个batch,目标函数为

$$ \mathcal{G}=\sum_{j=1}^{J} \frac{1}{2}\left[1+\log \left(\sigma_{i}^{2}\right)-\sigma_{i}^{2}-\mu_{i}^{2}\right]+\frac{1}{L} \sum_{l} E_{\sim q_{\theta}\left(z | x_{i}\right)}\left[\log p\left(x_{i} | z^{(i, l)}\right)\right] \tag{38} $$

其中 $J$ 是隐向量 $z$ 的维度,$L$ 是使用重参数技巧(re-parametrization trick)进行随机采样的个数。上面的目标函数是我们在训练过程中需要最大化的目标,为了得到损失函数,只需对 $\mathcal{G}$ 取负号即可:

$$ \mathcal{L}=-\sum_{j=1}^{J} \frac{1}{2}\left[1+\log \left(\sigma_{i}^{2}\right)-\sigma_{i}^{2}-\mu_{i}^{2}\right]-\frac{1}{L} \sum_{l} E_{\sim q_{\theta}\left(z | x_{i}\right)}\left[\log p\left(x_{i} | z^{(i, l)}\right)\right] \tag{39} $$

因此,VAE的训练过程就是找到最优的网络参数 $\left(\theta^{*}, \phi^{*}\right)$ 使得 $\mathcal{L}$ 最小:

$$ \left(\theta^{*}, \phi^{*}\right)=\operatorname{argmin}_{(\theta, \phi)} \mathcal{L}(\theta, \phi) \tag{40} $$

本文翻译自:Tutorial: Deriving the Standard Variational Autoencoder (VAE) Loss Function

Last modification:August 28th, 2019 at 03:18 pm
如果觉得我的文章对你有用,请随意赞赏