TPAMI 2024 | 将自注意力剪枝成单路径中的卷积层

共 44069字,需浏览 89分钟

 ·

2024-06-22 11:06

点击上方小白学视觉”,选择加"星标"或“置顶”

重磅干货,第一时间送达

题目:Pruning Self-Attentions Into Convolutional Layers in Single Path

将自注意力剪枝成单路径中的卷积层

作者:Haoyu He; Jianfei Cai; Jing Liu; Zizheng Pan; Jing Zhang; Dacheng Tao; Bohan Zhuang


摘要

视觉Transformer(ViTs)在各种计算机视觉任务中取得了令人印象深刻的性能。然而,使用多头自注意力(MSA)层建模全局关联会导致两个广为人知的问题:大量的计算资源消耗和缺乏对局部视觉模式进行建模的内在归纳偏差。为了解决这两个问题,我们提出了一种简单但有效的方法,名为单路径视觉Transformer剪枝(SPViT),以高效且自动地将预训练的ViTs压缩成具有适当局部性的紧凑模型。具体而言,我们首先提出了一种在MSA和卷积操作之间共享权重的新颖方案,提供了一个单路径空间来编码所有候选操作。通过这种方式,我们将操作搜索问题转化为寻找每个MSA层中要使用的参数子集,从而显著降低了计算成本和优化难度,并且可以使用预训练的MSA参数很好地初始化卷积核。依靠单路径空间,我们引入了可学习的二进制门来编码MSA层中的操作选择。同样地,我们进一步使用可学习的门来编码FFN层的细粒度MLP扩展比。通过这种方式,我们的SPViT优化了可学习的门,自动从庞大且统一的搜索空间中探索,并灵活调整每个密集模型的MSA-FFN剪枝比例。我们在两个代表性ViTs上进行了广泛的实验,表明我们的SPViT在ImageNet-1k上的剪枝达到了新的SOTA。例如,我们的SPViT可以将DeiT-B的FLOPs减少52.0%,同时获得令人印象深刻的0.6%的top-1精度提升。

关键词

  • 高效模型

  • 后训练剪枝

  • 视觉Transformer

I. 引言

视觉Transformer [1]、[2]、[3] 引起了大量研究兴趣,并成为各种图像识别任务(如分类 [2]、[3]、[4]、[5],分割 [6]、[7]、[8] 和检测 [9]、[10]、[11])中的主要骨干。

尽管最近的发展令人兴奋,但多头自注意力层 [12] 引入的两个局限性已被认识到。首先,MSA层的一个众所周知的问题是时间和内存复杂度的平方,这阻碍了ViTs的规模化开发和部署,尤其是在建模长序列时。为此,后训练剪枝方法 [13]、[14] 从预训练的ViT模型中剪除不太重要的ViT组件,以使ViTs适应更多资源受限的场景。另一个根本问题是MSA层缺乏一种局部信息建模机制,这对于处理低级图像模式(如边缘和形状)至关重要。在这种情况下,先前的艺术作品建议通过在ViTs中插入卷积层来引入归纳偏差,例如在前馈网络(FFNs)[15]、[16]中,在ViT编码器之前[17]或在MSA层之前[18]。因此,问题来了:如何在添加适当局部性的同时,自动将预训练的ViT模型剪枝成高效的模型?最近,几种神经架构搜索(NAS)方法 [19]、[20]、[21] 提议将候选的MSA和卷积操作分别纳入搜索空间,从而派生出具有局部性的高效ViT模型。在这种情况下,不同的操作作为不同的路径被维护,NAS问题通过从多路径空间中采样路径来解决,如图1(a)所示。然而,不同的操作是独立训练的,这在计算上是昂贵的,但不必要的。

在本文中,我们提出了一种简单但有效的方法,称为单路径视觉Transformer(SPViT),它自动将预训练的ViTs剪枝成具有适当局部性的紧凑模型。关键是,我们的SPViT具有低搜索成本。具体来说,受 [22] 中发现的启发,MSA层具有建模局部区域的能力,我们首先开发了MSA和卷积操作之间的权重共享方案,它用MSA参数的一个子集来表达卷积操作。通过这种权重共享方案,我们进一步设计了一个新颖的单路径搜索空间,将所有候选操作(卷积和MSA)封装到每层的单个MSA中,如图1(b)所示。因此,我们可以通过索引MSA中间结果直接派生卷积操作的输出,并使用预训练的MSA参数获得良好的卷积核初始化。而不是在多路径空间 [19]、[20]、[21] 中选择路径,我们的SPViT将搜索问题公式化为寻找每个MSA层中要使用的权重子集,从而显著降低了搜索成本和优化难度。

接下来,我们提出了一种使用NAS为ViTs搜索紧凑高效架构的新方法。我们的方法在统一的搜索空间内集成了MSA和FFN层的剪枝。具体来说,我们首先在每个MSA层中引入可学习的二进制门,以根据我们的单路径搜索空间编码操作选择(卷积和MSA)。此外,考虑到FFN层消耗了大量的计算(例如,处理224×224图像时,DeiT-B模型[4]的17.5 G Mult-Adds中有11.1 G),我们结合FFN层的剪枝。为了在统一的搜索空间中联合剪枝MSA和FFNs,我们还在每个FFN隐藏维度上引入了可学习的二进制门。我们的方法随后自动优化这些可学习的门,并探索广泛的高效架构,灵活调整每个密集模型的MSA-FFN剪枝比例。在搜索阶段之后,我们遵循剪枝文献[23]、[24]、[25] 对搜索到的架构进行微调,并交付紧凑模型。

我们的主要贡献可以总结如下:

  • 我们提出了一种在MSA和卷积操作之间的权重共享方案,使得能够在单路径搜索空间中将所有候选操作编码到一个MSA层中。然后我们将搜索问题转化为寻找MSA参数的子集,从而显著降低了搜索成本和优化难度。

  • 通过单路径搜索策略,我们提出了SPViT,它自动将昂贵的全局MSA操作剪枝成轻量级的局部卷积操作,同时在所需的效率约束下搜索细粒度的MLP扩展比。

  • 在ImageNet-1k上的广泛实验表明,我们的SPViT在剪枝DeiT[4]和Swin[3]模型时获得了有利的剪枝性能,达到了SOTA剪枝性能。例如,我们的SPViT变体通过知识蒸馏实现了23.1%、28.3%和52.0%的FLOPs减少,DeiT-Ti、DeiT-S和DeiT-B模型的top-1准确率分别比未压缩模型高1.0%、0.4%、0.6%。我们还在搜索的架构上做出了一些有趣的观察,揭示了剪枝ViTs时的架构偏好。

III. 方法

我们首先在第III-A节介绍MSA和卷积操作之间的权重共享方案。然后在第III-B节详细说明我们享受权重共享方案优势的单路径视觉Transformer剪枝方法。

A. MSA和卷积操作之间的权重共享

权重共享方案指在不同操作之间共享一部分参数,这已被广泛的NAS方法证明是有效的,例如在候选卷积操作之间共享权重大大降低了搜索成本[57]、[58]、[59]。本文中,我们引入了MSA和卷积操作之间的新颖权重共享方案,这本质上帮助我们获得了一个单路径空间,从而降低了搜索成本和优化难度。

1. 重新审视卷积和MSA层

卷积层是CNN的基本构建块。设 分别为输入特征和卷积核,卷积核大小为 ,其中 分别为空间宽度、高度、输入维度和输出维度。标准卷积层在局部感受野内聚合特征,该感受野在集合中定义。形式上,设 的宽度和高度索引集,则位置 的标准卷积层输出可以表示为
Transformer将MSA层作为主要构建块。然而,MSA层不仅聚合相邻特征,其感受野更大,覆盖整个序列。考虑与(1)中相同的输入 作为一组 维嵌入,并设 为任意关键特征嵌入的索引,则位置 的MSA层输出可以定义为
其中 表示MSA层的head数。给定 -th head的输出维度 ,我们将 定义为相应的可学习值、查询和关键线性投影。因此,表示 -th head的注意力图,带有 的索引, 变为标量。通过另一个可学习的线性投影 保持与 相同的维度。为简单起见,我们忽略了MSA层中的位置编码[60]和所有可学习偏差项,这些项伴随着卷积层和MSA层中的可学习投影。我们讨论了我们的权重共享方案不依赖于某种类型的位置编码(见第III-A3节)。

MSA层的计算复杂度为[2]。在这种情况下,当建模高分辨率特征图时, ,第二个平方项用于建模全局关联,主导了计算。相比之下,在输出相同宽度和高度特征图的标准卷积层中,卷积层的计算复杂度保持为

2. 权重共享方案

本节中,我们展示了MSA操作参数的一个子集如何在 时表达瓶颈卷积操作[61],如图2所示。在以下内容中,我们通过逐步消除MSA操作独有的参数来揭示两种操作的共享参数。
首先,由于卷积操作仅处理局部区域内的特征,我们将非局部关联的注意力分数固定为0,将局部关联的注意力分数固定为1,表示用MSA参数表达卷积操作。形式上
将其代入(2),我们得到
这里我们使用 来表示具有局部性限制的输出(3)。在(4)中,我们强制MSA操作的heads关注以 为中心的局部窗口大小为 的区域。
接下来,我们展示了可以进一步将MSA操作参数剖析成卷积核。类似于[22],一种方法是定义一个双射映射 ,它将选定的head子集 分配到局部窗口内的特定位置 。然后,我们可以在(4)中用 替代 ,即
在(5)中,每个head关注局部窗口中的某个位置。然而,许多ViT变体的 ,例如DeiT-Ti和DeiT-S [4],这将权重共享方案限制在仅大架构中。此外,在MSA层中,不同heads倾向于关注不同区域,如[22]、[62]中可视化的。因此,定义子集 和识别最适合局部位置的映射函数 是非平凡的。为了使权重共享方案适用于通用ViT架构以及简化定义 的难度,我们使用可学习参数 ,后跟softmax函数 ,首先将MSA heads组合成卷积核位置:
其中。在这种情况下,具有缩放因子 的MSA heads组合学习关注局部 窗口中的预定义位置。我们通过在 上应用softmax函数确保heads组合保持原始比例。
接下来,根据[22],可以通过用 替代 并应用结合律,将(6)中的 剖析成卷积核:
其中。相反,考虑到复杂性,我们选择形成瓶颈卷积操作[61]并修改(6)为
其中
注意,在这种形式中,卷积操作输出维度 ,我们通过 投影回 。可以很容易地发现(7)中的 和(10)中的 不大于heads的维度 。然而,前者的计算复杂度为 ,后者为 。由于 ,我们选择使用瓶颈卷积以提高效率。
我们进一步添加批归一化(BN)[63]和ReLU非线性[64],得到由 表示的瓶颈卷积操作:

3. 讨论

提出的权重共享方案鼓励MSA操作的heads建模全局区域,同时学习组合以处理局部区域特征。这种方法的直觉主要来自两个方面。首先,局部区域确实是全局区域的一部分,因此MSA heads的组合自然具有处理局部区域的能力。其次,由于ViT编码器中的一些注意力heads通常关注查询像素周围的局部区域[22]、[65],用这些heads表达卷积操作的行为可能会得到合理的输出。然而,考虑到添加局部性的最佳位置可能因不同ViT模型而异,因此我们引入了SPViT,自动学习选择最佳操作(见第III-B节)。
总而言之,与[22]相比,提出的权重共享方案有两个主要区别。首先,权重共享方案在ViT模型上有更广泛的应用,不受使用相对位置编码(二次编码)和 条件的限制。原因是[22]学习了head子集 和(5)中head与卷积核位置之间的双射映射 ,采用二次编码。而我们通过用 形成注意力heads组合来简化了这一过程。其次,我们提出使用具有相同秩上限但效率更高的瓶颈卷积操作,而不是[22]中的标准卷积操作。

B. 单路径视觉Transformer剪枝

我们介绍了我们的SPViT方法,通过NAS剪枝MSA和FFN层。为此,我们用统一的MSA(第III-B1节中的UMSA)和统一的FFN(第III-B2节中的UFFN)层替换了原始的MSA和FFN层,如图3左上所示。这两层通过可学习的二进制门编码架构配置选择,形成统一的搜索空间。我们用总体目标优化这些二进制门和模型权重(见第III-B3节),以探索广泛的高效架构。最后,我们对搜索到的架构进行微调,获得紧凑的剪枝ViTs(见第III-B4节)。

1. 搜索MSA或卷积层

我们首先搜索ViTs中将MSA层剪枝成卷积层的最佳位置,使用我们的UMSA层,如图3底部所示。具体来说,利用我们在第III-A节中的权重共享方案,我们设计了一个单路径搜索空间,通过根据(8)索引和缩放 直接获得预定义卷积核 的卷积输出,而不需要将每个候选卷积操作保留为单独的路径。接下来,通过(12)获得瓶颈卷积输出 。得益于单路径搜索空间,我们以低计算成本得到了瓶颈卷积和MSA的输出。按计算复杂度从低到高的顺序,我们可以将瓶颈卷积和MSA公式化为有序集 ,其中
接下来,按照[66]、[67],我们定义了一系列二进制门来编码操作选择。对于任何索引为 的操作,门 从具有可学习门参数 的Bernoulli分布中采样,该参数可以自动优化。具体来说,我们使用 确定选择第 个操作的概率:
出于效率考虑,我们每层只选择一个操作。通过用 替换 ,我们强制每个UMSA层最多只有一个打开的门: 。这本质上表明一旦选择了更复杂的门( ),所有较少复杂的门将被关闭。UMSA层的输出因此可以表示为
其中 是第 个操作的输出。每个UMSA层随后跟随残差连接和层归一化[68],如ViTs所做。当所有门关闭时,采用跳跃连接操作。因此,我们的UMSA搜索空间包括三种类型的候选操作:跳跃连接、不同候选卷积核大小的瓶颈卷积和MSA。

2. 搜索细粒度MLP扩展比

FFN层也消耗大量FLOPs。在这项工作中,我们将MSA和FFN剪枝配置联合纳入统一搜索空间,灵活调整每个特定密集模型的MSA-FFN比例。为此,我们剪枝不重要的FFN隐藏维度,以搜索细粒度MLP扩展比。
MLP扩展比控制FFN层的隐藏维度数量。之前的工作[19]、[69]、[70]将MLP扩展比的搜索空间定义为一系列粗粒度值,例如[69]中的[3, 3.5, 4]。相反,给定预训练模型定义的MLP扩展比 ,我们提出了一种简单但有效的方法,通过我们的UFFN层搜索细粒度MLP扩展比 ,满足 。UFFN层的概述如图3右上所示。给定输入特征 ,FFN层的输出可以表示为,其中。这里,我们用 表示预定义的MLP扩展比和FFN层的隐藏维度。 是可学习的全连接投影,前者将 的通道维度映射到 ,后者将通道维度投影回 。在两个可学习的全连接投影之间添加GeLU [71]激活。
如(13)所示,我们用二进制门 编码FFN隐藏维度选择。通过在搜索过程中对每个隐藏维度应用二进制门,我们定义UFFN层输出为
这使我们能够识别最重要的隐藏维度,从而搜索具有细粒度MLP扩展比的FFN层。

3. 搜索目标

为了获得具有期望效率约束的架构,我们用辅助计算复杂度损失来优化网络。具体来说,我们定义了一个查找表,包含候选操作和模块的计算复杂度。计算成本 定义为 ,其中 是当前网络的计算复杂度, 是目标计算复杂度。最终的总体搜索目标定义为 ,其中 是平衡 和交叉熵损失 的权衡超参数。受益于MSA和卷积操作之间的权重共享方案,搜索成本大大降低(见第IV-C1节)。
通过我们的统一搜索空间和全局损失,我们探索了广泛的高效架构,并灵活调整每个特定密集模型的MSA-FFN比例(见第IV-C6节),并对剪枝ViTs时的首选紧凑架构做出了一些有趣的观察(见第IV-D节)。

4. 微调

在微调过程中,我们将之前(13)(15)中的随机二进制门变为确定性门,即。与搜索相同,我们获得 ,并在微调过程中在UMSA层中仅选择最多一个操作。输出可以表示为
如果选择的操作是瓶颈卷积,我们可以将预训练的 剖析成卷积核,并直接执行卷积,如(10)所述并在图2中所示。其他MSA参数被剪枝。结合(12),我们可以在微调过程中定义瓶颈卷积层为,其中。类似地,对于UFFN层,我们可以自动找到隐藏维度的子集 ,表示选择的隐藏维度。因此, ,UFFN层在微调过程中输出可由 得出。
我们进一步将我们的SPViT与知识蒸馏结合[41],这是一种在模型压缩场景中常用的技术,以补偿精度损失[14]、[23]。具体来说,我们让剪枝后的网络从教师网络预测的硬标签中学习,类似于[41]。由于我们的SPViT向ViT模型引入了卷积归纳偏差,它促进了来自CNN教师的指导(见第IV-C2节)。因此,微调期间的最终目标是:
其中 是CNN教师和剪枝ViT学生之间的硬标签蒸馏损失, 是平衡权衡超参数。请注意,尽管我们使用了类似于DeiT [4]的CNN教师,但我们的SPViT不需要像DeiT那样额外的蒸馏token,从而在推理过程中不引入计算开销。

IV. 实验

在本节中,我们通过剪枝不同的ViT模型验证了SPViT的有效性。我们对两系列ViT模型进行了实验:DeiT [4] 和 Swin [3],它们分别是代表性的标准ViT和层次ViT [2]、[74]模型。接下来,我们首先在第IV-B节中与基线剪枝方法进行比较,然后在第IV-C节中对SPViT的重要组成部分进行消融实验,最后在第IV-D节中分析在变化目标计算复杂度时观察到的一些模式。

A. 实验设置

1. 实现细节

我们从预训练模型中进行搜索和微调。除非另有说明,否则我们SPViT的训练细节与DeiT [4] 和 Swin [3] 中介绍的相同。在搜索过程中,我们将门参数 和其他网络参数的初始学习率分别设置为 。门参数 的初始值设置为1.5,以启用较高的选择MSA操作的初始概率。在UMSA层中,我们在搜索空间中包括最流行的1×1和3×3瓶颈卷积。我们进行网格搜索,根据不同模型的冗余情况设置不同的超参数 进行剪枝。在(17)中将 设置为1.0。网络架构在ImageNet-1k上约10个epochs后收敛,然后我们按照[23]、[75]从搜索的架构中微调,学习率为 。对于ImageNet-1k [26]的实验,我们在8个NVIDIA V100 GPU上以1024的总批次大小进行搜索和微调,除了SPViT-DeiT-Ti在4个NVIDIA V100 GPU上训练。我们还使用相同的硬件规格评估我们的模型。

2. 比较方法

在剪枝DeiT [4]时,我们与先前的最先进的ViT模块剪枝方法进行了比较,包括VTP [13]、S^2 ViTE [24]、UVC [23] 和 MDC [38]。有关更多详细信息,请参阅第II-A节。我们还比较了基于显著性的结构化剪枝(SSP)[31]、[35],它通过利用权重、激活和梯度信息删除子模块(实现借鉴自[24])。此外,我们比较了具有代表性的token压缩方法。对于token剪枝方法[27]、[72]、[76],DynamicViT [76]和IA-RED^2 [27]逐步估计要剪除的无信息token,而EViT [72]剪除具有最低注意力得分的token。对于token合并方法[29]、[30],ToMe [29]合并最相似的token,而eTPS [30]和dTPS [30]将不重要的token合并到重要token中。由于我们是剪枝Swin [3]的先驱,我们将直通估计剪枝方法(STEP)[73]调整为比较基线。STEP通过可学习门参数优化MSA层和FFN层隐藏维度的注意力heads的重要性分数,保留重要性分数最高的前k个attention heads和隐藏维度。

B. 主要结果

我们通过表I中的SPViT与基线方法进行比较,调查我们方法的有效性。
首先,SPViT变体在模块剪枝方法中表现出色:例如,我们的SPViT-DeiT-B将DeiT-B的FLOPs压缩到原来的48.0%,同时仅损失0.3%的top-1准确率,超过了SOTA方法UVC-DeiT-B 0.9%的top-1准确率。我们推测,通过SPViT变体搜索到的混合架构可以学习到更多判别特征表示,从而提升性能。此外,由于插入的卷积在计算上是高效的,我们的SPViT在DeiT-B上表现出更高的FLOPs减少和更低的top-1准确率损失。在剪枝层次化Swin [3]模型时,SPViT-Swin变体也作为强基线。例如,SPViT-Swin-S节省了29.9%的FLOPs,仅损失0.6%的top-1准确率,超过STEP基线对照[73] 2.8%的top-1准确率。
其次,SPViT蒸馏变体在多个主干网络上取得了最强的结果,超过了SOTA方法:我们注意到SOTA方法(UVC和MDC)在剪枝DeiT-Ti和DeiT-S模型时比没有知识蒸馏的SPViT变体表现更好。原因是这些方法享有更大的剪枝空间,带来了更细粒度的剪枝(见第II-A节)。然而,我们的方法派生出的同时具有MSA和卷积的混合模型可以从CNN教师中进一步受益,并且表现更好。例如,我们的SPViT-DeiT-Ti、SPViT-DeiT-S和SPViT-DeiT-B在SOTA方法中实现了最高的top-1准确率,分别为73.2%、80.3%和82.4%。我们还实验证明,通过知识蒸馏,SPViT-Swin-Ti和SPViT-Swin-S取得了显著的性能提升,几乎没有top-1准确率损失。
此外,我们的SPViT-DeiT-B在更少的FLOPs下大幅超越了DeiT-B主干网络下的token压缩方法[28]、[30]、[72]。而在DeiT-S主干网络下,token压缩方法通常相对于模块剪枝方法表现强劲。值得注意的是,token压缩保持模型权重不变,甚至增加模型参数。我们推测,由于DeiT-S主干网络较小且更紧凑,使用模块剪枝剪除模型权重通常会影响性能。即便如此,我们的SPViT-DeiT-S仍然在FLOPs-准确率权衡上与SOTA token压缩方法dTPS [29]和ToMe [29]相当。

C. 消融实验

1. 单路径与多路径搜索

我们通过比较SPViT与多路径对照,调查了单路径搜索的有效性,见表II。具体来说,在多路径实现中,我们在搜索前随机初始化候选瓶颈卷积操作的权重,其他组件与单路径版本相同。我们观察到,与多路径对照相比,单路径搜索方法节省了超过50%的搜索成本,同时在相同计算复杂度下获得了更高性能和更少参数的更好架构。我们推测,单路径表述的优越性来自于优化难度的减轻和卷积核的良好初始化。

2. 每个模块的有效性

我们通过逐步添加SPViT的每个模块,在剪枝DeiT-Ti和DeiT-B时调查了它们的有效性。ImageNet-1k上的结果报告见表III。

3. 仅剪枝MSA或FFN层

我们在类似的FLOPs节省下,调查了仅剪枝MSA或FFN层的有效性。对于SPViT-DeiT-Ti和SPViT-DeiT-B,我们实验证明,仅剪枝MSA的性能略高于仅剪枝FFN,即top-1准确率高0.1%。然而,它表现出更多参数,因为FFN剪枝在参数上更有效,例如,节省0.04 G FLOPs分别剪除0.1 M和0.2 M模型参数的MSA(在我们SPViT和heads剪枝[31]中)和FFN剪枝。此外,在相同或更低FLOPs下,联合剪枝MSA和FFN层比仅剪枝MSA或FFN层性能相似或更好。我们推测,联合剪枝MSA和FFN层极大地扩大了搜索空间,有助于找到具有最佳MSA-FFN剪枝权衡的强大紧凑架构。我们还观察到,在相同FLOPs节省下,联合剪枝MSA和FFN层比仅剪枝FFN层有更多模型参数(SPViT-DeiT-Ti)。这是因为联合剪枝需要在剪枝MSA和FFN层之间进行平衡。然而,当剪枝具有更高冗余的DeiT-B时,我们可以在联合剪枝中寻求更积极的FLOPs节省,以在不牺牲太多准确率的情况下获得较少的参数。

4. 知识蒸馏

我们采用硬标签蒸馏以补偿精度损失,遵循最近的剪枝工作[23]、[24]。我们观察到,它比紧凑模型表现更好。为了进一步研究这一现象,我们在相同的SPViT-DeiT-Ti设置下报告了不同教师选择的结果,见表IV。为了更好地反映性能提升,我们以UVC-DeiT-Ti [23]为参考,它采用未压缩的DeiT-Ti模型作为教师并进行软标签蒸馏。我们发现,UVC-DeiT-Ti的蒸馏策略提供了显著的2.0% top-1准确率提升。然而,我们的SPViT-DeiT-Ti采用相同的蒸馏策略仅从SPViT-DeiT-Ti中提高了0.4%的top-1准确率。如第II节所介绍,卷积和MSA表现不同,学习不同的特征表示,我们推测,未压缩的DeiT-Ti教师的知识对我们SPViT搜索到的混合架构可能是噪声。然而,当切换到具有与DeiT-Ti相似容量的CNN教师RegNetY-1.6GF时,SPViT-DeiT-Ti实现了显著的性能飞跃,top-1准确率提高了2.1%,我们推测SPViT中搜索到的卷积层可以从CNN教师中获得更好的指导。最后,我们切换到具有更高容量的CNN教师RegNetY-16GF,显示性能可以进一步提高0.4%的top-1准确率,显著超过了具有RegNetY-16GF的UVC-DeiT-Ti。

5. 搜索策略的效果

为了验证我们对每个模块搜索策略的有效性,我们在表V中将SPViT-DeiT-Ti变体与随机搜索架构进行了比较。具体来说,我们在两种设置下进行了比较:1)保持FFN MLP扩展比固定,仅比较将MSA剪枝为卷积层与随机采样的MSA或卷积层;2)保持MSA层固定,仅比较将FFN扩展比剪枝为随机采样的FFN扩展比。对于每种设置,我们随机采样5个计算复杂度在目标FLOPs附近的架构,然后对这些架构进行微调并报告最高的top-1准确率。我们观察到,在相同的FLOPs下,我们的SPViT变体分别在两个设置中比随机搜索对照高0.5%和0.2%的top-1准确率。性能提升证明了我们搜索策略的有效性。

6. MSA和FFN层的剪枝比例

由于我们的SPViT可以在计算复杂度约束下自动搜索架构,我们在表VI中调查了MSA和FFN层的剪枝比例。我们观察到,不同模型之间的剪枝比例不同,例如,SPViT-Swin-Ti对MSA层的剪枝比例更高,而SPViT-Swin-S对FFN层的剪枝比例更高。这表明,在给定的目标效率约束下,SPViT可以灵活定制不同密集模型的适当剪枝比例。

7. 引入局部性的影响

我们的方法通过将自注意力剪枝为卷积层来压缩模型,同时引入局部性。为了调查引入局部性的影响,我们保持FFN层不变,在相同目标FLOPs下,将我们的方法与两个head剪枝基线进行了比较。一个基线随机选择剩余的heads,另一个基线使用与SPViT相同的NAS-based gating机制选择剩余的heads。结果见表VII。我们观察到,我们的方法在第2行和第3行的基线中表现出显著的优势。结果验证了我们在ViT剪枝中引入适当局部性的动机。

D. 搜索架构的观察

我们的SPViT可以在搜索过程中,在给定的效率约束下,自动将预训练的ViT模型剪枝为高效的模型,并适当插入局部性。这里,我们展示了使用SPViT对标准和层次化ViTs进行搜索架构的经验观察,在逐渐减少目标FLOPs时。如图4所示,我们可视化了目标FLOPs分别为3.9 G、3.3 G和2.4 G时的SPViT-DeiT-S,目标FLOPs分别为11.7 G、8.4 G和6.5 G时的SPViT-DeiT-B,以及目标FLOPs分别为3.9 G、3.4 G和2.7 G时的SPViT-Swin-Ti的搜索架构。

1. 在浅层块中鼓励局部性

如图4(a)-(i)所示,SPViT变体学习将浅层MSA层剪枝为瓶颈卷积层或跳跃连接(在SPViT-DeiT-S的前两个块中,在SPViT-DeiT-B的前六个块中,以及在SPViT-Swin-Ti的前三个块中)。这一观察与先前的工作[22]、[62]一致,表明浅层MSA层包含更多多余的全局关联。所有模型的最后两个块被观察到保持为MSA层,表明深层中全局操作的重要性。为了进一步支持这一观察,我们可视化了预训练DeiT-B模型的注意力概率,如图5所示。我们可以观察到,注意力概率在浅层和深层块之间显示出明显的不同模式。在浅层块(块1和2)中,许多MSA heads关注查询特征嵌入周围的局部区域,而几乎所有深层块(块11和12)中的heads关注全局区域。因此,注意力概率模式表明我们的SPViT鼓励将浅层MSA层剪枝为卷积层,同时保持深层MSA层不变是合理的。

2. 标准ViTs中深层FFN层具有更多冗余

对于预训练的DeiT模型,MLP扩展比在所有块中设置相同。如图4(a)-(f)所示,在剪枝DeiT模型时,网络优先剪除最后几个块中的隐藏维度。由于FFN层中更多的隐藏维度表明更高的容量,反之亦然[79]、[80],我们推测,最后几个FFN层需要的容量比其他层少,因为它们处理的注意力图较少多样化[81]、[82]、[83]。

3. 每个阶段内较浅的FFN层需要更高的容量在层次化ViTs中

与标准ViTs不同,Swin模型具有层次化架构,将块分成几个阶段。在每个阶段的开始,输入token合并到更高的嵌入维度。如图4(g)-(i)所示,在同一阶段内,较浅层保留更多隐藏维度。我们推测,建模由于token合并引起的特征空间变化需要更高的容量用于层次化ViTs。

V. 结论和未来工作

在本文中,我们介绍了一种在MSA和卷积操作之间的新颖权重共享方案,允许用MSA参数的一个子集表达卷积操作,并同时优化两种操作。基于权重共享方案,我们提出了SPViT,以减少ViTs的计算成本,并自动引入适当的归纳偏差,具有低搜索成本。具体来说,SPViT可以搜索是否在ViTs中将MSA剪枝为卷积层,并在所需的效率约束下搜索FFN层的最佳隐藏维度。通过将SPViT应用于两种流行的ViT变体,我们对ViT组件的重要性做出了一些有意义的观察。
SPViT的一个局限是将MSA剪枝为卷积层的粒度过粗,将前者剪枝为后者会由于巨大参数差距而导致判别能力的大幅损失。因此,我们将探索在SPViT中添加heads剪枝作为未来工作。另一个未来方向是将权重共享方案扩展到其他卷积操作变体,如深度卷积[84]和组卷积[85],提供更广泛的应用。

声明

本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与我们联系,我们将在第一时间回复并处理。
   
下载1:OpenCV-Contrib扩展模块中文版教程
在「小白学视觉」公众号后台回复:扩展模块中文教程即可下载全网第一份OpenCV扩展模块教程中文版,涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。
下载2:Python视觉实战项目52讲
小白学视觉公众号后台回复:Python视觉实战项目即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。
下载3:OpenCV实战项目20讲
小白学视觉公众号后台回复:OpenCV实战项目20讲即可下载含有20个基于OpenCV实现20个实战项目,实现OpenCV学习进阶。

交流群


欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉、传感器自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~


浏览 13
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报