[논문 리뷰] ZigMa: A DiT-style Zigzag Mamba Diffusion Model

Vincent Tao Hu, Stefan Andreas Baumann외 5명|arXiv (Cornell University)|2024. 03. 20.

Cellular Automata and Applications인용 수 4

한 줄 요약

ZigMa는 2D 및 3D 시각 생성에서 공간 연속성을 향상시키기 위해 주목사용 메커니즘을 새로운 '지그재그 맘바 블록'으로 대체하는 DiT 기반 확산 모델을 제안한다. 맘바의 스캔 경로를 지그재그 패턴으로 재정렬함으로써, 파rameter 오 overhead 가 없이도 초당 성능과 메모리 효율성이 뛰어나 1024×1024 FacesHQ 및 MS-COCO에서 최신 기준(FID) 성능을 달성한다.

ABSTRACT

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

연구 동기 및 목표

• 기존 맘바 기반 시각 방법에서 일반적으로 선형 행 또는 열 방향 스캔 순서를 사용함으로써 발생하는 공간 연속성 부족 문제를 해결한다.
• 경로 기반 스캔 전략을 통해 맘바 블록의 인덕티브 바이어스를 2D 및 3D 시각 데이터에 맞게 향상시킨다.
• 스토케스틱 인터폴란트 프레임워크 내에서 맘바 기반 아키텍처를 활용해 확장 가능하고 효율적이며 고해상도의 이미지 및 영상 생성을 가능하게 한다.
• FacesHQ 1024×1024 및 UCF101과 같은 대규모 데이터셋에서 맘바 기반 확산 모델의 실현 가능성과 우수성을 입증한다.
• 256×256 해상도를 초월해 1024×1024 및 영상 생성을 포함한 스케일러블한 확장 가능성에 대해 탐구한다.

제안 방법

• 맘바의 스캔 경로를 선형에서 지그재그로 재정렬함으로써 공간 국소성과 위치 인식 능력을 유지하는 지그재그 맘바 블록을 제안한다.
• 맘바 아키텍처를 수정하지 않고도 모든 레이어에 지그재그 스캔 패턴을 적용할 수 있는 플러그 앤 플레이 형식의 0 파라미터 기법을 도입한다.
• 공간과 시간 시퀀스를 분해함으로써 2D 지그재그 스캔을 3D로 확장하고, 시간적 冗 redundancy 를 활용하기 위해 공간 우선 스캔 체계를 사용한다.
• 주목사용 메커니즘 대신 맘바 블록을 사용하는 DiT 기반 확산 모델의 백본에 지그재그 맘바 블록을 통합한다.
• 스토케스틱 인터폴란트 프레임워크를 사용해 확산, 플로우 매칭, 노멀라이징 플로우 목표를 통합하여 훈련한다.
• 2×2 패치화된 잠재변수를 사용하는 잠재 확산 설정을 통해 1024×1024에서 4,096개 토큰으로 고해상도 생성을 가능하게 한다.

Figure 1 : Motivation. Our Zigzag Mamba method improves the network’s position-awareness by arranging and rearranging the scan path of Mamba in a heuristic manner.

실험 결과

연구 질문

• RQ1 표준 맘바 스캔 순서에서의 공간 연속성 부족은 시각 생성 성능에 어떤 영향을 미치는가?
• RQ2 지그재그 스캔과 같은 단순한 파라미터 없는 스캔 재정렬 전략이 맘바 기반 시각 모델의 성능을 크게 향상시킬 수 있는가?
• RQ3 지그재그 맘바는 1024×1024 FacesHQ 및 UCF101과 같은 고해상도 이미지 및 영상 생성 작업에 얼마나 잘 스케일링될 수 있는가?
• RQ4 스토케스틱 인터폴란트 프레임워크는 맘바 기반 확산 모델의 대규모 훈련을 효과적으로 가능하게 하는가?
• RQ5 3D 시퀀스를 2D 및 1D 지그재그 스캔으로 분해함으로써 영상 생성 품질과 효율성에 어떤 영향을 미치는가?

주요 결과

• 32개의 A100 GPU를 사용하여 FacesHQ 1024×1024에서 지그재그 맘바는 FID 점수 26.6을 달성했으며, 양방향 맘바(51.1 FID)를 능가하며 뛰어난 확장성과 성능을 보였다.
• MS-COCO 256×256에서 지그재그-8 버전은 FID 41.8을 기록했으며, 양방향 맘바(60.2 FID)와 지그재그-1(73.1 FID)를 크게 앞서며 뛰어난 성능을 보였다.
• UCF101 영상 생성에서 인과적 3D 지그재그 맘바는 16개의 A100 GPU를 사용해 프레임-FID 121.2, FVD 140.1을 기록했으며, 양방향 맘바(146.2 프레임-FID, 201.1 FVD)를 초월했다.
• FacesHQ 1024×1024 및 MultiModal-CelebA 512×512의 시각화 결과에서 다양한 해상도에서 높은 품질의 샘플 유지 능력을 입증했다.
• 제거 분석 결과, 다수의 스캔 전략을 통합함으로써 인덕티브 바이어스의 통합이 향상되고 일관된 성능 향상이 이루어짐을 확인했다.
• 높은 해상도에서도 트랜스포머 기반 베이스라인과 비교해 메모리 및 속도 효율성이 뛰어나다는 점을 입증했다.

Figure 2 : ZigMa. Our backbone is structured in L layers, mirroring the style of DiT [ 65 ] . We use the single-scan Mamba block as the primary reasoning module across different patches. To ensure the network is positionally aware, we’ve designed an arrange-rearrange scheme based on the single-scan