ZigMa: A DiT-Style Zigzag Mamba Diffusion Model
MCML Authors
Vincent Tao Hu
Dr.
* Former Member
Olga Grebenkova
Pingchuan Ma
Johannes Schusterbauer
Björn Ommer
Prof. Dr.
Core PI
Abstract
Vincent Tao Hu
Dr.
* Former Member
Olga Grebenkova
Pingchuan Ma
Johannes Schusterbauer
Björn Ommer
Prof. Dr.
Core PI
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 Zigzag Mamba, a simple, plug-and-play, minimal-parameter burden, DiT style solution, which outperforms Mamba-based baselines and demonstrates improved speed and memory utilization compared to transformer-based baselines, also this heterogeneous layerwise scan enables zero memory and speed burden when we consider more scan paths. 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 and UCF101, MultiModal-CelebA-HQ, and MS COCO .
inproceedings HBG+24
ECCV 2024
18th European Conference on Computer Vision. Milano, Italy, Sep 29-Oct 04, 2024.
Authors
V. T. Hu • S. A. Baumann • M. Gui • O. Grebenkova • P. Ma • J. Schusterbauer • B. OmmerLinks
DOI GitHubResearch Area
BibTeXKey: HBG+24