[논문 리뷰] Stand-Alone Self-Attention in Vision Models

Vision Transformer가 나오기 이전에도 self-attention을 computer vision 분야에 접목시키려는 많은 시도들이 있었다. 이번 논문은 그중에서 처음으로 self-attention만을 사용해서 모델을 만들었다.

https://arxiv.org/abs/1906.05909

Stand-Alone Self-Attention in Vision Models

---

1. Introduction

convolutional compute in CNN

최근은 Vision Transformer, MLP Mixer든 여러가지 형태의 네트워크들이 vision task에 등장하고 있지만, 그전까지는 CNN이 vision에서는 빼놓을 수 없는 dominant한 존재라는 것은 딥러닝을 연구하는 사람들이라면 알 것이다. CNN은 강한 inductive bias와 translation equivariance와 같은 특성으로 이미지 관련 task에서 엄청난 성능 향상을 불러일으켰는데, 아쉽게도 큰 receptive fields에 대한 scaling properties가 좋지 않아서 long range interactions이 어렵다는 단점이 있었다.

>> 아이러니하게도 주변 픽셀과의 relation이 높고 거리가 멀어질 수록 그 정도가 떨어진다는 "locality"특성을 잘 살린것이 CNN이면서 long range dependency를 못살리는 것이 단점이 되어버렸다.

이러한 long range interactions 문제를 해결하기 위해서 "attention" 이 등장하게 되었다. attention은 NLP에서 좋은 성능을 거두었는데, CNN에서도 이를 따라 attention을  함께 적용하여 성능을 끌어올리려는 연구를 하였다. 예로, Squeeze-Excite Network(SENet)와 spatially-aware attention등이 있다.

 

(위) SEBlock, (아래) attention augmented block

이렇게 Attention Mechanism은 convolutional models의 "add-on", 즉 additive method로써 사용이 되어왔는데, 저자는 한가지 의문을 제기한다.

Attention Mechainism을 주축으로 vision models을 만들면 좋은 성능을 낼 수 없을까?

 

논문에서는 이러한 질문의 대답으로써 attention layer로 모든 convolutional layer를 대체한 "stand-alone attention network"를 만들고, 이러한 모델이 baseline CNN 모델의 성능을 outperform하면서 parameter 및 compute efficient까지 좋다는 결과를 보여준다.

2. Stand-alone self-attention layer

논문에서는 Convolution layer를 attention으로 대체하기 위해서 "local attention layer"를 제안한다.

 

k=3 local attention

먼저 convolution과 비슷하게 self-attention을 할 windown size k를 정해준다. 논문에서는 이렇게 정한 block을 "memory block"이라고 부른다. 또한, 예시에서는 보기 쉽도록 2차원상에 그려놓았지만 실제로는 3차원 tensor형태를 가진다. Transformer의 self-attention 구조와 비슷하므로 빠른 이해를 위해 "attention is all you need" 리뷰를 읽고 오는 것을 권장한다.

https://aistudy9314.tistory.com/63

[논문리뷰] Attention is All you need

위 예시처럼 k=3일 때를 기준으로 알고리즘을 살펴보자!
pixel $x_{ij} \in \mathbb{R}^{d_{in}}$가 query, 그리고 그 주변 반경(${|a-i| \le k/2, |b-j| \le k/2}$)의 픽셀들을 key와 value라 할 것이다. 이 때 그 값들을 그대로 query, key, value로 사용하는 것이 아니라 weight layer를 통해 한번 transform을 해준다. 즉, queries $q_{ij} = W_Qx_{ij}$, keys $k_{ab} = W_Kx_{ab}$, values $v_{ab} = W_Vx_{ab}$이다.

>> 직관적으로 생각해보았을 때, 위와 같이 local한 window를 잡지 않고 모든 픽셀에 대해서 global attention을 하는 것이 Long-range dependency를 배우기 좋을 것이다. 하지만 이미지의 차원(픽셀의 개수)은 문장의 차원(token의 개수)보다 훨씬 크기 때문에 direct하게 global attention을 적용해버리면 computation측면에서 매우 큰 cost가 발생하게 된다.

Single-headed attention의 output은 다음과 같다.
$$y_{ij} = \sum_{a, b\in N_k(i, j)}softmax_{ab} (q^T_{ij} k_{ab}) v_{ab}$$
결국 self-attention이 하고 싶은 것은 query 픽셀과 key 픽셀들 간의 내적과 softmax를 취해 similarity matrix를 만들고 이를 values픽셀들과 내적하여 output을 냄으로써, 해당 픽셀과 주변 픽셀들 간의 relational information을 배우도록 유도하는 것이다.

또한 실제로는 single-head가 아닌 multe-head attention을 사용하고, 이는 head당 query, key, values의 depth를 나누어($x^n_{ij} \in \mathbb{R}^{d_{in} / N}$ 각각 다른 transforms($W^n_Q, W^n_K, W^n_V \in \mathbb{R}^{d_{out}/N \times d_{in} / N}$)을 적용함으로써 다양한 representation을 배우도록 한다.

추가로 local attention layer의 좋은 점은 window size와 파라메터의 수가 서로 independent하다는 것이다. 이는 $W_Q, W_K, W_V \in \mathbb{R}^{d_{out} \times d_{in}}$이고, keys와 values를 transform할 때 $d_{out} \times d_{in}$ vector를 element-wise로 곱하기 때문이다.

2.1 Positional Encoding

 

2D relative position embeddings

self-attention mechanism은 positional에 대한 정보를 알려주지 못하기 때문에 별도로 이러한 정보를 넣어주어야 한다. 대표적인 방법으로 "attention is all you need"에서 나온 sinusodial 함수를 사용한 embedding이 있는데, 여러 연구를 통해 이러한 absolute positional encoding보다는 relative positional embedding이 더 좋은 성능을 낸다는 것이 밝혀졌다.

논문에서도 "relative positional embedding"을 사용하고 이는 위 이미지와 같이 query기준으로 주변 픽셀들의 상대적인 위치 index들을 사용한다. 실제로는 위와 같이 (0, 0), (-1, 0) 등의 index value를 가지는 matrix가 아니고, learnable look-up table을 만들고서 해당 index의 값들을 가져오게 된다. 좀 더 자세히 설명하면, row offset과 column offset에 대해서 각 따로 learnable matrix가 look-up table로 존재하고 해당하는 offset의 벡터를 가져와서 이를 concat해준다.

>> row offset embedding과 column offset embedding 각 $\frac{1}{2} d_{out}$차원을 가지므로 합치면 $d_{out}$차원이 된다.

이를 식으로 나타내면 다음과 같다.
$$y_{ij} = \sum_{a,b \in N_k(i,j)} softmax_{ab} (q^T_{ij}k_{ab} + q^T_{ij}r_{a-i, b-j})v_{ab}$$

이렇게 positional information이 더해짐으로써 self-attention은 translation equivariance 특성을 가지게 된다.

3. Fully Attentional Vision Models

이제 local attention layer를 사용하여 모델을 구축할 시간이다.

먼저 저자들은 spatial convolution layer를 attention layer로 모두 대체하였다. 이 때, 1x1 conv는 spatial conv가 아니므로 제외되었다. 그리고 downsample이 필요할 때는 2x2 average pooling with stride 2를 사용하였다. 저자들 또한 이러한 transformation strategy이 suboptimal일 것이라 예상하고 있고, 추후 더 나은 architecture가 나오지 않을까라고 언급한다.

3.1 Replacing the convolutional stem

stem이란 보통 network의 매우 초반 layer를 말한다. 여러 실험을 통해서 edge와 corner같은 local features는 이러한 stem layer에서 학습이 되고, object의 detail(ex. texture)과 같은 global features는 후반 layer에서 학습이 된다고 알려져 있다.

저자들은 이러한 stem block에  local attention layer를 사용하면 성능이 떨어진다고 한다. 초반 input은 rgb 채널을 가진 이미지일 것이고 R, G, B 값들은 개별적으로는 아무런 정보를 담고 있지 않지만 spatially correlate한데 self-attention으로는 이러한 Input에서 좋은 정보를 뽑아내기 어렵다는 것이다.

>> local attention layer를 보면 채널 간 연산만 이루어지기 때문에 spatially correlation 정보를 배우는데 어려움이 있다.

따라서 저자들은 stem에 적합한 self-attention layer를 learnable weight $W_V$와 추가적인 position정보로 만들게 되는데, 자세한 내용은 논문의 appendix를 참고하길 바란다.

 

4. Experiments

4.1 ImageNet Classification

ImageNet 데이터셋에 대해서 성능을 비교한 논문이다. Baseline은 convolution을 사용한 모델, Conv-stem + attention은 stem만 convolution을 사용한 모델, full attention은 모든 layer를 attention으로 만든 모델이다.

표에서 볼 수 있듯이, conv-stem + attention이 가장 성능이 좋다.

>> 개인적인 생각으로 굳이 attention만을 사용하여서 모델을 만들어야하나 싶다. 많이 알려져 있는 사실로 CNN은 local features, VIT는 global feature를 많이 보는데 이를 위와 같이 알맞게 합치면 더 좋은 성능을 얻을 수 있지 않나 싶다.

 

두번째로 각 모델들의 parameters와 FLOP에 대한 accuracy를 살펴보았다. 마찬가지로 Conv-stem+Attention이 가장 efficient한 것을 볼 수 있다.

4.2. COCO Object Detection

Object Detection task에서 실험한 결과이다. Detection Head와 Backbone을 나누어 각 conv와 attention을 적용했을 때의 결과를 내었는데, 성능에서는 별로 차이가 안보이지만 FLOPS와 parameter 측면에서는 attention이 efficient한 것을 볼 수 있다.

4.3 Which Components are important?

self-attention을 사용한 모델을 만들 때의 고려사항 중 어떤 성분이 영향을 많이 미치는 지를 분석한 실험이다.

 

 

Spatial extent

먼저 window의 크기에 따른 결과이다(왼쪽). spatial extent가 커지면서 FLOPS가 증가하는 것을 볼 수 있고, 성능 또한 올라가는 것을 알 수 있다.

Positional encoding

positional encoding에 대한 결과이다(중간). positional encoding이 없을 때보다 있을 때가 더 성능이 좋고, absolute embedding보다는 relative가 더 좋은 성능을 낸다.

spatially-aware attention stem

이전에 설명했듯이, 그냥 self-attention을 stem으로 사용하게 되면 기존 conv를 사용했을 때보다 성능이 저하되는 것을 볼 수 있다. 이와 같은 문제를 해결하기 위해 spatially-aware attention stem을 만들었고, 이는 더 높은 성능을 보여준다.

>> spatially-aware attention을 stem뿐만 아니라 모든 layer에 사용한다면 어떤 결과가 나올 지 궁금하다.

 

---

 

처음으로 self-attention만을 사용해서 network를 만들었다는 데에 의의를 둘만한 논문이었다. 최근에 VIT에 관심이 많은데 한편으로는 VIT가 정말 좋은 것인가 싶기도 하다. 굳이 Attention을 쓰지 않고도 conv만을 사용하여 state-of-art 성능을 내는 논문이나 이러한 self-attention을 fourier transform으로 대체하는 논문 등도 많이 나왔기 때문에 같이 살펴보면서 비교할 필요성이 있다.