[2021 arXiv] MobileVIT: LIGHT-WEIGHT, GENERAL-PURPOSE, AND MOBILE-FRIENDLY VISION TRANSFORMER

[gusdlf93]

CNN은 적은 파라미터를 사용해서 네트워크를 만들 수 있지만, 지역 정보만을 사용한다는 단점이 있습니다.
본 논문에서는 CNN의 이러한 단점을 보완하고자 ViT를 사용해서 Global Representation을 학습합니다.

하지만 ViT만 사용해서 네트워크를 설계하면 너무 무거우므로, CNN과 ViT를 결합하여 연산량을 줄이면서도, 모바일 환경에서 적합한 네트워크를 설계합니다. 이러한 노력 끝에 동일 연산량 대비(6M parameters)기준 MobileNetv3(CNN 기반) 및 DeIT(ViT 기반)보다 3.2%, 6.2% 더 높은 정확도의 네트워크를 설계할 수 있었습니다.

MobileVit XXS : 1.3M parameters, 0.2 GFLOPs, Top-1 69%
MobileVit XS : 2.3M parameters, 0.6 GFLOPs, Top-1 74.8%
MobileVit S : 5.6M parameters, 1.1 GFLOPs, Top-1 78.4%
*Flops는 논문에 없으므로, Chris Ha가 임시로 구현한 모델 기준이라고 봐야 합니다.

모바일넷의 구현은 다음과 같습니다.

256x256의 Resolution을 가진 이미지를, 32x32크기의 Resolution의 feature map이 될 때까지 DownSampling을 진행합니다.
DownSampling은 Convolution과정에서 stride를 줌으로써, Maxpooling을 피합니다.
RegNet에서 주장했던 것처럼, 좋은 성능을 내는 네트워크들은 MaxPooling을 피하고, stride를 주로 사용하는 것으로 보입니다.
마찬가지로, 경량화 네트워크인데 Bottle Neck을 사용하지 않는 것도 이러한 이유가 아닌가? 싶습니다.

MobileVIT block입니다.
파란색 pixel은 Convolution을 통해서 local feature를 뽑고, 빨간색 픽셀은, 서로 다른 패치의 같은 위치의 정보로부터 정보를 추출합니다. 재밌는 방식으로 Attention을 경량화 한거 같네요.

이 논문에서 가장 핵심인 MobileVit block입니다.

이 부분이 사실 좀 난해했는데, 일반적으로는 Transformer를 적용할때, 2D Tensor가 나오면, 2D Tensor에서 Channel 방향으로 split을 해서 Multi Head를 구성함으로써 3D가 되는데, 여기서는 처음부터 3D가 나오드라구요. 어떤 특징정보가 Multi Head의 역할을 대체하는지 고민을 많이 했지만, Multi Head Self Attention을 구현하는 부분이 Channel을 나누는게 아닌거 같습니다.

**아직 코드가 공개가 되지 않아서 확답은 못하겠고, 추론으로 대체하겠습니다.

N은 각각의 이미지 패치들이고, d는 채널, P는 image patch(2x2)로, 패치 내 픽셀들을 의미합니다.
Transformer 들어가기전에 Convolution을 통해서 각 이미지 패치 내부의 local feature는 추출을 했으니 이번에는 image patch들간의 정보를 이용해야 할 차례입니다.

각각의 이미지는 P, N, d로 이루어져 있을겁니다. patch size * num of image path * dimension
attention 연산은 다음과 같이 이루어지게 될겁니다.

Q * K = P,N,d x P,d,N -> P, N, N
(Q*K)*V = P, N, N * P,N,d -> P,N,d가 만들어지게 될겁니다.

원래는 P위치에 multi head가 들어가야 하는데, 여기서 head를 pixel로 대입해서 문제를 해결했습니다.
각 픽셀들에 대해서, 이미지 패치들이 가진 채널정보를 이용해서 attention을 진행했다고 보시면 편할거 같습니다.
두 시간 정도 고민했는데 논문에서 제안한 설계 이념대로 구현한다면 저게 맞는거 같습니다.**

Image Classification 결과입니다.
가성비가 어마어마합니다.

object detection 결과입니다.
마찬가지로 성능이 꽤 좋습니다.

다만 단점은 latency입니다. parameter랑 accuracy는 매우 적지만, inference time은 CNN에 비해서 거의 8배이상 오래 걸린다고 볼 수 있습니다. mobile용으로 만든 거니, 추후 제품으로 나올 때에는 하드웨어 가속기를 사용해서 속도를 줄이지 않을까 예상됩니다.

Multi-Scale Sampler

이 부분은 Multi-Scale Training을 위해서 제안되었습니다. Yolo9000 같은 모델들을 보면, resolution에 강인하게 설계하기 위해서, 이미지의 크기를 여러개로 바꿔가면서 학습을 진행합니다.
epoch마다 다른 크기의 resolution을 넣어준것이죠.
우리는 모델을 학습하는 과정에서, pixel이 가진 정보량이 어떤지 알기가 상당히 어렵습니다.
이런 경우에 여러해상도의 이미지를 넣어서 학습을 진행하고, 어느정도 크기의 이미지가 우리 네트워크에 적합한지를 알아야 할 수 있습니다. 그렇기 때문에, 학습을 하는 과정에서 여러 scale의 이미지들을 넣어서 학습을 진행하고, gpu별로 이미지 크기가 다르니까 서로 다른 크기의 batch를 넣어 줬다는것으로 이해를 했습니다.

재밌는 논문이었습니다
Transformer가 일반적으로는 positional encoding을 학습해야 하기 때문에 많은 학습 데이터를 필요로 하는데, 이 부분을 원본 정보의 위치정보를 유지함으로써 해결하였습니다. 이것은 최근 트렌드를 잘 따라간거 같습니다만, 여전히 transformer가 가진 단점 중에 하나로, inference가 오래 걸린다는 문제는 해결하지는 못했습니다.
이런 점에서 보면 inference time을 생각한다면 아직은 CMT같은 구조가 더 좋은게 아닌가하는 생각이 들게 되네요.