ZigMa: Zigzag Mamba Diffusion Model
Authors: name: Vincent Tao Hu; name: Stefan Andreas Baumann; name: Ming Gui; name: Olga Grebenkova; name: Pingchuan Ma; name: Johannes Fischer; name: Bjorn Ommer
Published: 2024-03-20 17:59:14

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