[논문리뷰] Attention Augmented Convolutional Networks

1. Introduction

1.1 Convolutional Neural Network(CNN)

 

CNN은 computer vision에서 높은 성능을 보여주면서 dominant한 mechanism으로 자리잡았다. 이렇게 CNN이 이미지 task에최적화된 성능을 낼 수 있는 이유는 1) 일정 크기의 kernel을 통한 sliding 방식으로 receptive field를 제한하여 "locality"를 가지는 점과 2) weight sharing을 통한 translation equivariance 특성이 있다. 하지만 이러한 convolutional kernel의 locality로 인해서 이미지의 global contexts를 잡아내기 힘들다는 단점도 존재한다.

>> locality가 좋은 이유는 어떠한 픽셀이 그 주변의 픽셀들에 영향을 크게 받고 거리가 멀어질 수록 그 영향력이 떨어진다는 성질 때문이다. 하지만 경우에 따라서는 위 global한 특징을 잡아낼 필요성이 있다는 것을 이야기 하고 있다.

 

1.2 Self-attention

 

self-attention은 CNN과 반대로 long range interaction을 잘 학습할 수 있다는 장점을 가지고 있지만 적절한 inductive bias가 없기 때문에 locality와 translation equivariance 같은 특성을 학습하기 어렵다는 단점이 있다.

 

논문에서는 relative positional embedding을 적용하여, self-attention을 사용하면서 translation equivariance를 유지할 수 있는 relative self-attention을 제안한다. 이 mechanism을 사용하면 모든 convolution을 self-attention으로 대체했을 때에도 좋은 성능을 보여주지만 가장 높은 성능은 "Convolution"과 "Self-Attention"을 결합하였을 때 나타난다고 한다.

2. Traditional Attention Mechanism

이전에 vision task에서 CNN의 supplement 역할로써 사용되었던 attention mechanism들이 있다. 여기서 자세한 알고리즘을 살펴보진 않을 것이다.

 

(왼쪽) Squeeze-and-Excitation, (오른쪽) Gather-Exctite

 

먼저 전체 feature maps으로부터 global average pooling과 같은 method로 어떠한 채널이 더 큰 영향을 주는지에 대한 weight map을 뽑아서 feature channel을 reweight하는 방식의 Squeeze-and-Excitation과 Gather-Excite 방식이 있다.

 

(왼쪽) BAM, (오른쪽) CBAM

 

다음으로는 비슷한 방식이지만 channel과 spatial dimensions 두 방향으로 독립적이게 convolutional feature를 reweight하는 BAM과 CBAM method가 있다.

 

3. Methods

3.1 Notation

methods를 설명하기에 앞서 사용되는 Notation을 정리하고 가겠다.

• $H, \ W, \ F_{in}$은 activation map의 height, width 그리고 input filters의 개수
• $N_h, \ d_v, \ d_k$는 차례대로 head의 개수와 values의 depth, keys와 queries의 depth
• $d^h_v, \ d^h_k$는 각 depth를 head의 개수로 나눈 수, 즉 attention head 당 values와 queires, keys의 depth

3.2 Self-attention over images

이미지에 대해 self-attention을 어떻게 적용할지에 대한 내용이다. 

 

먼저 주어진 input tensor($H, \ W, \ F_{in}$)를 $X \in \mathbb{R}^{HW \times F_{in}}$ 크기로 flatten한다. 그리고 이를 그대로 Transformer architecture에서 사용되었던 multihead attention에 적용하게 된다. 

 

 

single head에 대한 attention map을 계산하는 식을 표현하면 다음과 같다.

$$O_h = Softmax(\frac{(XW_q)(XW_k)^T}{\sqrt{d^h_k}}(XW_v)$$

** $W_q, \ W_k \in \mathbb{R}^{F_{in} \times d^h_k}$이고 $W_v \in \mathbb{R}^{F_{in} \times d^h_v}$

 

간단하게 설명하면 input X를 각 다른 weight matrix를 통해 queries, keys, values vector로 embedding하고, queires와 keys간 matmul과 softmax를 통해 attention map을 구해 values vector에 최종적으로 곱해준다는 것이다. 이러한 방식이 이해가 안간다면 "attention is all you need" 리뷰를 한번 보고 오기를 추천한다.

 

그 다음으로 각 head에서 구한 output들을 서로 concat한 후에 한 번 더 projection해준다.

$$MHA(X) = Concat[O_1, \dots, O_{Nh}]W^O$$

**$W^O \in \mathbb{R}^{d_v \times d_v}$

 

마지막으로 이를 $(H, W, d_v)$ tensor shape로 reshape해주면 끝이다.

>> 저자들은 multi-head attention에서 complexity가 $O((HW)^2 d_k)$, 그리고 memory cost가 각 head 당  attention map을 저장하기 위해 $O((HW)^2 N_h)$만큼 필요하다고 한다. 이는 spatial dimension에 따라서 매우 커질 수 있기 때문에 문제점이라고 할 수 있다.

 

3.3 Two-dimensional Positional Embeddings

(왼쪽) sinusoidal encoding, (오른쪽) CoordConv

 

위와 같은 self-attention만을 사용하게 되면 input 성분(ex. pixel)들의 위치가 바뀌더라도 똑같은 결과가 나오는 "permutation equivariant" 현상이 발생한다. 이는 image와 같은 highly structured data를 모델링 할 때 좋지 않은 결과를 일으키는데 이러한 문제를 피하기 위해서 명시적으로 spatial information을 줄 수 있는 여러 positional encoding 방법들이 제안되었다. 몇가지 예로 transformer에서 사용된 sinusoidal waves를 이용한 absolute positional emcoding과 positional channel을 concat하는 CoordConv 방식이 있다.

 

그러나 이러한 encoding 방법들은 image classification이나 object detection과 같은 task에서 큰 효과를 보여주지 못하였는데, 저자들은 그러한 이유로 "해당 positional encoding들은 permutation equivariant하지만 translation equivariance는 만족시키지 못하기 때문이라고 가정하였다. 그리고 이를 해결하기 위해 relative position encoding을 2차원으로 확장해서 사용한다.

 

Relative Positional Embeddings

 

relative positional encoding은 말그대로 queries를 기준으로 상대적인 위치를 positional index로 사용하는 방법이다. 논문에서는 relative height information과 relative width information을 독립적으로 더해줌으로써 2-dimensional relative self-attention을 만든다. 

 

attention logit을 식으로 나타내보면 다음과 같다.

$$l_{i,j} = \frac{q^T_i}{\sqrt{d^h_k}}(k_j + r^W_{j_x - i_x} + r^H_{j_y-i_y})$$

**$q_i$는 pixel i에 해당하는 query vector, $k_j$는 pixel j에 해당하는 key vector

**$r^W_{j_x-i_x}$와 $r^H_{j_y-i_y}$는 relative width $j_x-i_x$와 relative height $j_y-i_y$에 해당하는 learned embeddings

 

head h의 output을 나타내보면 다음과 같다.

$$O_h = Softmax(\frac{QK^T+S^{rel}_H+S^{rel}_W}{\sqrt{d^h_k}})V$$

 

이 때 $S^{rel}_H, S^{rel}_W \in \mathbb{R}^{HW \times HW}$이고, 모든 relative embeddings $r_{ij}$를 저장하기 위해 $O((HW)^2 d^h_k)$만큼의 메모리가 필요하다. 이는 positional embedding을 사용하지 않았을 때의 memory cost인 $O((HW)^2 N_h)$와 비교했을 때 매우 큰데 따라서 mini-batch로 사용할 수 있는 이미지의 개수가 제한되는 문제가 발생한다. 즉, 공간적으로 너무 비효율적이라는 이야기이다. 

 

저자들은 unmasked relative self-attention 방법을 사용하여서 memory cost를 $O(HWd^h_k)$만 사용하여 이러한 문제를 해결한다. 

**relative positional embeddings $r^H$와 $r^W$는 heads간은 공유되고 layer간은 공유되지 않는다.

 

3.4 Attention Augmented Convolution

 

논문 저자가 제안한 Attention augmented convolution이 이전 method와 다른 점은 다음과 같다.

• attention mechanism을 사용하여 spatial과 feature subspaces를 jointly 참조할 수 있다.
• convolutional features를 refine(reweight)하는 것이 아니라 attention을 통해 새롭게 feature maps을 만들어낸다.

방법은 정말 간단하다. 위 이미지처럼 self-attention을 하는 path와 standard convolution을 하는 path 두 갈래로 나누어서 진행한 다음 마지막에 두 output을 Concatenate해주면 된다. 

 

$$AAConv(X) = Concat[Conv(X), MHA(X)]$$

 

개인적으로 위 이미지가 좀 난해했기 때문에 tensorflow code를 참조하여 single-head에 대한 attention augmented convolutional networks를 만들어보았다.

 

4. Experiments

마지막으로 여러 실험을 통해서 Attention Augmented Convolutional Networks의 우수성을 보여준다. 

4.1 CIFAR-100 image Classification

먼저 CIFAR-100 데이터셋에 대해서 다른 attention mechanism인 SEBlock및 GEBlock과의 성능을 비교한 논문이다. top-1 accuracy를 봤을 때 proposed network가 가장 높은 accuracy를 가지는 것을 볼 수 있다.

>> 근데 성능 차이가 그렇게 크지도 않고 심지어 top-5에서는 SEblock보다 더 떨어진다....SEblock논문이 2017년도에 나왔는데... 

 

4.2 ImageNet image classification with ResNet

이번에는 ImageNet 데이터셋에 대한 결과이다. 먼저 위 표는 다른 attention mechanism과 proposed method를 비교한 표인데, Attention Augmented method가 가장 적은 parameter를 가지면서 좋은 top-1 accuracy를 가진다.

** $\triangle$은 ResNet50과 비교했을 때 늘어나는 latency times를 말한다. 즉 낮을 수록 baseline과 비교했을 때 지연시간이 적다는 이야기이다.

 

 

두번째로 ResNet layer를 달리하면서 실험한 결과이다. 모든 layer에 대해서 proposed method가 적은 parameter를 가지면서(SEblock에 비해서) 좋은 accuracy를 가지는 것을 볼 수 있다. 

>> 이 실험 또한 accuracy 차이가 크지 않다. 특히 SEblock과 Attention Augmented block간의 차이는 커봤자 0.2...

 

4.3 ImageNet Classification with MnasNet

resource constrained setting에서 attention augmented method를 사용했을 때 어떠한 결과가 나오는지를 실험한 결과이다. 약간 parameter 수는 늘어나지만 performance가 상승하는 효과를 보여준다.

 

4.4 Object Detection with COCO dataset

Object Detection task에 적용했을 때 어떠한 결과를 보이는지에 대한 실험이다.

 

Squeeze-and-Excitiation operators는 오히려 detection 성능을 떨어트리는 반면에 Augmented Attention은 성능을 올려주는 효과를 보여준다. 저자들은 SEBlock이 localization에 필요한 spatial information을 없애기 때문에 이러한 결과가 나온다고 추측한다.

 

4.5 Ablation Study

Fully-attentional vision models

relative self-attention만을 사용해서 model을 만들었을 때 어떠한 결과가 나오는지에 대한 실험이다.

여기서 $\kappa$와 $v$ hyperparameter가 나오는데 각 $\kappa=\frac{d_k}{F_{out}}$, $v=\frac{d_v}{F_{out}}$를 의미한다.

 

conv연산에 사용되는 채널이 $F_{out}-d_k$가 되기 때문에 위 비율이 높을 수록 convolution을 덜 사용하는 것이라 볼 수 있다. 1.0일 때가 아예 convolution을 사용하지 않은 architecture인데, 비율이 높아질 수록 높은 memory cost 때문에 average pooling을 사용하여 spatial size를  줄이기 때문에 accuracy 저하가 일어난다고 한다. 그래도 더 적은 파라메터 수와 위와 같은 downsampling이 있었음에도 baseline과 2.5% 밖에 성능 차이가 없기 때문에 stand alone self-attention 모델이 충분히 powerful하다고 볼 수 있다.

 

Importance of position encodings

다른 position encodings을 사용했을 때 어떠한 결과가 나오는지에 대한 실험이다.

먼저 classification에 대해서 relative encoding을 사용했을 때 가장 성능이 높게 나온다.

>> 성능 차이가....거의 안난다...

 

 

두번째로 object detection에서의 실험인데 CoordConv의 경우 accuracy를 저하시키는 반면 relative positional encoding은 성능을 향상시킨다.