摘要

视觉Transformers（ViTs）在各种视觉任务上取得了领先的性能，但它们在计算复杂性上仍然很高，尤其是在语义分割这样的密集预测任务中。由于高分辨率的输入和输出通常意味着更多的token参与计算，直接移除注意力较少的token在图像分类任务中已被讨论，但不能直接扩展到语义分割，因为每个区域都需要密集的预测。为了解决这个问题，本文提出了一种基于token提前退出的动态token剪枝（Dynamic Token Pruning, DToP）方法，用于语义分割。该方法受到人类从粗糙到精细的分割过程的启发，将广泛采用的辅助损失网络架构自然地分成几个阶段，每个辅助块为每个token的难度等级进行分级。我们可以提前完成简单token的预测，而无需完成整个前向传播。此外，我们为每个语义类别保留k个最高置信度的token，以维持代表性的上下文信息。因此，计算复杂性将随着输入的难度而变化，类似于人类的分割方式。实验表明，所提出的DToP架构在不降低准确度的情况下，平均减少了20%到35%的计算成本。

拟解决的问题

论文旨在解决视觉变换器在语义分割任务中的高计算复杂性问题。由于高分辨率图像生成了大量输入token，导致计算负担加重，特别是在资源受限的环境中。

创新之处

• 提出了一种新的动态token剪枝范式，基于简单token的提前退出，用于语义分割变换器。
• 通过在每个阶段保留每个语义类别的k个最高置信度token，即使在非常简单的情况下，也保证了足够的上下文信息，从而提高了分割性能。
• 将DToP应用于主流的语义分割Transformer，并在三个具有挑战性的基准测试上进行了广泛的实验，结果表明DToP可以在不显著降低准确度的情况下减少高达35%的计算成本。

方法论

给定一个现有的普通视觉Transformer，我们使用固有的辅助磁头将其分为几个阶段。在第m阶段的最后一层(以为索引)，我们使用辅助块对所有令牌难度级别进行分级。我们在当前阶段完成高置信度容易令牌的预测，并在后续阶段处理其他低置信度硬令牌。为简单起见，没有给出每个语义类别保留的k个最高置信度令牌，以支持具有代表性的上下文信息。每个阶段的预测共同构成了最终结果。

动态Token剪枝（DToP）

动机：在语义分割中，并非所有图像区域（token）的识别难度都相同。一些区域可能很容易识别，而另一些则可能因为复杂性或细节而难以识别。DToP利用这一观察，通过在网络的不同阶段对token进行难度评估，并提前完成那些容易识别token的预测，从而减少后续层的计算负担。 

步骤：

1. 阶段划分：将视觉变换器网络划分为多个阶段（Stage），每个阶段的末尾都有一个辅助块（Auxiliary Block）。
2. 难度评估：在每个阶段的末尾，使用辅助块对所有token的难度进行评估。这是通过比较token的预测置信度与一个预设的阈值（例如0.9）来完成的。如果token的预测置信度高于这个阈值，它被认为是容易识别的。
3. 早期退出：对于那些被评估为容易识别的token，DToP将提前完成它们的预测，并且在后续的网络层中不再对它们进行进一步的处理。这意味着这些token的前向传播被“剪枝”或“停止”。
4. 保留上下文信息：为了防止在剪枝过程中丢失重要的上下文信息，DToP会为每个语义类别保留k个最高置信度的token。即使某些类别的大多数token都被剪枝了，这些保留的token仍然可以为网络提供必要的上下文信息。

查询匹配辅助块

动机：查询匹配辅助块（Query Matching Auxiliary Block, 简称QMAB）的目的是准确地估计每个token的难度等级，以便在DToP框架中有效地识别出哪些token是容易识别的，可以在网络的早期阶段就完成预测并停止进一步的前向传播。同时，QMAB需要保持轻量级的架构，以避免增加额外的计算负担。

实现：用于分级所有令牌难度级别的辅助块应该遵循两个原则:能够准确估计令牌难度级别，并具有轻量级架构。因此，采用最新的关注掩码模块(attention-to-mask module, ATM)来实现这一目标。具体来说，一系列可学习的类令牌使用Transformer解码器与编码器特征交换信息。输出类令牌用于获得类概率预测。每个类令牌的注意分数被用来组成一个掩码组。类概率和组掩码之间的点积产生最终的预测。

ATM模块来自于以下论文：

在DToP框架中，原始ATM模块的层数被减少，因为在DToP中观察到原始设置并没有显著的性能变化，同时这也保证了较低的计算开销。解耦多个级联ATM模块，并将它们用作单独的辅助分段头，每个模块都有单独的可学习类令牌。

维持上下文信息

动机：在DToP中，简单token在早期阶段就被剪枝，这可能导致某些语义类别的token全部被移除，从而丢失了这些类别的上下文信息。上下文信息对于语义分割的准确性至关重要，因为它有助于网络理解场景中各个对象之间的关系。

步骤：

1. 类别特定的保留：只考虑当前图像中出现的语义类别，而不是所有可能的类别。
2. 置信度阈值：对于每个特定的语义类别，找出所有预测置信度高于预设阈值​的token。
3. 保留k个最高置信度token：在这些高置信度token中，选择分数最高的k个token保留下来。如果某类别的高置信度token数量少于k，则保留所有这些token。
4. 联合计算：这些保留的token将与其他低置信度token一起参与后续层的计算，确保简单类别的语义信息得以保留，并用于类别间的信息交换。

结论

本文研究了如何减少基于纯视觉变换器的现有语义分割的计算成本问题。提出的DToP方法通过模仿人类的粗到细的分割过程，动态地减少了计算量，并通过保留简单类别的上下文信息来保证分割的准确性。广泛的实验结果表明，该方法在保持性能的同时显著减少了计算成本。尽管DToP在小批量计算效率方面还有待优化，但其为加速视觉变换器提供了一种有效的策略。