ZigMa: Модель диффузии Зигзаг Мамба

ZigMa: Zigzag Mamba Diffusion Model

March 20, 2024

Авторы: Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Fischer, Bjorn Ommer

cs.AI

Аннотация

Модель диффузии долгое время сталкивалась с проблемами масштабируемости и квадратичной сложности, особенно в структурах на основе трансформеров. В данном исследовании мы стремимся использовать возможность моделирования длинных последовательностей в модели пространства состояний под названием Mamba для расширения ее применимости к генерации визуальных данных. Во-первых, мы выявляем критическое упущение в большинстве существующих методов обработки изображений на основе Mamba, а именно отсутствие учета пространственной непрерывности в схеме сканирования Mamba. Во-вторых, опираясь на это открытие, мы представляем простой, готовый к использованию метод без параметров под названием Zigzag Mamba, который превосходит базовые методы на основе Mamba и демонстрирует улучшенную скорость и использование памяти по сравнению с базовыми методами на основе трансформеров. Наконец, мы интегрируем Zigzag Mamba с каркасом стохастического интерполянта для изучения масштабируемости модели на крупномасштабных визуальных наборах данных, таких как FacesHQ 1024 на 1024 и UCF101, MultiModal-CelebA-HQ и MS COCO 256 на 256. Код будет опубликован на 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/