DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents

要約

拡散モデル（DM）は生成学習に革命をもたらした。DMは拡散過程を利用して、データを単純なガウス分布にエンコードする。しかし、複雑で潜在的にマルチモーダルなデータ分布を単一の連続ガウス分布に符号化することは、不必要に困難な学習問題である。我々は、補完的な離散潜在変数を導入することで、このタスクを単純化する離散-連続潜在変数拡散モデル（DisCo-Diff）を提案する。我々は、学習可能な離散潜在変数でDMを補強し、エンコーダで推論し、DMとエンコーダをエンドツーエンドで訓練する。DisCo-Diffは事前に学習されたネットワークに依存しないため、このフレームワークは普遍的に適用可能である。離散レイテントは、DMの生成ODEの曲率を減少させることで、DMの複雑なノイズからデータへのマッピングの学習を大幅に単純化する。DisCo-Diffは小さなコードブックで少数の離散変数しか必要としないため、追加の自己回帰変換器が離散潜在量の分布をモデル化する。我々は、おもちゃのデータ、いくつかの画像合成タスク、および分子ドッキングにおいてDisCo-Diffを検証し、離散潜在量を導入することで一貫してモデルの性能が向上することを見出した。例えば、DisCo-DiffはODEサンプラーを用いたクラス条件付きImageNet-64/128データセットにおいて、最先端のFIDスコアを達成している。

要約(オリジナル)

Diffusion models (DMs) have revolutionized generative learning. They utilize a diffusion process to encode data into a simple Gaussian distribution. However, encoding a complex, potentially multimodal data distribution into a single continuous Gaussian distribution arguably represents an unnecessarily challenging learning problem. We propose Discrete-Continuous Latent Variable Diffusion Models (DisCo-Diff) to simplify this task by introducing complementary discrete latent variables. We augment DMs with learnable discrete latents, inferred with an encoder, and train DM and encoder end-to-end. DisCo-Diff does not rely on pre-trained networks, making the framework universally applicable. The discrete latents significantly simplify learning the DM’s complex noise-to-data mapping by reducing the curvature of the DM’s generative ODE. An additional autoregressive transformer models the distribution of the discrete latents, a simple step because DisCo-Diff requires only few discrete variables with small codebooks. We validate DisCo-Diff on toy data, several image synthesis tasks as well as molecular docking, and find that introducing discrete latents consistently improves model performance. For example, DisCo-Diff achieves state-of-the-art FID scores on class-conditioned ImageNet-64/128 datasets with ODE sampler.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL