单路径中将自注意力剪枝为卷积层
计算机视觉与模式识别
2024-01-17 v4 机器学习
摘要
视觉 Transformer（ViT）已在多种计算机视觉任务上取得令人印象深刻的性能。然而，利用多头自注意力（MSA）层建模全局相关性导致了两个广为人知的问题：巨大的计算资源消耗，以及缺乏建模局部视觉模式的内在归纳偏置。为解决这两个问题，我们设计了一种简单而有效的方法，称为单路径视觉 Transformer 剪枝（SPViT），以高效且自动地将预训练 ViT 压缩为添加了适当局部性的紧凑模型。具体而言，我们首先提出 MSA 与卷积操作间的一种新颖权重共享方案，提供编码所有候选操作的单路径空间。借此，我们将操作搜索问题转化为在每层 MSA 中寻找使用哪些参数子集，这显著降低了计算成本与优化难度，且卷积核可利用预训练 MSA 参数良好初始化。依托单路径空间，我们引入可学习二值门来编码 MSA 层中的操作选择。类似地，我们进一步采用可学习门来编码 FFN 层的细粒度 MLP 扩展比。如此，我们的 SPViT 优化可学习门，以从庞大统一的搜索空间中自动探索，并灵活调整每个独立稠密模型的 MSA-FFN 剪枝比例。我们在两个代表性 ViT 上进行了大量实验，表明我们的 SPViT 在 ImageNet-1k 上取得了剪枝的新 SOTA。例如，我们的 SPViT 可为 DeiT-B 削减 52.0% 的 FLOPs，同时获得显著的 0.6% top-1 准确率提升。源代码见 https://github.com/ziplab/SPViT。
引用
@article{arxiv.2111.11802,
title = {Pruning Self-attentions into Convolutional Layers in Single Path},
author = {Haoyu He and Jianfei Cai and Jing Liu and Zizheng Pan and Jing Zhang and Dacheng Tao and Bohan Zhuang},
journal= {arXiv preprint arXiv:2111.11802},
year = {2024}
}
备注
Accepted by TPAMI 2024