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를 Stochastic Interpolant 프레임워크와 통합하여 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/