ZigMa: A DiT-style Zigzag Mamba Diffusion Model

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 $1024\times 1024$ and UCF101, MultiModal-CelebA-HQ, and MS COCO $256\times 256$ . Code will be released at https://taohu.me/zigma/

한 줄 요약

ZigMa는 시각 데이터 생성을 위한 Mamba 기반 디퓨전 모델로, Zigzag 스캔 방식을 통해 성능과 메모리 효율성을 동시에 향상시킨다.

핵심 기여도

Zigzag Mamba는 기존 Mamba 기반 시각 모델에서 간과되었던 공간 연속성(Spatial Continuity) 문제를 해결한다.
Zigzag 스캔 방식은 추가 파라미터 없이도, 1~8개의 스캔 경로를 활용하여 512×512 이상의 고해상도 이미지에서 성능 향상(+1.6× 가속)을 보인다.
Zigzag Mamba는 1024×1024 해상도의 FacesHQ, UCF101, MultiModal-CelebA-HQ, MS COCO 256×256 데이터셋에서 Stochastic Interpolant 프레임워크와 결합하여 확장성을 검증한다.
3D 비디오 데이터에도 공간-시간 분해를 통해 Zigzag Mamba를 확장한다.

핵심 아이디어

기존 Mamba 기반 시각 모델은 2D 이미지에 적용할 때 행-열 기반 스캔 방식을 사용하여 공간 연속성을 무시하는 문제가 있었다. Zigzag Mamba는 이 문제를 해결하기 위해 2D 이미지의 공간적 연속성을 반영한 Zigzag 스캔 방식을 도입한다. 이 방식은 다양한 스캔 경로(1~8개)를 레이어 단위로 교차 적용하여, 시각 데이터의 인덕티브 바이어스를 최대한 활용한다. 특히, 스캔 경로 수가 증가할수록 512×512 이상의 고해상도 이미지에서 더 큰 성능 향상이 나타난다. Zigzag 스캔은 파라미터 증가 없이도 메모리와 속도 부담을 줄이며, 이는 기존의 병렬/양방향 Mamba와 비교해 큰 장점이다.

기술적 접근법

**Zigzag Mamba**: 2D 이미지에 적용된 Mamba 블록으로, 공간 연속성을 고려한 Zigzag 스캔 방식을 사용. 스캔 경로 수(Order Receptive Field)는 1~8개까지 조절 가능.
**Position Embedding**: 학습 가능한 임베딩이 Sinusoidal 임베딩보다 성능 향상에 기여함.
**3D 확장**: 공간-시간 분해를 통해 3D 비디오 데이터에도 적용 가능.
**Stochastic Interpolant**: 디퓨전, 플로우 매칭, 슈뢰딩거 브리지 등 다양한 생성 모델을 통합한 프레임워크로, 1024×1024 해상도의 이미지와 비디오에서 확장성을 검증.
**성능 비교**: U-ViT, DiT와 비교하여 FPS와 GPU 메모리 사용량에서 우수한 성능을 보임. 특히, 196×196 패치 크기에서도 메모리 초과 없이 처리 가능.

주요 결과

**MultiModal-CelebA-HQ**: Zigzag-8 스캔 방식에서 Zigzag-1 대비 성능 향상이 512×512 해상도에서 더 두드러짐.
**FPS/GPU-Memory**: 196×196 패치 크기에서 U-ViT는 메모리 초과 발생, DiT는 ZigMa와 유사한 GPU 사용량을 보임.
**Order Receptive Field**: Zigzag Mamba는 Order Receptive Field가 8일 때도 파라미터 증가 없이 1.6× 가속 효과를 보임.
**비교 성능**: Transformer 기반 DiT와 유사한 성능을 보이지만, FLOPS와 메모리 소모가 훨씬 적음.

의의 및 한계

Zigzag Mamba는 Mamba 기반 모델의 공간 연속성 문제를 해결하며, 고해상도 이미지와 3D 비디오 생성에 확장 가능한 구조를 제시한다. 특히, 추가 파라미터 없이도 성능과 메모리 효율성을 동시에 개선한 점이 학술적·실용적으로 의미가 크다. 그러나, Zigzag 스캔 방식의 최적 스캔 경로 수나 다양한 데이터셋에서의 일반화 가능성은 추가 연구가 필요하다. 또한, 3D 비디오 데이터에 대한 실험은 아직 초기 단계로, 더 많은 평가가 필요하다.

실용적 활용

Zigzag Mamba는 고해상도 이미지 생성, 3D 비디오 생성, 대규모 시각 데이터 처리 등에서 메모리 효율적인 디퓨전 모델로 활용 가능하다. 특히, 대규모 이미지 생성 시 GPU 메모리 제약이 큰 산업 현장에서 유용할 것으로 기대된다.