Twins: Revisiting the Design of Spatial Attention in Vision Transformers, 2021Q1
http://arxiv-export-lb.library.cornell.edu/abs/2104.13840
최근에 dense prediction task를 위한 vision transformer 아키텍쳐가 제안되었고 spatial attention 디자인의 중요성이 대두되고 있다. 이 연구에서, 간단하지만 세심하게 디자인된 spatial attention을 제안하며 이 성능은 현재의 SOTA만큼의 성능을 낸다. 결과적으로, 우리는 두개의 비전 트랜스포머 Twins-PCPVT와 Twins-SVT를 제안한다. 우리의 아키텍쳐는 최신 딥러닝 프레임워크에서 매우 최적화되어 있는 행렬곱으로만 이루어져 있어 매우 효율적이고 구현하기 쉽다. 제안된 아키텍쳐는 이미지레벨의 분류뿐 아니라 object detection같은 dense prediction에서 훌륭한 성능을 냈다.
최근 ViT는 CNN에 비해 1) long-range dependency를 학습하는데 효율적이고 2) 적은 inductive bias를 가정하고 3) 자연스럽게 이미지, 텍스트, 비디오 같은 다양한 modality를 처리할 수 있어 최근 많은 연구 관심을 받고 있다. ViT는 CNN을 대체하여 다음세대의 비전 연구의 가장 기본 요소가 될 것 같다.
vision task에 transformer를 적용할때의 가장 주요한 문제 중 하나는 트랜스포머내에 있는 self-attention으로 인한 $O(n^2)$의 계산 복잡도이다. (n은 인풋 이미지의 픽셀) 이를 해결하기 위해서 input이 윈도우로 묶이고 각 sub-window내에서 self-attention을 하는 locally-grouped self-attention(혹은 SwinTransformer에서는 non-overlapped windows)이 제안되었다. 이는 복잡도는 줄여주지만 다른 윈도우 사이의 연결을 막고, 제한된 시야(receptive field)를 가지게 한다. 많은 연구에서 충분히 큰 시야는 성능에 (특히 dense prediction에서) 중요하다. Swin은 이를 위해 local window의 바운더리가 레이어가 깊어질 수록 이동하는 Shifted Window 접근법을 제안했다. 이는 효과적이었지만, shifted window의 개수가 일정하지 않았다. 일정하지 않은 개수의 window는 ONNX나 TensorRT에서 deploy될 때에 문제가 있었다. 다른 접근 방식으로 Pyramid Transformer(PVT)가 제안되었다. 일반적인 self-attention operation과 달리, PVT는 인풋토큰의 subsample에 대해서만 어텐션을 계산했다. PVT 역시 이론상은 $O(n^2)$이지만, 실제로는 더 나은 효율성을 가졌다.
종합적으로 보면 언급한 ViT들은 어떻게 spatial한(인접한 픽셀끼리 묶어서 attention을 했으니 spatial하다고 표현하는듯) attention을 디자인했냐에 따라 달라졌다. 이 연구에서는 spatial attention의 디자인에 대해서 집중하려고 한다. 우리의 첫번째 발견은 PVT의 global sub-sampled attention이 매우 효율적이었고, applicable positional encoding(Conditional Position encoding ViT, CPVT)*과 함께, SOTA모델과 유사하거나 더 나은 성능을 보였다는 점이다. 이는 우리의 첫번째 아키텍쳐인 Twins-PCPVT를 제안하게 하였다. 그 위에, 우리는 단순하지만 섬세하게 디자인된 Spatially Seperable Self-Attention(SSSA)를 제안하였고, PVT를 더 효율적으로 만들었다. SSSA는 seperable depthwise convolution에 영감을 받았고, 두 종류의 attention연산으로 구성되어 있다: (1) locally-grouped self-attention(LSA) (2) global sub-sampled attention(GSA), LSA는 fine-grained, short-distance정보를 처리하고 GSA는 long-distance global 정보를 처리한다. 두번째로 제안된 Twins-SVT가 이러한 맥락에서 나오게 되었다.
*conditional positional encoding : https://arxiv.org/abs/2102.10882 *seperable depthwise convolution : https://coding-yoon.tistory.com/122
첫번째 방법은 PVT와 CPVT를 결합한 Twins-PCPVT이고, 두번째 방법은 우리의 SSSA에 기반한 Twins-SVT이다.
PVT는 dense prediction을 하기 위해 pyramid multi-stage 디자인을 도입했다. 이 모델에서 ViT나 DeiT처럼 absolute positional encoding을 사용했다. 모든 레이어는 글로벌 어텐션을 했고 spatial reduction을 통해 모든 시퀀스를 계산하는 비용을 줄이고자 했다. 최근에 제안된 Swin Transformer는 shifted local windows를 도입하여 PVT보다 (심지어 large receptive field가 더 중요한 dense prediction에서도) 좋은 성능을 냈다.
이 연구에서는, 덜 선호되는 PVT의 성능이 absolute positional encoding 때문임을 밝혀냈다. absolute positonal encoding은 이미지 크기가 달라질 때 문제를 해결하기 어렵고, translation invariance라는 룰을 깨기도 한다. SwinTransformer는 relative postional encoding을 사용하여 이러한 문제를 우회했고, 성능의 우위는 아키텍쳐가 아닌 적절한 포지셔널 인코딩때문임을 알 수 있었다.
*tranlation invariance : 이미지의 위치가 달라져도 output이 같아야 함 https://ganghee-lee.tistory.com/43
우리는 CPVT에서 제안된 CPE를 사용하여 absolute PE를 CPE로 바꾸었다. CPE는 인풋에 따라 임베딩이 달라져 위에서 언급한 absolute PE의 문제점을 보안한다. position wise generator(PEG)는 각 스테이미지의 첫번째 인코더 블록 뒤에 위치한다. 이미지 분류 태스크에서는, CPVT를 따라 우리는 class 토큰과 global average pooling(GAP)를 지웠다. 나머지의 비전 태스크에서는 PVT의 디자인을 따랐다. Swin에서도 CPE를 적용해봤지만 큰 성능 향상은 없었고, 이는 shifted window가 CPE와 잘 맞지 않기 때문으로 생각됐다.
high-resolution 인풋에 대한 dense prediction을 할때 ViT는 높은 계산 비용에 허덕였다. $H\times W$크기의 이미지에 $d$차원의 self-attention을 할 경우 $O(H^2W^2d)$의 복잡도를 가지고 있었다. 이를 해결하기 위해 LSA와 KSA가 있는 SSSA를 제안한다.