SlimFlow: Training Smaller One-Step Diffusion Models with Rectified Flow

要約

拡散モデルは高品質の生成に優れていますが、反復サンプリングによる推論の遅さに悩まされます。
最近の手法は拡散モデルをワンステップジェネレーターに変換することに成功しましたが、モデルサイズの削減が無視されているため、計算に制約のあるシナリオでの適用性が制限されています。
この論文は、推論ステップとモデル サイズの結合圧縮を調査することにより、強力な整流フレームワークに基づいた小型で効率的なワンステップ拡散モデルを開発することを目的としています。
整流フロー フレームワークは、リフローと蒸留の 2 つの操作を使用して 1 ステップの生成モデルをトレーニングします。
元のフレームワークと比較して、モデル サイズを縮小すると、次の 2 つの新たな課題が生じます。(1) リフロー中の大規模な教師と小規模な生徒の間の初期化の不一致。
(2) 小規模なスチューデント モデルでの単純な蒸留のパフォーマンスが低い。
これらの問題を克服するために、当社では、SlimFlow フレームワークを構成するアニーリング リフローとフローガイド蒸留を提案します。
私たちの新しいフレームワークを使用して、5.02 の FID と 1,570 万のパラメーターで 1 ステップ拡散モデルをトレーニングし、CIFAR10 上の以前の最先端の 1 ステップ拡散モデル (FID=6.47、1,940 万パラメーター) を上回るパフォーマンスを発揮します。
ImageNet 64$\times$64 および FFHQ 64$\times$64 では、私たちの方法は、より大きなモデルに匹敵する小さな 1 ステップ拡散モデルを生成し、コンパクトで効率的な 1 ステップ拡散モデルの作成における私たちの方法の有効性を示しています。

要約(オリジナル)

Diffusion models excel in high-quality generation but suffer from slow inference due to iterative sampling. While recent methods have successfully transformed diffusion models into one-step generators, they neglect model size reduction, limiting their applicability in compute-constrained scenarios. This paper aims to develop small, efficient one-step diffusion models based on the powerful rectified flow framework, by exploring joint compression of inference steps and model size. The rectified flow framework trains one-step generative models using two operations, reflow and distillation. Compared with the original framework, squeezing the model size brings two new challenges: (1) the initialization mismatch between large teachers and small students during reflow; (2) the underperformance of naive distillation on small student models. To overcome these issues, we propose Annealing Reflow and Flow-Guided Distillation, which together comprise our SlimFlow framework. With our novel framework, we train a one-step diffusion model with an FID of 5.02 and 15.7M parameters, outperforming the previous state-of-the-art one-step diffusion model (FID=6.47, 19.4M parameters) on CIFAR10. On ImageNet 64$\times$64 and FFHQ 64$\times$64, our method yields small one-step diffusion models that are comparable to larger models, showcasing the effectiveness of our method in creating compact, efficient one-step diffusion models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google