DiM: 고해상도 이미지 합성을 위한 효율적 확산 마바

DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

May 23, 2024

저자: Yao Teng, Yue Wu, Han Shi, Xuefei Ning, Guohao Dai, Yu Wang, Zhenguo Li, Xihui Liu

cs.AI

초록

디퓨전 모델은 이미지 생성 분야에서 큰 성공을 거두었으며, 그 백본이 U-Net에서 Vision Transformer로 진화해 왔습니다. 그러나 Transformer의 계산 비용은 토큰 수에 대해 2차적으로 증가하기 때문에, 고해상도 이미지를 다룰 때 상당한 어려움을 야기합니다. 본 연구에서는 State Space Model(SSM) 기반의 시퀀스 모델인 Mamba의 효율성과 디퓨전 모델의 표현력을 결합한 Diffusion Mamba(DiM)를 제안하여, 고해상도 이미지 합성을 효율적으로 수행합니다. Mamba가 2D 신호로 일반화되지 않는 문제를 해결하기 위해, 우리는 다방향 스캔, 각 행과 열 끝에 학습 가능한 패딩 토큰, 경량화된 지역 특징 강화 등 여러 아키텍처 설계를 도입했습니다. 우리의 DiM 아키텍처는 고해상도 이미지에 대한 추론 시간 효율성을 달성합니다. 또한, DiM을 사용한 고해상도 이미지 생성의 학습 효율성을 더욱 개선하기 위해, 저해상도 이미지(256×256)에서 DiM을 사전 학습한 후 고해상도 이미지(512×512)에서 미세 조정하는 "약한 학습에서 강한 학습으로(weak-to-strong)" 전략을 탐구했습니다. 더 나아가, 추가 미세 조정 없이 더 높은 해상도의 이미지(예: 1024×1024 및 1536×1536)를 생성할 수 있도록 학습 없는 업샘플링 전략을 탐구했습니다. 실험을 통해 우리의 DiM의 효과성과 효율성을 입증했습니다.

English

Diffusion models have achieved great success in image generation, with the backbone evolving from U-Net to Vision Transformers. However, the computational cost of Transformers is quadratic to the number of tokens, leading to significant challenges when dealing with high-resolution images. In this work, we propose Diffusion Mamba (DiM), which combines the efficiency of Mamba, a sequence model based on State Space Models (SSM), with the expressive power of diffusion models for efficient high-resolution image synthesis. To address the challenge that Mamba cannot generalize to 2D signals, we make several architecture designs including multi-directional scans, learnable padding tokens at the end of each row and column, and lightweight local feature enhancement. Our DiM architecture achieves inference-time efficiency for high-resolution images. In addition, to further improve training efficiency for high-resolution image generation with DiM, we investigate ``weak-to-strong'' training strategy that pretrains DiM on low-resolution images (256times 256) and then finetune it on high-resolution images (512 times 512). We further explore training-free upsampling strategies to enable the model to generate higher-resolution images (e.g., 1024times 1024 and 1536times 1536) without further fine-tuning. Experiments demonstrate the effectiveness and efficiency of our DiM.