NeurIPS 2022 | PDAE：利用预训练扩散模型进行表征学习

©PaperWeekly 原创 · 作者 | 张子健

单位 | 浙江大学

研究方向 | 计算机视觉生成模型

扩散模型（特指 DDPMs）使用一个马尔可夫链将数据分布逐步摧毁转化为标准高斯分布，并通过神经网络拟合其反向过程，得到一个从标准高斯分布到数据分布的马尔可夫生成链。扩散模型因为其稳定的训练过程和高质量的生成样本，得到了学术界和工业界的广泛关注。

作为一种隐变量生成模型（Latent Variable Generative Models），扩散模型天然包含可作为数据表征的隐变量，然而这些隐变量只是数据在原空间中的加噪结果，数量繁多，且缺乏数据的高级语义信息。

鉴于此，Diff-AE（https://github.com/phizaz/diffae）采用自编码器的范式，先使用编码器对数据进行编码，再将编码结果作为条件输入到扩散模型（作为解码器）中进行数据的重构，得到具有高级语义信息的数据表征，并将其成功应用于多个下游任务。

考虑到训练扩散模型会花费较长的时间，且现有大量预训练好的扩散模型，我们提出 Pre-trained DPMs AutoEncoder（PDAE），一种将预训练扩散模型改造适配为解码器并进行自编码器表征学习的方法。相比于 Diff-AE 的从零开始训练，PDAE 可以节约大约一半的表征学习训练时间，并且具有更好的性能表现、更高的灵活性和更广的应用场景。

论文标题：

Unsupervised Representation Learning from Pre-trained Diffusion Probabilistic Models

论文链接：

https://arxiv.org/abs/2212.12990

代码链接：

https://github.com/ckczzj/PDAE

研究背景

1.1 扩散模型

扩散模型（特指 DDPMs）通过一个预定义的马尔可夫链将数据分布 逐步摧毁转化为标准高斯分布 ：

其中 是一组固定的加噪参数。根据高斯分布加和原理，我们可以得到 与 的关系 ，其中 ，。

逆向过程使用高斯分布形式的参数化马尔可夫链：

其中 。最终的生成分布可以表示为 ，通过极大化似然 的变分下界优化模型参数，在进行了一些参数转换和简化后，最终的损失函数如下：

其中 是一个从 中预测噪声 的神经网络，一般采用 U-Net 架构。

对于训练好的 ，除了使用 DDPMs 的马尔可夫链采样法，还可以使用确定性的 DDIMs 采样迭代公式：

并且还可以使用跨步来实现加速采样。

1.2 分类器指导采样

扩散模型一个很好的性质是可以在模型之外额外训练一个分类器 来指导预训练好的无条件扩散模型进行条件采样，其原理是基于以下近似条件概率分布：

对于 DDIMs 采样法，可以使用如下修正后的估计器：

除了类别分类器外，任何评估任意条件 与 之间相似度的估计器都可以被用于指导采样，例如使用 Noisy-CLIP 进行文本指导采样。

1.3 Diff-AE

Diff-AE（https://github.com/phizaz/diffae）采用自编码器的范式，先使用编码器对数据进行编码，再将编码结果作为条件输入到扩散模型（作为解码器）中进行数据的重构。除了自编码器学到的包含高级语义信息的表征，Diff-AE 还使用反向 DDIM 对原始图像进行加噪，得到可以控制图像细节的噪声隐变量表征。此外，Diff-AE 还使用 Latent DDIM 建模自编码器学到的表征空间（post-hoc），以实现解码器的无条件采样。

研究思路

考虑在相同的实验环境下，分别训练无条件扩散模型 和条件扩散模型 去拟合相同的前向过程的后验分布 ，假设 都被设为常数，那么最终后者的损失函数会比前者低。下图为该结论在 MNIST 数据集上的验证（此时 为数字类别）：

这一结论表明， 比 更加接近于 ，所以 和 之间存在一个间隔（称为后验均值间隔），产生该间隔的本质原因是前向过程加噪导致的信息损失使得模型无法仅从 中预测 的所有信息，因此也就无法从 中复原 。而如果我们为扩散模型引入一些 的知识，例如这里的 ，那么这个间隔就会缩小，理论上， 包含 的信息越多，这个间隔就会越小。

从这个角度来看，上述的分类器指导采样方法可以被解释为，在 的基础上计算一个额外的均值偏移量 来填补这个间隔，使得 的信息在最终的样本中被重构。

受此启发，我们使用逆向思维，尝试从这个间隔中提取知识，也就是从数据中学习能填补这个间隔的表征，这个间隔被填补得越多，数据就被重构得越完整，学习到的表征中含有数据的信息就越多。为此，我们使用编码器 和梯度估计器 模拟 ，在预训练扩散模型 的基础上，训练均值偏移量 去尽可能地填补后验均值间隔 。

这等价于训练一个条件扩散模型 ，假设采用噪声预测的参数化方法，则最终的损失函数为：

注意我们使用的是预训练扩散模型， 在优化过程中保持不变。

对于梯度估计器 ，我们使用和 类似的 U-Net 架构，并且复用了一部分 的网络，减少了参数量，提高了训练效率。

我们还发现不同时间步的后验均值间隔中含有不同程度的数据信息，于是我们针对性地设计了权重 ，使得训练过程更加稳定高效，性能更好。

实验结果

3.1 训练效率

PDAE 每步的训练速度比 Diff-AE 快接近 10%，并且总的训练次数只需要 Diff-AE 的 。

3.2 填补间隔

使用在 FFHQ 数据集上训练的模型证明学习到的均值偏移量可以填补后验均值间隔。左图为对不同时间步 的 仅去噪一次得到的 对比，第一行为预训练扩散模型生成的样本，第二行为 PDAE 生成的样本，可以看到 越大， 噪声比例越大，后验均值间隔越大，但 PDAE 依然可以很好的填补间隔，生成和原图大致一样的样本。

右图为随机挑选 1000 个样本，在每个时间步 采样 并分别计算平均间隔 和 ，可以看到均值偏移量明显地填补并缩小了间隔。

3.3 图像重构

使用在 FFHQ 数据集上训练的模型进行图像重构实验。使用反向 DDIM 对原始图像进行加噪得到 inferred ，以它为起点可以得到较为完美的重构效果；而在引入随机性后，样本会出现一些细节的随机变化，例如头发、皮肤等。

3.4 图像插值

使用在 FFHQ 数据集上训练的模型进行图像插值实验。两行图像分别是使用 和 进行指导采样的结果，可以看到两者最终生成的样本接近， 可以被看作是对 的一个类似线性映射的函数。

3.5 图像操控

使用在 CelebA-HQ 数据集上训练的模型进行图像操控实验。利用 CelebA-HQ 数据集的属性标注，训练一个线性层对图像编码 进行分类，并使用属性对应的分类超平面的法向量作为移动方向，对 进行正负不同尺度的移动，再进行解码。

3.6 truncation-like effect

扩散模型分类器指导采样的一个特性是可以通过改变指导的尺度来实现 truncation-like effect。我们假设通过最小化后验均值间隔的方法可以训练出一个梯度估计器，为了进一步验证这一点，我们使用 ImageNet 数据集，用数据的类别标签 替换可学习的数据表征 ，只优化 去最小化后验均值间隔，并使用不同的指导尺度进行评估，同样可以得到 truncation-like effect：

3.7 无条件采样

正如前面所看到的，在 的指导下，仅仅去噪一次，PDAE 就可以生成较为不错的样本，所以我们使用一个 Latent DDIM 建模 空间并从中采样，得到的 可以用于解码器的无条件采样，在相同的迭代步数下，其采样结果远好于预训练扩散模型，并且超越了 Diff-AE。值得注意的是，得益于我们 和 的解耦，PDAE 可以作为预训练扩散模型的辅助提升器，而不用像 Diff-AE 那样从零开始训练。

参考文献

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33 (2020): 6840-6851.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. International Conference on Learning Representations. 2021.

Dhariwal, Prafulla, and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems 34 (2021): 8780-8794.

Preechakul, Konpat, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn. Diffusion autoencoders: Toward a meaningful and decodable representation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·