CAS-ViT: Convolutional Additive Self-attention Vision Transformers for Efficient Mobile Applications

TL;DR

CAS-ViT introduces Convolutional Additive Self-attention with the Convolutional Additive Token Mixer (CATM) to replace costly matrix multiplications and Softmax in Vision Transformers. By decomposing interactions into additive spatial and channel context maps, the method achieves linear-like complexity with respect to input size and enables efficient deployment on mobile hardware while preserving high accuracy. The paper presents a family of lightweight CAS-ViT models (XS/S/M/T) that excel across image classification, object detection/segmentation, and semantic segmentation, with strong throughput on GPU, ONNX, and iPhone Neural Engine. This work demonstrates a practical, scalable approach to edge-friendly transformers that balance performance, efficiency, and deployability.

Abstract

Vision Transformers (ViTs) mark a revolutionary advance in neural networks with their token mixer's powerful global context capability. However, the pairwise token affinity and complex matrix operations limit its deployment on resource-constrained scenarios and real-time applications, such as mobile devices, although considerable efforts have been made in previous works. In this paper, we introduce CAS-ViT: Convolutional Additive Self-attention Vision Transformers, to achieve a balance between efficiency and performance in mobile applications. Firstly, we argue that the capability of token mixers to obtain global contextual information hinges on multiple information interactions, such as spatial and channel domains. Subsequently, we propose Convolutional Additive Token Mixer (CATM) employing underlying spatial and channel attention as novel interaction forms. This module eliminates troublesome complex operations such as matrix multiplication and Softmax. We introduce Convolutional Additive Self-attention(CAS) block hybrid architecture and utilize CATM for each block. And further, we build a family of lightweight networks, which can be easily extended to various downstream tasks. Finally, we evaluate CAS-ViT across a variety of vision tasks, including image classification, object detection, instance segmentation, and semantic segmentation. Our M and T model achieves 83.0\%/84.1\% top-1 with only 12M/21M parameters on ImageNet-1K. Meanwhile, throughput evaluations on GPUs, ONNX, and iPhones also demonstrate superior results compared to other state-of-the-art backbones. Extensive experiments demonstrate that our approach achieves a better balance of performance, efficient inference and easy-to-deploy. Our code and model are available at: \url{https://github.com/Tianfang-Zhang/CAS-ViT}

Figures (6)

• Figure 1: Parameters v.s. Top-1 accuracy on ImageNet-1K deng2009imagenet. The circle size indicates Gflops and best viewed in color.
• Figure 2: Comparison of diverse self-attention mechanisms. (a) is the classical multi-head self-attention in ViT dosovitskiy2020image. (b) is the separable self-attention in MobileViTv2 mehta2022separable, which reduces the feature metric of a matrix to a vector. (c) is the swift self-attention in SwiftFormer shaker2023swiftformer, which achieves efficient feature association only with $\textbf{Q}$ and $\textbf{K}$. (d) is proposed convolutional additive self-attention.
• Figure 3: Upper: Illustration of the classification backbone. Four stages downsample the original image. Lower: Block architecture with $N_i$ blocks stacked in each stage.
• Figure 4: Heatmap visualization for the last layer of the backbone. From top to bottom are the original image, PoolFormer-S24 yu2022metaformer, SwiftFormer-L1 shaker2023swiftformer and CAS-ViT-M respectively. Our model accurately captures the global context and achieves a larger receptive field compared to other SOTAs, which is beneficial for dense prediction tasks.
• Figure 5: Visualization of object detection and instance segmentation results on COCO 2017.Upper: Prediction results of PoolFormer-S12 yu2022metaformer, Lower: Prediction results of CAS-ViT-M. Our method is able to detect and segment instances accurately.