[論文レビュー] Conffusion: Confidence Intervals for Diffusion Models

Eliahu Horwitz, Yedid Hoshen|arXiv (Cornell University)|Nov 17, 2022

Model Reduction and Neural Networks被引用数 6

ひとこと要約

Conffusionは、超解像や補完などの画像対画像タスクにおける拡散モデルのピクセル単位の信頼区間を生成する手法を導入し、統計的に保証された不確実性推定を可能にする。事前学習済み拡散モデルを分位数回帰損失で微調整することで、サンプリングベースのベースラインと比較して3桁以上高速な推論とより狭い信頼区間を達成しながら、理論的なカバレッジ保証を維持する。

ABSTRACT

Diffusion models have become the go-to method for many generative tasks, particularly for image-to-image generation tasks such as super-resolution and inpainting. Current diffusion-based methods do not provide statistical guarantees regarding the generated results, often preventing their use in high-stakes situations. To bridge this gap, we construct a confidence interval around each generated pixel such that the true value of the pixel is guaranteed to fall within the interval with a probability set by the user. Since diffusion models parametrize the data distribution, a straightforward way of constructing such intervals is by drawing multiple samples and calculating their bounds. However, this method has several drawbacks: i) slow sampling speeds ii) suboptimal bounds iii) requires training a diffusion model per task. To mitigate these shortcomings we propose Conffusion, wherein we fine-tune a pre-trained diffusion model to predict interval bounds in a single forward pass. We show that Conffusion outperforms the baseline method while being three orders of magnitude faster.

研究の動機と目的

• ユーザーが定めた確率で真のピクセル値が区間内に存在することを保証する、統計的に厳密な信頼区間を拡散モデルの出力に対して提供すること。
• 拡散モデルにおけるサンプリングベースの信頼区間推定の遅さを克服すること。
• 既存の分位数回帰手法と比較して、区間の狭さを向上させつつ理論的カバレッジ保証を維持すること。
• 1つの事前学習済み拡散モデルを用いて、複数のタスクやデータセットにわたる転移学習を可能にすること。
• 効率的な1パス推論を実現するため、ノイズ除去モデルを拡散プロセスから分離すること。

提案手法

• 1回の順伝播で各ピクセルの下限と上限を予測できるように、分位数回帰損失を用いて事前学習済み拡散モデル（例：ADM）を微調整する。
• ノイズ除去モデルを拡散プロセスから分離し、完全な拡散プロセスの再トレーニングなしに、再利用可能にすることで区間推定に活用できる。
• 適合予測の原則を用いて、真のピクセル値が予測された区間内に入る確率をユーザー指定で保証する。
• ImageNetからの強力な事前学習特徴を活用し、超解像や補完などのダウンストリームタスクへのゼロショット転移を可能にする。
• アモアタイズド推論を適用：1枚の画像あたり数千回の出力サンプリングではなく、1回の順伝播で境界を直接予測する。
• 複数の分位数（例：0.05 と 0.95）を同時に最適化できる1つのスカラー分位数損失を用いることで、区間境界を一括して最適化する。

実験結果

リサーチクエスチョン

• RQ1 画像対画像タスクにおける拡散モデル出力の統計的に妥当な信頼区間を構築できるか？
• RQ2 サンプリングベースのベースラインと比較して、信頼区間推定の推論速度を高速化できるか？
• RQ3 事前学習済み拡散モデルを区間予測用に微調整することで、既存の分位数回帰手法と比較してより狭く信頼性の高い境界が得られるか？
• RQ4 効果的な区間推定には拡散事前学習が不可欠か、それとも一般的な特徴抽出器で十分か？
• RQ5 1つの事前学習済みモデルが、超解像や補完といった多様なタスクに一般化できるか？

主な発見

• Conffusionは、サンプリングベースのベースラインと比較して3桁以上高速な推論を達成し、1000回以上の順伝播が必要な推論を1回の順伝播に削減する。
• 現在の最先端手法である ADM_{UQ} と比較して、理論的カバレッジ保証を維持しながらも、より狭い信頼区間を生成する。
• N-Conffusionは、サンプリングベースおよび分位数回帰ベースの手法を上回り、区間サイズと視覚的品質の両面で優れた性能を示し、現実的で明確な境界を生成する。
• 拡散事前学習は不可欠である：アブレーション実験では、非拡散モデル（例：ImageNet や Instagram で事前学習済みのエンコーダーを搭載した FPN）は Conffusion と比較して著しく広い区間を生成する。
• 本手法は、タスク固有の再トレーニングなしに、ImageNet で事前学習されたモデルから CelebA-HQ や他のデータセットへのタスクにわたり、強力な一般化性能を示す。
• Conffusionは、中間時間ステップでの信頼区間推定を可能にし、最終出力品質の高速なプロキシとしての早期停止を可能にする。