Scalable Diffusion Models with State Space Backbone

要約

この論文は、状態空間アーキテクチャに基づいて構築された拡散モデルのカテゴリーへの新たな探求を提示します。
私たちは、画像データの拡散モデルのトレーニングに努めています。このモデルでは、従来の U-Net バックボーンが状態空間バックボーンに取って代わられ、生のパッチまたは潜在空間で機能します。
拡散状態空間モデル (DiS) は、長距離の依存関係に対応する際の顕著な有効性を考慮して、時間、条件、ノイズの多い画像パッチを含むすべての入力をトークンとして扱うことで区別されます。
DiS の評価には、無条件画像生成シナリオとクラス条件付き画像生成シナリオの両方が含まれており、DiS が、相応のサイズの CNN ベースまたは Transformer ベースの U-Net アーキテクチャよりも優れているとは言えないまでも、同等のパフォーマンスを示すことが明らかになりました。
さらに、Gflops で定量化されたフォワード パスの複雑さによって測定される、DiS のスケーラビリティを分析します。
深さ/幅の増加または入力トークンの増加によって達成される、より高い Gflops を備えた DiS モデルは、一貫して低い FID を示します。
賞賛に値するスケーラビリティ特性を示すことに加えて、潜在空間の DiS-H/2 モデルは、クラス条件付き ImageNet ベンチマークで 256$\times$256 および 512$\times$512 の解像度で以前の拡散モデルと同様のパフォーマンス レベルを達成しながら、大幅にパフォーマンスを向上させます。
計算上の負担。
コードとモデルは https://github.com/feizc/DiS から入手できます。

要約(オリジナル)

This paper presents a new exploration into a category of diffusion models built upon state space architecture. We endeavor to train diffusion models for image data, wherein the traditional U-Net backbone is supplanted by a state space backbone, functioning on raw patches or latent space. Given its notable efficacy in accommodating long-range dependencies, Diffusion State Space Models (DiS) are distinguished by treating all inputs including time, condition, and noisy image patches as tokens. Our assessment of DiS encompasses both unconditional and class-conditional image generation scenarios, revealing that DiS exhibits comparable, if not superior, performance to CNN-based or Transformer-based U-Net architectures of commensurate size. Furthermore, we analyze the scalability of DiS, gauged by the forward pass complexity quantified in Gflops. DiS models with higher Gflops, achieved through augmentation of depth/width or augmentation of input tokens, consistently demonstrate lower FID. In addition to demonstrating commendable scalability characteristics, DiS-H/2 models in latent space achieve performance levels akin to prior diffusion models on class-conditional ImageNet benchmarks at the resolution of 256$\times$256 and 512$\times$512, while significantly reducing the computational burden. The code and models are available at: https://github.com/feizc/DiS.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google