CAS-ViT:用于高效移动应用的卷积加法自注意力视觉Transformer

近年来,Vision Transformer(ViT)在计算机视觉领域取得了巨大突破。然而ViT模型通常计算复杂度高,难以在资源受限的移动设备上部署。为了解决这个问题,研究人员提出了Convolutional Additive Self-attention Vision Transformers (CAS-ViT),这是一种轻量级的ViT变体,旨在在效率和性能之间取得平衡。

这是8月份再arxiv上发布的新论文，我们下面一起来介绍这篇论文的重要贡献

核心创新:卷积加法token混合器(CATM)

CAS-ViT的核心创新在于提出了一种新颖的加法相似度函数,称为卷积加法token混合器(CATM)。与传统ViT中的多头自注意力机制相比,CATM大大降低了计算复杂度。

让我们来看看CATM与之前工作的对比:

1. 传统ViT中的多头自注意力(图a):计算复杂度为O(N^2),其中N是序列长度。这在处理高分辨率图像时计算开销很大。
2. MobileViTv2中的可分离自注意力(图b):将矩阵特征度量简化为向量,降低了复杂度。
3. SwiftFormer中的swift自注意力(图c):将自注意力的键减少到两个(Q和K),进一步加速推理。
4. 论文提出的卷积加法自注意力(图d):定义了一个新的相似度函数,将Q和K的上下文分数相加。

具体来说,CATM中的相似度函数定义为:

Sim(Q, K) = Φ(Q) + Φ(K)

其中Φ(·)是上下文映射函数,由基于Sigmoid的通道注意力C(·)和空间注意力S(·)具体化。CATM的输出为:

O = Γ(Φ(Q) + Φ(K)) · V

由于CATM中的操作都用卷积表示,其复杂度为O(N),大大低于传统自注意力的O(N^2)。

网络架构

CAS-ViT的整体架构如下:

主要特点包括:

1. 4个编码阶段,每个阶段之间使用Patch Embedding进行2倍下采样。
2. 每个阶段包含Ni个堆叠的块。
3. 每个块包含三个部分:集成子网络、CATM和MLP,都使用残差连接。
4. 集成子网络由3个使用ReLU激活的深度可分离卷积层组成。

通过调整通道数Ci和块数Ni,构建了一系列轻量级CAS-ViT模型。

实验结果

CAS-ViT在多个计算机视觉任务上进行了评估,包括图像分类、目标检测和语义分割。

图像分类

在ImageNet-1K数据集上的分类和计算量的对比如下:

可以看出,CAS-ViT模型(红色点)在精度和计算效率之间取得了很好的平衡。特别是XS和S变体,在参数量较少的情况下达到了很高的Top-1准确率。

目标检测和实例分割

在COCO val2017数据集上的目标检测和实例分割结果如下:

CAS-ViT作为主干网络,在RetinaNet和Mask R-CNN框架下都取得了很好的性能,尤其是在计算效率方面具有优势。

语义分割

在ADE20K数据集上的语义分割结果如下:

CAS-ViT在mIoU和计算效率之间也达到了很好的平衡。

可视化分析

热力图可视化显示,CAS-ViT能够准确定位关键区域,并获得更大的感受野,这有利于后续的密集预测任务。

消融实验

消融实验表明,CATM中的每个组件都对模型性能有重要贡献。当移除或替换某个组件时,准确率都会下降。

总结

CAS-ViT通过创新的卷积加法自注意力机制,在保持高性能的同时大大降低了计算复杂度。这使得它非常适合部署在资源受限的移动设备上,为高效移动视觉应用开辟了新的可能性。未来的工作可以探索在更大规模的数据集和模型上应用CAS-ViT,以及进一步优化其在移动设备上的部署效率。