DevHwi

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) 리뷰\

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.

MobileViT 배경 설명

• MobileViT은 2022년 Apple에서 ICLR 2022에 제출한 논문이다. (Apple이여서, mobile에 대한 CNN이 더욱 필요했을 것이다.)
• CNN에서 mobilenet이 나왔듯, ViT에서도 light cost에 초점을 맞춘 논문이 등장하였다.

Abstract

• mobile 환경에서 구동 가능할 정도의 가벼운 vision task CNN 모델이 등장하였었다. (mobilenet) 
• 하지만, CNN과 달리 ViT는 최근 많은 vision task에 사용됨에도 불구하고, global representation들을 활용하기 위한 self-attention 구조를 사용하기 때문에, CNN에 비해 모델이 무겁다.
• 이 논문에서는 CNN과 ViT를 결합하여 mobile vision task에서 구동 가능한 가벼운 ViT 모델을 만들 수 있는가? 에 대한 답을 제시한다. 
• 이를 위해, 가볍고, 좋은 성능을 내고, 다양한 용도로 사용할 수 있는 MobileViT를 소개한다. 
• 실험에서는 MobileViT가 다양한 task와 dataset들에서 CNN과 ViT 기반의 모델보다 압도적인 성능을 보여준다.

Introduction

[배경]

• ViT의 소개 이후로, ViT의 연구 트랜드는 model의 parameter의 양을 늘리면서, 좋은 성능의 모델을 만들어내는 것이었다. (LLM처럼...)
• 하지만, 이런 모델 크기의 증가는 model의 용량 증가와 속도 저하등을 일으킨다. 
• 실생활에서는 mobile 기기 같은 제한된 H/W에서 구동되어야 하는 task들이 많이 존재하기 때문에, 가볍고 빠른 모델에 대한 요구사항이 있다. 
• 이를 위해, ViT의 parameter 개수를 줄이면, 오히려, 가벼운 CNN들보다도 좋지 않은 성능을 보여준다. 따라서, 어떻게좋은 성능을 유지하면서 가벼운 ViT를 만들 것인지에 대한 심도 깊은 고민이 필요하다.

[CNN + ViT]

• 일반적으로 ViT는 CNN에 많은 wieght를 필요로 하고, optimize가 어렵다. 예를 들어, ViT 기반의 segmentation network는 345 million의 parameter가 필요한데, 비슷한 성능의 CNN은 59 million으로, ViT 기반이 비효율적이다. 이것은 ViT가 CNN의 특성인 image specific inductive bias가 부족하기 때문이다.
• 좋은 성능의 ViT 모델을 위해서라면, convolution과 transformer의 하이브리드 방법을 채택하는 것이 효과적일 수 있다.
•  하지만, 기존 시도된 하이브리드 방법들은 아직 매우 큰 model size를 가지고 있고, data augmentation에 민감하다는 단점이 있다. 

[MobileViT]

• 이 논문에서는 제한된 H/W 리소스에서도 효과적으로 좋은 성능을 보이는 ViT 모델을 만들고자 한다. 
• 모델의 효율성 측정을 위해 보통 사용되는 FLops는 mobile 장비 등에서 성능 측정에는 충분하지 않다. 이는 FLOPs는 memory access나 parallelism 정도, platform 특성등을 고려하지 않기 때문이다. 
• 따라서, 논문에서는 FLOPs을 optimize 하기보다는, light weight, general-purpose, low latency에 초점을 맞췄다.  
• 이 논문에서는 CNN의 spatial inductive bias와 ViT의 global processing에 장점들을 결합가능한 MobileViT model을 소개한다.
• 특히, local 정보와 global 정보를 모두 encoding 할 수 있는 MobileViT block은 효과적인 representation을 포함한 feature를 만들 수 있게 해 준다.

MobileViT : A light-weight Transformer

• 일반적인 ViT model에서는 (H X W X C) 크기의 이미지를 P 사이즈의 patch들로 자르고, flatten 하고, projection 하여, (N X d) 크기의 inter patch representation으로 만든다.  이때, comnputational cost는 O(N*N*d)이다. (N: number of patch)
• 이러한 ViT 구조의 모델들은 CNN에 내재된 spatial inductive bias를 무시하기 때문에, CNN에 비해 더 많은 parameter를 필요로 한다. 
• 또한, 이런 구조들의 모델은 L2 regularization이나, overfitting을 피하기 위해 많은 data augmentation이 필요한 등 optimization이 어렵다. 

 
[MobileViT 아키텍처]
 
MiobileViT의 핵심 아이디어는 Transformer를 convolution으로 사용하여 global representation들을 학습하는 것이다. 
 

• MobileViT block 
  • MobileViT block은 model이 input의 local과 global 정보를 적은 paramter 안에 모두 담을 수 있도록 하는 것이 목적이다. 
  • 우선 MobileViT는 (n X n) 개의 convolutional layer와 (1 X 1) convolution을 수행해 준다. 이로 인해, local representation을 포함하게 된다. 
  • global representation을 학습하기 위해, 앞서 covolutional layer의 output을 flatten 하고 겹치지 않도록, patch로 나눠준다. 그러고 나서, 이 patch들을 input으로 Transformer에 넣어준다. 
  • 이로 인해, 아래 그림처럼, local 정보를 포함하면서, global 정보까지 포함할 수 있는 feature가 완성되게 된다.

• Light Weight
  • MobileViT는 image specific inductive bias를 사용할 수 있기 때문에, 일반 ViT에 비해 훨씬 더 가벼운 모델로 학습이 가능하다. (dimension 등을 줄일 수 있음)
• Computational Cost
  • MobileViT와 ViT의 computational cost를 비교하면, 각각 O(N*N*P*d)와 O(N*N*d)로 이론적으로는 MobileViT가 더 비효율적으로 보인다. 하지만, 실제로는 MobileViT가 DeIT ImageNet-1K image classification에서 2배 적은 FLOPs으로 1.8% 좋은 성능을 낼 정도로 더 효과적이다.
  • 이것은 더 가벼운 모델로 학습이 가능하기 때문이다. 

 
[Multi-scale sampler for training efficiency]

• 일반적으로 ViT 기반의 모델들은 fine-tuning 시에 다양한 스케일의 representation들을 학습한다. 이 과정에서 Positional embedding을 사용하는데, 이는 input size에 기반하여 interpolation이 필요하고, interpolation 방법에 따라 성능에 영향을 미치게 된다. 
• CNN과 마찬가지로, MobileViT에서도 이러한 positional embedding이 필요가 없다. 
• MobileViT에서 이러한 다양한 스케일의 representation들을 학습하기 위해, batch size를 변경시키는 방법을 사용했다. 다양한 spatial resolutions set을 만들어 놓고, 각 bach 시, resoluton을 무작위로 뽑고, 해당 resolution에 따라 batch size를 유동적으로 변경하여 학습했다. (small spatial resolution에는 많은 batch) 
• 이러한 학습 방법은 더 빠른 optimize가 가능하도록 하였다. 

실험 결과

• MobileViT 모델을 scratch부터 ImageNet-1K classification dataset으로 학습했다.
• multi-scale sampler로 각각 (160,160), (192,192), (256,256), (288,288), (320,320) 해상도를 사용했다. 
• 성능 측정은 top-1 accuracy로 진행했다. 

[CNN과 비교]

• MobileViT가 CNN 기반의 MobileNet 등의 light-weight CNN이나, ResNet 등의 일반적인 CNN의 성능을 압도했다.

[ViT 비교]

• ViT와 비교해서도 좋은 성능을 보인다. 특히, DataAugmentation에 따라 성능이 크게 바뀌는 다른 모델들과 달리, Basic Augmentation만으로도 다른 ViT보다 좋은 성능을 보인다. 

총평

• 매우 아이디어가 간단하고, 효과적인 것 같다.
• 논문 내용만 보면, 장점만 많아서, 진짜인지 직접 돌려보고 싶다. 

Reference

Mehta, S., & Rastegari, M. (2021). Mobilevit: light-weight, general-purpose, and mobile-friendly vision transformer. arXiv preprint arXiv:2110.02178.