Scalable Autoregressive Image Generation with Mamba
Authors: name: Haopeng Li; name: Jinyue Yang; name: Kexin Wang; name: Xuerui Qiu; name: Yuhong Chou; name: Xin Li; name: Guoqi Li
Published: 2024-08-22 09:27:49

Abstract: We introduce AiM, an autoregressive (AR) image generative model based on
Mamba architecture. AiM employs Mamba, a novel state-space model characterized
by its exceptional performance for long-sequence modeling with linear time
complexity, to supplant the commonly utilized Transformers in AR image
generation models, aiming to achieve both superior generation quality and
enhanced inference speed. Unlike existing methods that adapt Mamba to handle
two-dimensional signals via multi-directional scan, AiM directly utilizes the
next-token prediction paradigm for autoregressive image generation. This
approach circumvents the need for extensive modifications to enable Mamba to
learn 2D spatial representations. By implementing straightforward yet
strategically targeted modifications for visual generative tasks, we preserve
Mamba's core structure, fully exploiting its efficient long-sequence modeling
capabilities and scalability. We provide AiM models in various scales, with
parameter counts ranging from 148M to 1.3B. On the ImageNet1K 256*256
benchmark, our best AiM model achieves a FID of 2.21, surpassing all existing
AR models of comparable parameter counts and demonstrating significant
competitiveness against diffusion models, with 2 to 10 times faster inference
speed. Code is available at https://github.com/hp-l33/AiM