[论文解读] Unsupervised Representation Learning from Pre-trained Diffusion Probabilistic Models

Zijian Zhang, Zhou Zhao+1 more|arXiv (Cornell University)|Dec 26, 2022

Generative Adversarial Networks and Image Synthesis被引用 15

一句话总结

该论文提出PDAE，一种通过训练可学习的均值偏移预测器来填补前向过程中因信息丢失导致的后验均值差距，从而将预训练的扩散概率模型（DPMs）适配于无监督表示学习的自编码方法。PDAE在图像重建和无条件生成任务中表现优于Diff-AE，且训练效率更高，所需时间不足其一半。

ABSTRACT

Diffusion Probabilistic Models (DPMs) have shown a powerful capacity of generating high-quality image samples. Recently, diffusion autoencoders (Diff-AE) have been proposed to explore DPMs for representation learning via autoencoding. Their key idea is to jointly train an encoder for discovering meaningful representations from images and a conditional DPM as the decoder for reconstructing images. Considering that training DPMs from scratch will take a long time and there have existed numerous pre-trained DPMs, we propose extbf{P}re-trained extbf{D}PM extbf{A}uto extbf{E}ncoding ( extbf{PDAE}), a general method to adapt existing pre-trained DPMs to the decoders for image reconstruction, with better training efficiency and performance than Diff-AE. Specifically, we find that the reason that pre-trained DPMs fail to reconstruct an image from its latent variables is due to the information loss of forward process, which causes a gap between their predicted posterior mean and the true one. From this perspective, the classifier-guided sampling method can be explained as computing an extra mean shift to fill the gap, reconstructing the lost class information in samples. These imply that the gap corresponds to the lost information of the image, and we can reconstruct the image by filling the gap. Drawing inspiration from this, we employ a trainable model to predict a mean shift according to encoded representation and train it to fill as much gap as possible, in this way, the encoder is forced to learn as much information as possible from images to help the filling. By reusing a part of network of pre-trained DPMs and redesigning the weighting scheme of diffusion loss, PDAE can learn meaningful representations from images efficiently. Extensive experiments demonstrate the effectiveness, efficiency and flexibility of PDAE.

研究动机与目标

• 为解决从零开始训练扩散自编码器（Diff-AE）效率低下且性能欠佳的问题，通过利用现有的预训练DPMs。
• 克服预训练DPMs因前向过程中信息丢失而在从潜在变量重建图像时的失败。
• 开发一种可泛化、高效的通用方法，重用预训练DPM组件，并通过可训练的均值偏移预测器实现信息填补，从而提升表示学习。
• 通过最小化微调，使预训练DPM作为解码器实现高质量图像重建和无条件生成。

提出的方法

• PDAE引入一种可学习的梯度估计器，用于预测在反向扩散过程中填补预测后验均值与真实后验均值之间差距的均值偏移。
• 通过最大化潜在表示中的信息量，训练编码器使均值偏移预测器能够有效重建原始图像。
• 对扩散损失采用重新设计的加权方案，突出强调后验均值差距中信息更丰富的时刻步，提升表示学习效率。
• PDAE重用预训练DPM中的一部分U-Net架构，以加速收敛并降低训练成本。
• 利用分类器引导采样与后验均值差距之间的联系，将分类器的梯度解释为重建丢失信息的机制。
• 通过仅在采样步骤的最后30%应用均值偏移预测器，改善无条件采样，减少随机性并避免分布外样本的生成。

实验结果

研究问题

• RQ1 预训练的DPMs能否被有效重用于自编码框架中的无监督表示学习解码器？
• RQ2 为何预训练的DPMs在从潜在变量重建图像时会失败？后验均值差距的成因是什么？
• RQ3 能否将后验均值差距作为信号，用于训练编码器以学习更具信息量的表示？
• RQ4 如何设计均值偏移预测器，以在不重新训练DPM的情况下高效重建丢失的信息？
• RQ5 重用预训练DPM组件并优化扩散损失加权是否能提升训练效率和性能？

主要发现

• PDAE在图像重建质量上优于Diff-AE，且训练时间不足其一半，展现出更优的训练效率。
• 该方法在多个数据集（包括FFHQ128、Horse128、Bedroom128和CelebA64）上显著提升了无条件采样质量，FID得分表现更优。
• 仅在采样步骤的最后30%应用均值偏移预测器，通过减少随机性并避免分布外生成，提升了采样质量。
• 后验均值差距在不同时间步中所含信息量不一，为高效表示学习提供了重新设计损失加权方案的依据。
• PDAE可作为通用增强器应用于任意预训练DPM，无需重新训练DPM即可提升重建与生成质量。
• 即使在预训练DPM被冻结的情况下，该方法依然有效，证明其在实际部署中的兼容性与实用性。