[논문 리뷰] PVT: Pyramid Vision Transformer(2021)

이번에 리뷰할 논문은 Pyramid Vision Transformer 입니다.

 

 

 

 

https://paperswithcode.com/method/pvt

Papers with Code - PVT Explained

 

 

 

 

 

 

CNN 기반 모델들은 computer vision 분야에서 많은 성공을 이루었습니다. 하지만 2017년도에 나온 tranformer의 self-attention을 computer vision에 활용하기 위해 많은 연구들이 진행되었죠. 그 중 Vision transformer(ViT) 모델이 가장 대표적이라고 할 수 있을 것 같습니다. image를 patch단위로 split하여 sequence data로 변환한 후 self-attention을 적용하였고, 이는 기존의 CNN 모델들에 비해 더 좋은 performance를 보여주었습니다. 다만 ViT 모델의 경우 image classification 분야에 특화되었으며, dense prediction을 요하는 object detection, segmentation 분야에서는 한계를 보여주었죠.

 

 

ViT가 dense prediction에 한계를 가지는 이유는 위 figure를 통해 알 수 있습니다. 기존의 CNN 기반 모델의 경우 (a)와 같이feature map의 크기 즉, 해상도를 pooling layer를 통해 점차 줄여나감으로써 다양한 scale의 feature map을 출력할 수 있었지만, ViT의 경우 (b)와 같이 출력되는 feature map의 sclae이 모두 같은 single-scale임을 확인할 수 있습니다. 즉, attention을 수행하여 전역적인 feature를 추출한다는 장점은 있으나, dense prediction을 하기 위해 중요한 요소 중 하나인 multi-scale feature를 추출하지 못한다는 것이죠. 또한 object detection, segmentation과 같은 task의 경우 모델이 복잡하며, 이를 attention으로 수행할 때 많은 계산 비용이 든다는 단점도 존재합니다.

 

본 논문에서는 위와 같은 문제를 해결하고자 Pyramid Vision Transformer(PVT) 모델을 제안합니다. 위 figure의 (c)를 보면 낮은 해상도와 높은 연산량을 요하는 ViT와 달리 점진적으로 축소되는 피라미드를 사용하여 높은 해상도와 낮은 연산량을 가지도록 하였습니다. 두 번째로 PVT는 convolution 없이 다양한 vision task에 대한 통합된 backbone을 제공하며, 마지막으로 vision관련 많은 task에 대해 더 좋은 성능을 보여주었다고 말합니다.

 

 

PVT는 patch size를 4와 같이 아주 세밀한 patch를 사용하여 dense prediction에 필수적인 representation을 학습합니다. 또한 network가 깊어질 수록 점진적으로 축소되는 피라미드를 만들었기 때문에 계산 비용 역시 절약됩니다. 

 

 

 

본 논문에서는 다음 3가지를 PVT의 이점이라 설명합니다.

 

1. attention을 통해 항상 global receptive field를 생성하여, object detection, segmentation 분야에 적합하다.
2. pyramid 구조 덕에 많은 dense prediction pipeline에 통합 될 수 있다.
3. convolution이 없는 pipeline 구축이 가능하다.

 

 

 

 

 

 

 

 

 

Pyramid Vision Transformer (PVT)

우선 전체적으로 PVT 구조는 다음과 같습니다.

 

 

총 4개의 stage로 구성되어있으며, 각각 서로 다른 scale의 feature map을 생성합니다. 또한 모든 stage는 patch embedding layer 와 $ L_{i} $ Transformer encoder layer 로 구성되어 있습니다. 

 

첫 번째 stage에는 $ H \times W \times 3 $ 사이즈의 input image를 받아 $ \frac{HW} {4^{2}} $ 개의 patch들로 나눕니다. 즉, 각 patch의 size가 4x4x3 이라는 것이겠죠. 즉, 현재 size는 $ \frac{H} {4} \times \frac{W} {4} \times (4^{2} \times 3) $ 이 되며, 이를 flatten한 후 Linear projection을 통해 Patch embedding을 수행합니다. 해당 과정을 거치게 되어 embedding 되었을 때의 size는 $ \frac{HW} {4^{2}} \times C_{1} $ 가 됩니다. 

 

다음으로 Transformer encoder를 거쳐 생성된 feature map의 size는 $ \frac{H} {4} \times \frac{W} {4} \times C_{1} $이 됩니다. 여기까지는 ViT의 과정과 흡사합니다. 다만 PVT의 경우 다음에 오는 모든 stage에서 계속 Patch embedding을 취해줍니다. 이러한 방식때문에 stage를 거듭할수록 생성되는 feature map의 크기는 patch size에 비례하여 감소하게 되고 pyramid 형식을 띄며 전체적인 해상도는 감소하게 됩니다. 

 

 

 

 

Feature Pyramid for Transformer

다시 한번 Pyramid 구조에 대해 살펴보도록 합시다. 위에서 간단히 설명드렸던 것처럼 PVT는 patch embedding layer를 통해 feature map의 scale을 조절합니다. 

 

예를 들어 i 번째 stage에서 patch size를 $ P_{i} $라고 지정했을 때 input feature map의 사이즈가 $ (H_{i-1} \times W_{i-1} \times C_{i-1}) $ 라고 한다면, 총 $ \frac{H_{i-1}W_{i-1}} {P_{i}^{2}} $ 개의 patch가 만들어집니다. 다음으로 구분된 patch들을 flatten한 후 $ C_{i} $차원으로 projection을 수행하면 $ (\frac{H_{i-1}} {P_{i}} \times \frac{W_{i-1}} {P_{i}} \times C_{i}) $ 의 shape을 가지겠죠. 이를 통해 기존 해상도에 비해 $ P_{i} $ 만큼 줄었음을 확인할 수 있습니다.

 

 

 

 

 

Transformer Encoder

PVT의 경우 pyramid 구조를 위해 낮은 층에서 고해상도 feature map을 처리해야합니다. 위에서 예시로 들었던 것과 같이 patch size를 4로 하는 등의 과정을 거치기에 그만큼 연산량도 많아지겠죠. 이에 대해 본 논문은 해결책으로 spatial-reduction attention(SRA) 를 기존의 multi-head attention(MHA) 대신 사용하자고 제안합니다. SRA 역시 MHA와 비슷하게 Q(query), K(key), V(value) 를 받아 수행되지만 약간의 차이가 존재합니다.

 

 

위 figure를 보면 기존의 MHA와 달리 K, V의 해상도가 감소하여 Q와 적용되는 모습을 확인할 수 있습니다. 이 부분이 어떻게 수행되는지 자세히 살펴보도록 하죠.

 

 

위 (1) 식의 경우 multi-head를 concate를 한다고 보여주며, 이는 기존 transformer의 concate방식과 같습니다. 다음으로 각 head가 어떻게 계산되는 지 살펴보기 위해 (2) 식을 참고해야합니다. 기존 transformer의 방식과 같이 각 Q, K, V에 대해 다음과 같이 각 head의 차원으로 projection하는 linear projection을 수행하는 것은 동일합니다. 

 

 

물론 $ W^{O} $ 가중치 역시 차원의 변화 없이 linear layer를 적용시킵니다.

 

 

하지만 기존의 attention과 달리 K와 V에 SR이 적용된 것을 확인할 수 있습니다. 바로 위에서 언급했던 Spatial-Reduction을 수행하는 것인데 이에 대한 식은 다음과 같습니다.

 

 

우선 $ x $를 shape이 $ (H_{i}W_{i}) \times C_{i} $인 input sequence라고 했을 때, $ R_{i} $는 Stage i의 reduction ratio입니다. 다시 말해 $ Reshape(x, R_{i}) $는 $ x $를 $ \frac{H_{i}W_{i}} {R^{2}_{i}} \times (R^{2}_{i}C_{i}) $ 로 reshape 하는 것을 의미하죠. 

 

그 후에 $ C_{i} $로 projection하는 가중치 $ W_{S} $를 적용합니다. 

 

 

마지막으로 Norm(layer Nomalization)을 적용해주면 K, V의 shape은 $ \frac{H_{i}W_{i}} {R^{2}_{i}} \times C_{i} $ 가 됩니다. 즉, SR를 취해주는 것은 $ R_{i} $ 만큼 해상도를 감소시키는 작업을 해주는 것이 됩니다. 

 

이렇게 기존의 Q와 SR이 적용된 K, V 에 대해 각 head 차원인 $ d_{i} $ 차원으로 projection 한 후 attention을 수행하면 다음과 같이 계산이 진행됩니다. 우선 기존의 Attention 계산 과정은 다음과 같습니다.

 

 

위와 같이 Attention 계산을 감소된 K, V에 대해 적용하게 되면 다음과 같이 계산됩니다.

 

 

위와 같이 계산되기 때문에 Query의 해상도는 유지함과 동시에 각 위치의 정보를 $ R_{i}^{2} $ 만큼 요약함으로써 계산량을 줄일 수 있게 됩니다. 결국 attention score의 차원이 줄어든 것과 마찬가지이죠. 이와 같이 key, value의 scale을 감소하는 행위는 계산 비용을 크게 감소시킨다고 합니다. 

 

 

 

 

 

 

 

 

Hyper parameters

PVT 모델의 hyper parameter는 다음과 같습니다.

 

• $ P_{i} $ : Stage i 의 patch size
• $ C_{i} $ : Stage i 의 output channel 수
• $ L_{i} $ : Stage i 의 encoder layer 수
• $ R_{i} $ : Stage i 의 SRA reduction ratio
• $ N_{i} $ : Stage i 의 SRA head 개수
• $ E_{i} $ : Stage i 의 feed-forward layer의 expansion ratio