AnchorFormer: Differentiable Anchor Attention for Efficient Vision Transformer

TL;DR

AnchorFormer introduces differentiable anchor tokens to efficiently approximate global self-attention in vision transformers, significantly reducing computational complexity and improving performance across vision tasks.

Contribution

The paper proposes a novel differentiable anchor-based attention mechanism that reduces complexity and accelerates inference in vision transformers, applicable to multiple vision tasks.

Findings

Achieves up to 9.0% higher accuracy on ImageNet classification.

Reduces FLOPs by 46.7% compared to baselines.

Improves mAP by 81.3% on COCO detection.

Abstract

Recently, vision transformers (ViTs) have achieved excellent performance on vision tasks by measuring the global self-attention among the image patches. Given $n$ patches, they will have quadratic complexity such as $\mathcal{O}(n^2)$ and the time cost is high when splitting the input image with a small granularity. Meanwhile, the pivotal information is often randomly gathered in a few regions of an input image, some tokens may not be helpful for the downstream tasks. To handle this problem, we introduce an anchor-based efficient vision transformer (AnchorFormer), which employs the anchor tokens to learn the pivotal information and accelerate the inference. Firstly, by estimating the bipartite attention between the anchors and tokens, the complexity will be reduced from $\mathcal{O}(n^2)$ to $\mathcal{O}(mn)$, where $m$ is an anchor number and $m < n$. Notably, by representing the anchors with the neurons in a neural layer, we can differentiably learn these anchors and approximate global self-attention through the Markov process. It avoids the burden caused by non-differentiable operations and further speeds up the approximate attention. Moreover, we extend the proposed model to three downstream tasks including classification, detection, and segmentation. Extensive experiments show the effectiveness of our AnchorFormer, e.g., achieving up to a 9.0% higher accuracy or 46.7% FLOPs reduction on ImageNet classification, 81.3% higher mAP on COCO detection under comparable FLOPs, as compared to the current baselines.