ZigMa: Modello di Diffusione Mamba a Zigzag
ZigMa: Zigzag Mamba Diffusion Model
March 20, 2024
Autori: Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Fischer, Bjorn Ommer
cs.AI
Abstract
Il modello di diffusione è stato a lungo afflitto da problemi di scalabilità e complessità quadratica, specialmente all'interno di strutture basate su transformer. In questo studio, ci proponiamo di sfruttare la capacità di modellazione di sequenze lunghe di un modello State-Space chiamato Mamba per estenderne l'applicabilità alla generazione di dati visivi. In primo luogo, identifichiamo una critica lacuna nella maggior parte degli attuali metodi visivi basati su Mamba, ovvero la mancanza di considerazione per la continuità spaziale nello schema di scansione di Mamba. In secondo luogo, basandoci su questa intuizione, introduciamo un metodo semplice, plug-and-play e a zero parametri denominato Zigzag Mamba, che supera le baseline basate su Mamba e dimostra un miglioramento nella velocità e nell'utilizzo della memoria rispetto alle baseline basate su transformer. Infine, integriamo Zigzag Mamba con il framework Stochastic Interpolant per investigare la scalabilità del modello su dataset visivi ad alta risoluzione, come FacesHQ 1024x1024, UCF101, MultiModal-CelebA-HQ e MS COCO 256x256. Il codice sarà rilasciato all'indirizzo 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/