面向移动视觉 Transformer 的可分离自注意力

移动视觉 Transformer（MobileViT）能在包括分类与检测在内的若干移动视觉任务上取得最先进性能。尽管这些模型参数更少，但与基于卷积神经网络的模型相比具有较高延迟。MobileViT 中的主要效率瓶颈是 Transformer 中的多头自注意力（MHA），其相对于 token（或图像块）数 $k$ 需要 $O(k^2)$ 的时间复杂度。此外，MHA 需要昂贵操作（如批矩阵乘法）来计算自注意力，影响了资源受限设备上的延迟。本文引入一种具有线性复杂度即可 $O(k)$ 的可分离自注意力方法。所提方法一个简单而有效的特点是使用逐元素操作计算自注意力，使其成为资源受限设备的良好选择。改进模型 MobileViTv2 在包括 ImageNet 物体分类与 MS-COCO 物体检测在内的若干移动视觉任务上均为最先进。以约三百万参数，MobileViTv2 在 ImageNet 数据集上取得 75.6% 的 top-1 准确率，优于 MobileViT 约 1%，同时在移动设备上运行快 $3.2\times$。我们的源代码见：\url{https://github.com/apple/ml-cvnets}