Génération d'images autoregressive évolutive avec Mamba.

Scalable Autoregressive Image Generation with Mamba

August 22, 2024

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

cs.AI

Résumé

Nous présentons AiM, un modèle génératif d'images autoregressif (AR) basé sur l'architecture Mamba. AiM utilise Mamba, un modèle d'espace d'états novateur caractérisé par ses performances exceptionnelles pour la modélisation de longues séquences avec une complexité temporelle linéaire, pour remplacer les Transformers couramment utilisés dans les modèles de génération d'images AR, dans le but d'obtenir à la fois une qualité de génération supérieure et une vitesse d'inférence améliorée. Contrairement aux méthodes existantes qui adaptent Mamba pour traiter les signaux bidimensionnels via un balayage multidirectionnel, AiM utilise directement le paradigme de prédiction du jeton suivant pour la génération d'images autoregressive. Cette approche contourne le besoin de modifications étendues pour permettre à Mamba d'apprendre des représentations spatiales 2D. En mettant en œuvre des modifications simples mais ciblées pour les tâches génératives visuelles, nous préservons la structure de base de Mamba, exploitant pleinement ses capacités efficaces de modélisation de longues séquences et de mise à l'échelle. Nous fournissons des modèles AiM à différentes échelles, avec des comptes de paramètres allant de 148M à 1.3B. Sur le banc d'essai ImageNet1K 256*256, notre meilleur modèle AiM atteint un FID de 2.21, surpassant tous les modèles AR existants de comptes de paramètres comparables et démontrant une compétitivité significative par rapport aux modèles de diffusion, avec une vitesse d'inférence de 2 à 10 fois plus rapide. Le code est disponible sur 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