# ZigMa: A DiT-style Zigzag Mamba Diffusion Model

Preprint · Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Schusterbauer, Björn Ommer

- Archive id: arxiv:2403.13802
- Published: 2024-03-20
- arXiv: https://arxiv.org/abs/2403.13802
- Source URL: https://arxiv.org/abs/2403.13802
- PDF: https://arxiv.org/pdf/2403.13802
- Provenance: arxiv

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

_Abstract reproduced under CC0-1.0 (arXiv descriptive metadata)._

---
Space Kitty Archive · https://space-kitty.com · API: https://space-kitty.com/api/v1 · MCP: https://space-kitty.com/mcp · Premium: $0.20/request via x402 (Base Sepolia (testnet)). Metadata: arXiv (CC0), OpenAlex (CC0), Europe PMC (CC BY/CC0 works), ClinicalTrials.gov.