Exploring Plain Vision Transformer Backbones for Object Detection 리뷰

https://arxiv.org/abs/2203.16527

Exploring Plain Vision Transformer Backbones for Object Detection

배경

 

Abstract에서 첫 문장으로 We explore the plain, non-hierarchical Vision Transformer (ViT) as a backbone network for object detection.  라고 시작한다. 그리고, Masked Autoencoders (MAE) 기법도 보인다.

 

plain, non-hierarchical, ViT, MAE 는 이 논문의 핵심 키워드이다.

 

plain, non-hierarchical 이 무슨 뜻으로 쓰였을까? 

 

먼저 ViT를 봐보자

 

 

위 그림은 Original ViT 구조를 나타낸 그림이다. ViT구조를 봐보면 처음에 16x16 패치 단위로 이미지(3 x h x w)가 split된다. 그리고, reshape되어 (h/16 x w/16 x (16x16x3)) 꼴로 나타내어지고 Linearly embedding 된다. 이후에는 classification tocken 이 붙고 positional embedding이 더해진 뒤 Transformer Encoder를 쭉쭉 거치고 최종적으로 class 별 confidence score가 나오게 설계돼있다.

 

잘 생각해보면 (h, w) 축이 1/16 만큼 다운샘플링되고나서 (h, w) 축으로는 다운샘플링이 되지 않는 구조임을 알 수 있다. 이는 이전에 제안됐던 CNN 모델들이 대부분 (1/2, 1/4, 1/8, 1/16, 1/32) 단위로 (h, w)축이 점점 다운샘플링 되어지는 구조와는 확실히 다른 구조로 여길 수 있다.

 

이 논문에서는 Original ViT 의 구조처럼 피쳐맵의 스케일이 점진적으로 줄지 않고 동일한 스케일에서 피쳐를 추출하는 구조(single-scale property)를 plain/non-hierarchical 하다고 말하고 있다.

 

위 설명을 읽고 아래 그림을 보면 hierarchical 과 plain/non-hierarchical 차이를 이해할 수 있을 것이다. 우선 neck/head 쪽은 보지말고 backbone 쪽만 봐보자.

Object Detection 방법들을 봐왔다면 Feature Pyramid Networks(FPN) for Object Detection 논문을 안읽어봤을리가 없을 것이다. 이것도 Kaming He님이 저자로 들어가있는 논문인데, 인용수만 봐도 이 논문이 Object Detection 분야에 얼마만큼의 영향을 미쳤는지 파악이 가능하다. 

이 논문에서 제안된 방법은 backbone 모델에서 나오는 multi-scale의 피쳐맵을 서로 잘 지지고 볶고(aggregation) 이 지지고 볶인 피쳐맵들을 이용해서 object detection을 하면 검출 성능이 크게 올라간다는 것이였다.

 

위 그림에서 (d) 가 feature pyramid network 논문에서 나온 구조인데 그림만 봐도 hierarchical backbone에 의존적인 구조임을 알 수 있다. 그리고, 이 구조는 high performance object detection model을 설계하기 위한 하나의 inductive bias가 되어버렸다. 우리가 알고있는 대개 유명한 obejct detection 방법들(YOLOv3, YOLOv4, YOLOv5, PP-YOLO, FCOS, EfficientDet 등)은 대개 이 구조를 따라 디자인되었다. 2022년 05월 01일을 기준으로 SOTA object detector들은 Swin Transformer를 백본으로 하는 경우가 많은데 이 Swin Transformer도 hierarchical 구조로 FPN에 적합하다. 아예 Swin 논문에서 Swin은 Hierarchical 구조이므로 FPN 구조를 쉽게 적용가능하는 것이 장점인듯이 표현이 돼있는 걸 알 수 있다.

swin 논문에서 FPN 언급 부분

이 논문에서는 ViT를 object detection task의 백본으로 사용할때, 어떻게하면 plane, non-hierarchical 구조를 그대로 살리면서 높은 성능을 갖는 object detector를 설계할 수 있을지에 관하여 실험적으로 이 구조 저 구조를 시도해보며 적절히 현 SOTA 성능급과 견줄 수 있는 구조를 제안한다.

제안하는 방법

ViT를 Object detection 모델의 backbone으로 쓰는데 있어서 practical 하지 않은 부분이 존재한다. 바로 바로 Computational Cost!!! self-attention 연산이 차지하는 연산량이 너무너무 크다. self-attention 연산에서 query 와 key의 행렬곱에 필요한 연산량만 생각해봐도 O(n^2 x c) 다. 여기서, n은 패치수이고 c는 피쳐의 채널 수다. 입력 이미지의 해상도가 증가될수록 n이 증가되고, 연산량은 quadratic하게 증가된다.

 

그런데! Object detection 모델은 classification 모델에 비하여 상대적으로 고해상도의 입력 이미지를 인풋으로 받는다. 예를들어 이미지넷을 기준으로 설계된 대다수의 모델들이 224x224, 256x256, 384x384 정도의 해상도에 달하는 이미지를 입력으로 받는다고하면, YOLOX 같은 모델은 640x640 해상도에 달하는 이미지를 입력으로 받고, 이 논문에서 제안하는 object detector의 경우 1024x1024에 달하는 이미지를 입력으로 받는다. 그래서, ViT를 backbone으로 사용하고자 한다면, ViT의 self-attention 레이어를 손대지 않고서는 practical 한(=실제 필드에서 사용할만큼 빠르면서 고성능을 갖는) object detector를 설계하기란 쉬운 문제가 아니다.

 

저자들은 위 문제를 해결하기 위하여 ViT의 self-attention 레이어를 local(=window) self-attention 레이어로 교체하였다. 이는 attention을 global하게 수행하는 게 아니라 attention을 locally 수행하는 방법이다. 이에 대해서는 Swin 논문을 읽도록 하자. 이때 window 사이즈는 14x14 로 설정하였다.

 

why? 기존에 ViT가 224x224 이미지를 입력으로 받고 패치 임베딩될때 패치 사이즈를 16x16으로 설정하였는데 224/16 x 224/16 = 14 x 14가 된다. 그래서 14 x 14를 사용함. 

 

근데! 이렇게 self-attention 레이어를 local 로 바꿨다고 해서 다시 ImageNet 데이터갖고 pre-training을 시킨 게 아니다. object detector모델에 ViT모델을 Backbone으로 사용하여 Finetuning 시키는 단계에서만 기존에 global self-attention 레이어로 학습된 Original ViT 모델을 그대로 갖고와서 local self-attention 레이어가 되도록 수정을 가한것이다.

 

그런데...!! local self-attention 레이어는 패치간에 information propagation이 이뤄지지 않는다. Swin Transformer는 이 문제를 극복하기 위하여 패치를 shift 하는등의 트릭을 사용하였는데 이 과정이 은근 복잡하다(이건 그냥 내 생각이므로 무시하여도 됨). 

 

그래서, 패치간 information propagation이 되도록 두 가지 기법(모듈)을 제안한다.

 

먼저 ViT의 Transformer Encoder Layer수가 N이라고 하면 이를 균등하게 4등분하고 제안하는 모듈을 삽입/혹은 적용하여 information propagation이 되도록 모델을 수정하였다.

 

두 가지 방법은 Global propagation과 Convolutional propagation이다.

 

Global propagation은 일부 local self-attention layer만 global self-attention layer로 두는 방법이다.

 

Convolutional propagation은 convolutional layer로 구성된 block을 넣는 방법이다. 실험 과정에서 어떤 구조의 block이 좋은지 저자들이 실험을 해놨다.

 

그리고, 논문에서는 어떻게 neck 부분을 plain ViT에 적절하게 설계할 수 있는가에 대해서도 탐색하려는 시도를 했다. 최종적으로는 (c) 구조를 채택했다.

실험

실험은 아래와 같은 configuration 하에 진행됐다. 좀 굵직한 것만 써둔건데 논문을 읽어보면 디테일한 설정이 꽤있으므로 이 논문에 대한 관심이 커져 이 논문을 토대로 뭔가 연구를 해보고 싶으면 논문을 디테일하게 읽어보자.

 

Backbone: ViT-B, ViT-L, ViT-H

Detector head: Mask R-CNN, Cascade Mask R-CNN

Input Image Resolution: 1024x1024

Augmentation: Large-scale jittering

Finetuning Epochs: 100

Optimizer: Adam

 

Dataset: COCO17, LVIS

 

이 포스트에서는 COCO17에 대한 결과만 실겠다.

 

neck구조에 관한 실험이다. 기존 FPN보다 (c)의 성능이 더 좋다. 기존에 FPN이 multi-scale feature를 이용함으로써 성능이 올라가는 반면에 ViT는 single-scale 이고 어디서 떼어오든 receptive field의 차이가 없기 때문에 ViT를 Backbone으로 사용할때는 FPN이 큰 이득을 보지 못하는 것을 확인할 수 있다.

표2는 information progation 방법에 관한 성능을 비교해둔 표이다.

 

(a)를 보면, none은 information propagation 없이 window attention만 있을때의 결과이다. conv block을 썼을때 가장 성능이 좋은 것을 확인할 수 있다.

 

(b)를 보면, naive, basic, bottleneck이 있는데. 여기서 naive는 그냥 3x3 conv 하나를 두었을때의 결과이고 basic, bottleneck은 ResNet 논문에서 제안된 block 구조를 의미한다.

(c)는 information propagation을 특정 위치의 self-attention layer를 global self-attention layer로 세팅했을때의 성능을 비교한 표이다. global attention layer는 균일하게 두었을때와 마지막(last)에 몰아넣었을때 성능 향상폭이 큰 것을 확인할 수 있다. 앞쪽 self-attention layer를 global self-attention layer로 세팅하는것은 의미가 없어보인다.

 

(d)를 보면, global attention layer 개수에 따른 성능을 나타낸 표로 global attention layer의 개수가 늘어날수록 성능은 증가되는 걸 볼 수 있다. 그런데 24개를 썼을때와 4개를 썼을때 연산량 차이를 생각해보면 성능을 높이기 위해 global attention layer를 증가시키는 것은 좋지 않은 선택이 될 것이다.

 

표3은 Practical performance(latency 대비 performance)를 비교한 표이다. practical object detector를 설계하고 싶다면 information propagation 목적으로 conv layer를  사용하는 것이 유의미한 선택이 될 거 같다.

 

표4는 Backbone으로 사용되는 ViT의 pre-train 방법에 따른 detection 성능을 나타낸 표이다. MAE(Masked Autoecnoders) 기법을 사용해 학습된 ViT를 backbone으로 사용할때의 성능이 압도적으로 좋다. 심지어, ImageNet-21K(14M images) 를 사용하여 pre-train 된 ViT를 압살해버린다. pre-train을 하지 않은 모델이 지도학습으로 학습된 모델(IN-*K, supervised)과 비등한 성능을 보이는 것이 좀 눈에 띈다.

 

표5는 Hierarchical backbone 기반의 object detector간 성능을 비교한 표이다. 여기서, ViT의 position embedding 방법을 Swin, MViT와 마찬가지로 relative position embedding 방법을 사용하였다고한다.

 

Figure 3은 파라미터 사이즈, 연산량(FLOPs), 추론 시간(Test time)에 따라 Object detector들의 성능을 비교한 그림이다. 모델 사이즈가 커질수록(ViT-B -> ViT-L -> ViT-H) Plane ViT를 Backbone으로 쓸때, 타 Hierarchical Backbone보다 성능의 향상폭과 성능이 높은 것을 확인할 수 있다.

 

표6은 이전 실험에서 쓰인 head 보다 더 좋은 head인 Cascade Mask R-CNN을 쓰고, nms 기법으로는 soft-nms기법을 쓰고, 인풋 이미지 사이즈를 1024x1024 에서 1280x1280으로 세팅하였을때의 결과이다. 이로써 20220501 기준 SOTA급에 달하는 성능을 찍어버린다!