Efficient Diffusion Training via Min-SNR Weighting Strategy

要約

拡散モデルのノイズ除去は、画像生成の主流のアプローチでしたが、これらのモデルのトレーニングでは収束が遅くなることがよくあります。
この論文では、収束が遅いのは、タイムステップ間で最適化の方向が競合していることが一因であることを発見しました。
この問題に対処するために、拡散トレーニングをマルチタスク学習問題として扱い、Min-SNR-$\gamma$ と呼ばれるシンプルで効果的なアプローチを導入します。
この方法は、クランプされた信号対雑音比に基づいて時間ステップの損失重みを適応させ、時間ステップ間の競合を効果的にバランスさせます。
私たちの結果は、収束速度の大幅な改善を示しており、以前の重み付け戦略よりも 3.4$\times$ 高速です。
また、より効果的であり、以前の最先端技術で採用されていたものよりも小さなアーキテクチャを使用して、ImageNet $256\times256$ ベンチマークで 2.06 の新しい記録的な FID スコアを達成しています。
コードは https://github.com/TiankaiHang/Min-SNR-Diffusion-Training で入手できます。

要約(オリジナル)

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-$\gamma$. 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$\times$ faster than previous weighting strategies. It is also more effective, achieving a new record FID score of 2.06 on the ImageNet $256\times256$ 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google