Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions

Inc Lomin's avatar
Oct 29, 2021
Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
Vision Transformer (20.10) → PVT (21.02) → Swin Transformer (21.03) → Local VIT (21.04) → PVT v2 (21.06)
Vision Transformer 계열에서 처음으로 Dense precition에 활용 할 수 있도록 pyramid feature 구조를 가지는 네트워크를 제안했습니다.

Introduction

notion image
CNN에서는 Pooling layer가 존재하여, feature가 자연스럽게 pyramid 형태이기 때문에 FPN에 활용하기 쉽다. 하지만 Vision transformer의 경우 pyramid 형태가 아니게 되고, 이로 인해 dense prediction과 같은 task에 대해서는 backbone으로 활용하기 어렵습니다. 그렇기 때문에, 본 논문에서는 pyramid 형태를 가지는 vit backbone을 제시했습니다.

Proposed Method

notion image
Vision Transformer에서는 patch를 16x16를 자르지만, 해당 사이즈로 자르게되면 dense prediction에 사용하기에 feature 크기가 너무 작아집니다. 그렇기 때문에 본 논문에서는 4x4로 patch를 잘라서 각 patch를 token으로 만듭니다. 해당 feature 사이즈는 이미지의 4분의 1이기 때문에 feature pyramid network의 첫번째 feature로 쓰기 적절합니다.
하나의 stage가 끝나면 해당 feature를 다시 2x2 patch로 잘라서 각 patch를 token으로 사용하여 feature 사이즈를 절반으로 줄여줍니다. 2x2로 한번 자르게되면, 이미지 레벨에서는 8x8 patch를 쓴 효과입니다. 해당 작업을 32x32 patch를 쓸때까지 점차적으로 진행해줍니다.
 
하지만 이렇게 진행을 하게 되면, multi head attention에서의 연산이 너무나도 커지게 됩니다. 이미지 분류에서 가장 많이 사용하는 이미지 사이즈인 224x224를 기준으로 생각해보겠습니다.
첫번째 스테이지(4x4 patch)에서 사용되는 query, key, value 의 사이즈는 56x56이고 (56x56)^2 만큼의 연산량을 하게 됩니다. 이는 메모리나 연산량 측면에서 매우 손해이기 때문에 본 논문에서는 key와 value의 사이즈를 줄이는 trick을 사용했습니다.
notion image
Super-Resolution에서 자주 활용되는 pixel shuffling은 channel dimension을 줄여서 spatial dimension을 늘리는 방법입니다. inverse pixel shuffling을 시행하게 되면, spatial dimension을 줄여서, channel dimension이 늘어나게 됩니다. 본 논문에서는 inverse pixel shuffling 을 사용해서 spatial size를 줄이고, convolution을 통해 늘어난 channel dimension 을 원래의 channel dimension으로 맞추었습니다.
notion image
 
 

Experiment

Image Classification
notion image
Object Detection
notion image

Conclusions

 
 
Share article