[論文レビュー] ZigMa: A DiT-style Zigzag Mamba Diffusion Model

Vincent Tao Hu, Stefan Andreas Baumann+5 more|arXiv (Cornell University)|Mar 20, 2024

Cellular Automata and Applications被引用数 4

ひとこと要約

ZigMaは、注意メカニズムを置き換えることで2Dおよび3Dの視覚生成における空間的連続性を向上させる、DiTスタイルの拡散モデルを提案する。Mambaのスキャンパスをジグザグパターンに再順序することで、パrameterの追加コストゼロで優れた速度およびメモリ効率を達成し、1024×1024 FacesHQおよびMS-COCOで最先端のFIDスコアを達成した。

ABSTRACT

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 imes 1024$ and UCF101, MultiModal-CelebA-HQ, and MS COCO $256 imes 256$ . Code will be released at https://taohu.me/zigma/

研究の動機と目的

• 従来のMambaベースの視覚手法では通常、線形な行方向または列方向のスキャン順序が用いられるが、これにより空間的連続性に欠如している問題に対処する。
• ヒューリスティックなスキャンパスを用いて空間的連続性を組み込むことで、2Dおよび3Dの視覚データに対するMambaブロックのインダクティブバイアスを向上させる。
• Stochastic Interpolantフレームワーク内でのMambaベースのアーキテクチャを用いて、スケーラブルで効率的かつ高精度な画像および動画生成を実現する。
• FacesHQ 1024×1024やUCF101といった大規模データセットにおいて、Mambaベースの拡散モデルの実現可能性と優位性を示す。
• 256×256解像度を超えるスケーラビリティを検証し、1024×1024および動画生成を含む高解像度タスクへの応用を検討する。

提案手法

• Mambaのスキャンパスを線形からジグザグに再順序することで、空間的局所性を保持し、位置認識能力を向上させるジグザグMambaブロックを提案する。
• Mambaアーキテクチャを変更せずに、すべてのレイヤーにジグザグスキャンパターンを適用可能な、プラグアンドプレイ型でパrameterを追加しない技術を導入する。
• 2Dのジグザグスキャンを3Dに拡張するため、空間的および時間的シーケンスを分解し、空間優先のスキャン方式を用いて時間的冗長性を活用する。
• Zigzag MambaブロックをDiTスタイルの拡散モデルのバックボーンに統合し、注意メカニズムをMambaブロックに置き換える。
• Stochastic Interpolantフレームワークを用いてモデルを訓練し、拡散、フローマッチング、ノーマライジングフローの目的関数を統合的に訓練可能にする。
• 2×2のパッチ化された潜在変数を用いた潜在拡散設定を採用し、4,096トークンで1024×1024解像度の高解像度生成を実現する。

Figure 1 : Motivation. Our Zigzag Mamba method improves the network’s position-awareness by arranging and rearranging the scan path of Mamba in a heuristic manner.

実験結果

リサーチクエスチョン

• RQ1 標準的なMambaスキャン順序における空間的連続性の欠如が、視覚生成性能に与える影響は何か？
• RQ2 ジグザグスキャンのような単純でパrameterフリーのスキャン再順序戦略が、Mambaベースの視覚モデルを著しく改善できるか？
• RQ3 Zigzag Mambaは、1024×1024 FacesHQ やUCF101といった高解像度画像および動画生成タスクにどの程度スケーラブルに拡張可能か？
• RQ4 Stochastic Interpolantフレームワークは、大規模なMambaベースの拡散モデルの効果的訓練を可能にするか？
• RQ5 3Dシーケンスを2Dおよび1Dのジグザグスキャンに分解することで、動画生成の品質および効率にどのような影響を与えるか？

主な発見

• 32台のA100 GPUを用いて、FacesHQ 1024×1024でFIDスコア26.6を達成し、Bidirectional Mamba（51.1 FID）を上回り、スケーラビリティに優れた性能を示した。
• MS-COCO 256×256では、Zigzag-8バージョンがFID 41.8を達成し、Bidirectional Mamba（60.2 FID）およびZigzag-1（73.1 FID）を著しく上回った。
• UCF101動画生成では、Factorized 3D Zigzag Mambaが16台のA100 GPUでFrame-FID 121.2、FVD 140.1を達成し、Bidirectional Mamba（146.2 Frame-FID、201.1 FVD）を上回った。
• FacesHQ 1024×1024およびMultiModal-CelebA 512×512の可視化結果から、あらゆる解像度で高品質なサンプル再現性を維持していることが示された。
• アブレーションスタディにより、複数のスキャン方式を統合することでインダクティブバイアスの組み込みが向上し、一貫した性能向上が得られることを確認した。
• 大規模な解像度でも、トランスフォーマー基盤のベースラインと比較して、より優れたメモリおよび速度効率を示した。

Figure 2 : ZigMa. Our backbone is structured in L layers, mirroring the style of DiT [ 65 ] . We use the single-scan Mamba block as the primary reasoning module across different patches. To ensure the network is positionally aware, we’ve designed an arrange-rearrange scheme based on the single-scan