TokenFlow：用于多模态理解和生成的统一图像标记器

Paper是字节跳动发表在Arxiv 2024的工作

Paper Title：TokenFlow: Unified Image Tokenizer for Multimodal Understanding and Generation

Code地址

Abstract

我们提出了 TokenFlow，这是一种新颖的统一图像标记器，它弥合了多模态理解和生成之间长期存在的差距。先前的研究尝试使用单个重建目标矢量量化 (VQ) 编码器来统一这两个任务。我们观察到，理解和生成需要完全不同的视觉信息粒度。这导致了关键的权衡，尤其是在多模态理解任务中损害了性能。TokenFlow 通过创新的双码本架构解决了这一挑战，该架构将语义和像素级特征学习分离，同时通过共享映射机制保持它们的对齐。这种设计允许通过共享索引直接访问对理解任务至关重要的高级语义表示和对生成至关重要的细粒度视觉特征。我们广泛的实验证明了 TokenFlow 在多个维度上的优势。利用 TokenFlow，我们首次证明离散视觉输入在理解性能方面可以超越 LLaVA-1.5 13B，平均提高 7.2%。对于图像重建，我们在 384×384 分辨率下实现了 0.63 的强劲 FID 得分。此外，TokenFlow 在自回归图像生成中建立了最先进的性能，在 256×256 分辨率下 GenEval 得分为 0.55，实现了与 SDXL 相当的结果。

图 1. 使用 TokenFlow 实现的多模态理解结果。我们首次证明离散视觉输入在理解性能上可以超越 LLaVA-1.5 13B，平均提升 7.2%。

1. Introduction

大型语言模型 (LLM) 通过其统一的自回归框架彻底改变了自然语言处理，并展示了在不同任务中的卓越能力 [1, 2]。然而，在视觉和语言的多模态领域，感知和生成范式之间存在根本性的分歧。当前的方法通过不同的架构来解决这些问题：多模态理解模型利用视觉编码器和投影层将视觉表征与预训练的 LLM 对齐 [29, 52]，而视觉生成则依赖于基于扩散的方法 [39, 41] 或离散图像标记进行自回归生成 [38, 44, 51, 65]。这种分歧促使人们追求能够理解和生成的统一方法。

GPT-4o [59] 的出现极大地激发了人们对开发更通用的多模态模型的兴趣。 早期统一感知和生成能力的努力 [27, 46] 主要集中在为 LLM 配备扩散模型的功能。然而，这些方法引入了相当大的架构复杂性和计算开销，凸显了对更优雅的统一解决方案的需求。最近的努力探索了一个有希望的方向：使用单一变压器架构在下一个标记预测框架内统一视觉和文本信息 [48, 55]。这种方法依赖于 VQ 编码器将视觉输入转换为可以与文本一起处理的离散标记，从而提供了一个可能更简单、更高效的框架。通过将两种模态视为离散标记序列，该框架可以在单一架构内实现端到端训练。

然而，这种统一方法存在一个根本挑战。多模态理解需要丰富的语义表示来支持复杂的推理，而视觉生成则需要对空间结构和纹理细节进行精确编码。当前的方法主要采用以重建为目标的 VQ 编码器 [13, 73]，这些编码器主要针对重建保真度进行了优化。虽然这种优化使它们非常适合生成任务，但它可能会限制它们捕获对理解任务至关重要的高级语义特征的能力。虽然 Janus [57] 试图通过为理解和生成任务采用单独的编码器来解决这一冲突，但这增加了模型的复杂性，而没有从根本上解决潜在的表示差异。这些限制凸显了该领域的一个关键差距：缺乏能够有效服务于感知和生成目标的统一视觉编码机制。这激发了我们的核心研究问题：一个图像标记器能否得出适合多模态理解和生成的表示？

为了应对这一挑战，我们提出了 TokenFlow，这是一种新颖的统一图像标记器，它通过独特的双流设计弥合了理解和生成之间的差距。关键见解是将语义和像素级特征的学习解耦，同时通过共享索引映射保持它们的对齐。通过将具有语义和像素级相似性的补丁映射到相同的索引，量化的特征可以直接应用于自回归视觉生成和多模态理解。与在单个码本中约束不同特征级别的并发方法 [60] 不同，TokenFlow 的双码本设计支持专门学习，同时通过共享索引保持跨级别相关性。这项创新允许同时访问语义和像素级表示，而不会损害任何一个方面。具体来说，TokenFlow 采用双编码器架构和相应的专门码本。从 CLIP 风格的老师那里学习的语义编码器提供了强大的语义先验，而像素编码器则捕获详细的视觉信息。然后通过最小化语义和像素级距离的加权总和来量化提取的特征，从而创建联合表示空间。

我们的框架表现出了卓越的可扩展性，即使使用超过 130K 个条目的大规模码本，也能保持出色的码本利用率（95% 以上）——在容量和效率方面都大大超越了之前的方法 [13]。TokenFlow 还在 384×384 分辨率下实现了 0.63 的强劲 FID 分数。对于文本到图像的合成，我们在自回归范式中在 256×256 分辨率下建立了 0.55 的全新最佳 GenEval 分数，同时与 EMU3 [55] 和 LlamaGen [44] 等现有方法相比，所需的采样步骤明显减少。在多模态理解基准测试中，TokenFlow 以最小的训练开销实现了全新的最佳性能，平均超过 LLaVA-1.5 13B 7.2%——这是离散视觉输入首次能够超越这一强劲基线。这些结果验证了 TokenFlow 作为统一的视觉标记器的有效性，它弥合了理解和生成任务之间长期存在的差距。

图 2. 使用 TokenFlow 进行视觉生成的结果。我们展示了不同风格、主题和场景的多样化 256×256 结果。

2. Related Work

2.1. Tokenization for Visual Generation.

---

矢量量化 (VQ) 图像标记器在自回归图像生成的最新进展中发挥了至关重要的作用 [28、34、44、51、65]。 [54] 提出了 VQVAE，使用最近的码本条目量化块级特征，并通过重建损失使用编码器-解码器结构学习码本。 VQVAE-2 [40] 通过指数移动平均更新和分层多尺度方法改进了该框架。 VQGAN [13] 通过结合对抗性和感知损失进一步增强了架构，从而产生更精确和更详细的表示。 VQ 标记器的最新进展主要集中在三个方向：提高重建保真度和生成质量 [21、64、73]、提高码本利用率 [64、70、76]，以及探索新架构，如多尺度 VQVAE [25、51] 用于下一尺度图像预测。虽然这些方法在量化后有效地保留了局部细节，但它们往往难以捕捉语义级信息，限制了它们在自回归多模态图像理解任务中的有效性。我们提出的 TokenFlow 通过引入具有共享映射的双码本解决了这一限制，在自回归生成和多模态理解方面都实现了最先进的性能。

2.2. Tokenization for Unified Multimodal Understanding and Generation

---

最近出现了一些努力来弥合多模态理解和生成之间的差距 [23, 48, 55, 57, 60, 62]。Chameleon [48]、EMU3 [55] 和 Show-o [62] 等方法使用 VQ 标记器 [13, 66, 73] 为这两项任务编码图像。然而，这些方法通常需要从头开始进行多模态训练，并且由于其标记特征中的语义表示有限，在视觉感知任务中通常会遭受性能下降。 SEED-LLaMA [23] 引入了一种新颖的 VQ 标记器，结合了高级语义进行理解，并利用 SD [41] 作为生成解码器。Janus [57] 试图通过使用单独的标记器进行理解 [69] 和生成 [44] 来解决模态差距，尽管这会增加模型复杂性，而没有从根本上解决潜在的挑战。并行工作 [60] 提出了一个统一的视觉塔，在预训练期间将离散的视觉特征与文本对齐。然而，他们的方法将低级和高级表示限制在单个流程中，从而限制了下游性能的上限。相比之下，我们的工作认为统一理解和生成的关键在于学习通用映射。通过定义具有共享映射的双码本，TokenFlow 可以灵活地组合低级和高级特征，从而在所有下游任务中实现卓越的性能。

3. Method

3.1. Motivation

---

表 1. LLaVA-1.5 框架内各种视觉编码器对多模态理解的比较 [14, 23, 43]。VQKD 是从 CLIP ViT-B/14 中蒸馏出来的。“Sem.”指的是学习语义级表示的语义编码器，而“Pix.”则表示专注于低级视觉特征的像素级标记器。

将多模态理解和生成统一为一个有凝聚力的下一个标记预测范式需要一个 VQ 标记器来从输入图像中提取索引。虽然传统的 VQ 标记器 [13、54、66、76] 在像素级图像重建方面表现出色，但我们的研究表明，它们在图像理解能力方面存在很大局限性。 我们在 LLaVA-1.5 [29] 框架内利用这些标记器作为特征提取器进行了实验。如表 1 的实验 2-4 所示，这些离散标记器的性能始终落后于连续标记器 CLIP ViT-B/14 [37]。我们认为这种性能差距源于它们的预训练目标，这些目标主要优化以获得更好的低级重建质量。因此，提取的特征主要编码低级信息，缺乏对复杂视觉推理至关重要的语义级理解。

统一理解和生成的另一种直接解决方案是从预训练的 CLIP [8、37、45、69] 中提取离散标记，然后为其配备图像重建功能。如实验 5 所示，从 CLIP ViT-B/14 中提取的 VQKD 与其他离散标记器相比，性能差距显著缩小。我们进一步进行了一项实验，从 VQKD 提取的量化特征重建原始图像。重建后的图像表现出明显的模糊和明显的高频细节丢失，如图 8 所示。我们将此结果归因于 VQKD 编码器的性质，它将语义上接近的补丁映射到相同的码本索引中。如图 4 (a) 所示，它倾向于将具有相同语义含义的图像映射到相同的码本索引，而 VQGAN（图 4 (b)）倾向于将视觉上相似的图像映射到相同的码本索引，优先考虑低级特征而不是语义内容。因此，从 VQKD 聚合的低级不同块中重建细粒度细节变得极具挑战性。 这些观察结果强调了开发一种新颖的标记化方法的必要性，该方法可以有效地处理高级语义理解和低级视觉重建任务。

图 4. (a) VQKD [35]、(b) VQGAN [13] 和 © 我们的 TokenFlow 聚类的图像可视化。VQKD 聚类表现出语义相似性，而 VQGAN 聚类表现出低级相似性（即颜色）。我们的 TokenFlow 可以成功地结合语义和低级相似性。图像聚类的实现细节可以在附录 A.1 中找到。

图 3. TokenFlow 概览。我们将双编码器和码本与共享映射结合起来，从而实现高级语义和低级像素细节的联合优化。对于给定的输入图像，距离 $d_{\text{sem }}$ 和 $d_{\text{pix }}$ 分别从像素级和语义级码本计算得出，最终的码本索引和特征通过最小化加权和 $d_{\text{sem }}+w_{\text{dis }}\cdot d_{\text{pix }}$ 确定。对得到的量化特征进行独立解码以进行语义对齐和图像重建训练，然后将它们连接起来，为理解和生成的下游任务提供统一的表示。

3.2. Unified Image Tokenizer

---

为了弥补这一差距，我们提出了 TokenFlow（图 3），这是一种新颖的统一图像标记器，可在语义和像素级别实现联合表示学习。我们发现统一理解和生成的关键在于学习通用映射。如果标记器可以将高级和低级的块映射到相同的码本索引，那么量化特征就可以轻松解码并直接应用于自回归视觉生成任务和多模态理解任务。

编码器。与以前使用单个编码器提取低级图像信息的方法不同，我们提出了一种双编码器架构，包括语义编码器 ${\mathcal{E}}_{\text{sem }}$ 和像素编码器 ${\mathcal{E}}_{\text{pix }}$。这种设计可以提取两种不同类型的图像特征。对于语义编码器，我们使用预先训练的文本对齐视觉编码器（例如 CLIP ViT-B/14）对其进行初始化。这种初始化策略有助于更好地学习语义码本中的高级文本对齐嵌入，最终增强模型的多模态理解能力。为简洁起见，我们省略了特征表示的空间索引，其中 ${\hat{z}}_{\text{sem }}={\mathcal{E}}_{\text{sem }}(x)\in {\mathbb{R}}^{d_{\mathrm{sem}}}$ 和 ${\hat{z}}_{\text{pix }}={\mathcal{E}}_{\text{pix }}(x)\in {\mathbb{R}}^{d_{\mathrm{pix}}}$ 是来自语义和像素编码器的编码特征。

量化。我们引入了一种采用双码本的创新量化方法：语义级嵌入 Z sem  = { z sem  , i } i = 1 K ∈ R K × d sem  \mathbf{Z}_{\text {sem }}=\left\{z_{\text {sem }, i}\right\}_{i=1}^K \in \mathbb{R}^{K \times d_{\text {sem }}} Zsem ​={zsem ,i​}i=1K​∈RK×dsem ​ 和像素级嵌入 Z p i x = { z p i x , i } i = 1 K ∈ R K × d p i x \mathbf{Z}_{\mathrm{pix}}=\left\{z_{\mathrm{pix}, i}\right\}_{i=1}^K \in \mathbb{R}^{K \times d_{\mathrm{pix}}} Zpix​={zpix,i​}i=1K​∈RK×dpix​，其中 $K$ 是码本条目数。这两个码本共享统一的映射，从而能够在量化过程中同时考虑高级语义信息和低级像素细节。给定编码特征表示 ${\hat{z}}_{\text{sem }}$ 和 ${\hat{z}}_{\text{pix }}$，我们计算它们在 $l_2$-norm [64] 之后到各自码本嵌入的距离：

$$\begin{array}{rl}{d}_{\mathrm{sem},i}& ={\Vert {\hat{z}}_{\mathrm{sem}}-z_{\mathrm{sem},i}\Vert }_2^2,\text{ for }i=1,\dots ,K\\ d_{\mathrm{pix},i}& ={\Vert {\hat{z}}_{\mathrm{pix}}-z_{\mathrm{pix},i}\Vert }_2^2,\text{ for }i=1,\dots ,K\\ i^{\ast }& =\underset{i}{\mathrm{argmin}}\left(d_{\mathrm{sem},i}+w_{\mathrm{dis}}\cdot d_{\mathrm{pix},i}\right)\end{array}$$

最佳量化指标 $i^{\ast }$ 由最小化这两个距离的加权和确定，其中 $w_{\text{dis }}$ 是距离平衡权重，如公式 (3) 所示。这种联合优化方法与以前的 VQ 方法有很大不同，以前的 VQ 方法通常侧重于学习单一特征类型的分布。我们进一步采用多尺度 VQ (MSVQ) 结构 [51] 来增强码本表示的丰富性。我们的共享映射策略使码本能够学习高级语义和低级特征的联合分布，从而带来几个关键优势：

❶ 可扩展性：随着码本大小的增加，我们的方法在生成和理解任务中都表现出一致的性能改进，因为较大的码本大小提供了更多的高级和低级特征组合可能性。在扩展码本大小为 131,072 的情况下，它仍然可以保持超过 95% 的极高利用率，同时实现最佳图像重建质量和多模态理解性能。

❷ 多任务能力：通过学习语义和像素级特征的联合分布，我们的方法弥合了生成任务和理解任务之间的差距。这种统一的表示使单个标记器能够在两个领域中表现出色。这种设计还允许无缝集成更多代码本以嵌入其他类型的特征表示，从而能够在不修改架构的情况下扩展到更多下游任务。

解码器和训练目标。我们的架构包含两个不同的解码器，包括语义解码器 ${\mathcal{D}}_{\text{sem }}$ 和像素解码器 ${\mathcal{D}}_{\text{pix }}$，用于重建语义特征和原始图像。我们使用教师模型 [35]（与语义编码器的初始化相同）进行目标特征提取。语义损失 ${\mathcal{L}}_{\text{sem }}$ 计算为解码特征和教师提取特征之间的 $l_2$ 距离。重建损失公式为：

$${\mathcal{L}}_{\mathrm{pix}}={\ell }_2(x,\hat{x})+{\mathcal{L}}_{\mathrm{P}}(x,\hat{x})+{\lambda }_{\mathrm{G}}{\mathcal{L}}_{\mathrm{G}}(\hat{x})$$

其中 $\hat{x}={\mathcal{D}}_{\text{pix }}(z)，{\ell }_2$ 表示逐像素重建损失，${\mathcal{L}}_{\mathrm{P}}(\cdot )$ 表示使用 LPIPS 的感知损失，${\mathcal{L}}_{\mathrm{G}}(\cdot )$
表示以 ${\lambda }_{\mathrm{G}}$ 为权重系数的对抗损失。按照矢量量化惯例，我们采用直通梯度估计器：$z=\mathrm{sg}[z-\hat{z}]+\hat{z}$，其中 $\mathrm{sg}[\cdot ]$ 表示停止梯度操作。码本学习目标是：${\mathcal{L}}_{\mathrm{VQ}}=\Vert \mathrm{sg}[\hat{z}]-z{\Vert }_2^2+\beta \Vert \hat{z}-\mathrm{sg}[z]{\Vert }_2^2$，其中第二项表示具有平衡因子 $\beta$ 的承诺损失。总训练目标是所有损失的总和：${\mathcal{L}}_{\text{total }}={\mathcal{L}}_{\text{sem }}+{\mathcal{L}}_{\mathrm{VQ}}+{\mathcal{L}}_{\text{pix }}$。

图 5. 我们框架中不同采样策略的定性比较。（a）单次 top-k（k=1200）和 top-p（p=0.8）采样表现出不一致的模式和伪影。（b）我们提出的多步采样策略产生了更连贯、更具视觉吸引力的结果。放大后效果更佳。

3.3. Visual Generation with TokenFlow

---

TokenFlow 帮助我们使用下一代预测范式在自回归文本到图像生成中实现 SOTA 性能。下面，我们详细介绍了高质量图像合成的训练和推理策略。

训练策略。我们的视觉生成架构建立在预先训练的 LLM 模型 [53] 之上。对于文本编码，我们利用模型的本机 BPE 标记器将输入文本转换为离散标记序列并提取特征表示。原始词汇表通过专门的视觉标记进行扩展。我们使用 TokenFlow 提取图像标记，将其传递给 MLP，并将其与文本标记连接起来进行训练。考虑到模型的自回归性质，我们采用专门在图像标记上计算的交叉熵损失。为了在推理过程中启用无分类器指导 [17]，我们在训练期间以概率 $p_{\text{drop }}=0.1$ 随机将条件文本替换为空字符串。遵循 [11, 48, 56]，我们结合 QKnormalization 和规范重新排序来增强训练稳定性并防止损失峰值。

推理策略。我们观察到，传统的 topk-top-p 采样策略在 nextscale 范式中使用时，经常会导致图像崩溃和重复的局部模式。这可以归因于交叉熵训练目标，它主要与 top-1 预测建立基于注意力的关系。在推理过程中对每个 token 进行独立的 top-k 采样可能会导致 token 缺乏直接相关性，从而导致不一致或重复的模式，只能通过后续尺度的注意力来部分补救。这个问题在推理步骤有限的情况下变得更加严重。

为了解决这一根本限制，我们提出了一种新颖的多步采样方法：（i）初始采样：使用参数 $k_1$ 和 $p_1$ 执行 top- $k$ top- $p$ 采样。（ii）细化：使用采样输出作为第二轮采样的输入，以相同的比例减少参数 $k_2<k_1$ 和 $p_2<p_1$。这种采样空间的逐步缩小保持了创造性的多样性，同时通过细化步骤实现了一致性。与单次采样方法相比，实证结果表明，生成的结果更加连贯且更具视觉吸引力（参见图 5 和附录 B.1 中的详细消融）。

3.4. Multimodal Understanding with TokenFlow

---

TokenFlow 充当多尺度 VQ 标记器，其中量化的多尺度特征可以直接输入到预先训练的 LLM 中，以遵循 LLaVA-1.5 [29] 范式进行多模态理解训练。来自双流的联合特征表示作为模型的输入。我们验证了多种特征输入策略：（i）来自所有尺度的特征（ii）仅最终尺度的特征（iii）来自所有尺度的残差特征。我们发现来自最终尺度的特征实现了最佳的整体性能，如附录 B.1 中所述。这表明最终尺度捕获了与多模态理解最相关的语义信息，而额外的尺度特征或残差特征可能会引入影响性能的噪音。我们的模型比现有的离散多模态方法有显着的改进。值得注意的是，可以以最小的计算开销实现性能提升，使用 LLaVA 1.5 训练数据在 8×A100 GPU 上进行不到 24 小时的训练。