TiTok：将图像标记为32个Tokens以实现图像重建和生成

最近在生成模型方面的进展强调了图像标记(Tokenization)在高效生成高分辨率图像中的关键作用。图像标记通过将图像转换为潜在表示，减少了相对于直接处理像素的计算需求，并提高了生成过程的效果和效率。之前的方法，例如VQ-GAN，通常使用具有固定下采样因子的二维(2D)潜在网格。然而，这些2D标记在处理图像中固有的冗余时面临挑战，因为相邻区域经常显示相似性。为了解决这个问题，作者引入了基于Transformer的一维图像标记器：TiTok（Transformer-based 1-Dimensional Tokenizer），这是一种将图像标记为一维(1D)潜在序列的创新方法。TiTok提供了一种更紧凑的潜在表示，比传统技术产生的表示更高效和有效。例如，尺寸为256×256×3的单张图像可以减少到仅32个离散标记(tokens)，相比于之前方法获得的256或1024个标记有显著减少。尽管其表示紧凑，TiTok的性能仍与最先进的方法相媲美。具体来说，使用相同的生成框架，TiTok在ImageNet 256×256基准测试中达到了1.97的gFID，显著超越了MaskGIT基线4.21的分数。TiTok在更高分辨率下的优势更加明显。在ImageNet 512×512基准测试中，TiTok不仅在gFID（2.74对比3.04）上超越了最先进的扩散模型DiT-XL/2，而且将图像标记减少了64倍，使生成过程加快了410倍。作者提出的最佳变体在生成高质量样本的同时，显著超越了DiT-XL/2（gFID 2.13对比3.04），速度快了74倍。

题目：An Image is Worth 32 Tokens for Reconstruction and Generation
作者：Qihang Yu, Mark Weber, Xueqing Deng, Xiaohui Shen, Daniel Creme rs, Liang-Chieh Chen （来自字节跳动、慕尼黑工业大学）
文章地址：https://arxiv.org/abs/2406.07550
代码地址：https://github.com/bytedance/1d-tokenizer
内容整理：张俸玺

引言

近年来，由于Transformer和扩散模型的重大进展，图像生成取得了显著进步。类似于生成语言模型的发展趋势，许多当代图像生成模型的架构都包含标准的图像标记器(tokenizer)和解标记器(de-tokenizer)。这些模型使用标记化的图像表示形式，将连续的原始图像像素转换为离散的潜在空间。潜在空间（例如32 × 32）比原始图像空间（256×256×3）要紧凑得多。它提供了一种压缩但富有表现力的表示形式，从而不仅促进了生成模型的高效训练和推理，还为模型规模的扩大铺平了道路。

虽然图像标记器(Tokenizer)在图像生成工作流程中取得了巨大的成功，但它们在固有设计上遇到了一个基本的限制。这些标记器基于一个假设，即潜在空间应该保留二维结构，以维持潜在分词和图像块之间位置的直接映射。例如，左上角的潜在分词直接对应于左上角的图像块。这限制了标记器有效利用图像中固有的冗余性来构建更压缩的潜在空间的能力。

退一步讲，作者提出一个问题：“二维结构对于图像标记是必要的吗？”为了回答这个问题，作者从多个基于从输入图像中提取高层次信息的图像理解任务中获得灵感——例如图像分类、目标检测、分割和多模态大语言模型。这些任务不需要解标记器(de-tokenizer)，因为它们的输出通常以图像以外的特定结构呈现。换句话说，它们通常将更高层次的一维序列格式化为输出，但仍然能捕捉到所有与任务相关的信息。先前的研究，如目标查询或感知器重采样器，将图像编码为预定数量的标记(Tokens)的一维序列（例如，64个Tokens）。这些Tokens有助于生成诸如边界框或标题的输出。这些方法的成功激励作者在图像重构和生成的背景下，探索更紧凑的一维序列作为图像潜在表示。值得注意的是，高层次和低层次信息的综合对于生成高质量的图像至关重要，这为极其紧凑的潜在表示带来了挑战。

在本文中，作者介绍了一种基于Transformer的框架，该框架设计用于将图像标记为一维离散序列，之后可以通过解标记器(de-tokenizer)解码回图像空间。具体来说，作者提出了基于Transformer的一维标记器：TiTok，它由一个视觉Transformer（ViT）编码器、一个ViT解码器和一个遵循典型向量量化（VQ）模型设计的向量量化器组成。在标记阶段，图像被分割并展平成一系列的图像块，然后与一维序列的潜在标记连接。在ViT编码器的特征编码过程之后，这些潜在分词构建了图像的潜在表示。在向量量化步骤之后，使用ViT解码器从掩码分词序列中重构输入图像。

基于TiTok，作者进行了广泛的实验，以探究一维图像标记的动态。作者的研究考察了潜在空间大小、模型大小、重构保真度和生成质量之间的相互关系。从这一探索中，作者得出了一些见解：

• 增加表示图像的潜在标记数量可以持续改善重构性能，但在超过128个标记后，收益变得微不足道。有趣的是，32个标记就足以实现合理的图像重构。
• 扩大标记器模型的规模显著提高了重构和生成的性能，特别是在标记数量有限（例如32或64个）时，展示了在潜在空间中实现紧凑图像表示的有前途的途径。
• 一维标记打破了先前二维图像标记器中的网格限制，这不仅使每个潜在标记能够重构超出固定图像网格的区域，从而导致更灵活的标记器设计，还学习到了更多高级和语义丰富的图像信息，尤其是在紧凑的潜在空间中。
• 一维标记在生成训练中表现出优越的性能，不仅显著加快了训练和推理速度，还在使用更少标记的情况下，获得了与典型二维标记器相比具有竞争力的FID分数。

鉴于这些发现，作者引入了TiTok系列，包括各种模型规模和潜在空间大小的模型，能够实现高度紧凑的标记，最少只需32个Tokens。作者进一步通过MaskGIT框架确认了该模型在图像生成中的有效性。TiTok被证明能够在图像生成方面实现最先进的性能，同时所需的潜在空间比传统方法小8到64倍，从而在训练和推理阶段显著加速。它还可以生成质量相似或更高的图像，速度却比最先进的扩散模型如DiT快多达410倍，如图2所示。

方法

VQ-VAE

图像标记器通过在潜在空间中提供紧凑的图像表示，在促进生成模型方面起到了关键作用。在本文的讨论范围内，主要关注向量量化（VQ）标记器，因为它在包括但不限于图像和视频生成、大规模预训练和多模态模型等各个领域具有广泛的适用性。

尽管在VQ-VAE上进行了许多改进（例如损失函数、模型架构和量化/码书策略），但其基本工作流程（例如基于二维网格的潜在表示）基本上没有改变。

从2D到1D的Tokenization

尽管现有的VQ模型已经取得了显著成就，但标准工作流程中存在一个显著的限制：潜在表示Z_2D通常被设想为一个静态的二维网格。这种配置固有地假设潜在网格和原始图像块之间存在严格的一对一映射。这一假设限制了VQ模型充分利用图像中存在的冗余（例如相邻块之间的相似性）的能力。此外，这种方法限制了选择潜在大小的灵活性，最常见的配置是f=4、f=8或f=16，对于尺寸为256×256×3的图像，分别产生4096、1024或256个标记。受1D序列表示在解决广泛计算机视觉问题中的成功启发，作者建议使用1D序列作为图像重建和生成的高效和有效的潜在表示，而不是在二维标记化中潜在表示和图像块之间的固定对应关系。

利用TiTok进行图像重建

作者建立了一个基于Transformer的一维图像标记器的新框架: TiTok，利用Vision Transformer(ViT)将图像标记化为一维潜在标记，然后从这些一维潜在标记中重建原始图像。如图3所示，TiTok在标记化和去标记化过程中都使用了标准的ViT（编码器和解码器都是ViT）。

尽管这一概念相对简单，作者强调紧凑型1D图像标记在现有文献中仍然未被充分探索。因此，所提出的TiTok作为一个基础平台，用于探索1D标记和解码在自然图像中的潜力。值得注意的是，虽然可以将二维网格的潜在表示展开成一维序列，但这与作者提出的一维标记器有显著不同，因为隐含的二维网格映射限制仍然存在。

利用TiTok进行图像生成

除了标记器训练所针对的图像重建任务外，作者还评估了其在图像生成中的有效性，并遵循典型的流程。具体来说，作者采用MaskGIT作为生成框架，因为它简单且有效，这使作者可以通过将其VQGAN标记器替换为TiTok来训练MaskGIT模型。作者没有对MaskGIT进行任何其他具体修改，但为了完整性，·作者简要描述了TiTok在整个生成过程中的应用。

图像被预先标记为一维离散标记。在每个训练步骤中，随机比例的潜在标记会被掩码标记替换。然后，双向Transformer将掩码标记序列作为输入，预测这些掩码标记对应的离散标记ID。推理过程包括多个采样步骤，在每个步骤中，Transformer对掩码标记的预测将根据预测置信度进行采样，然后用于更新掩码图像。通过这种方式，图像从充满掩码标记的序列逐步生成到具有生成标记的图像，之后可以解码回像素空间。与自回归模型相比，MaskGIT框架在生成过程中的速度显著提升。

TiTok两阶段训练

现有的VQ模型训练策略：尽管大多数VQ模型遵循简单的公式，它们对训练过程却非常敏感，模型的性能在很大程度上受到采用更有效训练范式的影响。例如，与DALL-E的dVAE相比，VQGAN在ImageNet验证集上的重建FID（rFID）取得了显著提升。这一改进归因于感知损失和对抗损失的进步。此外，MaskGIT现代实现的VQGAN通过精炼的训练技术，在没有架构改进的情况下进一步提升了性能。值得注意的是，这些改进大多在训练阶段（即通过辅助损失）专门应用，并显著影响了模型的效能。考虑到损失函数的复杂性、所涉及的超参数的大量调整，以及最重要的是缺乏公开可用的代码库以供参考或再现，为提出的TiTok建立一个最佳实验设置提出了巨大的挑战，尤其是当目标是先前文献中很少研究的紧凑型1D标记化时。

两阶段训练解决方案：尽管在典型的Taming-VQGAN设置下训练TiTok是可行的，但作者引入了一个两阶段训练范式以进一步提高性能。两阶段训练策略包括“预热”和“解码器微调”阶段。具体而言，在第一个“预热”阶段，作者建议使用现成的MaskGIT-VQGAN模型生成的离散码（称为代理码，Proxy Codes）来训练1D的VQ模型，而不是像现有方法那样直接回归RGB值并采用各种损失函数。这种方法使作者能够绕过复杂的损失函数和GAN架构，从而专注于优化1D标记设置。重要的是，这一修改并不影响TiTok中的标记器和量化器的功能，它们仍然可以完全用于图像标记和解标记；主要的调整只是涉及TiTok解码器输出的处理。具体来说，这一输出包含一组代理码，随后被输入到相同的现成VQGAN解码器中以生成最终的RGB输出。值得注意的是，代理码的引入不同于简单的蒸馏。作者的实验验证了TiTok比MaskGIT-VQGAN显著更好的生成性能。

在使用代理码完成第一阶段训练后，可以选择进行第二阶段“解码器微调”，这一阶段旨在提高重建质量。具体来说，作者保持编码器和量化器冻结，只针对像素空间训练解码器，采用典型的VQGAN训练配方。作者观察到，这种两阶段训练策略显著提高了训练的稳定性和重建图像的质量。

实验

初步实验

实验设置

实验对象：实验使用了分辨率为256×256的图像进行训练，主要通过开源的MaskGIT-VQGAN模型提供的代码进行训练。

图像分块和编码：在图像分块和编码过程中，图像被分割成大小为16×16的块，并与32个潜在Tokens拼接，输入到Vision Transformer（ViT）编码器中。

潜在Tokens数量：实验探讨了不同潜在Tokens数量（从16到256）的影响，主要关注三种模型大小：小型（TiTok-S）、基础（TiTok-B）和大型（TiTok-L）。

实验结论

• 潜在Token数量对重建性能的影响：增加潜在Token数量会提高图像重建性能，但当数量超过128时，性能提升变得不明显。令人惊讶的是，仅32个潜在Token就足以实现合理的图像重建。
• 模型大小对性能的影响：增大编码器和解码器的模型大小显著提升了重建和生成性能，尤其是在潜在Token数量有限（如32或64）的情况下。
• 1D图像标记化的优越性：相比于传统的2D标记方法，1D图像标记不仅打破了2D网格限制，使每个潜在标记能够重建超出固定图像网格的区域，还学习到了更多高层次和语义丰富的图像信息。
• 生成训练的性能：1D标记化在生成训练中表现出色，不仅显著加快了训练和推理速度，还在使用更少Token的情况下取得了与典型2D标记化相当甚至更好的生成性能。

主要实验

本文所提出的TiTok通过牺牲较大的模型尺寸来换取更紧凑的图像潜在尺寸。本文的实验主要关注与现有SOTA技术相比的ImageNet生成基准测试，并在生成性的MaskGIT框架中评估TiTok作为一维图像标记器的表现。

实现细节

文章研究所涉及的TiTok变体有：TiTok-S-128（即包含128个标记的小模型）、TiTok-B-64（即包含64个标记的基础模型）和TiTok-L-32（即包含32个标记的大模型），每个变体的设计目标都是在模型尺寸增加的同时将潜在空间大小减半。对于分辨率为512的情况，作者将潜在尺寸加倍，以确保在更高分辨率下保留更多细节，从而得到TiTok-L-64和TiTok-B-128。在TiTok的最终训练设置中，码本配置为N=4096，训练时长延长至100万次迭代（200个周期）。

作者还采用了“解码器微调”阶段来进一步提升模型性能，此阶段中编码器和量化器保持冻结状态，解码器微调500,000次迭代。对于生成模型的训练，作者采用了MaskGIT框架，除了采用了arccos掩码调度之外没有进行任何特定修改。所有其他参数与之前的设置相同。

主要结果

表1和表2中总结了分辨率为256×256和512×512的ImageNet-1K生成基准测试结果。

对于表1中的ImageNet(256×256)结果，TiTok在使用比其他VQ模型更少数量的潜在标记情况下，可以实现类似水平的重建FID (rFID)。具体来说，仅使用32个标记，TiTok-L-32达到了2.21的rFID，与经过良好训练的MaskGIT的VQGAN（rFID为2.28）相当，而其潜在表示大小小了8倍。此外，使用相同的生成框架和相同的采样步骤，TiTok-L-32在gFID上大幅优于MaskGIT（从6.18提高到2.77），展示了紧凑1D标记在更有效生成器训练中的优势。与其他基于扩散的生成模型相比，TiTok在享受超过100倍采样速度提升的同时也能实现具有竞争力的性能。

具体来说，TiTok-L-32的gFID优于LDM-4（2.77对比3.60），生成图像速度快了254倍（101.6样本/秒对比0.4样本/秒）。表现最好的变体TiTok-S-128在gFID上优于最先进的扩散方法DiT-XL/2（1.97对比2.27），速度提升了13倍。

对于表2中的ImageNet(512×512)结果，TiTok在使用更少数量的潜在标记情况下仍能保持较好的rFID，特别是在高压缩比的情况下。与基线MaskGIT相比，所有TiTok变体的表现大幅优于其基线。与基于扩散的模型相比，TiTok-L-64的表现优于DiT-XL/2（2.74对比3.04），速度快了410倍。表现最好的变体TiTok-B-128在gFID上显著优于DiT-XL/2（2.13对比3.04），同时生成高质量样本的速度快了74倍。

消融实验

表3中报告了关于最终模型设计的消融研究。

具体来说，在表3(a)中对标记器设计在图像重建上的效果进行了消融研究。从基线模型TiTok-L-32开始，其rFID为6.59。使用更大的码本尺寸将rFID提高了0.74，而将训练迭代次数从100个周期增加到200个周期则进一步将rFID提高了0.37。在此基础上，“解码器微调”（第二阶段训练策略）可以显著提高整体重建性能，rFID达到2.21。

在表3(b)中研究了不同掩码调度对TiTok在MaskGIT中的效果。与原始MaskGIT设置经验发现余弦掩码调度显著优于其他调度不同，作者观察到配备TiTok的MaskGIT更倾向于arccos或线性调度。此外，与报告的根掩码调度表现远不如其他调度不同，作者发现TiTok对不同的掩码调度具有相当的鲁棒性。作者将这些观察归因于TiTok相较于2D-VQGAN能提供更紧凑且更具语义意义的标记，因为相较于余弦掩码调度，线性和arccos调度在早期步骤中具有较低的掩码比例。这与通常冗余信号的掩码比例较高（例如图像中的75%掩码比例）而语义意义输入的掩码比例相对较低（例如语言中的15%掩码比例）的观察结果一致。

在表3(c)中研究了训练范式的影响。从Taming-VQGAN的训练设置开始，其中TiTok-B-64在相同训练设置下获得了5.15的rFID，优于原始2D Taming-VQGAN的7.94 rFID。通过构建TiTok-B64的2D变体展示了1D标记的必要性，其中架构保持不变，只是使用图像块而不是潜在标记作为图像表示。结果显示，2D变体的性能大大下降（15.58对比5.15的rFID），因为2D标记中的固定对应关系限制了在紧凑潜在空间下的合理重建。这一结果证明了所提出的1D标记的有效性，尤其是在更紧凑的潜在空间下。

尽管TiTok在简单的单阶段训练下可以取得相当好的性能，但由于缺乏强大的训练配方（没有公开参考或访问），与MaskGIT-VQGAN相比仍存在性能差距。因此，作者采用了代理码的两阶段训练，这被证明是有效的，并且可以超越MaskGIT-VQGAN（1.70对比2.28的rFID）。值得注意的是，两阶段训练对于获得合理的1D标记器并不是至关重要的，作者相信TiTok在简单的单阶段Taming-VQGAN训练设置下，通过在更大规模的数据集上训练，也能受益。

结论

在本文中，作者探索了一种用于重建和生成自然图像的紧凑一维标记化方法TiTok。不同于现有将图像潜在空间视为二维网格的VQ模型，作者提出了一种更紧凑的方案，将图像标记为一维潜在序列。所提出的TiTok可以用比常用的二维标记器少8到64倍的标记来表示图像。此外，紧凑的一维标记不仅显著提高了生成模型的训练和推理吞吐量，还在ImageNet基准测试中取得了具有竞争力的FID分数。