ZigMa: Modelo de Difusión Mamba Zigzag

ZigMa: Zigzag Mamba Diffusion Model

March 20, 2024

Autores: Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Fischer, Bjorn Ommer

cs.AI

Resumen

El modelo de difusión ha enfrentado durante mucho tiempo problemas de escalabilidad y complejidad cuadrática, especialmente dentro de estructuras basadas en transformadores. En este estudio, buscamos aprovechar la capacidad de modelado de secuencias largas de un Modelo de Espacio de Estados llamado Mamba para extender su aplicabilidad a la generación de datos visuales. En primer lugar, identificamos una omisión crítica en la mayoría de los métodos visuales actuales basados en Mamba, específicamente la falta de consideración de la continuidad espacial en el esquema de escaneo de Mamba. En segundo lugar, basándonos en esta observación, introducimos un método simple, plug-and-play y sin parámetros llamado Zigzag Mamba, que supera a las líneas base basadas en Mamba y demuestra una mejora en la velocidad y la utilización de memoria en comparación con las líneas base basadas en transformadores. Por último, integramos Zigzag Mamba con el marco de Interpolante Estocástico para investigar la escalabilidad del modelo en conjuntos de datos visuales de alta resolución, como FacesHQ 1024x1024, UCF101, MultiModal-CelebA-HQ y MS COCO 256x256. El código se publicará en https://taohu.me/zigma/.

English

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/