端到端视觉 token 化器调优-CSDN博客

本文链接：https://blog.csdn.net/yorkhunter/article/details/148356940

25年5月来自自动化所、中科院大学、北京智源、大连理工和清华的论文“End-to-End Vision Tokenizer Tuning”。

现有的视觉 token 化方法将视觉 token 化器的优化与下游训练分离开来，隐式地假设视觉 token 可以推广到各种任务，例如图像生成和视觉问答。针对低级重建而优化的视觉 token 化器与需要各种表示和语义的下游任务无关。这种解耦范式引入了一个关键的错位：视觉 token 化的丢失可能成为目标任务的表示瓶颈。例如，图像中文本的token化错误会导致识别或生成文本时的结果不佳。为了解决这个问题，ETT，是一种端到端的视觉 token 化器调优方法，可实现视觉 token 化和目标自回归任务之间的联合优化。与之前仅使用来自冻结视觉 token 化器的离散索引的自回归模型不同，ETT 利用 token 化器码本的视觉嵌入，并针对重建和字幕目标对视觉 token 化器进行端到端优化。ETT 易于实现和集成，无需调整大语言模型的原始码本或架构。端到端视觉 token 化器调优方法显著提升了性能，例如，与冻结的 token 化器基线相比，在多模态理解和视觉生成任务中实现 2-6% 的性能提升，同时保留原有的重建能力。

近年来，大语言模型 (LLM) 和多模态预训练的快速发展，推动了自回归 (AR) 模型超越其在自然语言处理领域的主导地位，将其影响力扩展到视觉和多模态任务。在自回归模型的下一个 token 预测 (NTP) 范式下，多模态学习通常将图像和文本等多模态数据编码为紧凑的离散 token，以实现统一的序列建模。例如，最近的研究 Emu3 [54] 将文本、图像和视频token为离散token，并在统一的 token 空间中执行下一个token预测。随后的大量研究 [37, 61, 69, 57] 进一步推进了这一方向，在视觉生成和感知方面取得了更高的性能。这种统一的下一个token预测范式为大规模训练和推理提供了一个灵活的多模态学习框架。

token 化在自回归框架中起着关键作用。对于图像和视频等模态，训练一个高效通用的 token 化器面临着巨大的挑战，因为 token 化器很难同时兼顾紧凑和无损。token 化器的损失可能会成为目标任务的瓶颈。例如，图像中文本的 token 错误会导致识别或生成效果不佳。然而，现有的 token 化方法忽略这种偏差，单独训练视觉 token 化器，并直接将其集成到下游训练中，假设获得的视觉 token 化器能够很好地泛化到各个任务中。针对自编码优化的视觉token 化器与需要各种表示和语义的下游任务无关。例如，大多数token 化器专注于低级像素级重构。学习到的表示质量本质上受到矢量量化 (VQ) 造成的信息损失的限制，导致其在视觉理解任务中的表现不如使用连续高级表示的模型，例如 CLIP [38]。此外，现有的自回归流程通常仅使用视觉 token化器中的离散索引，这与大语言模型中视觉嵌入的随机初始化相关，以实现下游任务。这对学习视觉表征和视觉语言对齐提出了挑战，而这两者对于多模态学习至关重要。

本文提出一种端到端的视觉 token 化器调优方法 (ETT)，该方法能够联合优化视觉 token 化和目标自回归下游任务。受近期视觉-语言模型 [47, 46, 27, 26] 的启发，这些模型在训练过程中更新其连续视觉编码器，以优化视觉表征和视觉语言对齐，从而实现更好的视觉理解。其提出以端到端的方式调优离散视觉token 化器，并利用大语言模型作为视觉 token 化器的视觉助手。

视觉 token 化器

初步工作。基于矢量量化 (VQ) 的视觉token化器包含编码器、量化器和解码器。编码器将输入图像 I 投影到特征图 f，其中 D 为特征维度，h = H/s，w = W/s，s 为下采样因子。量化器将 f 中的每个特征映射到码本 B 中最近的代码，以获得量化的嵌入 z，其中 K 为码本大小。最后，解码器将量化的嵌入重建回原始图像。传统的矢量量化 (VQ) 模型在扩展码本大小和代码维度方面往往存在困难。作为先驱研究之一，IBQ [43] 首次尝试将代码维度扩展到 256，同时通过在每个训练步骤同时更新整个码本，保持极大的码本大小（即 262,144）。

ETT 中的视觉token化器。主要采用 IBQ [43] 框架进行图像 token 化，使用下采样因子 s = 16。码本中每个离散 token 的维度为 D = 256。在原始 IBQ 的基础上，将码本大小调整为 131,072。用于token化器训练的损失函数 L_vq 为：L_vq = L_rec + L_quant + L_lpips + λ_G · L_GAN + λ_E · L_entropy。其中，L_rec 为像素重建损失，L_quant 为量化嵌入与编码特征之间的量化损失，L_lpips 为 LPIPS [67] 的感知损失，L_GAN 为 PatchGAN [18] 的对抗损失，L_entropy 为熵损失 [65]。λ_G 和 λ_E 分别为对抗损失和熵损失的权重。

端到端视觉 token 化器调优

离散索引到码本嵌入。Emu3 [54] 等方法以及 [48] 等类似方法在下游任务中仅使用视觉 token 化器离散索引，从而舍弃了视觉 token 化器嵌入丰富的表征能力。由于仅依赖离散码本索引，这些方法会阻止梯度传播，从而无法进行端到端训练。为了突破这一限制，ETT 将视觉 token 化器中的码本嵌入直接连接到 LLM，从而有效利用视觉 token 化器中编码的更丰富的特征表征，同时实现端到端训练。

LLM 桥接端到端调优。具体而言，如图所示，给定输入图像 I，首先从token 化器的码本中获取其量化嵌入 z。为了确保与预训练的 LLM 兼容，采用带有 GeLU 激活函数的多层感知器作为轻量级投影器。该投影层将量化的视觉嵌入 z 映射到 x^I，其中 C 表示大语言模型的隐维度大小。由于整个计算图（包括预训练的 LLM 和视觉 token 化器）保持可微分，因此可以使用基于梯度的优化对整个结构进行端到端训练。对于文本输入 T，利用预训练 LLM 中的 token 化器和文本嵌入层将其转换为文本 token 嵌入 x^T。

请添加图片描述

保留重构能力。虽然端到端训练增强了视觉 token 化器的表示，但保持其重构能力以确保高保真图像合成至关重要。为此，将总体训练目标设定为字幕损失 L_cap 和 VQ 损失 L_vq 的组合。具体而言，将图像 token 嵌入 x^I 和文本 token 嵌入 x^T 同时输入到 LLM 中。对于文本 token，应用交叉熵 (CE) 损失：

请添加图片描述

此外，直接复用损失函数 L_vq 进行视觉重建。因此，端到端视觉 token 化器调优目标变为：L=L_cap + α·L_vq，其中 α 是损失权重，用于控制多模态感知和视觉重建之间的权衡。通过使用 LLM 联合训练 token 化器编码器和解码器，该方法在保持模型重建能力的同时，确保学习的视觉 token 在语义上有意义，并且对于多模态理解和生成有效。

多模态生成与理解的训练方案

借鉴前人的研究 [7, 8]，下游多模态感知和生成的整个训练过程遵循三个连续的训练阶段。所使用的训练数据包括公开可用的图像数据集，并补充了用于理解和生成的各种指导数据，如表所示。

请添加图片描述

第一阶段：对齐学习。第一个训练阶段是有效地建立视觉-语言对齐。用预训练的大语言模型和视觉 token 化器，并将它们冻结，仅使用图像转文本字幕损失函数 L_cap 训练视觉投影层。这种设置使 LLM 能够直接从 token 化器获取视觉概念和实体，从而有效地连接视觉和语言模态。具体而言，从构建的数据集 SOL-recap 中挑选出一个包含 1200 万张图像的子集，该子集包含 3200 万张来自公开数据集（例如 SA-1B [20]、OpenImages [22] 和 LAION [42]）的图像-文本对。需要注意的是，所有图像均使用改进的字幕引擎重新生成字幕，该引擎遵循 [8] 进行了改进。此阶段的高质量数据可以增强训练的稳定性和跨模态对齐。

第二阶段：语义学习。第二阶段是整个训练流程中最关键的部分，实现了端到端的视觉 token 化调优。在此阶段，解冻 LLM、投影器和视觉 token 化器的权重，并用字幕损失 L_cap 和重建损失 L_vq 联合优化它们。来自 SOL-recap 的高质量子集（即 1200 万个图文对）用于多模态理解和重建学习。此阶段能够高效学习视觉 token 化器的感知能力，支持视觉重建和理解。精心设计的第二阶段可以增强视觉 token 化器与下游任务之间的一致性，同时保留其原有的重建能力。

第三阶段：后训练。通过提出的端到端视觉 token 化器调优获得增强的视觉 token 化器后，遵循标准的后训练流程，实现多模态理解和生成。在此训练阶段，进一步对两个专业模型（即 ETT-Chat 和 ETT-Gen）进行后期训练，方法是冻结视觉 token 化器部分，并调整视觉投影器以及大语言模型层，以分别增强多模态理解和文本-到-图像生成的指令遵循能力。对于多模态理解，收集一系列多样化的高质量、多源指令数据，包括 SOL-recap、LLaVA-OneVision [27] 和 Infinity-MM [15]。对于视觉生成，用 Flux 模型 [23] 构建 1400 万个 AI 生成样本，并从开源网络数据 [12, 3] 中额外筛选出 1600 万个图文对，并根据图像分辨率和 LAION 美学评分进行筛选 [24]。

数据准备

(1) 视觉语言预训练和视觉 token 化器数据集。采用预处理流程 [8] 对 SA-1B [20]、OpenImages [22] 和 LAION [42] 进行预处理，分别生成 1100 万、700 万和 1400 万张图像。利用 [8] 之后的字幕引擎生成 3200 万条高质量字幕。

(2) 监督式微调数据集。为了理解数据集，从 Infinity-MM [15] 中提取 3180 万个多任务样本，并从 LLaVA-OneVision [27] 中提取 350 万个指令数据，并优先考虑复杂的对话结构。为了生成数据集，使用 Flux 模型 [23] 生成 1400 万个 AI 生成的样本，并进一步从开源网络数据 [12, 3] 中筛选出 1600 万个图文对，并根据图像分辨率和美观度评分进行筛选 [24]。

实施细节

用 Adam 优化器 [19] 在 8-A100 节点上训练 ETT。第 1、2 和 3 阶段的批次大小分别设置为 1024、1024 和 1024，最大学习率分别为 4 × 10−5、4 × 10−5 和 2 × 10−5。采用 0.03 比率的预热策略，并在所有阶段使用余弦衰减调度程序。除非另有说明，图像处理分辨率为 5122，并在第 3 阶段使用 LLaVA-mix-665K [31] 进行消融研究。在所有实验中，采用 Qwen2.5-1.5B [50] 作为多模态序列建模的大语言模型。

对于 ETT 中的视觉 token 化器，采用 Adam 优化器 [19]，固定学习率为 1 × 10−4，β1 = 0.5，β2 = 0.9。token 化器训练 500,000 步，全局批次大小为 256，输入分辨率为 256 × 256。对抗损失权重 λG 设置为 0.1，熵损失权重 λE 设置为 0.05。还采用 LeCAM 正则化 [52] 进行鉴别器训练，以提高训练稳定性。