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

📄 arXiv: 2504.17789v2 📥 PDF

作者: 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.CV

发布日期: 2025-04-24 (更新: 2025-04-27)

备注: Project Page: https://ma-xu.github.io/token-shuffle/ Add related works

💡 一句话要点

提出Token-Shuffle，提升自回归模型在高分辨率图像生成中的效率与质量。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 自回归模型 图像生成 高分辨率 Transformer 多模态学习 Token-Shuffle 效率优化

📋 核心要点

  1. 自回归图像生成模型受限于图像tokens数量，导致训练和推理效率低，难以生成高分辨率图像。
  2. Token-Shuffle通过token-shuffle和token-unshuffle操作，在Transformer中减少图像tokens数量，提升效率。
  3. 实验表明，Token-Shuffle在2048x2048分辨率下实现了优秀的图像生成效果，并在GenAI-benchmark中超越了其他模型。

📝 摘要（中文）

自回归（AR）模型在语言生成领域占据主导地位，并逐渐应用于图像合成，但通常被认为不如基于扩散的模型。一个主要的限制是AR模型需要大量的图像tokens，这限制了训练和推理效率以及图像分辨率。为了解决这个问题，我们提出了一种新颖而简单的方法Token-Shuffle，它减少了Transformer中的图像tokens数量。我们的关键见解是多模态大型语言模型（MLLM）中视觉词汇的维度冗余，其中来自视觉编码器的低维视觉代码直接映射到高维语言词汇。利用这一点，我们考虑两个关键操作：token-shuffle，它沿通道维度合并空间局部tokens以减少输入token数量；token-unshuffle，它在Transformer块之后解开推断的tokens，以恢复输出的空间排列。与文本提示联合训练，我们的策略不需要额外的预训练文本编码器，并使MLLM能够以统一的下一个token预测方式支持极高分辨率的图像合成，同时保持高效的训练和推理。我们首次将AR文本到图像生成的边界推向2048x2048的分辨率，并获得了令人满意的生成性能。在GenAI-benchmark中，我们的27亿参数模型在困难提示上实现了0.77的总体得分，优于AR模型LlamaGen 0.18和扩散模型LDM 0.15。大规模人工评估也证明了我们在文本对齐、视觉缺陷和视觉外观方面的卓越图像生成能力。我们希望Token-Shuffle可以作为MLLM中高效高分辨率图像生成的基础设计。

🔬 方法详解

问题定义：现有自回归图像生成模型在处理高分辨率图像时，由于需要大量的图像tokens，导致计算复杂度高，训练和推理效率低下。这限制了模型生成高分辨率图像的能力，并且难以应用于实际场景。现有方法难以在效率和图像质量之间取得平衡。

核心思路：论文的核心思路是利用多模态大型语言模型（MLLM）中视觉词汇的维度冗余性，通过减少图像tokens的数量来提高效率。具体来说，通过token-shuffle操作合并空间局部tokens，减少输入Transformer的tokens数量，从而降低计算复杂度。

技术框架：整体框架包括视觉编码器、Token-Shuffle模块、Transformer和Token-Unshuffle模块。首先，视觉编码器将输入图像编码为低维视觉代码。然后，Token-Shuffle模块合并空间局部tokens，减少tokens数量。接着，Transformer处理压缩后的tokens。最后，Token-Unshuffle模块解开tokens，恢复空间排列，生成最终图像。整个过程以统一的下一个token预测方式进行。

关键创新：最重要的技术创新点是Token-Shuffle模块，它通过合并空间局部tokens来减少输入Transformer的tokens数量，从而降低计算复杂度，提高效率。与现有方法不同，Token-Shuffle不需要额外的预训练文本编码器，并且能够支持极高分辨率的图像合成。

关键设计：Token-Shuffle模块包含token-shuffle和token-unshuffle两个操作。Token-shuffle操作沿通道维度合并空间局部tokens，减少tokens数量。Token-unshuffle操作在Transformer块之后解开推断的tokens，恢复空间排列。论文采用联合训练策略，同时训练视觉编码器、Transformer和Token-Shuffle模块。损失函数包括标准的交叉熵损失，用于下一个token预测。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Token-Shuffle在GenAI-benchmark中，27亿参数模型在困难提示上实现了0.77的总体得分，优于AR模型LlamaGen 0.18和扩散模型LDM 0.15。大规模人工评估也证明了其在文本对齐、视觉缺陷和视觉外观方面的卓越图像生成能力。该模型首次将AR文本到图像生成的边界推向2048x2048的分辨率，并获得了令人满意的生成性能。

🎯 应用场景

Token-Shuffle具有广泛的应用前景，包括高质量图像生成、图像编辑、艺术创作等领域。该技术可以应用于生成逼真的高分辨率图像，例如用于游戏开发、电影制作和广告设计。此外，Token-Shuffle还可以用于图像修复和增强，提高图像质量。未来，该技术有望推动多模态内容生成的发展，实现更智能、更高效的图像处理。

📄 摘要（原文）

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.