[論文レビュー] Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation

Dong‐Jun Kim, Seungjae Shin+3 more|arXiv (Cornell University)|Jun 10, 2021

Advanced Neuroimaging Techniques and Applications被引用数 10

ひとこと要約

本稿では、固定された打ち切りハイパーパrameterの代わりに確率的変数を導入することで、すべての拡散時間におけるスコア推定を向上させる、スコアベースの拡散モデル向けの普遍的な訓練技術「ソフト打ち切り」を提案する。各訓練ステップで最小拡散時間を確率的にサンプリングすることで、CIFAR-10、CelebA、CelebA-HQ、STL-10において、負の対数尤度（NLL）とフレシェ・インセプション距離（FID）の両面で最先端の性能を達成し、密度推定とサンプル品質の間の長年のトレードオフを解消する。

ABSTRACT

Recent advances in diffusion models bring state-of-the-art performance on image generation tasks. However, empirical results from previous research in diffusion models imply an inverse correlation between density estimation and sample generation performances. This paper investigates with sufficient empirical evidence that such inverse correlation happens because density estimation is significantly contributed by small diffusion time, whereas sample generation mainly depends on large diffusion time. However, training a score network well across the entire diffusion time is demanding because the loss scale is significantly imbalanced at each diffusion time. For successful training, therefore, we introduce Soft Truncation, a universally applicable training technique for diffusion models, that softens the fixed and static truncation hyperparameter into a random variable. In experiments, Soft Truncation achieves state-of-the-art performance on CIFAR-10, CelebA, CelebA-HQ 256x256, and STL-10 datasets.

研究の動機と目的

• スコアベースの拡散モデルにおける密度推定（NLL）とサンプル生成品質（FID）の逆相関を解消すること。
• 特にNLLが最も感受性が高い低ノイズ領域において、全拡散時間範囲にわたるスコアネットワークの訓練を改善すること。
• FIDやNLLの両方を最適化するための別個のモデル設定を不要にする、統一された訓練技術を導入すること。
• アーキテクチャの変更なしに、尤度と生成性能の両方を向上させる理論的裏付けがあり、普遍的に適用可能な手法を提供すること。

提案手法

• 固定されたεではなく、各ミニバッチごとに再サンプリングされる確率的打ち切り変数τを導入し、スコア推定の最小拡散時間を定義する。
• スコアネットワークを[τ, T]の区間でのみ訓練し、τ未満のスコアは無視することで、確率的サンプリングにより全時間範囲をカバーする。
• この方法は、拡散時間における重み関数を暗黙的に導入し、損失における一般化された期待値に相当する。これにより、NLLとFIDの両方をバランスよく向上させる理由が説明される。
• ソフト打ち切りは、任意のスコアベースの拡散モデルおよびSDE定式化と互換性があり、モデルアーキテクチャやサンプリングプロセスの変更を一切不要としない。
• 訓練目的は、拡散モデルに特化したMLEの一般化形である「最大摂動尤度推定（MPLE）」として定式化される。
• 本手法は、VE、VP、ODEに基づく定式化を含む、複数のデータセットおよびSDEにおいて実証的に検証されている。

実験結果

リサーチクエスチョン

• RQ1 単一の訓練技術が、拡散モデルにおける密度推定（NLL）とサンプル品質（FID）の両方を同時に改善できるか？
• RQ2 実際の応用において、固定された打ち切りハイパーパrameter εがなぜNLLとFIDの間のトレードオフを引き起こすのか？
• RQ3 訓練中に最小拡散時間を確率化することで、全時間範囲にわたるより良いスコア推定が達成できるか？
• RQ4 ソフト打ち切りは理論的に重み付き尤度目的と同等であり、その重みはどのように導出されるか？
• RQ5 ソフト打ち切りは、異なるモデルアーキテクチャ、SDE定式化、データセットに一般化可能か？

主な発見

• CIFAR-10（VP SDE、DDPM++）において、ソフト打ち切りはFID 3.45を達成し、FIDに有利なモデルと同等の性能を発揮しながら、NLL 3.03を維持しており、NLLに有利なモデルと同等の水準を保つ。
• CelebA-HQ 256では、ソフト打ち切りがFID 3.96、NLL 3.01を達成し、FID最適化型およびNLL最適化型のベースラインを上回る。
• アブレーションスタディでは、τ ∼ P₁を用いたソフト打ち切りは、固定ε=10⁻⁵と比較してCIFAR-10（RVE、UNCSN++）でFIDを2.74ポイント改善し、NLLは低く保ったままにした。
• 固定打ち切りと比較して、CIFAR-10ではFIDを最大2.74ポイント、CelebAでは1.26ポイント改善し、NLLに劣化は認められなかった。
• P₁（[0.1, 1.0]上での一様分布）を用いたソフト打ち切りは最良のバランスを達成し、CIFAR-10（VP、DDPM++）ではFIDを6.70から3.96に低下させた一方、NLLは安定した。
• 本手法はアーキテクチャにかかわらず頑健である。CIFAR-10ではFIDが1.26〜2.74ポイント改善され、CelebAでは0.7〜1.2ポイント改善され、NLL性能は一貫して維持された。