MobileViT v2 논문 리뷰

GPU 없는 환경에서 Image Classification을 해야 할 상황이 되었다. MobileNet으로 실험할까 하다, 성능도 어느 정도 챙기고 싶어, MobileViT 논문을 읽게 되었다. 

MobileViT v2 배경 설명

• MobileVIT v2는 2022년 Apple에서 발표한 "Separable Self-attention for Mobile Vision Transformers" 논문에서 소개된 모델이다.
• 기존 MobieViT가 model의 parameter 경량화에 성공하면서, CNN 기반의 mobilenet보다 좋은 성능을 보였지만, inference 속도가 느리다는 점을, separable self-attention 개념으로 해결한 논문이다.

 

Abstract

• MobileVit가 적은 parameter로 좋은 성능을 보였지만, CNN 기반의 모델에 비해 수행시간이 오래 걸리는 문제가 있다.
• 수행 시간 감소에 가장 큰 bottleneck은 MobileViT의 multi-headed self-attention(MHA)이다. 이 연산은 patch의 개수 k에 대해 O(k^2)의 연산 복잡도를 갖기 때문이다. 또한, MHA는 batch 단위의 matrix multiplication 같은 cost가 큰 연산을 포함하여, 수행시간이 길어지게 된다.
• 이 논문에서는 선형복잡도를 가지는 "separable self-attention"을 제안한다. separable self-attention을 위해, element-wise 연산을 수행하여, 제한된 device 환경에서도 효율적으로 사용가능하다.
• separable self-attention을 적용한 모델인 MobiileViTv2는 mobile vision task의 다양한 분야에서 SOTA 성능을 보인다. (ImageNet classification, MS-COCO OD)
• 3 million parameter만으로 ImageNet top-1 accuracy 75.6%를 달성하였다. 이것은 MobileViT v1보다, 3.2배 빠르면서, 1% 정도의 높은 성능을 보여준 것이다.

 

Introduction

[배경]

• MobileViT의 등장으로, ViT를 제한된 device 환경에서도 활용할 수 있게 되었지만, 이러한 모델은 속도 향상을 막는 주요 원인은 global representation을 학습하기 위해 꼭 필요한, multi-headed self-attention(MHA)이다. 
• MHA 계산을 위해, token(ViT에서는 patch)의 제곱에 해당하는 연산 복잡도(O(k^2))를 가진 연산이 필요하다. 
• 또한, MHA 계산에는 복잡한 matrix 연산이 필요한데, 이러한 cost가 큰 연산들은 제한된 메모리 환경과 power 환경에서 제한된 성능을 가지게 된다.
• 따라서, 이 논문은 제한된 device 환경에서 사용할 수 있도록, transformer의 self-attention을 효율화할 수 있을 것인가에 대한 결과이다. 

[관련 연구]

• ViT에 한정하지 않더라도, 몇 가지 연구들이 self-attention 연산을 최적화하기 위한 시도를 하였다.
• 그들 중, 대다수는 self-attention layer의 sparsity를 도입하여, 각 toekn이 다른 token 전체를 계산하는 것이 아니라, 일부를 참조하도록 바꾸는 방향이다. 이러한 시도들은 연산복잡도를, O(k*sqrt(k))나 O(k*logk)까지 낮췄다. 하지만, 그에 따라 성능도 감소했다.
• 다른 방법으로는 low-rank approximation을 이용한 연산이다. Linformer(아직 안 읽어봄)는 linear projection을 통해, self-attention 연산을 여러 개의 작은 self-attention 연산으로 decompose 하였다. 이로 인해, 연산복잡도는 O(k)까지 떨어졌다. 그러나, Linformer는 여전히 cost가 큰 batch 단위의 matrix multiplication 등을 유지하고 있다.

[소개]

• 이 논문에서는 separable self-attention 이라는 O(k)의 연산 복잡도를 가진 새로운 self-attention 방법을 소개한다.
• 효율적인 inference를 위해, 복잡한 연산들을 element-wise operation(덧셈, 곱셈)으로 대체했다. 

 

Model : MobileViT V2

• MobileViT는 CNN과 ViT를 결합한 네트워크이다. MobileViT는 transformer를 convolution으로 보는 방식으로, convolution과 transformer의 장범을 모두 활용했다. MobileViT는 적은 parameter만으로도, CNN 기반의 MobileNet보다 좋은 성능을 보였으나, 속도가 느리다는 단점이 있고, 그 bottleneck은 multi-headed attention(MHA)이다.
• MHA는 각 patch들간의 context 관계 연산을 위해, scaled dot-product attention을 사용한다. 하지만, MHA는 O(k^2)의 복잡 연산도를 가지기에, k가 큰 경우(큰 이미지나, patch를 작은 단위로 나눈 경우)에는 매우 연산량이 많아진다. 
• 또한, MHA는 연산과 memory의 소모가 큰 batch 단위의 matrix multiplication 등의 연산을 사용하는데, 이것은 제한된 device 환경에서 성능 저하의 요인이 될 수 있다. 
• 이를 해결하기 위해, 논문에서는 선형 복잡도를 가지는 separable self-attention을 제시한다.
• separable self-attention의 핵심 아이디어는 latent token L에 대한 context score를 계산하는 것이다. 이 socre를 이용하여 input token을 re-weight 하고, global information을 담은 context vector도 만들어낸다. 

Overview of MHA

• MHA는 transformer가 token 간의 관계를 encoding 할 수 있도록 하는 모듈이다.
• Transformer의 MHA에 대한 자세한 내용은 아래 글 참조 바란다. 

2023.05.08 - [NLP 논문] - Transformer (Attention Is All You Need) - (1) 리뷰\

Transformer (Attention Is All You Need) - (1) 리뷰

 

Separable self-attention

• separable self-attention은 MHA의 구조를 따랐다. MHA와 비슷하게, input 값은 Query, Key , Value의 3개의 branch를 통해 처리된다.

[Query]

• ViT에서 Query는 이미지의 patch를 linear layer에 통과하여 얻은 d dimension의 feature이다. 이를 d X 1 linear layer에 통과시키고, softmax를 취해주어, context score를 생성한다. (기존 Transformer에서 연산량을 크게 줄인 부분, 과거에는 feature 단위로 attention을 계산하였지만, latent vector 단위로 attention을 계산하겠다.)

[Key]

• 한편, Key에서는 input이 d X d linear layer를 거쳐서, k X d의 key feature를 만들고, 앞서 구한, context score와 key feature의 weighted sum으로 context vector를 구한다. 이렇게 구해진 context vector는 가볍지만, input x에 대한 모든 token 정보를 지니고 있다.

[Value]

• 마지막으로, value에서는 input이 d X d linear layer와 ReLU를 거쳐서, k X d의 feature를 만들고, 앞서 구한 context vector를 Value feature에 적용해 준다. 이렇게 구해진 feature를 최종 d X d linear layer에 거쳐서, self-attention 값을 생성한다. 

MobileViTv2 아키텍처

• separable self-attention의 효과를 입증하기 위해, 기본 MobileViT 구조에서 attention만 바꿨다. 
• 추가적으로 기본 MobileViT에서 skip connection과 fusion block을 제거하여, 성능을 약간 향상했다. 

 

실험 결과

[실험 셋팅]

ImageNet-1k 학습

• batch size 1024로 300 epoch을 scratch부터 학습시켰다.
• optimizer로는 AdamW를 사용하였고, ImageNet-1k를 1.28 million과 5000개로 나누어, train set과 validation set을 구성했다. 

ImageNet-21k-P로 Pretraining & ImageNet-1k finetuning

• ImageNet-21k(13 mullion, 19000 classes)를 이용하여 pretraining을 진행했다. 
• ImageNet-21k의 validation set과 ImageNet-1k validationset은 중복되지 않는다. 

 

[기존 모델들과 비교]

• 성능과 속도 측면, parameter 양에서 종합적으로 고려해 보았을 때, MobileViT는 좋은 성능을 보인다. 

 

[다른 Task]

• 각각 semantic segmentation과 OD에서도 종합적으로 좋은 성능을 보인다. 

 

 

총평 

• 간단한 아이디어인데, 논문에서 비슷한 내용이 계속 반복되는 느낌이 든다. 
• 해결한 방법보다, 현재 transformer의 구조에서 속도가 느린 점을 명확히 파악한 것이 더 흥미롭다. 
• 가벼운 모델 하면, EfficientNet과 MobileNet 정도 생각했었는데, 사용 가능한 선택지가 하나 추가된 것 같다. 

 

Reference

MEHTA, Sachin; RASTEGARI, Mohammad. Separable self-attention for mobile vision transformers. arXiv preprint arXiv:2206.02680, 2022.