논문 링크: Attention Augmented Convolutional Networks
Attention 구조는 기존 RNN와는 다르게 Parallelization이 가능해져서 많은 것들을 개선할 수 있다. 특히 Sequence 에 두 feature가 서로 interact 하기 까지의 연산량을 대폭 줄일수 있다.
RNN은 Timestep 마다 implicit하게 이전 step t-1에서 t로 position에 대한 정보를 넘겨주었기 때문에 위치 정보에 대한 별다른 조치가 필요하지 않다. 하지만 Attention 구조에서는 parallelization + 이전 상태에 대한 정보 (memory cell)이 없기 때문에 position의 대한 정보를 별도로 제공해주어야 한다.
따라서 position의 대한 정보를 제공하기 위해 positional embedding을 추가해 주어야 한다. 이 논문에서는 Width와 Height에 대한 정보를 relative positional embedding에 추가한다. 이 두개의 embedding은 trainable이 가능하기 때문에 train해야하는 parameter가 약간 더 늘어나게 된다.
먼저 Self-attention이기 때문에 Query, Key, Value를 위해 attention maps이 각각 split된다.
(그러기 위해서는 Input convolution Layer의 output channel dimension이 $d_{q} + d_{k}+ d_{v}$ 이 되게끔 해주어야 한다. 이때 $d_{q} = d_{k}$이다).
다음으로 split된 Channel dimension이 Head의 갯수로 짝을 지어 dimension / Nh 만큼 split된다. multi-head attention Layer후 모든 Head output이 Concat되고 linear projection(1x1 conv)을 거쳐 최종 attention feature가 나오게 된다. 최종적으로 attention feature와 standard convolution을 거친 feature가 concat이 된다.
참고로 앞서 설명했듯이 multi-head attention 에서 Key+query과 Value가 interaction할때 추가적으로 relative positional embedding term이 2개가 붙게 된다.
$$O_{h} = Softmax (\frac{(QK^{T} + S_{rel}^{H} + S_{rel}^W)} {\sqrt{d^{h}_{k}}}) * V $$
이때 $S_{rel}^{H} = \gamma^{H} * q^{T}_{i}$ , $S_{rel}^{W} = \gamma^{W} * q^{T}_{i}$ 이다.
$q^{T}_{i}$ 는 Q 의 i th row 이고 $\gamma^{H, W}$는 learned positional embedding이다.
이 논문에서는 기존 Transformer에 Convolution을 추가해서 image classification이 가능하다. 논문에서 AAConv모델과 다른 기존의 CNN 모델들과 비교했는데 비슷하거나 약간 나은 결과를 보여주었다. 결론적으로 Attention 기반 model이 기존의 CNN을 충분히 대체 가능하다는게 이 논문의 핵심인듯 하다.
위 논문의 appendix에 tensorflow implemetation 코드가 있어 이해하는데 도움이 많이 된다...