ZigMa : Modèle de diffusion Mamba en zigzag

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

Résumé

Le modèle de diffusion a longtemps été confronté à des problèmes d'évolutivité et de complexité quadratique, en particulier au sein des architectures basées sur les transformers. Dans cette étude, nous cherchons à exploiter la capacité de modélisation de longues séquences d'un modèle d'état appelé Mamba pour étendre son applicabilité à la génération de données visuelles. Tout d'abord, nous identifions une lacune critique dans la plupart des méthodes visuelles actuelles basées sur Mamba, à savoir l'absence de prise en compte de la continuité spatiale dans le schéma de balayage de Mamba. Ensuite, en nous appuyant sur cette observation, nous introduisons une méthode simple, prête à l'emploi et sans paramètre, nommée Zigzag Mamba, qui surpasse les modèles de référence basés sur Mamba et démontre une amélioration de la vitesse et de l'utilisation de la mémoire par rapport aux modèles de référence basés sur les transformers. Enfin, nous intégrons Zigzag Mamba avec le cadre d'interpolation stochastique pour étudier l'évolutivité du modèle sur des ensembles de données visuelles à haute résolution, tels que FacesHQ 1024 × 1024, UCF101, MultiModal-CelebA-HQ et MS COCO 256 × 256. Le code sera disponible à l'adresse suivante : 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/