[Paper Review] FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization

 

Paper : https://arxiv.org/abs/2303.14189

FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization

GitHub : https://github.com/apple/ml-fastvit

GitHub - apple/ml-fastvit: This repository contains the official implementation of the research paper, "FastViT: A Fast Hybrid V

 

 

 

 

 

Overview

오늘 리뷰할 논문은 애플에서 공개한 "FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization" 라는 논문입니다. 본 논문에서는 제안하는 FastViT는 the state-of-the-art latency-accuracy trade-off를 만족하는 hybrid vision transformer 구조입니다. 여기서 Reparameterization 기법을 이용하여 신경망에서 Skip-connection을 제거하고, memory access 비용을 낮추는 새로운 RepMixer(a novel token mixing operator)를 제안합니다. 현재 hybrid transformer architecture인 CMT 보다 3.5배 빠르고, EfficiendNet 보다 4.9배 빠르고, 모바일 장치에서 ConvNeXt 보다 1.9배 빠르다고 합니다. 그리고 유사한 latency 에서 본 논문에서 제안하는 모델이 ImageNet에서 MobileOne 모델보다 4.2% 더 나은 정확도를 보여준다고 하네요. 무엇보다도 모바일 장치와 데스크톱 GPU 환경에서 latency가 크게 개선되었다고 합니다. 

 

 

위 지표를 보시면, iPhone 12 Pro, NVIDIA RTX-2080Ti 에서 측정되었으며 , FastViT 모델이 다른 모델에 비해 빠르고 정확도도 높기 때문에 Mobile Latency, GPU Latency에서 우수하다는 것을 볼 수 있습니다. 

 

요즘 연구에서는 정확도를 유지하면서 효율적으로 네트워크를 설계하기 위해 vision transformer 구조에서 conv 구조와 transformer의 구조를 병합하는 하이브리드 방식이 대세인 듯 하네요! 최근 공개된 Metaformer 아키텍처는 skip connection이 존재하는 token mixer와 다른 skip connection이 있는 Feed Forward Network(FFN)으로 구성되곤 하는데 이러한 skip connection 구조는 상당한 오버헤드를 차지한다고 합니다. 왜냐면 메모리 엑세스 비용이 필요하기 때문이죠. 그래서 이러한 오버헤드를 해결하기 위해 Reparameterization을 사용하는겁니다. 

 

그래서 본 논문에서는 이러한 skip connection 구조를 제거하여 메모리 비용을 아낄 수 있는 Reparameterization 기법인 RepMixer를 소개합니다. 이렇게 Reparameterization 하는 방식은 이전 연구에서도 많이 연구가 진행되었는데, 하이브리드 구조에서 적용된 사례는 처음보는 듯 합니다. 앞으로 이런 연구가 활발하게 진행되지 않을까 싶네요. 

 

아무튼, FastViT 아키텍처의 주요 설계 원칙은 다음과 같다고 합니다. contribution이라고 할 수 있겠네요. 

 

1. RepMixer 블록을 사용하여 skip connection 제거

2. 정확도 향상을 위한 linear train-time overparameterization 사용

3. early stage에서 self-attention 레이어를 대체하기 위해 large conv kernel 사용 

 

 

Related Work

 

Hybrid Vision Transformers

이전 연구에서 하이브리드 아키텍처에서 일부 설계는 patchify stem 를 convolutional layers로 대체하거나 early conv stage를 도입하거나 window attention을 통해 암시적으로 하이브리드화 했다고 합니다. 또한 최근에는 token(or patch)간의 더 나은 정보 교환을 위해 명시적으로 하이브리드 구조를 구축한다고 합니다. 대부분의 하이브리드 구조에서 token mixer는 self-attention 기반이라고 하네요. 최근 MetaFormer 구조에서는 간단하고 효율적인 token mixing candidate인 Pooling을 도입했습니다. 

 

Structural Reparameterization

최근에는 skip-connection을 reparameterization 하여 메모리 액세스 비용을 낮추는 이점을 보여주었으며, 본 논문에서는 inference 과정 중에 reparam 작업을 할 수 있는 RepMixer를 소개합니다. 더 나은 효율성을 위해 1x1 pointwise conv가 뒤따르는 depthwise or grouped conv을 사용하여 factorized kxk conv를 도입합니다. 이러한 방식은 모델의 전반적인 효율성을 개선하는데 매우 효과적이지만 매개변수 수가 적으면 용량이 줄어들 수 있습니다. 이러한 모델 용량을 개선하기 위해 overparameterization에 대한 linear train-time이 도입되었습니다. 모델에서 분해된 k x k conv를 사용하고 linear train-time overparameterization을 사용하여 이러한 레이어의 용량을 높입니다. 아직까지는 skip-connection을 제거하기 위한 structural reparameterization 방법과 linear overparameterization 방법은 hybrid transformer architecture에서 사용된적이 없다고 합니다. 

 

 

 

Architecture

Overview

 

FastViT

Reparameterizing Skip Connections

1. RepMixer 

Convolutional mixing은 ConvMixer에서 처음 소개되었다고 합니다. 입력 텐서의 X의 경우 mixing block에서 다음과 같이 구현됩니다. 

 

여기서 σ는 non-linear activation function이고 BN은 batch norm, DWConv는 depthwixe convolutional layer를 뜻합니다. 이 블록이 효과적인걸로 나타났지만 RepMixer에서는 아래와 같이 σ 를 제거합니다. 

 

 

본 논문에서 RepMixer 설계의 주요 이점은 추론 시간에 하나의 depthwise conv로 다시 매개변수화 할 수 있다는 것입니다. 

 

위 그림에서 이 부분에 해당합니다. 

 

 

2. Positional Encodings 

본 논문에서는 positional encodings을 사용합니다. 이는 입력 토큰의 로컬 이웃을 조건으로 삼게 되고, 이러한 인코딩은 depthwise conv operator의 결과로 생성되며, patch embedding에 추가됩니다. 이 작업 그룹에서는 non-linearity가 없습니다. 이 block은 다시 매개변수화가 됩니다. 아래 그림 2a에 해당합니다. 

 

3. Empirical Analysis

reparameterizing skip connections의 이점을 확인하기 위해 MetaFormer S12 아키텍처에서 Pooling 또는 RepMixer를 사용하지 않고 실험해봅니다. 제거된 두 모델 모두 ~1.8G FLOP을 갖고, iPhone 12 Pro에서 224x224 해상도에서 1024x1024 해상도 까지 다양한 입력 해상도에 대해 모델 시간을 측정합니다. 그림 3에서 RepMixer가 더 높은 해상도에서 Pooling 보다 정확도가 더 크게 향상되었음을 볼 수 있습니다. 

 

 

 

Linear Train-time Overparameterization

parameter count, FLOPs,latency의 효율성을 위해 모든 dense k×k conv 레이어를 factorized version으로 대체합니다. 즉 k x k depthwise 레이어 뒤에 1 x 1 pointwise conv를 붙입니다. 하지만 factorization으로 인해 매개변수 수가 적으면 모델 용량이 줄어들 수 있다고 하네요. 분해된 레이어의 용량을 늘리기 위해 MobileOne에서 사용하고 있는 방법인 linear train-time overparameterization을 수행합니다. 이는 성능향상에 도움이 됩니다. 표3에서 train-time overparameterization 방법은 정확도를 향상시킵니다. 그러나 이는 추가된 분기의 계산 오버헤드로 인해 학습시간이 증가하게 됩니다. 그래서 본 논문의 아키텍처에서는 dense k x k conv를 factorized form으로 대체하는 레이어만 overparameterization 시킵니다. 

 

 

 

Large Kernel Convolutions

RepMixer의 receptive field는 self-attention token mixer에 비해 local 인데, self-attention 기반 token mixer는 계산 비용이 많이 듭니다. self-attention을 사용하지 않는 초기 단계의 receptive field를 개선하기 위해서는 depthwise large kernel convolution을 통합하는 것입니다. FFN 및 patch embedding layer에 depthwise large kernel conv를 도입합니다. 다음 표 4에서 이러한 방법이 self-attention을 사용하는 방법에 비해 굉장히 경제적일 수 있음을 보여줍니다. 여기서 RM은 RepMixer-FFM을 뜻하고 SA는 Self Attention-FFM 을 뜻합니다. 

 

 

 

 

Experiments 

Image Classification

image classification에서의 결과는 다음과 같습니다. FastViT-T8 모델이 MobileOne-S0 보다 GPU latency가 더 빠르고, 모바일에서의 latency는 똑같네요. 하지만 정확도는 4.2% 정도 좋다는 것을 볼 수 있습니다. 그 외 논문들도 같이 살펴보면 좋겠네요. 여기서 GPU 결과를 측정할 때는 TensorRT를 사용했고, 모바일에서는 CoreML을 사용했다고 합니다. 변환이 안되는 것은 - 표시해두었다고 하네요. 

 

다음 결과는 KD를 이용했을 때의 결과입니다.  

 

Robustness Evaluation 

다음은 Robustness benchmark dataset으로 평가한 결과입니다. 

 

3D Hand mesh estimation 

3D hand mesh estimation 분야에서도 사용된 결과입니다. 

모델은 다음과 같이 구성했다고 하네요. 

 

정성 결과도 있습니다. 

 

 

 

Semantic Segmentation and Object Detection 

정말 다양한 분야에서 평가를 해놓았네요! ResNet을 사용했을 때 보다 훨씬 좋은걸 볼 수 있습니다.