在单路径中将自注意力机制修剪为卷积层

视觉变换器（ViTs）在各种计算机视觉任务中取得了令人印象深刻的性能。然而，使用多头自注意力（MSA）层建模全局相关性导致了两个广为人知的问题：巨大的计算资源消耗和缺乏对局部视觉模式建模的内在归纳偏置。为了解决这两个问题，我们设计了一种简单而有效的方法，称为单路径视觉变换器剪枝（SPViT），以高效且自动地将预训练的ViTs压缩为具有适当局部性的紧凑模型。具体而言，我们首先提出了一种新颖的MSA与卷积操作之间的权重共享方案，提供了一个单一路径空间来编码所有候选操作。通过这种方式，我们将操作搜索问题转化为在每个MSA层中寻找要使用的参数子集的问题，这显著降低了计算成本和优化难度，并且可以利用预训练的MSA参数对卷积核进行良好的初始化。基于这一单一路径空间，我们引入了可学习的二值门控机制来编码MSA层中的操作选择。同样地，我们进一步使用可学习的门控机制来编码前馈网络（FFN）层中细粒度的MLP扩展比例。这样，我们的SPViT通过优化这些可学习的门控机制，从一个庞大且统一的搜索空间中自动探索，并灵活调整每个独立密集模型的MSA-FFN剪枝比例。我们在两个代表性的ViTs上进行了广泛的实验，结果表明我们的SPViT在ImageNet-1k数据集上的剪枝效果达到了新的最先进水平（SOTA）。例如，我们的SPViT可以将DeiT-B模型的浮点运算量减少52.0%，同时获得0.6%的Top-1精度提升。源代码可在https://github.com/ziplab/SPViT获取。