ZigMa: Modelo de Difusão 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
Resumo
O modelo de difusão tem sido historicamente afetado por problemas de escalabilidade e complexidade quadrática, especialmente em estruturas baseadas em transformers. Neste estudo, buscamos aproveitar a capacidade de modelagem de sequências longas de um Modelo de Espaço de Estados chamado Mamba para estender sua aplicabilidade à geração de dados visuais. Primeiramente, identificamos uma falha crítica na maioria dos métodos visuais atuais baseados em Mamba, especificamente a falta de consideração pela continuidade espacial no esquema de varredura do Mamba. Em segundo lugar, com base nessa percepção, introduzimos um método simples, plug-and-play e sem parâmetros, denominado Zigzag Mamba, que supera as linhas de base baseadas em Mamba e demonstra melhorias na velocidade e utilização de memória em comparação com as linhas de base baseadas em transformers. Por fim, integramos o Zigzag Mamba ao framework Stochastic Interpolant para investigar a escalabilidade do modelo em conjuntos de dados visuais de alta resolução, como FacesHQ 1024x1024, UCF101, MultiModal-CelebA-HQ e MS COCO 256x256. O código será disponibilizado em 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/