Token-Shuffle: В сторону генерации изображений высокого разрешения с использованием авторегрессивных моделей

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

April 24, 2025

Авторы: 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

Аннотация

Авторегрессионные (AR) модели, долгое время доминировавшие в генерации текста, всё чаще применяются для синтеза изображений, однако зачастую считаются менее конкурентоспособными по сравнению с моделями на основе диффузии. Основное ограничение заключается в значительном количестве токенов изображения, требуемых для AR-моделей, что ограничивает как эффективность обучения и вывода, так и разрешение изображений. Для решения этой проблемы мы представляем Token-Shuffle — новый, но простой метод, который сокращает количество токенов изображения в Transformer. Наше ключевое наблюдение заключается в избыточности размерности визуальных словарей в мультимодальных больших языковых моделях (MLLMs), где низкоразмерные визуальные коды из визуального кодера напрямую отображаются на высокоразмерные языковые словари. Используя это, мы рассматриваем две ключевые операции: token-shuffle, которая объединяет локальные в пространстве токены по канальному измерению для уменьшения количества входных токенов, и token-unshuffle, которая разделяет выведенные токены после блоков Transformer для восстановления пространственной структуры на выходе. Совместное обучение с текстовыми подсказками позволяет нашей стратегии обходиться без дополнительного предобученного текстового кодера и позволяет MLLMs поддерживать синтез изображений с чрезвычайно высоким разрешением в едином подходе предсказания следующего токена, сохраняя при этом эффективность обучения и вывода. Впервые мы расширяем границы AR-генерации текста в изображения до разрешения 2048x2048 с впечатляющими результатами генерации. В бенчмарке GenAI наша модель с 2.7 миллиардами параметров достигает общего балла 0.77 на сложных запросах, превосходя AR-модель LlamaGen на 0.18 и диффузионную модель LDM на 0.15. Масштабные человеческие оценки также демонстрируют нашу выдающуюся способность к генерации изображений с точки зрения соответствия тексту, визуальных дефектов и внешнего вида. Мы надеемся, что Token-Shuffle может стать основополагающим подходом для эффективной генерации изображений высокого разрешения в рамках 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.