O(n2)复杂度没了？谷歌重磅Token Cropr：极快视觉Transformer！4倍加速ViT！

3DCV · 公众号 · · 2024-12-07 00:00

正文

点击下方卡片，关注 「3DCV」 公众号
选择星标，干货第一时间送达

来源：3DCV

添加小助理：cv3d001，备注：方向+学校/公司+昵称，拉你入群。文末附3D视觉行业细分群。

扫描下方二维码，加入 「3D视觉从入门到精通」知识星球 ( 点开有惊喜 ) ，星球内凝聚了众多3D视觉实战问题，以及各个模块的学习资料：近20门独家秘制视频课程、最新顶会论文、计算机视觉书籍、优质3D视觉算法源码等。想要入门3D视觉、做项目、搞科研，欢迎扫码加入！

0. 论文信息

标题：Token Cropr: Faster ViTs for Quite a Few Tasks

作者：Benjamin Bergner, Christoph Lippert, Aravindh Mahendran

机构：Hasso Plattner Institute for Digital Engineering, University of Potsdam、Hasso Plattner Institute for Digital Health at the Icahn School of Medicine at Mount Sinai、Google DeepMind

原文链接：https://arxiv.org/abs/2412.00965

代码链接：https://github.com/benbergner/cropr

1. 导读

在资源受限的应用中采用视觉转换器(vit)需要提高推理吞吐量。为此，已经提出了几种令牌修剪和合并方法，它们通过连续减少令牌的数量来提高效率。然而，设计一种快速、保持高性能并适用于各种视觉任务的记号简化方法仍然是一个公开的问题。在这项工作中，我们提出了一个令牌剪枝器，它使用辅助预测头来学习基于任务相关性端到端地选择令牌。这些辅助头可以在训练后被移除，导致吞吐量接近随机修剪机的吞吐量。我们在图像分类、语义分割、对象检测和实例分割上对我们的方法进行了评估，结果显示速度提高了1.5到4倍，性能略有下降。作为一个最好的例子，在ADE20k语义分割基准上，我们观察到相对于无修剪基线的2倍加速，而跨5个种子的0.1中值mIoU的性能损失可以忽略不计。

2. 引言

视觉Transformer（Vision Transformer，ViT）是图像分类、分割和目标检测[29]等计算机视觉任务中广泛使用的架构。ViT将图像表示为一系列块（patch）标记（token）的序列，并使用多头自注意力（Multi-Head Self-Attention，MHSA）Transformer块对其进行处理。自注意力机制计算所有标记之间的成对点积，导致相对于序列长度的二次时间和空间复杂度，即O(n²)。对于需要低延迟或较小计算预算的实际应用，序列长度因此成为负担，尤其是在模型尺寸、图像分辨率和更精细的标记化不断增加的情况下。

图像在空间上存在冗余，包含不显著的背景和重复模式。这表明可以使用较少的Transformer块处理多个块，从而有机会修剪掉无信息量的标记、减少高层中的序列长度，进而提高计算效率。然而，这引发了一个核心问题：如何准确高效地评估给定任务中单个标记的重要性？

最近的标记修剪方法依赖于启发式方法，如自注意力分数，来识别信息性标记。其他方法通过合并相似标记来减少标记数量。然而，这些方法没有明确地为给定任务建模标记的重要性，这可能导致任务性能显著下降。相比之下，归因方法如显著性（Saliency）、遮挡（Occlusion）和注意力展开（Attention Rollout）估计输入对预测的贡献，但需要完整的正向传递，由于相关开销，这不是一个可行的选择。

计算机视觉中任务类型的多样性进一步复杂化了任务相关性估计的问题。图像分类是最简单的视觉任务，一直是许多先前工作的重点等。然而，诸如语义分割等密集任务为标记修剪带来了新的挑战，因为它们需要在像素级别进行预测，这与修剪标记的想法本质上相冲突。

我们通过跨注意力修剪（Cross-attention pruning，Cropr）解决了这些问题，这是一种简单的ViT标记修剪方法，能够高效地估计每个标记的任务相关性，同时适用于各种视觉任务。Cropr模块在中间层用于标记修剪。每个模块由一个基于跨注意力的聚合机制和一个辅助预测头组成。后者学习解决任务，而前者根据任务相关性对标记进行排序，仅将最相关的标记转发到更深层。训练后可以丢弃辅助头，从而最大限度地减少开销并使标记修剪高效。最后，在网络的后续部分，通过一种称为最后一层融合（Last Layer Fusion）的技巧重新引入修剪掉的标记，以支持密集任务。

3. 算法概述

给定一系列块标记，我们的目标是通过在网络中传播时逐步减少标记数量来提高ViT的推理效率。为此，我们在ViT块之上添加了Cropr模块，每个模块选择最具辨别力的标记，同时修剪掉信息量最少的标记。通过这种方式，减少了后续层的计算量，同时保留了相关信息，从而最大限度地减少了修剪对任务性能的影响。为了选择最具辨别力的标记，Cropr模块使用基于跨注意力的路由和聚合机制，该机制从辅助头接收任务特定的训练信号。通过稍微定制这些组件，我们的方法可以应用于各种视觉任务，如图像分类、分割和目标检测。特别是，对于语义分割等密集任务，模型需要进行像素级预测，因此需要来自所有标记的信息。我们提出了最后一层融合（Last Layer Fusion，LLF）作为从修剪掉的标记中恢复信息的简单而有效的方法。在推理期间，可以引入进一步的优化来精简我们的模块并提高吞吐量。我们以一个现实世界的例子结束本节，该例子说明了我们的修剪计划。Cropr的实现可在以下网址找到：https://github.com/benbergner/cropr。推荐课程： Transformer如何在自动驾驶领域一统江湖！

交叉注意力修剪(Cropr)模块连续地修剪不太相关的记号，只保留最有区别的记号用于更深的层。我们的方法在保持高性能的同时加速ViTs，并且适用于许多视觉任务，从分类到分割和检测。示例城堡图像说明了修剪过程。热图可视化了在网络中的每个块1到L处哪些令牌被剪除。

培训期间的Cropr模块。路由器对显著的保留标记进行评分，并将其与要删除的无信息标记分开。评分者的注意力矩阵A在聚合器中重用，聚合器的输出用于进行中间预测。用红色虚线表示的梯度流反馈到计分器和查询中。

推理过程中的Cropr模块。聚合函数和辅助头被移除，所有查询都被聚合为一个查询。(b)这些优化加快了Cropr的速度，吞吐量与随机选择器相当结果展示了语义分割。

4. 实验结果

我们对图像分类、语义分割、目标检测和实例分割评估了我们的方法。即使在积极的修剪计划下，我们也展示了强大的性能。例如，当微调EVA-02主干时，我们能够在ImageNet-1k上保持89.7%的top-1准确率，与未修剪模型相比仅下降了0.2个百分点，同时实现了2.1倍的速度提升。我们还评估了标记修剪对不同编码器容量和图像分辨率的影响，表明我们的方法在大规模情况下表现尤为出色。消融研究为我们的设计选择提供了实证支持，定性评估则提供了对修剪过程的见解。