ZigMa: Zigzag Mamba Diffusion Model

要約

拡散モデルは、特にトランスベースの構造内で、スケーラビリティと二次複雑性の問題に長い間悩まされてきました。
この研究では、Mamba と呼ばれる状態空間モデルの長いシーケンス モデリング機能を活用して、その適用可能性をビジュアル データ生成に拡張することを目的としています。
まず、現在のほとんどの Mamba ベースの視覚手法における重大な見落とし、つまり Mamba のスキャン スキームにおける空間的連続性の考慮の欠如を特定します。
次に、この洞察に基づいて、Zigzag Mamba という名前のシンプルなプラグアンドプレイのゼロパラメーターメソッドを導入します。これは、Mamba ベースのベースラインを上回り、トランスフォーマーベースのベースラインと比較して速度とメモリ使用率の向上を示します。
最後に、Zigzag Mamba を Stochastic Interpolant フレームワークと統合して、FacesHQ $1024\times 1024$ や UCF101、MultiModal-CelebA-HQ、MS COCO $256\times 256$ などの大解像度のビジュアル データセットでのモデルのスケーラビリティを調査します。
。
コードは https://taohu.me/zigma/ で公開されます。

要約(オリジナル)

The diffusion model has long been plagued by scalability and quadratic complexity issues, especially within transformer-based structures. In this study, we aim to leverage the long sequence modeling capability of a State-Space Model called Mamba to extend its applicability to visual data generation. Firstly, we identify a critical oversight in most current Mamba-based vision methods, namely the lack of consideration for spatial continuity in the scan scheme of Mamba. Secondly, building upon this insight, we introduce a simple, plug-and-play, zero-parameter method named Zigzag Mamba, which outperforms Mamba-based baselines and demonstrates improved speed and memory utilization compared to transformer-based baselines. Lastly, we integrate Zigzag Mamba with the Stochastic Interpolant framework to investigate the scalability of the model on large-resolution visual datasets, such as FacesHQ $1024\times 1024$ and UCF101, MultiModal-CelebA-HQ, and MS COCO $256\times 256$. Code will be released at https://taohu.me/zigma/

arxiv情報

提供元, 利用サービス

arxiv.jp, Google