[論文レビュー] Improving Diffusion Model Efficiency Through Patching

Troy Luhman, Eric Luhman|arXiv (Cornell University)|Jul 9, 2022

Generative Adversarial Networks and Image Synthesis被引用数 5

ひとこと要約

この論文では、空間的パッチをチャネル次元に変換することで、拡散モデルにおける計算コストとメモリ使用量を削減する、Patched Diffusion Models (PDM) を提案する。ViT風のパッチングを適用することで、生成品質を維持したまま、特にImageNet 256² や FFHQ 1024² などの高解像度画像において、推論時間とメモリ使用量を顕著に削減できる。

ABSTRACT

Diffusion models are a powerful class of generative models that iteratively denoise samples to produce data. While many works have focused on the number of iterations in this sampling procedure, few have focused on the cost of each iteration. We find that adding a simple ViT-style patching transformation can considerably reduce a diffusion model's sampling time and memory usage. We justify our approach both through an analysis of the diffusion model objective, and through empirical experiments on LSUN Church, ImageNet 256, and FFHQ 1024. We provide implementations in Tensorflow and Pytorch.

研究の動機と目的

• 高解像度画像における拡散モデルのサンプリングにかかる高い計算コストとメモリ使用量を低減すること。
• 最適再構成において、ぼんやりとした特徴が生じるため、ほとんどのタイムステップで高周波数の畳み込み層が冗長であることを特定すること。
• パッチングにより空間的特徴をチャネル次元に再編成することで、シンプルで即座に統合可能な効率化手法を開発すること。
• 計算オーバーヘッドを低減しつつ、拡散プロセスの統計的同等性を保持すること。
• スケーラブルなパッチベースアーキテクチャにより、ImageNet 256² などの困難なデータセットでも効率的かつ高精度な画像生成を可能にすること。

提案手法

• 空間的画像パッチをチャネル次元に変換するため、ViT風のパッチングを適用し、空間解像度を低下させるとともにチャネル数を増加させる。
• パッチ化された潜在表現上で動作するように拡散モデルを再定式化し、同じ目的関数と生成プロセスを維持する。
• 訓練の安定性を向上させ、脆弱性を低減するため、ノイズ予測の代わりにデータ予測ヘッド (x_θ) を使用する。
• ImageNet 256² のような複雑なデータセットにスケーラブルに適用できるよう、マルチパートネットワーク分割戦略を導入する。
• 低レベル特徴の学習を優先するために、損失関数を γₜ = √(αₜ / (1−αₜ)) で最適化する。
• 広範なアクセシビリティと再現可能性を実現するため、TensorFlow および PyTorch 両方のフレームワークで手法を実装する。

実験結果

リサーチクエスチョン

• RQ1 パッチングによって、サンプル品質を損なわずに、拡散モデルのサンプリングにおける計算コストとメモリ使用量をどの程度削減できるか？
• RQ2 高周波数の畳み込み層の冗長性は、拡散モデルにおける最適再構成にどのように影響するか？
• RQ3 パッチング下で、データ予測 (x_θ) がノイズ予測 (ε_θ) よりも訓練の安定性とサンプル忠実度において優れているか？
• RQ4 パッチド拡散モデルは、ImageNet 256² のような困難なデータセットにおいても高精度な画像合成を達成できるか？
• RQ5 パッチング変換は、拡散プロセスの統計的同等性と尤度目的関数にどのように影響を与えるか？

主な発見

• Patched Diffusion Models は、特に FFHQ 1024² や ImageNet 256² のような高解像度画像において、サンプリング時間とメモリ使用量を顕著に削減する。
• 最適再構成 x*(zₜ,t) は低次元多様体上に位置し、ぼんやりとした外観を示しており、ほとんどのタイムステップで高周波数成分が重要でないことが示唆される。
• 高周波数の畳み込み層は、特に後段のタイムステップで最適再構成にほとんど寄与しないため、拡散プロセスにおいて冗長である。
• データ予測 (x_θ) はノイズ予測 (ε_θ) よりも脆弱性が低く、パッチド環境下でもより安定した訓練とより優れたサンプル品質を実現する。
• 提案手法は、分割ネットワークアーキテクチャを用いて ImageNet 256² において高精度な画像合成を達成し、複雑なデータセットへのスケーラビリティを示した。
• パッチング操作は拡散プロセスの統計構造を保持しており、任意の既存の拡散モデルに即座に統合可能な「プラグアンドプレイ」性を有する。