17日前

DisCo-Diff:離散潜在変数を用いた連続拡散モデルの強化

Yilun Xu, Gabriele Corso, Tommi Jaakkola, Arash Vahdat, Karsten Kreis
DisCo-Diff:離散潜在変数を用いた連続拡散モデルの強化
要約

拡散モデル(Diffusion Models, DMs)は生成学習の分野を革命的に変革した。これらのモデルは、データを単純なガウス分布に符号化する拡散プロセスを利用する。しかし、複雑で多モーダルなデータ分布を単一の連続的ガウス分布に符号化することは、本質的に過度に困難な学習課題であると指摘できる。本研究では、補完的な離散潜在変数を導入することで、この課題を簡素化する「離散・連続潜在変数拡散モデル」(Discrete-Continuous Latent Variable Diffusion Models, DisCo-Diff)を提案する。DisCo-Diffは、エンコーダによって推論される学習可能な離散潜在変数をDMに追加し、DMとエンコーダをエンドツーエンドで同時学習する。このアプローチは事前学習済みネットワークに依存しないため、汎用的なフレームワークとして適用可能である。離散潜在変数の導入により、DMの生成的ODEにおける曲率が大幅に低減され、複雑なノイズからデータへの写像を学習するプロセスが著しく簡素化される。さらに、離散潜在変数の分布をモデル化するための自己回帰型トランスフォーマーを追加するが、DisCo-Diffでは少量の離散変数と小さなコードブックで十分であるため、このステップは非常にシンプルである。本手法は玩具データ、複数の画像合成タスク、および分子ドッキングの実験において検証され、離散潜在変数の導入が一貫してモデル性能の向上をもたらすことが確認された。例えば、ODEサンプラを用いた条件下で、クラス条件付きImageNet-64/128データセットにおいて、最先端のFIDスコアを達成した。