从Autoencoder到beta-VAE
https://lilianweng.github.io/posts/2018-08-12-vae/#reparameterization-trick 中文翻译
从Autoencoder到beta-VAE
Autoencoder是一个为了使用中间有一个狭窄的bottleneck的神经网络模型重建高维数据而发明的（对于Varaitional Autoencoder可能不是这样的，我们将在后面的章节中详细研究它）。一个很好的副产品是降维：bottleneck层捕获了压缩的潜在编码。这样的低维表示可以用作各种应用（即搜索中的嵌入向量），帮助数据压缩，或者揭示潜在的数据生成因素。
Notions
| 符号 | 意义 |
|---|---|
| \(\mathcal{D}\) | 数据集，\(\mathcal{D}=\left\{\mathbf{x}^{(1)}, \mathbf{x}^{(2)}, \ldots, \mathbf{x}^{(n)}\right\}\),包括\(n\)个数据样本；\(\mid \mathcal{D}\mid=n\). |
| \(\mathbf{x^{i}}\) | 每个数据点是一个\(d\)维的向量，\(\mathbf{x}^{(i)}=\left[x_1^{(i)}, x_2^{(i)}, \ldots, x_d^{(i)}\right]\). |
| \(\mathbf{x}\) | 数据集中的一个样本，\(\mathbf{x}\in\mathcal{D}\). |
| \(\mathbf{x^{\prime}}\) | \(\mathbf{x}\)的重建版本 |
| \(\tilde{\mathbf{x}}\) | \(\mathbf{x}\)的损坏版本 |
| \(\mathbf{z}\) | bottleneck层学到的压缩编码 |
| \(a_j^{(l)}\) | 第 \(l\) 个隐藏层中第 \(j\)个神经元的激活函数 |
| \(g_\phi(.)\) | 以\(\phi\)为参数的编码函数 |
| \(f_\theta(.)\) | 以\(\theta\)为参数的解码函数 |
| \(q_{\phi}\mathbf{(z\mid x)}\) | 估计后验概率函数，也称为概率编码器. |
| \(p_{\theta}\mathbf{(z\mid x})\) | 给定潜在代码生成真实数据样本的可能性，也称为概率解码器. |
Autoencoder
Autoencoder是一种神经网络，旨在以无监督的方式学习一个恒等函数，以重建原始输入，同时在这个过程中压缩数据，从而发现一个更有效和压缩的表示。这个想法起源于20世纪80年代，后来被Hinton&Salakhutdinov, 2006的开创性论文所推广。
它由两个网络组成：
Encoder network：它将原始的高维输入转变成成潜在的低维编码。输入的尺寸比输出的尺寸大。
Decoder network：解码器网络从编码中恢复数据，可能会有越来越大的输出层。
编码器网络本质上完成了降维，就像我们如何使用主成分分析（PCA）或矩阵分解（MF)一样。此外，自动编码器针对编码的数据重建进行了显式优化。一个好的中间表示不仅可以捕获潜在变量，而且有利于全面的解压过程。
该模型包含以\(\phi\)为参数的编码器函数\(g(.)\)，并有以\(\theta\)为参数的解码器函数\(f(.)\)。在bottleneck层为输入\(\mathbf{x}\)学习的低维编码是\(\mathbf{z}=g_\phi(\mathbf{x})\)，重建的输入是\(\mathbf{x}^{\prime}=f_\theta\left(g_\phi(\mathbf{x})\right)\)。
参数\((\theta, \phi)\)被一起学习，以输出一个与原始输入相同的重建数据样本，\(\mathbf{x} \approx f_\theta\left(g_\phi(\mathbf{x})\right)\)，或者换句话说，学习恒等函数。有各种指标来量化两个向量之间的差异，例如激活函数为sigmoid时的交叉熵，或者简单的MSE损失： \[ L_{\mathrm{AE}}(\theta, \phi)=\frac{1}{n} \sum_{i=1}^n\left(\mathbf{x}^{(i)}-f_\theta\left(g_\phi\left(\mathbf{x}^{(i)}\right)\right)\right)^2 \]
Denoising Autoencoder
由于autoencoder学习恒等函数，当网络参数多于数据点时，我们面临“过度拟合”的风险。
为了避免过度拟合并提高鲁棒性，Denoising Autoencoder(Vincent et al. 2008) 提出了对基本autoencoder的修改。通过以随机方式向输入向量的某些值添加噪声或掩盖输入向量的某些值来部分破坏输入，\(\tilde{\mathbf{x}} \sim \mathcal{M}_{\mathcal{D}}(\tilde{\mathbf{x}} \mid \mathbf{x})\)。然后训练模型以恢复原始输入（注意：不是损坏的输入）。 \[ \begin{aligned} \tilde{\mathbf{x}}^{(i)} & \sim \mathcal{M}_{\mathcal{D}}\left(\tilde{\mathbf{x}}^{(i)} \mid \mathbf{x}^{(i)}\right) \\ L_{\mathrm{DAE}}(\theta, \phi) & =\frac{1}{n} \sum_{i=1}^n\left(\mathbf{x}^{(i)}-f_\theta\left(g_\phi\left(\tilde{\mathbf{x}}^{(i)}\right)\right)\right)^2 \end{aligned} \] 其中\(\mathcal{M}_{\mathcal{D}}\)定义了从真实数据样本到噪声或者损坏样本的映射。
这种设计的动机是人类可以很容易地识别一个物体或场景，即使视图被部分遮挡或损坏。为了"修复 "部分被破坏的输入，denoising autoencoder码器必须发现和捕捉输入的各个维度之间的关系，以便推断出丢失的部分。
对于具有高冗余度的高维输入，如图像，该模型可能依赖于从许多输入维度的组合中收集的根据来恢复去噪版本，而不是过度拟合一个维度。这为学习稳健的潜在表征奠定了良好的基础。
噪声是由一个随机映射控制的\(\mathcal{M}_\mathcal{D}(\mathbf{\tilde{x}|x})\)，而且它不是特定于某一类型的破坏过程（即掩蔽噪声、高斯噪声、盐和胡椒噪声等）。当然，破坏过程可以配备先验知识。
在原始DAE论文的实验中，噪声是这样应用的：随机选择固定比例的输入维度，并将其值强制为0。 听起来很像dropout，对吗？嗯，去噪自动编码器是在2008年提出的，比dropout论文早4年（Hinton, et al. 2012）。
Sparse Autoencoder
Sparse Autoencoder对隐藏单元的激活应用了一个“稀疏”约束，以避免过度拟合并提高鲁棒性。它迫使模型只有少量的隐藏单元同时被激活，或者换句话说，一个隐藏神经元应该在大部分时间内不被激活。
回想一下，常见的激活函数包括 sigmoid、tanh、relu、leaky relu 等。当值接近 1 时激活神经元，当值接近 0 时神经元失活。
假设在第\(l\)个隐藏层中有\(s_l\)个神经元，该层中第\(j\)个神经元的激活函数被标记为\(a_j^{l}(.),j=1,...,s_l\)。该神经元的激活比例\(\hat{\rho_j}\)应该是一个小数\(\rho\)，被称为稀疏参数；常见的配置是\(\rho=0.05\)。
这个公式是用来计算第\(l\)个隐藏层中第\(j\)个神经元的激活比例\(\hat{\rho_j}\)的。它是通过对\(n\)个输入样本的激活函数\(a_j^{(l)}\)进行求平均得到的。其中，\(x^{(i)}\)是第\(i\)个输入样本，\(n\)是输入样本的数量。由于这个平均值通常会近似于稀疏参数\(\rho\)，因此可以用这个平均值作为\(\hat{\rho_j}\)的近似值。 \[ \hat{\rho}_j^l=\frac{1}{n}\sum_{i=1}^n[a_j^{(l)}x(^{(i)})]\approx\rho \] 这个约束是通过在损失函数中加入一个惩罚项来实现的。KL-散度\(D_{KL}\)测量两个伯努利分布之间的差异，一个是平均值\(\rho\)，另一个是平均值\(\hat{\rho}_j^{(l)}\)。超参数\(\beta\)控制控制我们想要对稀疏性损失应用的惩罚强度。 \[ \begin{aligned} L_{\mathrm{SAE}}(\theta) & =L(\theta)+\beta \sum_{l=1}^L \sum_{j=1}^{s_l} D_{\mathrm{KL}}\left(\rho \| \hat{\rho}_j^{(l)}\right) \\ & =L(\theta)+\beta \sum_{l=1}^L \sum_{j=1}^{s_l} \rho \log \frac{\rho}{\hat{\rho}_j^{(l)}}+(1-\rho) \log \frac{1-\rho}{1-\hat{\rho}_j^{(l)}} \end{aligned} \]
\[ \text{图3. 均值为}\rho=0.5\text{的伯努利分布与均值为}0\le\hat{\rho}\le1\text{的伯努利分布之间的KL散度} \]
\(k-\)Sparse Autoencoder
在\(k-\)Sparse Autoencoder（Makhzani和Frey，2013）中，通过在具有线性激活函数的bottleneck层中只保留前\(k\)个最高的激活来强制执行稀疏性。首先，我们通过编码器网络运行前馈，得到压缩后的代码： \(\mathbf{z=g(x)}\)。对代码向量\(\mathbf{z}\)中的值进行排序，只保留最大的\(k\)个值，而其他神经元被设置为0。现在我们有了一个稀疏化的编码： \(\mathbf{z'{=Sparsify(z)}}\)。计算输出和稀疏化代码的损失、 \(\mathbf{L=|x-f(z')|_2^2}\)。而且，反向传播只经过前\(k\)个被激活的隐藏单元！
Contractive Autoencoder
与sparse autoencoder类似，Contractive Autoencoder（Rifai, et al, 2011）鼓励学习到的表征保持在收缩空间中，以获得更好的鲁棒性。
它在损失函数中增加了一个项，以惩罚对输入过于敏感的表示，从而提高对训练数据点周围的小扰动的稳健性。敏感性是由编码器激活的雅各布矩阵对输入的Frobenius准则来衡量的： \[ \left\|J_f(\mathbf{x})\right\|_F^2=\sum_{i j}\left(\frac{\partial h_j(\mathbf{x})}{\partial x_i}\right)^2 \] 其中\(h_j\)是压缩代码 \(z=f(x)\) 中的一个单元输出。
这个惩罚项是所学编码相对于输入维度的所有偏导数的平方之和。作者声称，根据经验，这个惩罚项被发现可以刻画出一个对应于低维非线性流形的表征，同时对流形正交的多数方向保持不变。
VAE: Variational Autoencoder
变量自动编码器（Kingma & Welling, 2014）的想法，简称VAE，实际上与上述所有自动编码器模型不太相似，但深深扎根于变分贝叶斯和图模型的方法。
我们不是把输入映射成一个固定的向量，而是要把它映射成一个分布。让我们把这个分布标记为\(p_\theta\)，以\(\theta\)为参数。数据输入\(x\)和潜在编码向量\(z\)之间的关系可以完全定义为：
先验 \(p_\theta(z)\)
似然 \(p_\theta(x|z)\)
后验 \(p_\theta(z|x)\)
假设我们知道这个分布的真实参数\(θ^*\)。为了生成一个看起来像真实数据点\(x^{(i)}\)的样本，我们遵循以下步骤：
- 首先，从先验分布\(p_{\theta^*}(z)\)中采样 \(z^{(i)}\)。
- 然后从条件分布 \(p_{θ^∗}(x|z=z^{(i)})\) 生成值 \(x^{(i)}\)。
最佳参数\(θ^∗\)是使生成真实数据样本的概率最大化的参数： \[ \theta^*=\arg \max _\theta \prod_{i=1}^n p_\theta\left(\mathbf{x}^{(i)}\right) \] 通常我们使用对数概率将 RHS 上的乘积转换为总和： \[ \theta^*=\arg \max _\theta \sum_{i=1}^n \log p_\theta\left(\mathbf{x}^{(i)}\right) \] \(p_\theta\left(\mathbf{x}^{(i)}\right)\) 是在参数 \(\theta\) 下生成数据点 \(\mathbf{x}^{(i)}\) 的概率密度函数值。在 VAE 中，我们希望最大化所有数据点的生成概率密度函数的乘积，从而找到最优的参数 \(\theta^*\)。所以，该公式就是要求最大化所有数据点的生成概率密度函数的乘积，以找到最优的参数 \(\theta^*\)。
现在让我们更新这个方程，以更好地展示数据生成过程，从而涉及编码矢量。 \[ p_\theta\left(\mathbf{x}^{(i)}\right)=\int p_\theta\left(\mathbf{x}^{(i)} \mid \mathbf{z}\right) p_\theta(\mathbf{z}) d \mathbf{z} \] 不幸的是，按照这种方式计算 \(p_\theta\left(\mathbf{x}^{(i)}\right)\) 是非常昂贵的，因为要检查所有可能的 \(\mathbf{z}\) 值并将它们相加。为了缩小值空间以便更快地进行搜索，我们想引入一个新的近似函数，用于输出给定输入 \(\mathbf{x}\) 的可能编码，即 \(q_\phi(\mathbf{z} \mid \mathbf{x})\)，由参数 \(\phi\) 参数化。
\[ \text{图5. 变分自编码器中涉及的图形模型。实线表示生成分布}\:p_\theta(.)，\text{虚线表示后验分布}\:q_\phi(\mathbf{z}\mid \mathbf{x})\text{以近似难以计算的后验}\:p_\theta(\mathbf{z} \mid \mathbf{x})。 \] 现在这个结构看起来很像一个autoencoder：
条件概率\(p_\theta(\mathbf{x}|\mathbf{z})\)定义了一个生成模型，类似于上面介绍的解码器\(f_θ(\mathbf{x|z})\)。 \(p_\theta(\mathbf{x}|\mathbf{z})\)也被称为概率解码器。
近似函数\(q_{\phi}\mathbf{(z|x)}\)是概率编码器，扮演的角色与上述\(g_\phi(\mathbf{z|x})\)类似。
Loss Function: ELBO
估计的后验\(q_\phi(\mathbf{z|x})\)应该非常接近实际的\(p_\theta(\mathbf{z|x})\)。我们可以用Kullback-Leibler散度来量化这两个分布之间的距离。KL散度\(D_{KL}(X|Y)\)衡量如果用分布Y来表示X，会损失多少信息。
在我们的例子中，我们想要最小化关于\(\phi\) 的\(D_{KL}(q_\phi(\mathbf{z|x})\mid p_\theta(\mathbf{z|x}))\)。
但为什么要使用\(D_{KL}(q_\phi|p_\theta)\)（反向KL）而不是\(D_{KL}(p_\theta|q_\phi)\)（前向KL）？Eric Jang在他关于贝叶斯变异方法的文章中有一个很好的解释。快速回顾一下:
前向 \(\mathrm{KL}\) 散度: \(D_{\mathrm{KL}}(P \mid Q)=\mathbb{E}_{z \sim P(z)} \log \frac{P(z)}{Q(z)}\); 我们必须确保 \(\mathrm{P}(\mathrm{z})>0\)时\(\mathrm{Q}(\mathrm{z})>0\)。优化后的变分分布\(q(z)\)必须覆盖整个 \(p(z)\)。
反向\(\mathrm{KL}\) 散度: \(D_{\mathrm{KL}}(Q \mid P)=\mathbb{E}_{z \sim Q(z)} \log \frac{Q(z)}{P(z)} ;\) 最小化反向 KL 散度会让 \(Q(z)\) 的分布逐渐靠近并受限于 \(P(z)\) 的分布
现在让我们扩展方程： \[ \begin{aligned} & D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) \mid p_\theta(\mathbf{z}\vert\mathbf{x}) ) & \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z} \vert \mathbf{x})} d\mathbf{z} & \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})p_\theta(\mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} d\mathbf{z} & \scriptstyle{\text{; Because }p(z \vert x) = p(z, x) / p(x)} \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \big( \log p_\theta(\mathbf{x}) + \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} \big) d\mathbf{z} & \\ &=\log p_\theta(\mathbf{x}) + \int q_\phi(\mathbf{z} \vert \mathbf{x})\log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} d\mathbf{z} & \scriptstyle{\text{; Because }\int q(z \vert x) dz = 1}\\ &=\log p_\theta(\mathbf{x}) + \int q_\phi(\mathbf{z} \vert \mathbf{x})\log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{x}\vert\mathbf{z})p_\theta(\mathbf{z})} d\mathbf{z} & \scriptstyle{\text{; Because }p(z, x) = p(x \vert z) p(z)} \\ &=\log p_\theta(\mathbf{x}) + \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z} \vert \mathbf{x})}[\log \frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z})} - \log p_\theta(\mathbf{x} \vert \mathbf{z})] &\\ &=\log p_\theta(\mathbf{x}) + D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z})) - \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x})}\log p_\theta(\mathbf{x}\vert\mathbf{z}) & \end{aligned} \] 方程的LHS正是我们在学习真实分布时想要最大化的东西：我们想要最大化产生真实数据的（对数）可能性（也就是\(logp_{\theta}(x)\)），同时最小化真实分布和估计后验分布之间的差异（术语\(D_{KL}\)的作用就像一个正则器）。请注意，\(p_{θ}(x)\)相对于\(q_\phi\)是固定的。
上述的负数定义了我们的损失函数： \[ \begin{aligned} L_{\mathrm{VAE}}(\theta, \phi) & =-\log p_\theta(\mathbf{x})+D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) \| p_\theta(\mathbf{z} \mid \mathbf{x})\right) \\ & =-\mathbb{E}_{\mathbf{z} \sim q \phi(\mathbf{z} \mid \mathbf{x})} \log p_\theta(\mathbf{x} \mid \mathbf{z})+D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) \| p_\theta(\mathbf{z})\right) \\ \theta^*, \phi^* & =\arg \min _{\theta, \phi} L_{\mathrm{VAE}} \end{aligned} \] 在变分贝叶斯方法中，这种损失函数被称为变分下界或evidence下界。名称中的“下界”部分来自 KL 散度始终为非负的事实，因此$ −L_{VAE}$ 是 \(logp_θ(x)\) 的下界。
\[ -L_{\mathrm{VAE}}=\log p_\theta(\mathbf{x})-D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) \| p_\theta(\mathbf{z} \mid \mathbf{x})\right) \leq \log p_\theta(\mathbf{x}) \] 这句话中，"variational lower bound" 指的是变分下界或evidence下界，即能够被用来估计模型生成的样本与真实样本之间差异的指标。"KL divergence is always non-negative" 表示 KL 散度总是非负数。因此，\(-L_{\mathrm{VAE}}\) 的值就是 \(\log p_\theta(\mathbf{x})\) 的下界，这意味着我们可以使用 \(-L_{\mathrm{VAE}}\) 来估计 \(\log p_\theta(\mathbf{x})\)，并且保证得到的结果不会低于实际值。
Reparameterization Trick
损失函数中的期望项涉及从\(\mathbf{z} \sim q_\phi(\mathbf{z} \mid \mathbf{x})\)生成样本 。采样是一种随机过程，因此我们无法反向传播梯度。为了使其可训练，引入了重参数化技巧：通常可以将随机变量 \(\mathbf{z}\) 表示为一个确定性变量 \(\mathbf{z}=\mathcal{T}_\phi(\mathbf{x}, \boldsymbol{\epsilon})\)，其中 \(\boldsymbol{\epsilon}\) 是一个辅助的独立随机变量，由 \(\phi\) 参数化的变换函数 \(\mathcal{T}_\phi\) 将 \(\boldsymbol{\epsilon}\) 转换为 \(\mathbf{z}\)。
例如，\(q_\phi(\mathbf{z} \mid \mathbf{x})\) 的常见选择是具有对角协方差结构的多元高斯分布： \[ \begin{align} &\mathbf{z} \sim q_\phi\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)=\mathcal{N}\left(\mathbf{z} ; \boldsymbol{\mu}^{(i)}, \boldsymbol{\sigma}^{2(i)} \boldsymbol{I}\right)\\ &\mathbf{z}=\boldsymbol{\mu}+\boldsymbol{\sigma} \odot \boldsymbol{\epsilon}, where\:\boldsymbol{\epsilon} \sim \mathcal{N}(0, \boldsymbol{I}) \quad ; Reparameterization\:trick. \end{align} \] 其中 \(\odot\) 表示逐元素相乘。
\[
\text{图7. 说明了重新参数化技巧如何使得}\mathbf{z}\text{采样过程可训练。(图片来源：Kingma在NIPS 2015研讨会上的幻灯片第12页)}
\] 
\[
\text{图8. 具有多元高斯假设的variational autoencoder模型的图示}
\]
编码步骤：
- 将 \(\mathbf{x}\) 通过编码器 \(q_{\phi}(\mathbf{z} \mid \mathbf{x})\) 转化为潜在空间中的分布参数，即 \(\mu_{\phi}(\mathbf{x})\) 和 \(\log \sigma^2_{\phi}(\mathbf{x})\)。
- 从标准正态分布\(\mathcal{N}(0, \boldsymbol{I})\)中采样噪声 \(\mathbf{\epsilon}\)。
- 使用重参数技巧，将采样的噪声 \(\mathbf{\epsilon}\) 与分布参数 \(\mu_{\phi}(\mathbf{x})\) 和 \(\sigma_{\phi}(\mathbf{x})\) 结合，得到 \(\mathbf{z} = \mu_{\phi}(\mathbf{x}) + \sigma_{\phi}(\mathbf{x}) \odot \mathbf{\epsilon}\)，其中 \(\odot\) 表示元素级别的乘法。
- 最终得到的 \(\mathbf{z}\) 即为从数据 \(\mathbf{x}\) 中提取的隐变量。
需要注意的是，在第三步使用重参数技巧是为了将采样过程转化为可微分的操作，这样可以通过反向传播算法优化模型参数。
在编码器中，通常会使用神经网络将输入 \(\mathbf{x}\) 映射到潜在空间中的分布参数 \(\mu_{\phi}(\mathbf{x})\) 和 \(\log \sigma^2_{\phi}(\mathbf{x})\)。具体来说，输入 \(\mathbf{x}\) 会经过多层神经网络的计算，最终得到均值向量 \(\mu_{\phi}(\mathbf{x})\) 和方差向量 \(\sigma^2_{\phi}(\mathbf{x})\)。由于 \(\sigma^2_{\phi}(\mathbf{x})\) 需要保证是正数，因此一般会将神经网络输出的 \(\sigma^2_{\phi}(\mathbf{x})\) 进行非负化处理，例如通过取指数等方式，再加上一个很小的正数，以保证不会出现零方差。最终得到的是 \(\log \sigma^2_{\phi}(\mathbf{x})\)。这样就可以得到潜在分布的均值和方差，从而生成随机样本 \(\mathbf{z}\)。
Beta-VAE
如果推断得到的潜在表示 \(\mathbf{z}\) 中的每个变量只对一个单独的生成因子敏感，而对其他因子相对不变，那么我们会说这个表示是分离的或因子化的。分离的表示通常具有良好的可解释性，并且易于推广到各种任务中。
例如，对于人脸照片训练的模型可能会在单独的维度中捕捉到肤色、头发颜色、头发长度、情绪、是否戴眼镜等许多相对独立的因素。这样的分离表示对于人脸图像生成非常有益。
\(\beta\)-VAE (Higgins et al., 2017) 是变分自编码器的一种修改，特别强调发现分离的潜在因子。与 \(\mathrm{VAE}\) 中的相同激励一样，我们希望最大化生成真实数据的概率，同时保持真实后验分布与估计后验分布之间的距离很小（例如，小于一个常量 \(\delta\)）： \[ \begin{aligned} & \max _{\phi, \theta} \mathbb{E}_{\mathbf{x} \sim \mathcal{D}}\left[\mathbb{E}_{\mathbf{z} \sim q \phi(\mathbf{z} \mid \mathbf{x})} \log p_\theta(\mathbf{x} \mid \mathbf{z})\right] \\ & \text { subject to } D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) \| p_\theta(\mathbf{z})\right)<\delta \end{aligned} \] 我们可以在 KKT 条件下使用拉格朗日乘子 \(\beta\) 重新表述它。只有一个不等式约束的上述优化问题等价于最大化以下方程 \(\mathcal{F}(\theta, \phi, \beta)\)： \[ \begin{array}{rlr} \mathcal{F}(\theta, \phi, \beta) & =\mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z} \mid \mathbf{x})} \log p_\theta(\mathbf{x} \mid \mathbf{z})-\beta\left(D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) \| p_\theta(\mathbf{z})\right)-\delta\right) \\ & =\mathbb{E}_{\mathbf{z} \sim q \phi(\mathbf{z} \mid \mathbf{x})} \log p_\theta(\mathbf{x} \mid \mathbf{z})-\beta D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) \| p_\theta(\mathbf{z})\right)+\beta \delta & \\ & \geq \mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z} \mid \mathbf{x})} \log p_\theta(\mathbf{x} \mid \mathbf{z})-\beta D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) \| p_\theta(\mathbf{z})\right) \quad ; \text { Because } \beta, \delta \geq 0 \end{array} \] \(\beta\)-VAE 的损失函数定义如下： \[ L_{\mathrm{BETA}}(\phi, \beta)=-\mathbb{E}*{\mathbf{z} \sim q*\phi(\mathbf{z} \mid \mathbf{x})} \log p_\theta(\mathbf{x} \mid \mathbf{z})+\beta D_{\mathrm{KL}}\left(q_\phi(\mathbf{z} \mid \mathbf{x}) | p_\theta(\mathbf{z})\right) \] 在这里，拉格朗日乘子\(\beta\)被视为超参数。由于\(L_{\mathrm{BETA}}(\phi, \beta)\)的否定是拉格朗日函数\(\mathcal{F}(\theta, \phi, \beta)\)的下界，因此最小化损失等价于最大化拉格朗日函数，从而适用于我们最初的优化问题。当\(\beta=1\)时，它与VAE相同。当\(\beta>1\)时，它对潜在bottleneck施加了更强的约束，并限制了\(\mathbf{z}\)的表示能力。对于某些条件独立的生成因子，保持它们分离是最有效的表示方法。因此，较高的\(\beta\)鼓励更有效的潜在编码，并进一步鼓励解耦。同时，较高的\(\beta\)可能会在重构质量和解耦程度之间产生权衡。Burgess等人（2017）深入讨论了\(\beta\)-VAE中的解耦，并受信息瓶颈理论的启发，进一步提出了一种修改方法以更好地控制编码表示能力。
VQ-VAE and VQ-VAE-2
VQ-VAE（"Vector Quantised-Variational AutoEncoder"）模型通过编码器学习离散的潜在变量，因为对于语言、语音、推理等问题，离散表示可能更自然地适用。
向量量化（VQ）是一种将\(K\)维向量映射到有限的“码”向量集合中的方法。该过程非常类似于 \({KNN}\) 算法。样本应映射到的最优质心码向量是欧几里得距离最小的向量。
设 \(\mathbf{e} \in \mathbb{R}^{K \times D}\)，\(i=1, \ldots, K\) 为 VQ-VAE 中的潜在嵌入空间（也称为“码本”），其中 \(K\) 是潜在变量类别的数量，\(D\) 是嵌入大小。单个嵌入向量是 \(\mathbf{e}_i \in \mathbb{R}^D\)，\(i=1, \ldots, K\)。
编码器的输出 \(E(\mathbf{x})=\mathbf{z}_e\) 经过最近邻查找以匹配 \(K\) 个嵌入向量之一，然后匹配的码向量成为解码器的输入 \(D(\).\() :\)
\(\mathbf{z}_q(\mathbf{x})=\) Quantize \((E(\mathbf{x}))=\mathbf{e}_k\) ，\(k=\arg \min _i\left\|E(\mathbf{x})-\mathbf{e}_i\right\|_2\)
请注意，在不同的应用程序中，离散潜在变量的形状可能不同；例如，语音为1D，图像为2D，视频为3D。
由于 argmin() 在离散空间上是不可微的，因此从解码器输入 \(\mathbf{z}_q\) 到编码器输出 \(\mathbf{z}_e\) 的梯度 \(\nabla_z L\) 被复制。除了重构损失之外，VQ-VAE 还优化以下内容：
- \(V Q\) 损失：嵌入空间和编码器输出之间的 L2 误差。
- 承诺损失：一种鼓励编码器输出保持接近嵌入空间的度量，以防止其从一个码向量频繁波动到另一个码向量。
\(L=\underbrace{\left\|\mathbf{x}-D\left(\mathbf{e}_k\right)\right\|_2^2}_{\text {reconstruction loss }}+\underbrace{\| \operatorname{sg}[E(\mathbf{x})]}_{\text {vQ loss }}-\mathbf{e}_k \|_2^2+\underbrace{\beta\left\|(\mathbf{x})-\operatorname{sg}\left[\mathbf{e}_k\right]\right\|_2^2}_{\text {commitment loss }}\)
其中 \(\mathrm{sq}[.]\) 是 stop_gradient运算符。
codebook 中的embedding vectors通过 EMA（指数移动平均）进行更新。给定一个code vector\(\mathbf{e}_i\)，假设我们有 \(n_i\) 个编码器输出向量\(\left\{\mathbf{z}_{i, j}\right\}_{j=1}^{n_i}\)，它们被量化为 \(\mathbf{e}_i\)： \[ N_i^{(t)}=\gamma N_i^{(t-1)}+(1-\gamma) n_i^{(t)} \quad \mathbf{m}_i^{(t)}=\gamma \mathbf{m}_i^{(t-1)}+(1-\gamma) \sum_{j=1}^{n_i^{(t)}} \mathbf{z}_{i, j}^{(t)} \quad \mathbf{e}_i^{(t)}=\mathbf{m}_i^{(t)} / N_i^{(t)} \] 其中 \((t)\) 指时间序列批次。\(N_i\) 和 \(\mathbf{m}_i\) 分别是累积向量计数和体积。
VQ-VAE-2（Ali Razavi 等，2019）是一个两级分层的 VQ-VAE，与自注意力自回归模型结合
- 第 1 阶段是训练分层 VQ-VAE：分层潜变量的设计旨在将局部模式（例如纹理）与全局信息（例如对象形状）分离。更大的底层codebook 的训练受到更小的顶层码的限制，这样就不必从头开始学习所有内容。
- 第2阶段是学习一个先验分布，使我们可以从中采样并生成图像。这样，解码器可以接收从训练中采样的类似分布的输入向量。使用一个强大的自回归模型，增强了多头自注意力层来捕捉先验分布（例如PixelSNAIL；Chen等人，2017）。
考虑到VQ-VAE-2取决于在简单的分层设置中配置的离散潜在变量，它生成的图像质量非常惊人。
TD-VAE
TD-VAE (“Temporal Difference VAE”; Gregor et al., 2019)适用于连续的数据。它依赖于三个主要想法，描述如下。
1. 状态空间模型
在（潜在的）状态空间模型中，一系列未观察到的隐藏状态\(\mathbf{z}=\left(z_1, \ldots, z_T\right)\)确定了观测状态\(\mathbf{x}=\left(x_1, \ldots, x_T\right)\)。在图13中的马尔可夫链模型中，与图6中类似，可以通过将难以处理的后验\(p(z \mid x)\)近似为一个函数\(q(z \mid x)\)来训练每个时间步。
2. 信念状态
一个智能体应该学会编码所有过去的状态以便推断未来，这被称为信念状态，\(b_t=\operatorname{belief}\left(x_1, \ldots, x_t\right)=\) belief \(f\left(b_{t-1}, x_t\right)\)。有了这个，将未来状态的分布条件化于过去就可以写成\(p\left(x_{t+1}, \ldots, x_T \mid x_1, \ldots, x_t\right) \approx p\left(x_{t+1}, \ldots, x_T \mid b_t\right)\)。在递归策略中，隐藏状态被用作TD-VAE中智能体的信念状态，因此\(b_t=\operatorname{RNN}\left(b_{t-1}, x_t\right)\)。
3. 跳跃式预测
此外，智能体应该能够基于迄今为止收集的所有信息来想象远期，这表明具有跳跃式预测能力，即预测未来几步的状态。
回顾以上方差下界的内容： \[ \begin{aligned} \log p(x) & \geq \log p(x)-D_{\mathrm{KL}}(q(z \mid x) \| p(z \mid x)) \\ & =\mathbb{E}_{z \sim q} \log p(x \mid z)-D_{\mathrm{KL}}(q(z \mid x) \| p(z)) \\ & =\mathbb{E}_{z \sim q} \log p(x \mid z)-\mathbb{E}_{z \sim q} \log \frac{q(z \mid x)}{p(z)} \\ & =\mathbb{E}_{z \sim q}[\log p(x \mid z)-\log q(z \mid x)+\log p(z)] \\ & =\mathbb{E}_{z \sim q}[\log p(x, z)-\log q(z \mid x)] \\ \log p(x) & \geq \mathbb{E}_{z \sim q}[\log p(x, z)-\log q(z \mid x)] \end{aligned} \] 现在让我们将状态\(x_t\)的分布建模为一个概率函数，条件为所有过去的状态\(x_{<t}\)和当前时间步和上一步的两个潜在变量\(z_t\)和\(z_{t-1}\)： \[ \log p\left(x_t \mid x_{<t}\right) \geq \mathbb{E}_{(z t-1, z t) \sim q}\left[\log p\left(x_t, z_{t-1}, z_t \mid x_{<t}\right)-\log q\left(z_{t-1}, z_t \mid x_{\leq t}\right)\right] \] 继续扩展方程： \[ \begin{aligned} & \log p\left(x_t \mid x_{<t}\right) \\ & \geq \mathbb{E}_{(z t-1, z t) \sim q}\left[\log p\left(x_t, z_{t-1}, z_t \mid x_{<t}\right)-\log q\left(z_{t-1}, z_t \mid x_{\leq t}\right)\right] \\ & \geq \mathbb{E}_{\left(z_{t-1}, z_t\right) \sim q}\left[\log p\left(x_t \mid z_{t-1}, z_t, x_{<t}\right)+\log p\left(z_{t-1}, z_t \mid x_{<t}\right)-\log q\left(z_{t-1}, z_t \mid x_{\leq t}\right)\right] \\ & \geq \mathbb{E}_{(z t-1, z t) \sim q}\left[\log p\left(x_t \mid z_t\right)+\log p\left(z_{t-1} \mid x_{<t}\right)+\log p\left(z_t \mid z_{t-1}\right)-\log q\left(z_{t-1}, z_t \mid x_{\leq t}\right)\right] \\ & \geq \mathbb{E}_{(z t-1, z t) \sim q}\left[\log p\left(x_t \mid z_t\right)+\log p\left(z_{t-1} \mid x_{<t}\right)+\log p\left(z_t \mid z_{t-1}\right)-\log q\left(z_t \mid x_{\leq t}\right)-\log q\left(z_{t-1} \mid z_t, x_{\leq t}\right)\right] \end{aligned} \] 注意三点：
根据马尔可夫假设，红色的项可以忽略。
根据马尔可夫假设，蓝色蓝色的项进行了展开。
绿色的项进行了展开，包括了向过去进行一个步长预测的平滑分布。
准确地说，有四种类型的分布需要学习：
- \(p_D(\).) 是decoder distribution：
- \(p\left(x_t \mid z_t\right)\) 是通常定义下的编码器；
- \(p\left(x_t \mid z_t\right) \rightarrow p_D\left(x_t \mid z_t\right)\)；
- \(p_T(\).) 是transition distribution：
- \(p\left(z_t \mid z_{t-1}\right)\) 捕捉潜变量之间的时序依赖关系；
- \(p\left(z_t \mid z_{t-1}\right) \rightarrow p_T\left(z_t \mid z_{t-1}\right)\)；
- \(p_B(\).$) 是belief distribution：
- \(p\left(z_{t-1} \mid x_{<t}\right)\) 和 \(q\left(z_t \mid x_{\leq t}\right)\) 都可以使用置信状态来预测潜变量；
- \(p\left(z_{t-1} \mid x_{<t}\right) \rightarrow p_B\left(z_{t-1} \mid b_{t-1}\right)\)；
- \(q\left(z_t \mid x_{\leq t}\right) \rightarrow p_B\left(z_t \mid b_t\right)\)；
- \(p_S(\).$) 是smoothing distribution：
- 可以将回溯到过去的平滑项 \(q\left(z_{t-1} \mid z_t, x_{\leq t}\right)\) 重新写成依赖于belief状态的形式；
- \(q\left(z_{t-1} \mid z_t, x_{\leq t}\right) \rightarrow p_S\left(z_{t-1} \mid z_t, b_{t-1}, b_t\right)\)；
为了结合跳跃预测的思想，顺序 ELBO 不仅要对$ t\(,\)t+1$起作用，还要对两个相距较远的时间戳 $t_1t_2 $起作用。这是要最大化的最终 TD-VAE 目标函数： \[ J_{t_1, t_2}=\mathbb{E}\left[\log p_D\left(x_{t_2} \mid z_{t_2}\right)+\log p_B\left(z_{t_1} \mid b_{t_1}\right)+\log p_T\left(z_{t_2} \mid z_{t_1}\right)-\log p_B\left(z_{t_2} \mid b_{t_2}\right)-\log p_S\left(z_{t_1} \mid z_{t_2}, b_{t_1}, b_{t_2}\right)\right] \]
