点击下方卡片,关注「集智书童」公众号
想要了解更多:
前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉(分割、检测)、多模态、目标跟踪、NerF」
行业技术方案 👉「AI安防、AI医疗、AI自动驾驶」
AI模型部署落地实战 👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」
欢迎扫描上方二维码,加入「 集智书童-知识星球 」,日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑,期待交流!
免责声明
凡本公众号注明“来源:XXX(非集智书童)”的作品,均转载自其它媒体,版权归原作者所有,如有侵权请联系我们删除,谢谢。
在资源受限的应用中采用视觉 Transformer (ViTs)需要提高推理吞吐量。为此,已经提出了一些基于Token剪枝和合并的方法,通过逐步减少Token数量来提升效率。
然而,设计一种既快速又能在各种视觉任务上维持高性能的Token减少方法仍然是一个开放问题。在这项工作中,作者提出了一种使用辅助预测 Head 的Token剪枝方法,这些辅助 Head 能够端到端地基于任务相关性学习选择Token。
训练完成后,可以移除这些辅助 Head ,从而使得吞吐量接近随机剪枝方法的水平。作者在图像分类、语义分割、目标检测和实例分割任务上评估了该方法,并展示了1.5-4倍的速度提升,同时性能下降较小。
作为最佳案例,在ADE20k语义分割基准测试中,与无剪枝 Baseline 相比,作者观察到了2倍的速度提升,同时性能下降几乎可以忽略不计,5次种子实验中的中值mIoU损失不到0.1分。
- Introduction
Vision Transformer [17] 是用于图像分类、分割和目标检测等计算机视觉任务的一种广泛使用的架构。ViTs 将图像表示为每个片段的序列 Token,并使用多头自注意力(MHSA)Transformer 块来处理这些 Token。自注意力机制计算所有 Token 对之间的点积,这导致其时间和空间复杂度随着序列长度平方级增长,即 (O(n^2)) [61]。对于需要低延迟或较小计算预算的实际应用而言,序列长度成为一个负担,尤其是在模型规模增加 [13]、图像分辨率提高 [3] 和更精细 Token 化 [4] 的背景下。
图像,如图1中所示的例子,具有空间冗余性,包含非显著的背景和重复的模式。这表明,可以使用较少的Transformer块来处理某些patches,从而有机会修剪不具信息性的Token,减少高层的序列长度,并因此提高计算效率。
然而,这引发了一个核心问题:作者如何能够准确且高效地评估个体Token在给定任务中的重要性?近年来的Token修剪方法依赖于启发式方法,如Self-Attention分数,来识别信息性Token[20, 35]。另一种方法通过合并相似Token来减少Token数量[6, 18, 42]。
然而,这些方法并未明确建模一个Token对于给定任务的重要性,这可能导致任务性能显著下降。相比之下,归因方法,如可解释性[53]、遮挡[72]和注意力展开[1]可以估计输入对预测的贡献,但需要完整的前向传递,由于相关的开销,这不是可行的选择。
估计任务相关性的难题进一步被计算机视觉中任务类型的多样性所复杂化。图像分类,最简单的视觉任务,一直是许多先前研究的关注点。然而,密集型任务如语义分割为Token裁剪带来了新的挑战,因为这类任务需要在像素级进行预测,这与Token裁剪的基本理念存在内在冲突。
作者通过跨注意力剪枝(Cropr)来解决这些问题,Cropr是一种针对ViTs的简单Token剪枝方法,能够高效地估计每个Token的任务相关性,并适用于各种视觉任务。Cropr模块在中间层用于Token剪枝,如图1所示。每个模块由基于跨注意力的聚合机制和一个辅助预测头组成。后者负责学习解决任务,而前者则根据任务相关性对Tokens进行排名,只将最具相关性的Tokens传递到更深的层。训练完成后可以丢弃辅助头,从而最小化开销并使Token剪枝变得高效。最后,剪枝后的Tokens会在网络的后期通过一种称为最后一层融合的技术重新引入,以支持密集型任务。作者在第3节中详细介绍了作者的方法。
作者在第4章中评估了本方法在图像分类、语义分割、目标检测和实例分割任务上的表现。即使在严格的剪枝计划下,作者仍表现出色。例如,在对EVA-02 [19] 主干网络进行微调时,作者的方法在ImageNet-1k上维持了89.7%的Top-1准确率,与未剪枝模型相比,准确率仅下降了0.2个百分点,同时实现了2.1倍的速度提升。作者还评估了不同编码器容量和图像分辨率下的Token剪枝效果,表明作者的方法在大规模应用中表现尤为突出。消融研究提供了作者设计选择的经验支持,而定性的评估则揭示了剪枝过程中的关键见解。作者在第5章总结了作者的发现并提出了未来的工作方向。
- Related Work
几种方法已被提出,通过剪枝/合并Token来缩短视觉Transformer中的序列长度。
基于类 Token (CLS)的注意力分数进行分类的Token修剪。一种常见策略是将类 Token 的注意力分数作为从上而下的线索,应用于图像Token的修剪 [20, 35, 39, 68]。具有较低注意力分数的Token被认为不那么重要,并被裁剪掉。值得注意的是,Haurum等人 [23] 表明,简单的Top-K选择器是一个强大的基准。然而,现代融合 Kernel 实现 [12, 47] 经常限制直接访问注意矩阵,因此需要采用替代策略。相反,作者采取从上而下的方法,利用辅助头信号保留相关的任务Token。
其他方法采用参数化模块来预测应保留哪些Token [32, 49],但引入了额外的层和损失,这可能会干扰主要任务。Cropr模块通过应用stop-gradient避免梯度干扰,并在推理时仅限制额外参数到单个 Query Token。另一种常见的设计选择是使Token剪枝具有适应性,在简单输入下剪枝更多的Token 。相比之下,作者使用了一种针对吞吐量优化的静态方法,该方法固定剪枝数量的Token以实现不同输入间的批量处理。
超越分类的Token修剪。很少有工作将Token修剪应用于超出分类的领域:
刘等人[38]和唐等人[56]通过添加辅助头,在基于信心值修剪Token方面将其扩展到了语义分割。作者在实验中与唐等人[56]进行了对比。刘等人[37]在目标检测和实例分割中使用两层MLP进行Token修剪,小型网络可获得高达34%的适度加速。相比之下,作者省去了推理时的额外层,并将Cropr应用于较大的ViT,实现了1.9倍的加速和63.0 AP^box。作者相信,Cropr通过其快速性、高性能以及适用于各种视觉任务的特点,在Token修剪领域取得了重大进展。
Token 合并。Token 合并背后的假设是相似的 Token 表示贡献了冗余信息,可以因此进行合并。硬合并方法通过聚类或二分匹配 [6] 等方式将相似的 Token 组合成不重叠的组。相比之下,软合并方法通过学习空间 Token 的凸组合来生成摘要 Token [22, 28, 50, 77]。例如,Renggli 等人 [50] 和 Jaegle 等人 [28] 使用可学习的 Query 进行跨注意力操作来实现这一目的。作者同样使用可学习的 Query 进行跨注意力操作,但是用于 Token 选择而不是合并,其中注意力分数反映任务相关性,并且聚合后的 Token 仅用于辅助头中。
剪枝和合并。最近,剪枝和合并的概念也被联合应用于模型中[7, 30, 66]。许多剪枝方法还将剪枝后的Token聚合为一个或几个新Token。Cropr则通过在最终Transformer块之前将激活的剪枝后Token与保留下来的Token进行 ConCat 来重新激活这些Token,而无需进行任何Token总结。
- Methods
给定一系列每个Patch的Token,作者的目标是通过在网络中逐步减少Token的数量来提高ViT的推理效率。为此,作者基于ViT模块之上添加了Cropr模块,每个模块选择最具鉴别性的Token并修剪最不具信息性的Token。
这样,在后续层中的计算量得以减少,同时保留相关的信息,从而最小化修剪对任务性能的影响。为了选择最具鉴别性的Token,Cropr模块采用了基于交叉注意力的路由和聚合机制,该机制接收来自辅助Head的任务特定训练信号。
通过对这些组件进行略微定制,作者的方法可以应用于各种视觉任务,例如图像分类、分割和目标检测(第3.2节)。特别是对于密集型任务如语义分割来说,这些模型需要所有Token的信息来进行像素级预测。作者提出了一种简单且有效的方法——最后一层融合(LLF),用于从修剪后的Token中恢复信息(第3.3节)。
在推理过程中,可以通过进一步优化来减小作者的模块并提高吞吐量(第3.4节)。本节以一个真实的例子结束,展示了作者的修剪计划(第3.5节)。
Cropr的实现可以在以下链接获取:https://github.com/benbergner/cropr。
3.1. Module description
Cropr模块在图2中有所展示。每个模块接收输入Token ,并输出一组不相交的“保留”和“裁剪”Token,分别为和,其中且为裁剪率。每个模块包含四个组成部分:评分器与选择器,两者共同构成了路由器,以及聚合器和任务头。
评分器为Token集合分配分数。然后这些分数被传递给选择器,选择器保留得分最高的个Token,并剔除剩余的个Token。
其中 (\backslash) 表示集合差运算。
Cross-attn模块通常使用线性Query、Key和Value投影、多个注意头以及LayerNorm(LN)[17, 61]。作者发现,在作者的设置中,这些组成部分都不是实现高任务性能所必需的。这使得作者可以简化模块的同时增加吞吐量(表4a)。作者将A映射到a,通过在Query轴上求和注意矩阵来实现。对于
这完成了路由器的设计。
为了学习反映一个词元对预测贡献的分数,聚合器利用注意力矩阵A计算输入词元的加权平均值,然后将这些值传递给辅助头。因此,在训练过程中,评分器会赋予更具有区分性的词元更多的权重,并且这些词元将被保留以供后续Transformer块处理。作者发现,在聚合器中增加Transformer块前馈模块的容量并加入LayerNorm和带有残差连接的MLP(见表4c),是有益的。因此,对于,
聚合输出由任务特定的Head进行处理,以生成中间预测,这些中间预测进而为训练聚合器和评分器提供梯度。
最后,需要注意的是,在评分和聚合模块之前应用了stop-gradient操作。从概念上讲,这有一个优势,即可以隔离辅助头与主干网络。因此,编码器不会受到来自辅助损失的矛盾梯度的影响。此外,在训练过程中这也更加高效,因为Cropr组件的梯度不会反向传播到编码器。
3.2. Task-specific designs
作者的评分器和聚合器采用灵活的 Query 机制,类似于Perceiver IO [28]的设计,从而支持任意输出形状,并且易于适应各种任务。在本节中,通过调整可学习的 Query 数量、设计辅助头和损失函数,作者为每个视觉任务实例化Cropr,如下所示。
图像分类。评分器使用单个可学习的Query,。然后聚合器输出一个单一的Token,该Token经过LayerNorm和线性投影处理,类似于最终分类头的操作。后者输出所有类别的logits。使用Softmax交叉熵损失函数。
语义分割。主头和辅助头均采用Segmenter [54] 的线性头。评分器通过一个可学习的 Query ,对每个patch token,获得网格结构化的表示形式。聚合器输出经过LayerNorm和线性投影处理,类似于图像分类,但在每个patch位置上独立进行,随后是一个基于每个patch的softmax交叉熵损失。为了在辅助头中减少计算复杂度,与Segmenter不同,不将logits上采样至输入分辨率,而是将标签下采样到特征图分辨率。下采样的标签可以被多个Cropr模块重用。
联合检测和实例分割。作者将Cropr应用于Cascade Mask R-CNN [8, 24] 来完成这一任务。但由于这种多阶段检测器计算成本较高,在辅助头的应用上不太实际。作者提出了一种 Agent 辅助头及其损失,这些辅助头和损失能够为两个任务——多标签分类——提供强烈的信号。其核心思路是目标检测和实例分割都要求识别出图像中存在的所有类别。具体而言,ground truth标签以二进制向量形式编码,其中每个维度对应一个类别的存在。评分器使用类似于图像分类中的单一可学习 Query ,。然后,LayerNorm和线性投影将聚合的token映射为与数据集中类别数量相同的logits。在此多标签设置中,使用sigmoid激活函数和二元交叉熵损失。
3.3. Last Layer Fusion (LLF)
在密集任务中,例如语义分割,预测是在像素 Level 进行的。然而,在输入中有较大比例的数据被丢弃时,这会变得困难。此外,许多任务 Head 需要空间特征图来进行上采样,在剪枝过程中这些特征图并未得到维持。
作者通过LLF解决这些挑战,这是一种高效且有效的方法,可以重新激活剪枝的Token并保留来自所有图像Patch的信息。具体来说,来自所有Cropr模块的剪枝Token会在其相应的位置与倒数第二个ViT块输出的保留Token一起插入(见图1)。换句话说,这些剪枝Token会被传递到最终的ViT块,而不是完全丢弃。最终的ViT块处理这个合并序列,使之前被剪枝的Token能够注意保留Token的深层特征。作者在附录F中通过tSNE [60] 图来可视化其效果。作者在最终的ViT块中禁用DropPath [33],以确保Token融合。
LLF 不引入额外参数却能在与其他融合方法相比时表现出色(参见表5)。请注意,LLF 并不限于 Cropr;实际上,作者在实验中为几种 Baseline 方法配备了 LLF。
3.4. Efficient inference
Crop跨注意力Transformer块和辅助头构成了重要的计算开销。然而,图2中黄色标注的这些组件仅在训练评分器时需要,在推理阶段可以安全地丢弃,仅保留用于Token选择的路由器,如图3a所示。
评分器仍按 的复杂度进行计算,在 Query 数量 较大的情况下(如语义分割中, 随图像分辨率增加),这种复杂度是代价高昂的。但由于聚合器已被丢弃,因此不再需要构建跨注意力矩阵。只需要求和后的分数 即可。利用矩阵乘法的分配律,可以很容易地证明,式 (4) 可简化为向量-矩阵乘法的形式,时间复杂度为 。
其中, 是一个聚合 Query ,可以预先计算。每个 Cropr 模块因此可以简化为一个路由器,该路由器包含一个高效的评分函数和一个 Top 选择器。借助这些改进,Cropr 的吞吐量与随机选择器相当(图3b)。
3.5. Pruning schedule
作者通过一个工作示例来解释剪枝计划。考虑一个具有24个块的ViT-L模型,输入图像尺寸为,patch大小为16,因此得到196个patch token。除非另有说明,作者会在每个块之后插入Cropr模块,这是一种按块进行的计划,每次剪枝个token。作者的目标是在网络的末尾去除大部分token。
不采用LLF时,在每层块(除了最后一层)之后都会进行剪枝。在作者的例子中,当设置时,作者剪枝个Token,留下12个Token,总的剪枝比例(TPR)为。采用LLF时,在除了最后两层之外的每层块之后都进行剪枝,结果是输出20个Token,TPR为。在这种情况下,不会在倒数第二层之后进行剪枝,因为被剪枝的Token会立即重新插入。
作者观察到,在高分辨率图像中,保持保持Token的数量为8的倍数可以提高吞吐量(附录E)。由于Patch序列长度通常能够被8整除,因此在可能的情况下作者将设置为8的倍数。此外,由于ViT模型通常采用分类(CLS)Token,因此作者在第一个模块中将增加1。遵循常见的做法[6, 20, 35],CLS Token永远不会被剪枝。
- Experiments
作者在四种不同的ViT架构、网络容量和图像分辨率下,评估了Cropr在四个视觉任务上的表现。为了证明Cropr能够选择与任务相关的Token,作者将其与具有挑战性的 Baseline 方法进行了比较:(1) 不剪枝(上限基准),(2) 随机剪枝,(3) 方差剪枝[43],基于RGB通道的每个Patch像素方差进行排序,以及(4) Attn Top-K,该方法根据自注意力分数选择Token,并已被证明是表现最好的方法之一[23]。为了进行公平的比较,作者在(2)、(3) 和(4) 中使用了LLF。除了任务特定的指标外,作者还报告了单次推理前向传播时的FLOPs/吞吐量(跨批量大小优化),使用自动混合精度(AMP)和NVIDIA A100 GPU。超参数详见附录C。
4.1. Image classification
与 Baseline 方法及先前工作对比。作者使用预先训练好的 Mask 自编码器(MAE)对ViTL在ImageNet-1k [51]上进行微调,并按照He等人的设置[25]应用作者的工作示例中的剪枝方案(第3.5节)。表1显示了作者的方法在三种不同场景下的全面评估结果。首先,作者将作者的方法与 Baseline 方法(2)-(4)进行比较。Cropr在同等吞吐量的情况下优于所有剪枝 Baseline 。作者还包含了非显著性选择器的结果,该选择器通过剪枝最相关Token的方式反转了Cropr。正如预期,这种方法的表现比随机剪枝要差。
接下来,在表1中部,作者比较了裁剪无LLF的方法与先前的工作。作者的方法在性能和吞吐量方面具有竞争力。后者尤其是在不同方法之间差异显著,其中K-Medoids [42] 和 ATS [20] 的速度慢于未剪枝的基础模型。
最后,作者观察到一些方法在作者的块状剪枝计划下并未收敛。因此,在表1底部,作者展示了使用更轻量级三阶段计划的结果,其中在第6、12和18个块之后剪除50个Token,最终剩下46个Token(TPR为77%)。在这种情况下,Cropr表现最佳,并显示出竞争力的吞吐量。
与未剪枝的基础模型相比,Cropr在准确率上仅下降了0.3-0.7个百分点,但在速度上实现了1.6-1.9倍的加速。在附录E中,作者展示了使用更轻的调度策略不会影响性能。
裁剪规模效应。在Cropr模型与未剪枝 Baseline 模型相比,在网络容量对性能和吞吐量的影响方面有何不同?鉴于倾向于更大模型的趋势[13],这一问题尤为重要。为了研究这一点,作者使用LLF对ViT-B/16、L/16和(分别包含12、24和32个块)进行了Cropr处理,并将这些模型大小的设置为每块16、8和8个Token,从而分别获得TPR值为82%、90%和94%。图4显示,随着模型规模的增加,Cropr的相对性能惩罚减小,从ViT-B的减少到ViT-H的,尽管TPR更高。这一观察可能归因于更深的模型中,剪枝被分布在更多的层上,导致早期丢弃更少的Token。此外,由于更多层受益于减少的Token计数,Cropr的加速效果随规模增加而改善,从ViT-B的增加到ViT-H的。当作者扩展图像分辨率时(附录D),也观察到了类似的效果。
将最先进的模型应用于实验。作者使用EVA-02-L,这是一个最新的开源Vision Transformer [19]。作者从一个IN-21K微调预训练权重开始训练,并将图像调整大小为,设置patch大小为14。表2展示了实验结果。首先,作者在不剪枝的情况下训练EVA-02,得到的准确率为,这与Fang等人在[19]中报告的相比较为接近。然后,作者训练了一个裁剪后的版本Cropr,其中作者将并启用LLF,得到的TPR为。作者观察到的准确率为,只下降了0.2个百分点,但速度提高了倍,且FLOPs减少了约。
此外,作者还报告了在不使用LLF的情况下更为激进的剪枝计划的结果( Token 为↓),在这个计划中,在第3个区块之后应用一个Cropr模块,可以剪枝掉总Token数量的80%,即825个Token。与未剪枝的 Baseline 相比,这导致性能下降了1.1个百分点,但提供了约76%的FLOP减少和4.1倍的速度提升。表2的其余部分列出了其他一些先进模型的选择结果。经过剪枝后,EVA-02仍然保持在准确性排名第3位,但速度加快了一倍。使用更激进的计划,作者的模型在速度上有较大优势,同时仍优于一些其他模型。
4.2. Semantic segmentation
作者在ADE20k数据集[76]上进行了实验,并使用线性解码头对Segmenter[54]进行了微调(详见第3.2节)。编码器被替换为EVA-02-L,遵循Fang等人[19]的设置。图像以512×512的分辨率处理,patch大小为16,从而产生了1024个patch。模型训练了64个epoch。在评估时,作者将最大边调整为512像素,并保持纵横比的同时填充较短的边。这种1-shot评估方法注重效率,相比常用的单尺度评估设置更具挑战性,后者通过滑动窗口平均预测结果。
在这一设置中,未经修剪的模型在5次种子实验中实现了56.7%mIoU,超过Seg-L-Mask/16 (54% mIoU [51.8%]),尽管Seg-L-Mask/16采用更高的分辨率为640×640,并使用更复杂的 Mask Transformer 解码器,且采用单一尺度评估设置。作者将其归因于作者使用了预训练的EVA-02主干网络。
在应用Cropr时,作者激活LLF,并在前22个块中的每一个之后剪枝个Token,结果得到一个TPR为。为了促进学习,在前32个epoch中使用了一个关于的递增课程,线性地将从1增加到40。
与 Baseline 方法的比较。作者将Cropr与(l)-(4)中的 Baseline 方法进行了比较。剪枝 Baseline 使用LLF以适用于分割场景。此外,Attn Top-K现在使用平均自注意力矩阵来评分片,因为在 Head 未使用CLS Token 。每种模型分别运行五次,使用不同的随机种子,结果汇总在图5中。Cropr在中间mIoU上得分为56.6%,这比无剪枝 Baseline 仅低0.1个百分点,同时速度快2.0倍。此外,作者的模型在吞吐量相似的情况下,相比所有剪枝 Baseline 表现出更高的中间性能。有趣的是,作者发现即使是随机剪枝器,通过利用LLF也能取得不错的性能。
与先前工作的对比。作者重新实现了DToP的日志融合方法[56],并通过使用与作者方法相同设置来进行公平比较。DToP使用辅助头根据预测置信度选择Token,并将剪枝和保留Token的日志概率进行拼接以获得最终预测。与此类似,作者也在辅助头中使用LN和线性输出投影。不同的是,LLF融合的是特征而不是日志概率,并且Cropr组件的梯度不通过编码器反向传播。表5显示,使用LLF的Cropr明显优于DToP的日志融合方法。
刘等[38]也使用了辅助头,但在任务头之前,他们将剪枝和保留的Token特征进行了拼接。作者在表5中与这种融合方法进行了比较,结果显示LLF的表现优于“Token-Concat”。此外,他们在基于ViT-S的Segmenter上应用了Token剪枝,并在ADE20k数据集上报告了一个mIoU下降0.35的值,相对无剪枝 Baseline 减少了18%的FLOPs。作者通过减少41%的FLOPs实现了仅有0.1 mIoU的中位数下降。定性评价。图6展示了室内和室外场景的剪枝热力图。在这个任务中,每个像素的输出都会贡献到评估指标,因此确定需要剪裁哪些信息是有挑战性的。作者发现注意力主要集中在显著目标上;然而,在后续层中仍保留了一些背景片段,这可能是因为它们对任务的整体相关性。此外,即使进行了剪枝,作者仍然观察到对于一些更小、难以分割的目标,预测结果是一致的。这可能是由于LLF的作用,它使得早期剪裁的Token能够关注邻近Token的更深表示,从而促进了预测的一致性。
4.3. Object detection and instance segmentation
作者在COCO数据集[36]上使用EVA-02-L Backbone对Cropr进行基准测试,初始化来自Objects365[52]的微调预训练权重。作者遵循Fang等人的做法[19],使用Cascade Mask R-CNN作为任务头以支持检测和分割。图像被调整为像素,patch大小为16,从而产生个patch。如同Fang等人的做法[19],全局注意力与窗口注意力交替使用。窗口大小设为16,这生成了一个初始的窗口网格。为了在保持窗口大小的同时实现剪枝,作者应用了五阶段的剪枝计划。在每一阶段,token的数量减少到,最终实现了97%的TPR。剪枝发生在第5、8、11、14和20层之后,即全局注意力层之前。使用LLF方法,允许任务头在不作修改的情况下使用。
表3显示,Cropr在检测和分割任务上均优于 Baseline 模型(2)-(4)。与未剪枝模型相比,Cropr的性能差距适中,这符合其较高的TPR预期。尽管采用了优化的窗口注意力结构,Cropr仍实现了54%的FLOPs减少(未剪枝 Baseline :2790 GFlops,Cropr:1273 GFlops),同时编码器的速度提高了2.4倍,整体模型的速度提高了1.9倍。
图7表明,Cropr模块专注于与目标目标相应的任务相关图像区域。有趣的是,即使随机剪枝器也可以在结合LLF后作为有效的检测器使用,尽管可能会出现更多错误。
4.4. Ablation study
Crop模块设计。在表4a中,作者将作者简化版的跨注意力设计(单头,不包含QKV和头投影,不包含LN)与更复杂的多头自注意力(MHA)设计(16个头,包含QKV和头投影,包含LN)进行了比较。简化的设计在效率和性能指标上均优于MHA。在表4b中,作者评估了一种替代的选择方法,该方法从跨注意力分布中无放回地进行采样。这种采样方法的效果不如Top-K方法。表4c显示,将带有LN和残差连接的MLP集成到聚合模块中可以提高Token的选择效果。重要的是,这种修改不会影响推理时间的效率指标,因为在训练完成后移除了聚合器。最后,在表4d中,作者展示了在Crop中阻止梯度 Stream 可以带来更好的结果,这很可能是因为梯度干扰被防止了。
Token融合。作者将LLF与表5中的几种替代方案进行了比较:
Cross-Attn 应用了一个具有格子状学习 Query 的交叉注意块,这些 Query 跨注意力关注最后一个层产生的保留Token,也就是说,裁剪掉的Token不会被重新激活。进一步说明的是,这个交叉注意块是从头开始进行训练的。
Token Concat 通过在最后一个层之后将裁剪掉的Token连接起来来重新激活它们。“Cross-Attn + Concat”结合了这两种方法,在最后一个层之后跨注意力关注连接起来的Token。
MHSA + Concat 与此类似,但使用的是一个从头开始训练的完整自注意力 Transformer 块。
最后, DToP是第四章第二节中讨论的logit融合方法。所有方法都在ADE20k数据集上用于语义分割的评估。
不重新激活剪枝的Token,如在“Cross-Attn”中表现非常差。"Token Concat"、“Cross-Attn + Concat”和“DToP”会重新激活剪枝的Token,但不支持剪枝和保留Token之间的自注意力机制,因此表现不佳。
相比之下,“MHSA+6 Concat”和LLF允许Token之间进行注意交互,从而提高mIoU。值得注意的是,与未经剪枝的基准相比,LLF在不引入任何额外参数的情况下表现出色。
- Conclusion
实验表明,通过剪枝特定任务中最具信息量的Token之外的Token,ViTs可以在较小的性能损失下得到加速。作者通过将其应用于分类之外的语义分割、实例分割以及目标检测等多种任务,展示了该方法的灵活性。
未来的工作可以将Cropr扩展到其他视觉任务中,通过适应辅助头来进行。此外,作者方法基于Token的特性表明它在其他模态(如语言和音频)中也具有更广泛的应用潜力。
参考
[0]. Token Cropr: Faster ViTs for Quite a Few Tasks .
扫码加入👉「集智书童」交流群
(备注: 方向+学校/公司+昵称 )
点击下方“ 阅读原文 ”,
了解更多AI学习路上的 「武功秘籍」
