ZigMa: Zigzag-Mamba-Diffusionsmodell

ZigMa: Zigzag Mamba Diffusion Model

March 20, 2024

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

cs.AI

Zusammenfassung

Das Diffusionsmodell wurde lange Zeit von Skalierbarkeits- und quadratischen Komplexitätsproblemen geplagt, insbesondere innerhalb von auf Transformatoren basierenden Strukturen. In dieser Studie zielen wir darauf ab, die Fähigkeit zur Modellierung langer Sequenzen eines Zustandsraummodells namens Mamba zu nutzen, um seine Anwendbarkeit auf die Generierung visueller Daten zu erweitern. Zunächst identifizieren wir ein entscheidendes Versäumnis in den meisten aktuellen Mamba-basierten Methoden für die Bildverarbeitung, nämlich das Fehlen der Berücksichtigung räumlicher Kontinuität im Scan-Schema von Mamba. Zweitens bauen wir auf dieser Erkenntnis auf und führen eine einfache, Plug-and-Play-, nullparameter Methode namens Zigzag Mamba ein, die Mamba-basierte Baselines übertrifft und im Vergleich zu auf Transformatoren basierenden Baselines eine verbesserte Geschwindigkeit und Speichernutzung aufweist. Schließlich integrieren wir Zigzag Mamba in das Stochastic Interpolant Framework, um die Skalierbarkeit des Modells auf großen Auflösungsdatensätzen zu untersuchen, wie z.B. FacesHQ 1024mal 1024 und UCF101, MultiModal-CelebA-HQ und MS COCO 256mal 256. Der Code wird unter https://taohu.me/zigma/ veröffentlicht.

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/