Token Compensator：无需重新微调即可改变视觉Transformer的推理成本

Token Compensator: Altering Inference Cost of Vision Transformer without Re-Tuning

https://arxiv.org/pdf/2408.06798

引言

研究背景：

随着ViTs规模的快速增长，其不断增加的计算成本已成为一个紧迫的问题。因此，大量研究正致力于加速ViTs的训练和推理过程。ViTs的一个特点是它们能够容纳可变数量的输入标记（tokens）。因此，除了卷积神经网络中广泛使用的模型剪枝、量化和蒸馏等传统技术外，最近的研究还提出了通过标记压缩来加速ViTs的方法，如剪除不重要的标记或合并相似的标记。

标记压缩技术具有显著的优势。与剪枝和蒸馏等技术相比，一些标记压缩方法（如ToMe[2]）可以零样本地应用于现成的模型，或用于加速训练；与量化不同，标记压缩方法不需要低精度操作器的支持。此外，标记压缩方法与上述其他技术正交，使得它们在ViTs中具有广泛的应用前景。

提出问题：

然而，当应用于下游任务时，观察到标记压缩存在以下缺点：

首先，尽管一些标记压缩技术可以应用于现成的模型，但它们往往会导致显著的性能下降。

其次，即使仅在训练过程中应用标记压缩以加速训练过程，而在推理过程中不压缩标记，模型的性能仍然低于未应用标记压缩训练的模型。

当训练和推理阶段的标记压缩程度不一致时，模型的性能会次优，如图1所示。这一限制限制了其在各种场景中的应用。例如，如果旨在根据服务器负载动态调整部署模型的计算成本，需要为每个候选计算成本和每个下游任务训练一个模型以实现最佳性能，这将导致显著的训练和存储开销。

另外，如果打算在最小化模型推理性能损失的情况下加速训练，则仅在训练时间应用标记压缩可能会导致明显的性能下降。

本文指出在不同标记压缩程度下微调后的模型在参数层面存在一定的差距，这导致了在推理过程中改变压缩程度时的性能下降。

还观察到，这种差距可以在不同的下游数据集之间传递。

本文方法：

基于此，提出了Token Compensator（ToCom），一个预训练的插件，旨在解耦训练和推理之间的标记压缩程度，以解决上述挑战。

ToCom是一个参数高效的模块，仅包含可忽略不计的参数数量，它描述了不同压缩程度模型之间的差距。

为了获得ToCom，在预训练数据集上通过不同压缩程度之间的快速自蒸馏过程对其进行训练。

具体来说，教师模型和学生模型都是相同的冻结预训练模型，而学生模型则集成了ToCom。

在每个步骤中，随机为教师模型和学生模型分配不同的压缩程度，而ToCom则通过蒸馏学习它们之间的差距。

此外，为不同的压缩程度对分配ToCom参数的不同子集，使ToCom能够通过单次训练过程适应各种压缩程度对。

在推理过程中：

直接将ToCom集成到在下游任务上微调过的现成模型中，而无需任何进一步的训练。

通过选择ToCom参数的子集，微调后的模型可以直接应用于各种标记压缩程度，并实现与训练和推理压缩程度一致时相当的性能。

重要的是，ToCom只需预训练一次，并可以应用于在任意下游数据集上以任意标记压缩程度微调过的模型，从而使任何单个现成的模型都能够处理动态延迟约束而无需修改参数。

实验结果表明：

在超过20个数据集上，作为即插即用模块的ToCom可以有效地解耦训练和推理过程中的标记压缩程度。例如，在VTAB-1k基准测试中，ToCom可以将DeiT-B的平均性能相对于ToMe提高多达2.0%，如图1所示。ToCom还可以应用于不同规模的模型或不同对象预训练的模型，或用于增强包括标记合并和标记剪枝在内的各种标记压缩方法。

Delve into Token Compression

压缩度的影响

首先，对ViT中的Token压缩方法进行公式化。ViT的单层由两个块组成，即多头自注意力（MHSA）和多层感知机（MLP）。该层可以形式化为：

其中 Xl ∈ R^{ N×d }是第 l 层的输入，其长度为N且维度为d，LN 表示层归一化。

本文主要关注一种代表性的、最新的无需训练的Token压缩方法，即ToMe [2]，然后将其推广到其他方法。ToMe在MHSA和MLP块之间操作，利用patch tokens的键来评估它们的相似性，并通过二分图软匹配（Bipartite Soft Matching）合并 r 个相似的tokens。

在ToMe中，每层合并的tokens数量被视为超参数，用于调整ViT的吞吐量，这通常根据训练前的推理需求来确定。合并的tokens越多，模型在训练和推理时就越快，如图3所示。

然而，在现实场景中，训练时的压缩度（称为源压缩度）和推理时的压缩度（称为目标压缩度）可能并不相等。即，一个在某个压缩度下训练的现成模型可能会被应用在不同的压缩度下，而无需重新训练。这种场景具有实际意义，例如在利用下载的检查点进行推理时。

压缩度不一致的模型性能差距

当训练和推理的压缩度不匹配时，模型性能会受到显著影响。为了分析这种影响，进行了实验，研究了不同压缩度下模型性能的差距。实验结果表明，即使在源压缩度下训练的模型在源压缩度下表现良好，但当目标压缩度改变时，性能也会大幅下降。这限制了Token压缩方法在实际应用中的灵活性。

压缩度不一致的模型差距可转移性

进一步地，发现不同压缩度下模型之间的性能差距可以在不同的下游数据集之间转移。具体来说，我们在源数据集上训练了不同压缩度的模型，并将这些模型之间的差距应用到另一个目标数据集上。实验结果表明，通过这种差距转移，可以在一定程度上恢复目标数据集上的模型性能。这为提出Token Compensator（ToCom）提供了动机。

Token Compensator的提出

基于上述观察，提出了Token Compensator（ToCom），一个参数高效的插件，用于在训练和推理阶段之间解耦Token压缩度。ToCom通过在不同压缩度之间进行快速自蒸馏训练得到，能够描述不同压缩度模型之间的差距。在推理阶段，ToCom可以直接插入到任何现成的模型中，无需进一步训练，即可在不同压缩度下获得一致的性能。 

Token Compensator

本文提出了一种名为Token Compensator（ToCom）的框架，旨在解决Vision Transformers（ViTs）在训练和推理阶段压缩度不一致时性能下降的问题。ToCom是一个轻量级的插件，能够在不重新训练的情况下，直接插入到任意预训练的下游模型中，以弥补不同压缩度之间的性能差距。

Token Compensator设计

ToCom被设计为一个参数高效的模块，其包含的参数数量极少。它的主要目标是描述具有不同压缩度模型之间的性能差距。为了获得ToCom，在预训练数据集上进行了快速的自蒸馏过程，其中教师模型和学生模型都是相同的预训练模型，但学生模型中集成了ToCom。在每一步中，教师模型和学生模型被随机分配不同的压缩度，而ToCom则通过学习两者之间的差距来进行蒸馏。

进一步地，为不同的压缩度对分配了ToCom参数的不同子集，这使得ToCom能够通过单次训练过程适应各种压缩度对。这样，ToCom就能够在推理阶段被直接插入到任意预训练的下游模型中，以匹配不同的压缩度，而无需重新训练。

训练和推理过程

在训练阶段，我们首先在预训练数据集上使用自蒸馏方法来训练ToCom。在这个过程中，利用冻结的预训练模型作为教师模型，并在其基础上添加ToCom以形成学生模型。通过随机分配不同的压缩度给教师模型和学生模型，ToCom学习到了两者之间的性能差距。

在推理阶段，ToCom被直接插入到预训练的下游模型中，而无需任何额外的训练。通过选择ToCom中对应的参数子集，我们可以使得下游模型能够适应不同的推理压缩度，而无需重新调整其参数。这样，单个预训练模型就能够处理不同的延迟约束，而无需进行参数修改。