Castling-ViT：通过在 Vision Transformer 推理时切换至线性-角度注意力来压缩自注意力
计算机视觉与模式识别
2024-07-26 v5
摘要
Vision Transformers（ViTs）已展现出令人印象深刻的性能，但与卷积神经网络（CNNs）相比仍需要较高的计算成本，原因之一在于 ViTs 的注意力度量全局相似性，从而随输入 token 数量呈二次复杂度。现有高效 ViTs 采用局部注意力（如 Swin）或线性注意力（如 Performer），牺牲了 ViTs 捕获全局或局部上下文的能力。在本工作中，我们提出一个重要研究问题：ViTs 能否在学习全局与局部上下文的同时，在推理时更高效？为此，我们提出名为 Castling-ViT 的框架，其使用线性-角度注意力和基于掩码 softmax 的二次注意力训练 ViTs，但在 ViT 推理时仅切换为具有线性角度注意力。我们的 Castling-ViT 利用角度核通过谱角度度量查询与键之间的相似性。我们进一步用两种技术简化它：（1）一种新颖的线性-角度注意力机制：我们将角度核分解为线性项和高阶残差，仅保留线性项；（2）我们采用两个参数化模块来近似高阶残差：一个深度可分离卷积和一个辅助掩码 softmax 注意力，以帮助学习全局和局部信息，其中 softmax 注意力的掩码被正则化以逐渐变为零，从而在 ViT 推理时不产生额外开销。在三个任务上的大量实验和消融研究一致验证了所提 Castling-ViT 的有效性，例如，在与具有原始 softmax 注意力的 ViTs 相比 FLOPs 相当时，在 ImageNet 分类上实现高达 1.8% 的更高准确率或 40% 的 MACs 降低，在 COCO 检测上 mAP 提高 1.2。
引用
@article{arxiv.2211.10526,
title = {Castling-ViT: Compressing Self-Attention via Switching Towards Linear-Angular Attention at Vision Transformer Inference},
author = {Haoran You and Yunyang Xiong and Xiaoliang Dai and Bichen Wu and Peizhao Zhang and Haoqi Fan and Peter Vajda and Yingyan Celine Lin},
journal= {arXiv preprint arXiv:2211.10526},
year = {2024}
}
备注
CVPR 2023 Camera Ready