[論文レビュー] 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

ひとこと要約

本論文は、前もって学習された拡散確率モデル（DPM）を、自己符号化を通じて教師なし表現学習に適応させる手法PDAEを提案する。情報損失が前向きプロセスに起因する事後平均ギャップを埋めるために、学習可能な平均シフト予測器を訓練することで実現される。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.

研究の動機と目的

• 事前学習済みDPMを活用することで、拡散自己符号化器（Diff-AE）を再訓練する際の非効率さと最適でない性能を是正すること。
• 前向きプロセスにおける情報損失により、事前学習済みDPMが潜在変数から画像を再構成できない問題を克服すること。
• 事前学習済みDPMのコンponentsを再利用し、学習可能な平均シフト予測器によるギャップ補填によって表現学習を向上させる汎用的で効率的な手法を開発すること。
• 最小限のファインチューニングで、事前学習済みDPMをデコーダーとして用いて高品質な画像再構成と無条件生成を可能にすること。

提案手法

• PDAEは、逆方向拡散プロセスにおける予測事後平均と真の事後平均のギャップを埋めるために、学習可能な勾配推定器（平均シフト予測器）を導入する。
• エンコーダーを訓練して、潜在表現に含まれる情報量を最大化することで、平均シフト予測器が元の画像を効果的に再構成できるようにする。
• 拡散損失の重み付けを再設計し、事後平均ギャップに豊富な情報が含まれる時間ステップに重点を置くことで、表現学習の効率を向上させる。
• 収束を早め、訓練コストを削減するために、事前学習済みDPMのU-Netアーキテクチャの一部を再利用する。
• 分類器ガイドドサンプリングと事後平均ギャップの関係を活用し、分類器の勾配を失われた情報を再構成するメカニズムとして解釈する。
• 無条件サンプリングの品質を向上させるために、平均シフト予測器をサンプリングステップの最後30%にのみ適用することで、確率的変動を低減し、ドメイン外のサンプルを回避する。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みDPMは、教師なし表現学習の自己符号化フレームワークにおいて、デコーダーとして効果的に再利用可能か？
• RQ2 事前学習済みDPMが潜在変数から画像を再構成できない理由は何か？また、事後平均ギャップの原因は何か？
• RQ3 事後平均ギャップを、より情報の豊かな表現を学習するためのエンコーダー訓練の信号として活用できるか？
• RQ4 DPMの再学習なしに、失われた情報を効率的に再構成できるように、平均シフト予測器をどのように設計できるか？
• RQ5 事前学習済みDPMのコンponentsを再利用し、拡散損失の重み付けを最適化することで、訓練効率と性能が向上するか？

主な発見

• PDAEは、Diff-AEよりも優れた画像再構成品質を達成しながら、訓練時間の半分未満で実現し、優れた訓練効率を示した。
• FFHQ128、Horse128、Bedroom128、CelebA64の複数のデータセットで、FIDスコアを用いて測定したところ、無条件サンプリング品質が著しく向上した。
• 平均シフト予測器をサンプリングステップの最後30%にのみ適用することで、確率的変動を低減し、ドメイン外の生成を回避でき、サンプル品質が向上した。
• 事後平均ギャップには、時間ステップごとに情報量のばらつきが存在することが示され、効率的な表現学習のための再設計された損失重み付けの妥当性が裏付けられた。
• PDAEは、任意の事前学習済みDPMに一般化可能なブースターとして適用可能であり、DPMの再学習なしに再構成品質および生成品質を向上させることができる。
• 事前学習済みDPMを固定した状態でも効果を発揮するため、実世界の展開においても互換性と実用性を示した。