Home

Awesome

Fcaformer: Forward Cross Attention in Hybrid Vision Transformer

Official PyTorch implementation of FcaFormer


Comparison with SOTA models. The latency is measured on a single NVIDIA RTX 3090 GPU with batchsize=64.

<p align="center"> <img src="https://s1.ax1x.com/2023/03/20/pptbZ7Q.png" width=100% height=100% class="center"> </p>

Architecture

<p align="center"> <img src="https://s1.ax1x.com/2023/03/20/pptbmkj.png" width=100% height=100% class="center"> </p> A FcaFormer stage which consists of L FcaFormer blocks. Compared with standard ViT blocks, we add token merge and enhancement (TME) part, which uses long stride large kernel convolutions to merge tokens spatially as cross tokens, and small kernel convolutions to further enhance tokens for channel mixing (FFN). The cross tokens are then used in later blocks as extra tokens for multi-head attention, after being calibrated by multiplying them with learned scaling factors. <p align="center"> <img src="https://s1.ax1x.com/2023/03/20/pptbnts.png" width=60% height=60% class="center"> </p> The overall structure of FcaFormer Mobels for image classification tasks. (a) Plain FcaFormer model. This model is directly modified from the DeiT structure. (b) Hybrid FcaFormer model. The ConvNet stages are composed of ConvNext blocks.

Introduction

Currently, one main research line in designing a more efficient vision transformer is reducing the computational cost of self attention modules by adopting sparse attention or using local attention windows. In contrast, we propose a different approach that aims to improve the performance of transformer-based architectures by densifying the attention pattern. Specifically, we proposed forward cross attention for hybrid vision transformer (FcaFormer), where tokens from previous blocks in the same stage are secondary used. To achieve this, the FcaFormer leverages two innovative components: learnable scale factors (LSFs) and a token merge and enhancement module (TME). The LSFs enable efficient processing of cross tokens, while the TME generates representative cross tokens. By integrating these components, the proposed FcaFormer enhances the interactions of tokens across blocks with potentially different semantics, and encourages more information flows to the lower levels. Based on the forward cross attention (Fca), we have designed a series of FcaFormer models that achieve the best trade-off between model size, computational cost, memory cost, and accuracy. For example, without the need for knowledge distillation to strengthen training, our FcaFormer achieves 83.1% top-1 accuracy on Imagenet with only 16.3 million parameters and about 3.6 billion MACs. This saves almost half of the parameters and a few computational costs while achieving 0.7% higher accuracy compared to distilled EfficientFormer.

Image classification on ImageNet 1k

Models#params (M)MACs(M)Top1 acc
Deit-T5.5-72.2
FcaFormer-L05.9-74.3
Swin-1G6.31.578.4
FcaFormer-L16.21.480.3
ConvNext-Tiny294.582.1
Swin-Tiny294.581.3
FcaFormer-L216.33.683.1

Semantic segmentation on ADE20K

MethodBackbonemIOU#paramsMACs
DNLResNet-10146.0691249
OCRNetResNet-10145.356923
UperNetResNet-10144.9861029
UperNetDeiT III (ViT-S)46.842588
UperNetSwin-T45.860945
UperNetConNext-T46.760939
UperNetFcaFormer-L247.646730

Object detection on COCO

Backbone#params.MACsAPboxAPbox50APbox75APmaskAPmask50APmask75
Mask-RCNN 3 × schedule
Swin-T4826746.068.150.341.665.144.9
ConvNext-T4826246.267.950.841.765.044.9
FcaFormer-L23724947.068.951.842.165.745.4
Cascade Mask-RCNN 3 × schedule
X101-6414097248.366.452.341.764.045.1
Swin-T8674550.469.254.743.766.647.3
ConvNext-T8674550.469.154.843.766.547.3
FcaFormer-L27472851.069.455.543.967.047.4

Test on edge device

batch size=1, image size=224, four threads. ARM:Quad Core Cortex-A17

Models#params.MACsLatency (ms)Memeory (M)Acc (%)pretrained weights
ConvNext-Tiny294.587512982.1-
ConvNext-Small508.7161821183.1-
ConvNext-Base8915.4270836483.8-
ConvNext-Large19834.4560476484.3-
Swin-Tiny294.558813981.3-
Swin-Small508.7157622283.0-
Swin-Base8815.4262437883.5-
FcaFormer-L1(Micro)6.21.43124280.3-
FcaFormer-L2(Tiny)163.67289583.1FcaFormer-L2 Access Code:1234
FcaFormer-L3(Small)286.7134414884.2-
FcaFormer-L4(Base)6614.5262432884.9-