ZigMa: Zigzag Mamba Diffusiemodel
ZigMa: Zigzag Mamba Diffusion Model
March 20, 2024
Auteurs: Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Fischer, Bjorn Ommer
cs.AI
Samenvatting
Het diffusiemodel wordt al lang geplaagd door schaalbaarheidsproblemen en kwadratische complexiteit, vooral binnen transformer-gebaseerde structuren. In deze studie streven we ernaar om de mogelijkheid van langere sequentiemodellering van een State-Space Model genaamd Mamba te benutten om de toepasbaarheid ervan uit te breiden naar visuele datageneratie. Ten eerste identificeren we een kritisch overzicht in de meeste huidige Mamba-gebaseerde visiemethoden, namelijk het gebrek aan aandacht voor ruimtelijke continuïteit in het scanschema van Mamba. Ten tweede, voortbouwend op dit inzicht, introduceren we een eenvoudige, plug-and-play, nul-parameter methode genaamd Zigzag Mamba, die de Mamba-gebaseerde basislijnen overtreft en verbeterde snelheid en geheugenbenutting laat zien in vergelijking met transformer-gebaseerde basislijnen. Tot slot integreren we Zigzag Mamba met het Stochastic Interpolant-framework om de schaalbaarheid van het model te onderzoeken op grote-resolutie visuele datasets, zoals FacesHQ 1024x1024 en UCF101, MultiModal-CelebA-HQ, en MS COCO 256x256. De code zal worden vrijgegeven op 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/