Unified Latents: 扩散模型的框架下，令生成与编码相互配合

本文介绍 google deepmind 团队提出的 Unified Latents 架构，在扩散模型的架构下统一了生成和表示压缩，并且给出了优雅的信息界。

生成与编码的权衡

在生成模型设定中（详细见此文 ） ，首先我们有一个在高维空间中的分布

历史上， 变分自编码器（Variational Autoencoder, VAE）引入了一个潜变量 ，并且得出了一个关于 的下界。然后考虑真实的条件概率 与其变分近似 :

即：先从原始数据中压缩出潜变量 ，然后从潜变量 中还原出 。并且让编码器 靠近标准正态分布 以得到优良的潜空间性质。

在这里，编码器 从原始数据中提取出潜变量 ，而生成器 从潜变量 中还原出原始数据 。这个由信息瓶颈（Information Bottleneck）引入的潜变量 连接了生成和编码两个过程。

信息瓶颈

信息瓶颈描述了在编码过程中，如何在保留有用信息的同时压缩数据。具体来说，信息瓶颈试图找到一个潜变量 ，使得它能够最大程度地保留关于输入数据 的有用信息，同时尽可能地压缩 的表示。

数学上，信息瓶颈可以通过以下优化问题来描述： 其中， 是输入数据 和潜变量 之间的互信息， 是潜变量 和目标变量 之间的互信息， 是一个权衡参数。

较小的 会使得模型更倾向于最小化 ，从而得到一个更保留输入数据 的潜变量 ，而较大的 则会使得模型更倾向于最小化 ，从而得到一个更压缩的潜变量 。

变分自编码器通过调控 KL 散度项 来实现信息瓶颈的效果。较大的 KL 散度会使得潜变量 更接近于标准正态分布，从而令 较小，得到一个更压缩的潜变量 ；反之亦然。

问题在于，变分自编码器中，我们无法显式地控制 ，或者说潜变量 中保留了多少关于输入数据 的信息，只能通过调节损失函数里面的 KL 散度权重来间接控制。这就导致了复杂的调试过程，而且因模型而异。

Unified Latents 架构

首先，Unified Latents 架构使用了一个确定性的编码器 ，将输入数据 映射到一个潜空间中的表示 。

然后，对这个确定性的编码加噪得到带噪编码 ，其中 是标准高斯噪声。

接下来，使用两个独立的扩散模型，一个先验扩散模型（diffusion prior）负责从加噪编码 生成原始编码。另一个从一个设定好的固定程度加噪（，见下文）编码 作为条件生成数据 。

Unified Latents 架构

信噪比为 同时定义 ， 大意味着信号强、噪声弱（数据清晰）； 小意味着噪声主导。

好处是可以方便地表达信号和噪声的比例：

注：

手画的演示

损失函数推导 与 ELBO 重加权

对于扩散模型，文中使用了以下的损失：

逻辑是，首先最小化负对数似然展开成 ELBO。

在扩散模型的设定里， 是标准高斯噪声 。

但是不同于 VAE，扩散模型中的 代表的是一整个时间步的噪声过程，也即是路径空间 上的作为随机变量。然后这个路径由前向和后向过程 Itô SDE 定义。

根据 Girsanov 定理，路径空间上的 KL 散度可以被重写为一个时间积分： 其中 是模型的得分函数， 是真实的得分函数。

而 和 的关系由扩散模型的定义给出：

然后由 SDE 的对应关系给出，。

带入化简一下就可以得到文中使用的损失函数。

而这里只降噪到 之类的 ，所以实际上 time integral 的下界不是 0。然后还需要加上 端点的 KL 散度：

而 reweight 就是在路径空间上对不同时间步的损失进行加权，根据 Jensen’s inequality 可以证明，重加权后的损失函数仍然是 ELBO 的上界。

TODO: 有时间写一个详细的推导过程。

双阶段训练

TODO.