Geração de Imagens Autoregressiva Escalável com Mamba

Scalable Autoregressive Image Generation with Mamba

August 22, 2024

Autores: Haopeng Li, Jinyue Yang, Kexin Wang, Xuerui Qiu, Yuhong Chou, Xin Li, Guoqi Li

cs.AI

Resumo

Apresentamos AiM, um modelo generativo de imagens autoregressivo (AR) baseado na arquitetura Mamba. AiM utiliza Mamba, um modelo de espaço de estados inovador caracterizado por seu desempenho excepcional para modelagem de sequências longas com complexidade de tempo linear, para substituir os Transformadores comumente utilizados em modelos de geração de imagens AR, com o objetivo de alcançar tanto uma qualidade de geração superior quanto uma velocidade de inferência aprimorada. Ao contrário dos métodos existentes que adaptam o Mamba para lidar com sinais bidimensionais por meio de varredura multidirecional, AiM utiliza diretamente o paradigma de previsão do próximo token para geração de imagens autoregressiva. Esta abordagem contorna a necessidade de extensas modificações para permitir que o Mamba aprenda representações espaciais 2D. Ao implementar modificações simples, porém estrategicamente direcionadas para tarefas generativas visuais, preservamos a estrutura central do Mamba, explorando totalmente suas eficientes capacidades de modelagem de sequências longas e escalabilidade. Fornecemos modelos AiM em várias escalas, com contagens de parâmetros variando de 148M a 1.3B. No benchmark ImageNet1K 256*256, nosso melhor modelo AiM alcança um FID de 2.21, superando todos os modelos AR existentes com contagens de parâmetros comparáveis e demonstrando uma competitividade significativa em relação aos modelos de difusão, com uma velocidade de inferência de 2 a 10 vezes mais rápida. O código está disponível em https://github.com/hp-l33/AiM

English

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