[論文レビュー] Maximum Likelihood Training of Implicit Nonlinear Diffusion Models

Dongjun Kim, Byeonghu Na+4 more|arXiv (Cornell University)|May 27, 2022

Generative Adversarial Networks and Image Synthesis被引用数 10

ひとこと要約

本稿では、潜在空間における正規化フローを学習することで、データに適応した非線形拡散過程を学習する、暗黙的非線形拡散モデル（INDM）を提案する。この手法により、データ空間における非線形のドリフト項とボラティリティを暗黙的にモデル化できる。INDMはCelebAで1.75のSOTA FIDを達成し、最大尤度推定（MLE）による学習を実施することで、DDPM++などの線形ベースラインより優れたサンプル品質と学習安定性を実現した。

ABSTRACT

Whereas diverse variations of diffusion models exist, extending the linear diffusion into a nonlinear diffusion process is investigated by very few works. The nonlinearity effect has been hardly understood, but intuitively, there would be promising diffusion patterns to efficiently train the generative distribution towards the data distribution. This paper introduces a data-adaptive nonlinear diffusion process for score-based diffusion models. The proposed Implicit Nonlinear Diffusion Model (INDM) learns by combining a normalizing flow and a diffusion process. Specifically, INDM implicitly constructs a nonlinear diffusion on the extit{data space} by leveraging a linear diffusion on the extit{latent space} through a flow network. This flow network is key to forming a nonlinear diffusion, as the nonlinearity depends on the flow network. This flexible nonlinearity improves the learning curve of INDM to nearly Maximum Likelihood Estimation (MLE) against the non-MLE curve of DDPM++, which turns out to be an inflexible version of INDM with the flow fixed as an identity mapping. Also, the discretization of INDM shows the sampling robustness. In experiments, INDM achieves the state-of-the-art FID of 1.75 on CelebA. We release our code at https://github.com/byeonghu-na/INDM.

研究の動機と目的

• 既存の拡散モデルが固定された線形拡散過程に依存するという制限を解決すること。これは表現力の制限と学習効率の低下を引き起こす可能性がある。
• データ空間における学習可能な非線形前向きSDE（確率微分方程式）—潜在空間における学習可能な正規化フローを介して実現—が、生成モデルの性能を向上させることを調査すること。
• フローに基づく変換により、前向きSDEのドリフト項とボラティリティ係数を暗黙的にモデル化することで、拡散モデルにおける最大尤度推定（MLE）学習を可能にすること。
• 非線形SDEの離散化が異なるステップサイズに対しても安定するように保証することで、サンプリングのロバスト性を向上させること。
• 適切に訓練された非線形拡散が、尤度とサンプル品質の両方において線形拡散ベースラインを上回ることを実証すること。

提案手法

• INDMは、潜在空間における線形拡散を、学習可能な正規化フローを用いて変換することで、データ空間に非線形拡散過程を構築する。
• 正規化フローは可逆的かつ微分可能であり、前向きおよび逆向きのプロセスにわたるバックプロパゲーションを可能にし、非線形SDEの学習に不可欠である。
• データ空間における前向きSDEは、フロー写像によって暗黙的に定義される：$\mathbf{x}_t = \mathbf{f}_{\bm{\phi}}(\mathbf{z}_t)$、ここで$\mathbf{z}_t$は線形SDEに従う。
• スコアネットワーク$\mathbf{s}_{\bm{\theta}}$は、真のデータスコアを近似するノイズ除去スコア損失を最小化することで学習され、変分下界を介してMLE学習が実現される。
• 尤度の下界を保証するため、重み関数$\lambda(t) = g^2(t)$を用いる。
• サンプリングは、事前分布$\mathbf{x}_T \sim \mathcal{N}(0, I)$から始まり、逆向きSDEを時間反転して解くことで実行され、学習済みスコアネットワークとフローに基づくドリフト補正が用いられる。

実験結果

リサーチクエスチョン

• RQ1 データに適応した学習可能な非線形拡散プロセスが、固定された線形拡散と比較して生成モデルの性能を向上させられるか？
• RQ2 潜在空間における正規化フローを介して前向きSDEを暗黙的にモデル化することで、標準的な線形拡散と比較してドリフト項およびボラティリティ係数の学習がより効果的に行えるか？
• RQ3 提案手法は拡散モデルにおける最大尤度推定（MLE）学習を達成できるか？また、スコアベース学習と比較して尤度とサンプル品質の両面でどのように差がつくか？
• RQ4 非線形SDEの離散化ロバスト性は、DDPM++のような線形拡散モデルと比較してどの程度優れているか？
• RQ5 フローネットワークが導入する非線形性は、CIFAR-10 や CelebA といったベンチマークデータセットにおけるFIDおよびNLL指標をどの程度改善するか？

主な発見

• INDMはCelebA $64 \times 64$ で1.75のSOTA FIDを達成し、DDPM++ や NCSN++ を含むすべての先行拡散モデルを上回った。
• CIFAR-10では、補正後NLLが2.97、NELBOとNLLの差が0.02にとどまり、ほぼ最適な尤度推定が達成された。
• CIFAR-10（VE）ではFIDが2.28、VPでは2.90にまで低下し、DDPM++（FID 2.64）およびNCSN++（FID 2.20）を両方の指標で上回った。
• サンプリングの離散化に対して強いロバスト性を示し、粗い時間ステップでも高品質な生成が維持された。これは線形ベースラインとは対照的であった。
• INDMはほぼMLE性能を達成しており、CIFAR-10におけるNELBOとNLLの差がわずか0.02にとどまり、尤度最適化が効果的に行われた。
• アブレーションスタディにより、フローネットワークによる非線形拡散が不可欠であることが確認された。フローを恒等写像（すなわちDDPM++に類似）に固定した場合、著しく劣化した学習曲線と高いFIDが観察された。