Token-Shuffle : Vers la génération d'images haute résolution avec des modèles autorégressifs

Token-Shuffle: Towards High-Resolution Image Generation with Autoregressive Models

April 24, 2025

Auteurs: Xu Ma, Peize Sun, Haoyu Ma, Hao Tang, Chih-Yao Ma, Jialiang Wang, Kunpeng Li, Xiaoliang Dai, Yujun Shi, Xuan Ju, Yushi Hu, Artsiom Sanakoyeu, Felix Juefei-Xu, Ji Hou, Junjiao Tian, Tao Xu, Tingbo Hou, Yen-Cheng Liu, Zecheng He, Zijian He, Matt Feiszli, Peizhao Zhang, Peter Vajda, Sam Tsai, Yun Fu

cs.AI

Résumé

Les modèles autorégressifs (AR), longtemps dominants dans la génération de langage, sont de plus en plus appliqués à la synthèse d'images mais sont souvent considérés comme moins compétitifs que les modèles basés sur la diffusion. Une limitation majeure réside dans le nombre substantiel de tokens d'image requis par les modèles AR, ce qui contraint à la fois l'efficacité de l'entraînement et de l'inférence, ainsi que la résolution des images. Pour remédier à cela, nous présentons Token-Shuffle, une méthode novatrice mais simple qui réduit le nombre de tokens d'image dans les Transformers. Notre idée clé repose sur la redondance dimensionnelle des vocabulaires visuels dans les Modèles de Langage Multimodaux de Grande Taille (MLLMs), où les codes visuels de faible dimension issus de l'encodeur visuel sont directement mappés sur des vocabulaires linguistiques de haute dimension. En exploitant cela, nous considérons deux opérations clés : le token-shuffle, qui fusionne les tokens locaux spatialement le long de la dimension des canaux pour diminuer le nombre de tokens en entrée, et le token-unshuffle, qui démêle les tokens inférés après les blocs Transformer pour restaurer l'arrangement spatial en sortie. Entraîné conjointement avec des prompts textuels, notre stratégie ne nécessite aucun encodeur de texte pré-entraîné supplémentaire et permet aux MLLMs de supporter la synthèse d'images à très haute résolution de manière unifiée via la prédiction de tokens suivants, tout en maintenant un entraînement et une inférence efficaces. Pour la première fois, nous repoussons les limites de la génération d'images à partir de texte par AR à une résolution de 2048x2048 avec des performances de génération satisfaisantes. Dans le benchmark GenAI, notre modèle de 2,7 milliards de paramètres atteint un score global de 0,77 sur les prompts difficiles, surpassant les modèles AR LlamaGen de 0,18 et les modèles de diffusion LDM de 0,15. Des évaluations humaines exhaustives à grande échelle démontrent également notre capacité prééminente en génération d'images en termes d'alignement textuel, de défauts visuels et d'apparence visuelle. Nous espérons que Token-Shuffle pourra servir de conception de base pour une génération d'images haute résolution efficace au sein des MLLMs.

English

Autoregressive (AR) models, long dominant in language generation, are increasingly applied to image synthesis but are often considered less competitive than Diffusion-based models. A primary limitation is the substantial number of image tokens required for AR models, which constrains both training and inference efficiency, as well as image resolution. To address this, we present Token-Shuffle, a novel yet simple method that reduces the number of image tokens in Transformer. Our key insight is the dimensional redundancy of visual vocabularies in Multimodal Large Language Models (MLLMs), where low-dimensional visual codes from visual encoder are directly mapped to high-dimensional language vocabularies. Leveraging this, we consider two key operations: token-shuffle, which merges spatially local tokens along channel dimension to decrease the input token number, and token-unshuffle, which untangles the inferred tokens after Transformer blocks to restore the spatial arrangement for output. Jointly training with textual prompts, our strategy requires no additional pretrained text-encoder and enables MLLMs to support extremely high-resolution image synthesis in a unified next-token prediction way while maintaining efficient training and inference. For the first time, we push the boundary of AR text-to-image generation to a resolution of 2048x2048 with gratifying generation performance. In GenAI-benchmark, our 2.7B model achieves 0.77 overall score on hard prompts, outperforming AR models LlamaGen by 0.18 and diffusion models LDM by 0.15. Exhaustive large-scale human evaluations also demonstrate our prominent image generation ability in terms of text-alignment, visual flaw, and visual appearance. We hope that Token-Shuffle can serve as a foundational design for efficient high-resolution image generation within MLLMs.