[论文解读] An Image is Worth 32 Tokens for Reconstruction and Generation

Qihang Yu, Mark Weber+4 more|arXiv (Cornell University)|Jun 11, 2024

Anatomy and Medical Technology被引用 4

一句话总结

该论文提出 TiTok，一种基于 1D Transformer 的图像分词器，通过利用区域冗余性，仅用 32 个离散标记即可表示 $256\times256$ 的图像，与 2D VQ 基方法相比显著减少了标记数量。其在 ImageNet-256 上实现了 1.97 gFID 的最先进性能，并在 512×512 分辨率下实现高达 74 倍的生成速度提升，证明了紧凑的 1D 潜在序列可实现高效且高保真度的图像重建与生成。

ABSTRACT

Recent advancements in generative models have highlighted the crucial role of image tokenization in the efficient synthesis of high-resolution images. Tokenization, which transforms images into latent representations, reduces computational demands compared to directly processing pixels and enhances the effectiveness and efficiency of the generation process. Prior methods, such as VQGAN, typically utilize 2D latent grids with fixed downsampling factors. However, these 2D tokenizations face challenges in managing the inherent redundancies present in images, where adjacent regions frequently display similarities. To overcome this issue, we introduce Transformer-based 1-Dimensional Tokenizer (TiTok), an innovative approach that tokenizes images into 1D latent sequences. TiTok provides a more compact latent representation, yielding substantially more efficient and effective representations than conventional techniques. For example, a 256 x 256 x 3 image can be reduced to just 32 discrete tokens, a significant reduction from the 256 or 1024 tokens obtained by prior methods. Despite its compact nature, TiTok achieves competitive performance to state-of-the-art approaches. Specifically, using the same generator framework, TiTok attains 1.97 gFID, outperforming MaskGIT baseline significantly by 4.21 at ImageNet 256 x 256 benchmark. The advantages of TiTok become even more significant when it comes to higher resolution. At ImageNet 512 x 512 benchmark, TiTok not only outperforms state-of-the-art diffusion model DiT-XL/2 (gFID 2.74 vs. 3.04), but also reduces the image tokens by 64x, leading to 410x faster generation process. Our best-performing variant can significantly surpasses DiT-XL/2 (gFID 2.13 vs. 3.04) while still generating high-quality samples 74x faster.

研究动机与目标

• 探究图像分词中的 2D 结构约束是否对高质量图像表征是必要的。
• 通过提出 1D 潜在序列表示，解决 2D 分词器在处理图像冗余方面的低效问题。
• 实现更紧凑、高效且语义丰富的图像表征，以支持高分辨率图像生成。
• 证明 1D 分词在重建保真度和生成速度方面可超越 2D 对应方法，尤其在低标记数量下表现更优。
• 通过消融实验与最先进模型（如 MaskGIT 和 DiT-XL/2）的对比，验证 1D 分词的有效性。

提出的方法

• TiTok 使用视觉 Transformer（ViT）编码器处理将图像块展平后与可学习的 1D 潜在标记拼接的结果。
• 向量量化器（VQ）将编码特征映射为离散标记，形成紧凑的 1D 离散潜在代码序列。
• ViT 解码器从掩码后的 1D 标记序列重建原始图像，支持端到端训练与重建损失。
• 模型采用两阶段训练策略，引入代理代码以提升标记质量，尤其在标记数量有限时表现更优。
• 在生成过程中应用掩码策略以模拟自回归生成过程，消融实验表明其对线性和反余弦掩码策略均具有更强鲁棒性。
• 该框架打破了 2D 网格约束，使每个潜在标记能够关注非局部的、语义相关的图像区域，从而提升表征效率。

实验结果

研究问题

• RQ1 1D 潜在序列表征是否能与基于 2D 网格的分词器在图像重建与生成性能上相媲美？
• RQ2 1D 分词是否能有效利用图像冗余，在不损失保真度的前提下减少标记数量？
• RQ3 1D 分词的性能如何随模型规模和标记数量变化，尤其在低标记数量（如 32 或 64）时？
• RQ4 与 2D 对应方法相比，1D 分词在生成速度和吞吐量方面提升程度如何？
• RQ5 两阶段训练中使用代理代码是否对实现 1D 分词的最先进性能至关重要？

主要发现

• TiTok 将 $256\times256$ 图像压缩为仅 32 个离散标记，相比 2D VQGAN 方法减少了 64 倍。
• 在 ImageNet-256 上，TiTok 实现 1.97 gFID，比 MaskGIT 基线提升 4.21 gFID。
• 在 ImageNet-512 上，TiTok 实现 2.13 gFID，优于 DiT-XL/2（3.04 gFID），同时生成速度提升 74 倍。
• 通过解码器微调和更大的码本，TiTok-L-32 在重建任务上达到 2.21 rFID，相比基线提升 5.8 倍。
• TiTok-B-64 的 1D 变体在紧凑潜在空间下优于其 2D 对应模型 10.43 rFID（5.15 vs. 15.58），证明了 1D 设计的优越性。
• 采用代理代码的两阶段训练使 TiTok 实现 1.70 rFID，优于 MaskGIT-VQGAN（2.28 rFID），证明了该训练策略的有效性。