基于 Mamba 的可扩展自回归图像生成
Scalable Autoregressive Image Generation with Mamba
AI总结:
本文提出 AiM,一种基于 Mamba 架构的自回归图像生成模型,通过直接采用 next-token prediction 范式并保留 Mamba 核心结构,在 ImageNet1K 基准上以更快推理速度超越了同等参数量的现有 AR 模型。
摘要:
我们引入了 AiM,一种基于 Mamba 架构的自回归(AR)图像生成模型。AiM 采用 Mamba(一种新型状态空间模型,以其在长序列建模中线性时间复杂度的卓越性能为特征),来取代 AR 图像生成模型中常用的 Transformers,旨在实现更优的生成质量和更快的推理速度。与现有方法通过多向扫描使 Mamba 适应二维信号不同,AiM 直接利用 next-token prediction 范式进行自回归图像生成。这种方法避免了对 Mamba 进行大量修改以使其学习 2D 空间表示的需求。通过对视觉生成任务实施简单但具有战略针对性的修改,我们保留了 Mamba 的核心结构,充分利用其高效的长序列建模能力和可扩展性。我们提供了各种规模的 AiM 模型,参数量从 148M 到 1.3B 不等。在 ImageNet1K 256*256 基准上,我们最佳的 AiM 模型达到了 2.21 的 FID,超越了所有同等参数量的现有 AR 模型,并展现出与扩散模型相比的显著竞争力,且推理速度快 2 到 10 倍。代码可在 https://github.com/hp-l33/AiM 获取。
英文摘要:
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
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。