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层
2. 权重共享方案
3. 讨论
B. 单路径视觉Transformer剪枝
1. 搜索MSA或卷积层
2. 搜索细粒度MLP扩展比
3. 搜索目标
4. 微调
IV. 实验
A. 实验设置
1. 实现细节
2. 比较方法
B. 主要结果
C. 消融实验
1. 单路径与多路径搜索
2. 每个模块的有效性
3. 仅剪枝MSA或FFN层
4. 知识蒸馏
5. 搜索策略的效果
6. MSA和FFN层的剪枝比例
7. 引入局部性的影响
D. 搜索架构的观察
1. 在浅层块中鼓励局部性
2. 标准ViTs中深层FFN层具有更多冗余
3. 每个阶段内较浅的FFN层需要更高的容量在层次化ViTs中
V. 结论和未来工作
声明
下载1:OpenCV-Contrib扩展模块中文版教程
在「小白学视觉」公众号后台回复:扩展模块中文教程,即可下载全网第一份OpenCV扩展模块教程中文版,涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。 下载2:Python视觉实战项目52讲 在「小白学视觉」公众号后台回复:Python视觉实战项目,即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。 下载3:OpenCV实战项目20讲 在「小白学视觉」公众号后台回复:OpenCV实战项目20讲,即可下载含有20个基于OpenCV实现20个实战项目,实现OpenCV学习进阶。
交流群
欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~
