[Paper] Next-ViT: Next Generation Vision Transformer for Efficient Deployment inRealistic Industrial Scenarios

Instroduction

최근 ViT는 다양한 computer vision 분야에서 주목을 받고 많은 성공을 이루어냈다. 하지만 deployment 관점의 ViT는 CNN보다 훨씬 느리기때문에 여전히 CNN이 압도하고 있다.

 

ViT의 높은 latency를 해결하기위해 많은 연구가 이루어졌다. 예를 들어 Swin Transformer, PVT(Pyramid Vision Transformer)는 MHSA(Multi-Head Self Attention)의 quadratic하게 증가하는 연산량을 개선시켰다. 이외에도 MobileViT, Efficientformer, Coatnet 등은 convolution block들과 Transformer block들을 결합하여 accuracy와 latency의 trade-off를 더 개선시킨 CNN-Transformer hybrid architecture이다. 또한 이런 hybrid architecture는 낮은 stage에는 convolution block, 마지막 몇 stage에는 Transformer block을 쌓아놓는다. 하지만 hybrid strategy와 같은 것은 downstream에서 saturation이 잘 발생한다는 것을 발견했다고 한다. 그리고 convolution block과 Transformer block 모두 성능과 효율성을 동시에 가질 수 없다고 한다.

 

위 문제들을 개선하기위해 세 가지 중요한 요소들을 개발하였다.

 

1. Next Convolution Block (NCB)

deploy에 친화적인 MHCA(Multi-Head Convolutional Attention)를 사용하여 short-term dependency information을 캡처할 수 있다.

2. Next Transformer Block (NTB)

NTB는 long-term dependency information을 캡쳐할 수 있을 뿐만 아니라 가볍고 모델링 능력을 향상시키기위한 high-and-low frequency signal mixer를 사용한다.

3. Next Hybrid Strategy (NHS)

각 stage에 NCB와 NTB를 쌓는 hybrid paradigm이다. 각 stage에 Transformer block의 비율은 줄이고 다양한 downstream task에서 ViT의 높은 정확도를 유지킨다.

Next-ViT overall architecture

Method

Next Convolution Block (NCB)

• NCB는 BottleNeck block의 deployment면에서의 이점과 Transformer block의 성능을 가졌다.
• NCB는 MetaFormer의 architecture를 따른다.
• efficient token mixer로써 Multi-Head Convolutional Attention(MHCA)를 고안했고 deployment에 최적화된 convolution operation을 가졌다. (efficient token mixer는 보통 self-attention이라고 생각하면 될 것 같다. ref)
• 수식으로 표현하면 다음과 같다.

$z^{l-1}$: $l-1$ block의 input

$\tilde{z}^l, z^l$: 각각 MHCA과 $l$ NCB의 output

 

Multi-Head Convolutional Attention (MHCA)

• 효율적인 convolution 연산을 위해 Convolutional Attention (CA)를 고안했다.
• CA는 MHSA의 multi-head에서 아이디어를 얻었으며 효과적인 local representation learning을 위해 다른 position에서의 representation 정보도 같이 attention하는 구조로 빌드했다.

 

• MHSA를 수식으로 간략히 정의하면 다음과 같다.

• MHCA는 parallel representation subspace $h$에서 information을 캡처한다.
• input feature인 z를 channel dimension에서 multi-head form에 맞게 나눈것을 $z = [z_1, z_2,...,z_h]$로 표현한다.
• MHCA에는 head들끼리 information interaction을 위해 projection layer($W^P$)를 붙인다.
• CA는 다음과 같이 정의된 single-head convolutional attention이다.

$T_m, T_n$: input feature $z$에 서로 인접한 토큰들

$O$: 학습 가능한 parameter $W$와 input tokens $T_{{m, n}}$의 내적

• CA는 parameter $W$를 학습시키는 것을 통해 local receptive field에서 다른 토큰들 사이의 친화성(?) 유사성(?)을 학습할 수 있다.
• Figure 3의 (f)를 보면 MHCA는 group convolution(multi-head convolution)과 point-wise convolution(Conv 1x1)으로 구현되어있는 것을 확인할 수 있다.
• TensorRT에 다양한 data type에서 빠른 inference를 위해 모든 MHCA의 head dim은 32로 지정했다고 한다.
• 전통적인 Transformer block의 LayerNorm(LN)과 GELU보다 inference 가속을 위해 Batch Norm(BN)과 ReLU를 적용했다고 한다.

Next Transformer Block (NTB)

Transformer block은 global information을 캡쳐할 수 있는 능력이 강하다. 하지만 Transformer block은 high-frequency (local) information을 악화시킨다고 한다.

 

• NTB는 multi-frequency signals mixer로써 작동하여 전체적인 모델링 능력을 향상시킨다.
• NTB는 우선 Efficient Multi-Head Self Attention(E-MHSA)로 low-frequency signal을 캡쳐하며 다음과 같이 수식으로 표현할 수 있다.

$z = [z_1, z_2,...,z_h]$: input feature인 z를 channel dimension에서 multi-head form에 맞게 나눈것

SA: PvTv2의 Linear SRA와 비슷한 spatial reduction self-attention

Attention: 일반적은 attention 연산

$W^Q, W^K, W^V$: context encoding을 위한 linear layer

$P_s$: 연산량을 줄이기위해 attention operation 전에 downsampling을 위한 avg-pool

 

E-MHSA의 시간 소비가 channel 수에 상당히 영향을 받는 것을 확인했다고 한다. 그래서 E-MHSA module 전에 point-wise convolution(Conv 1x1)으로 channel dimension reduction을 수행하였다. (Fig 2 확인) 또한 효율적인 deployment를 극대화하기위해 E-MHSA module에 Batch Normalization을 적용하였다.

 

수축률 $r$은 channel reduction을 위한 값이다.

 

• NTB는 E-MHSA와 MHCA를 같이 사용함으로써 multi-frequency signal을 캡쳐할 수 있게 하였다. 그 후 E-MHSA와 MHCA의 output feature는 high-low-frequency 정보를 섞기위해 concat한다.
• 마지막으로 MLP layer을 통해 더 본질적이고 뚜렷한 feature들을 extraction한다.

 

$\tilde{z}^l, \hat{z}^l, z^l$: 각각 E-MHSA, MHCA, NTB의 output

Proj: point-wise convolution layer

 

• NTB 또한 BN과 ReLU를 적용하였다.

Next Hybrid Strategy (NHS)

효율적인 deployment를 위한 CNN과 Transformer 결합을 위해 많은 노력이 있었다. Figure 4의 (b), (c)를 보면 보통 대부분이 shallow stage에서는 Convolution Block을 사용하고 마지막 한, 두 stage에는 Transformer Block을 사용한다. 하지만 이런 traditional hybrid strategy는 downstream task(e.g. segmentation, detection)에서 saturation 성능이 안좋게 나온다고 한다. 그 이유는 classification task의 경우 마지막 stage에서의 output을 사용할 뿐이지만 downstream task(e.g. segmentation, detection)에서는 각 stage에서의 feature들에 의존하여 더 나은 결과를 얻을 수 있다는 것이다.

 

NCB와 NTB를 $(N + 1) * L$로 쌓는 새로운 개념, NHS를 제안한다.

 

• shallow stage에 global information을 캡쳐할 수 있는 능력을 부여하기위해 Figure 4의 (d)와 같이 $(NCB * N + NTB * 1)$로 NCB와 NTB를 쌓는다. 구체적으로, shallow layer에서 global representation을 학습하기위해 NTB를 각 stage 끝에 위치시킨다.

 

Detail

• shirink ratio $r$ = 0.75 (uniformly)
• spatial reduction ratio $s$ = [8, 4, 2, 1] (E-MHSA)
• MLP layer의 expansion ratio = NCB: 3, NTB:2
• E-MHSA, MHCA head dim = 32
• BN, ReLU 적용