ToCom：一次训练随意使用，华为提出通用的ViT标记压缩器 | ECCV 2024

标记压缩通过减少冗余标记的数量（例如，修剪不重要的标记或合并相似的标记）来加快视觉变换器（ViTs）的训练和推理。然而，当这些方法应用于下游任务时，如果训练和推理阶段的压缩程度不匹配，会导致显著的性能下降，这限制了标记压缩在现成训练模型上的应用。因此提出了标记补偿器（ToCom），以解耦两个阶段之间的压缩程度。该插件通过在预训练模型上额外执行了一个快速参数高效的自蒸馏阶段获得一个小型插件，描述了不同压缩程度下模型之间的差距。在推理过程中，ToCom可以直接插入到任何下游现成模型中，无论训练和推理的压缩程度是否匹配，都能获得通用的性能提升，而无需进一步训练。

来源：晓飞的算法工程笔记 公众号，转载请注明出处

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

• 论文地址：https://arxiv.org/abs/2408.06798
• 论文代码：https://github.com/JieShibo/ToCom

Introduction

---

  视觉变换器（ViTs）在计算机视觉的多个领域取得了显著成功，包括图像分类、目标检测、语义分割等。然而，随着ViTs规模的快速增长，计算成本的增加已成为一个迫切问题。因此，大量研究工作集中在加速ViTs的训练和推理上。ViTs的特点在于能够处理可变数量的输入标记，除了卷积神经网络中广泛使用的传统技术，如模型剪枝、量化和蒸馏，近期的研究通过标记压缩来加速ViTs，例如修剪不重要的标记或合并相似的标记。

  与剪枝和蒸馏等技术相比，标记压缩技术具有明显的优势。一些标记压缩方法（例如，ToMe）可以以零样本的方式应用于现成模型或用于加速训练。与量化不同，标记压缩方法不需要对低精度操作符的支持。此外，标记压缩方法与上述其他技术在操作上是正交的，使其在ViTs中具有广泛的适用性。

  然而，当标记压缩应用于下游任务时，论文观察到如图1所示的以下缺点：

1. 尽管一些标记压缩技术可以应用于现成模型，但通常会导致显著的性能下降。
2. 即使标记压缩仅在训练期间应用以加速过程，而在推理期间不进行压缩，模型的性能仍然低于未经标记压缩训练的模型。
3. 总之，当训练和推理阶段的标记压缩程度不一致时，模型的性能表现不佳。

  论文指出，经过不同标记压缩程度微调的模型在参数层面存在一定的差距，这导致在推理过程中改变压缩程度时性能下降。此外，这一差距可以在不同的下游数据集之间转移。基于此，论文提出标记补偿器（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

---

Impact of Compression Degrees

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

$$\begin{array}{c}{\tilde{\mathbf{X}}}^l={\mathbf{X}}^l+\text{MHSA}(\text{LN}({\mathbf{X}}^l)),{\mathbf{X}}^{l+1}={\tilde{\mathbf{X}}}^l+\text{MLP}(\text{LN}({\tilde{\mathbf{X}}}^l)),\end{array}$$

  其中， ${\mathbf{X}}^l\in {\mathbb{R}}^{N\times d}$ 是第 $l$ 层的输入，具有长度 $N$ 和维度 $d$ ，LN表示层归一化。

  论文主要关注一种具有代表性且最先进的无训练标记压缩方法ToMe，并进一步推广到其他方法。ToMe在MHSA和MLP模块之间操作，利用图像块标记的键来评估它们的相似性，并通过二分软匹配将 $r$ 个相似的标记进行合并。

  在ToMe中，每层合并的标记数量被视为超参数，以调整ViTs的吞吐量，这通常在训练之前根据推理需求来确定。合并的标记越多，模型在训练和推理时的速度就越快，如图3所示。然而，在实际场景中，训练期间的压缩程度（称为源压缩程度）和推理期间的压缩程度（称为目标压缩程度）可能不一定相等。也就是说，一个在某一压缩程度下训练好的现成模型，可能会在没重新训练的情况下进行不同的压缩程度下的应用。这种情况具有实际意义，例如，使用下载的checkpoint而无法访问训练数据或重新训练资源时，或根据服务器负载动态调整推理期间的压缩程度。此外，在现有计算资源有限的情况下，可能需要在训练期间使用较高的压缩程度以减少内存和时间开销，但在推理期间恢复到较低的压缩程度以确保性能。

  为了研究标记压缩方法在源压缩程度与目标压缩程度不一致时的性能，论文在五个下游数据集上进行了实验。如图2所示，论文对DeiT-B进行了ToMe的 $r=0$ 和 $16$ 的微调，并报告了在推理期间使用 $r=0,2,4,\dots ,16$ 的性能。可以看到，对于特定的目标压缩程度，当源压缩程度与其匹配时，模型的表现更好。源压缩程度与目标压缩程度之间的差距越大，性能下降的程度就越显著。

  然而，由于在较低压缩程度下训练的模型在训练期间接触到了更多的标记，这意味着它们遇到的信息范围比在较高压缩程度下训练的模型更广泛，因此，前者在各种目标压缩程度下理应优于后者。这表明，不同源压缩程度下训练的模型之间存在差距，使得在不同压缩程度之间的迁移效果较差。

Transfer across Tasks

  对于具有不同源压缩程度的模型之间的差距，是否可以在不同任务之间转移？更具体地说，令 ${\mathcal{M}}_m^{{\mathcal{D}}_{\mathcal{A}}}$ 和 ${\mathcal{M}}_n^{{\mathcal{D}}_{\mathcal{A}}}$ 表示在数据集 ${\mathcal{D}}_{\mathcal{A}}$ 上以压缩程度 $m$ 和 $n$ 训练的模型， ${\mathcal{M}}_m^{{\mathcal{D}}_{\mathcal{B}}}$ 和 $\mathcal{M}$