[論文レビュー] Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models

Zhendong Wang, Yifan Jiang+6 more|arXiv (Cornell University)|Apr 25, 2023

Advanced Neuroimaging Techniques and Applications被引用数 32

ひとこと要約

Patch Diffusion は、座標条件付き入力とランダムなパッチサイズスケジューリングを用いたパッチ単位の条件付きスコアマッチングを導入し、拡散モデルのトレーニングを加速しデータ効率を向上させる。サンプル品質を維持または向上させつつ、少なくとも2x のトレーニング速度向上を達成する。

ABSTRACT

Diffusion models are powerful, but they require a lot of time and data to train. We propose Patch Diffusion, a generic patch-wise training framework, to significantly reduce the training time costs while improving data efficiency, which thus helps democratize diffusion model training to broader users. At the core of our innovations is a new conditional score function at the patch level, where the patch location in the original image is included as additional coordinate channels, while the patch size is randomized and diversified throughout training to encode the cross-region dependency at multiple scales. Sampling with our method is as easy as in the original diffusion model. Through Patch Diffusion, we could achieve $\mathbf{\ge 2 imes}$ faster training, while maintaining comparable or better generation quality. Patch Diffusion meanwhile improves the performance of diffusion models trained on relatively small datasets, $e.g.$, as few as 5,000 images to train from scratch. We achieve outstanding FID scores in line with state-of-the-art benchmarks: 1.77 on CelebA-64$ imes$64, 1.93 on AFHQv2-Wild-64$ imes$64, and 2.72 on ImageNet-256$ imes$256. We share our code and pre-trained models at https://github.com/Zhendong-Wang/Patch-Diffusion.

研究の動機と目的

• 拡散モデルのトレーニングにかかる時間とデータ要件を低減して民主化を促進する。
• UNet バックボーンとサンプリング手法に依存しないパッチ単位のトレーニングフレームワークを提案する。
• パッチレベルの座標条件付けとパッチサイズスケジューリングを開発し、領域間の依存関係を捉える。
• パッチ拡散がトレーニング時間を半減しつつ生成品質を維持または向上させられることを示す。特に小規模データセットで。

提案手法

• ランダムな画像パッチ上でデノイジングスコアマッチングを用いて訓練されるパッチ-wise デノイザー D_theta(x; sigma_t) を定義する。
• パッチのスコアをパッチ位置 (i, j) とパッチサイズ s に条件付けし、小さなパッチと大きなパッチの混合からサンプルする。
• 入力に2つの座標チャネルを追加してピクセルレベルの座標系を導入し、パッチ位置をエンコードする。
• トレーニング中にパッチサイズを混合するためのパッチサイズスケジューリング（確率的または段階的）を使用し、 controllable ratio p で全サイズの画像を含める。
• 各逆ステップで全画像の座標を提供して標準的な拡散モデルと同様にサンプリングを維持する。
• 評価時に CFG をオプションとして使用しつつ、EDM-DDPM++ および Latent Diffusion のバックボーン上でパッチ拡散を実証する。

実験結果

リサーチクエスチョン

• RQ1 パッチごとのスコアマッチングが位置とサイズの条件付けと組み合わさることで、拡散モデルのデータ効率とトレーニング時間を改善できるか。
• RQ2 パッチサイズスケジューリング（確率的 vs 段階的）が生成品質とトレーニングコストにどう影響するか。
• RQ3 パッチベースの訓練を用いつつ、 explicit なパッチ結合なしで標準的な拡散モデルと同様のサンプリングを行えるか。
• RQ4 小規模データセットや事前学習済み拡散モデルのファインチューニング時にパッチ拡散は有利か。

主な発見

• Patch Diffusion はベンチマークデータセットで少なくとも2x のトレーニング速度向上を達成しつつ、生成品質は同等またはそれ以上を維持。
• 小さなパッチと大きなパッチの混合と座標条件付けにより領域間の一貫性とFIDスコアが改善。
• p ≈ 0.5（トレーニング中の全サイズ画像の比率）で効率と品質のバランスが最適化される特性。
• Patch Diffusion はデータ効率を改善し、約5k程度の比較的小規模データセットでも競合的な結果を実現。
• 大規模事前学習モデル（例：ControlNet）を Patch Diffusion でファインチューニングするとトレーニング効率が2倍になっても性能を損なわない。
• 制限サイズのデータセットでは、Patch Diffusion は基準法より一貫してFIDで優れており、約2x のトレーニング時間削減を実現。