ZigMa: A DiT-style Zigzag Mamba Diffusion Model
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866913584179576832 |
|---|---|
| author | Hu, Vincent Tao Baumann, Stefan Andreas Gui, Ming Grebenkova, Olga Ma, Pingchuan Schusterbauer, Johannes Ommer, Björn |
| author_facet | Hu, Vincent Tao Baumann, Stefan Andreas Gui, Ming Grebenkova, Olga Ma, Pingchuan Schusterbauer, Johannes Ommer, Björn |
| contents | 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/ |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2403_13802 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | ZigMa: A DiT-style Zigzag Mamba Diffusion Model Hu, Vincent Tao Baumann, Stefan Andreas Gui, Ming Grebenkova, Olga Ma, Pingchuan Schusterbauer, Johannes Ommer, Björn Computer Vision and Pattern Recognition Artificial Intelligence Computation and Language Machine Learning 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/ |
| title | ZigMa: A DiT-style Zigzag Mamba Diffusion Model |
| topic | Computer Vision and Pattern Recognition Artificial Intelligence Computation and Language Machine Learning |
| url | https://arxiv.org/abs/2403.13802 |