[論文レビュー] Efficient Diffusion Training via Min-SNR Weighting Strategy

Tiankai Hang, Shuyang Gu+6 more|arXiv (Cornell University)|Mar 16, 2023

Domain Adaptation and Few-Shot Learning被引用数 5

ひとこと要約

本稿では、スナップショット信号対ノイズ比（SNR）に基づくタイムステップ固有の重みを割り当てることで、拡散学習をマルチタスク学習問題として扱うシンプルでありながら効果的な損失重み戦略、Min-SNR-γを提案する。この戦略により、収束が著しく加速され、より小さなモデルでもImageNet 256×256で2.06という新しいSOTA FIDスコアを達成した。

ABSTRACT

Denoising diffusion models have been a mainstream approach for image generation, however, training these models often suffers from slow convergence. In this paper, we discovered that the slow convergence is partly due to conflicting optimization directions between timesteps. To address this issue, we treat the diffusion training as a multi-task learning problem, and introduce a simple yet effective approach referred to as Min-SNR-$γ$. This method adapts loss weights of timesteps based on clamped signal-to-noise ratios, which effectively balances the conflicts among timesteps. Our results demonstrate a significant improvement in converging speed, 3.4$ imes$ faster than previous weighting strategies. It is also more effective, achieving a new record FID score of 2.06 on the ImageNet $256 imes256$ benchmark using smaller architectures than that employed in previous state-of-the-art. The code is available at https://github.com/TiankaiHang/Min-SNR-Diffusion-Training.

研究の動機と目的

• タイムステップ間で最適化方向が相反することに起因する、ノイズ除去拡散モデル学習の収束遅延を解消すること。
• 各タイムステップを個別のタスクとして扱うマルチタスク学習フレームワークとして拡散学習を再定式化すること。
• 適応的パレート最適化の不安定性と計算コストを回避する、安定的で効率的かつスケーラブルな損失重み戦略を開発すること。
• アーキテクチャの変更なしに、学習効率と生成品質を向上させ、ImageNetベンチマークで最先端のパフォーマンスを達成すること。

提案手法

• マルチタスク学習フレームワークにおいて、各ノイズ除去タイムステップを個別のタスクとして扱う。
• スナップショット信号対ノイズ比（SNR）を重み要因として用いる、事前に定義されたグローバル損失重み戦略であるMin-SNR-γを導入する。
• 極端な重みを防ぐためにSNRにクラムプ関数を適用し、学習の安定化と勾配の相反を軽減する。
• クラムプ強度を制御するハイパーパrameter γを用い、さまざまな学習設定においてもロバスト性を確保する。
• 実行時における適応的重み付けを避けるために、固定で段階的な損失重みスケジュールを採用し、効率性と安定性を向上させる。
• アーキテクチャの変更なしに、標準的なアーキテクチャ（UNet、ViT）を用いて新しい損失重み戦略でモデルを学習する。

実験結果

リサーチクエスチョン

• RQ1 なぜ、タイムステップ間で共有されるモデル重みを持つにもかかわらず、拡散学習は収束が遅いのか？
• RQ2 タイムステップ間の相反する最適化方向を、原理的損失重み戦略によって軽減できるか？
• RQ3 適応的マルチタスク学習手法（例：パレート最適化）と比較して、グローバルかつ事前に定義された損失重み戦略は拡散学習においてどのように機能するか？
• RQ4 シンプルなSNRベースの重み戦略は、従来のベースラインと比較して収束速度とFIDスコアの両面で優れているか？
• RQ5 Min-SNR-γは、さまざまなアーキテクチャや生成ターゲット（例：ε予測 vs. x₀予測）に一般化可能か？

主な発見

• Min-SNR-γは、従来の重み戦略と比較して3.4倍の高速な収束を達成し、学習を著しく加速した。
• 本手法は、ImageNet 256×256ベンチマークで2.06という新しいSOTA FIDスコアを達成し、先行SOTA手法を上回った。
• より小さなモデル（例：395Mパラメータ）でも、ImageNet 256×256でFID 2.81を達成し、DiT-XL-2のようなより大きなモデルを上回った。
• Min-SNR-5を用いて学習したViT-XLモデルは、たった210万イテレーションでFID 2.08を達成し、DiTの3.3倍速で収束した。
• 本手法はUNetおよびViTバックボーンの両方で性能向上を示し、ε予測およびx₀予測ターゲットの両方で効果的に機能した。
• アブレーションスタディの結果、ハイパーパrameter γの変更に対しても性能の低下が最小限に抑えられ、ロバスト性が確認された。