CNN에서 long range에 대한 대응을 위해 self-attention을 CNN에 도입한 SASA(Stand-Alone Self-attention)가 있었습니다. SASA는 Attention layer로 Convolution layer를 대체하면서 sota를 달성했는데요. 여기서 학습된 attention layer가 Convolution layer와 비슷하게 동작하는지에 관한 의문이 생깁니다.
본 논문은 실험과 증명을 통해서 attention layer가 convolution을 수행할 수 있음을 보여줍니다. multi-head self-attention layer가 충분한 multi-head를 가질 때 적어도 convolution layer만큼의 표현력을 가지게 됨을 확인합니다.
Introduction
self-attention의 도입
최근 자연어처리에서의 성과는 Transformer에 기반합니다. Transformer가 CNN이나 RNN과의 가장 큰 차이점은 입력이 동시에 모든 input sequence에 영향을 주는 점입니다. Transformer의 self-attention은 input sequence 내의 두 단어간의 유사도를 attention score를 기반으로 관계의 거리를 계산합니다.각 단어의 representation은 Transformer layer를 지나면서 attention score가 높은 단어들을 기반으로 업데이트가 되게 됩니다.
Vision에서의 self-attention활용
단어간의 상호관계를 학습하는 능력에서 영감을 받아 self-attention을 활용하기 시작했습니다. 특히 SASA는 attention과 convolution을 결합하여 sota를 달성하였고 self-attention only architecture로 경쟁력있는 image classification accuracy를 보였습니다. 여기서 생기는 의문은 self-attention layer가 image를 다룰 때 convolution layer와 비슷한 방식으로 처리하는지에 대한 궁금증이 생깁니다. Perez et al. (2019)에는 attention 기반의 multi-layer architecture에 추가적인 positional encoding을 더하면 이론적으로 unbounded precision arithmetic와 같은 문제에 대해 Turing complete하다고 하나 일반적인 task에 대해서는 그 능력이 있음만을 실험으로 보여줍니다.
Contribution
본 논문에서는 이론과 실험으로 self-attention이 convolution과 유사하게 동작한다는 증거를 제시합니다.
- 이론적으로 self-attention layer가 convolution layer를 표현할 수 있음을 증명합니다.
- SASA의 attention-only architecture에서 앞단의 레이어들은 위의 이론과 같이 convolution처럼 동작을 하며 query pixel 주변의 grid-pattern에 대해 학습합니다.
위를 통해서 이전 work들에서 귀납적으로 유도된 앞단의 레이어를 convolution으로 구성하는 것이 옳음을 알 수 있습니다.
Background
Multi-head self-attention layer
X는 input matrix로서 D_in dimension을 가지느 T token으로 구성됩니다. NLP에서는 token은 각 단어의 임베딩이 되지만 이러한 방식은 모든 sequence에 적용이 가능합니다. self-attention layer는 token t를 D_in에서 D_out으로 아래의 식을 통해 mapping합니다.
positional encoding
위의 self-attention model에서는 input token의 순서가 반영되지 않습니다. 따라서 positional encoding P를 input matrix에 더해서 사용하게 됩니다.
Multi-head를 self-attention에 적용하는 것은 각각 다른 query, key, value를 사용함에 따라서 각 head가 서로 다른 부분에 집중 할 수 있도록 합니다. 위에서 본 Self-attention의 각 head들의 output은 concat되고 Linear layer를 통해서 projection되면 Multi-head self-attention(MHSA)의 결과를 얻을 수 있습니다.
![W_out: [D_head, D_out] b_out[D_out]](https://inblog.ai/_next/image?url=https%3A%2F%2Fwww.notion.so%2Fimage%2Fhttps%253A%252F%252Fs3-us-west-2.amazonaws.com%252Fsecure.notion-static.com%252Fd56bea25-118a-4768-864d-2673ff1c8c38%252FUntitled.png%3Ftable%3Dblock%26id%3D5acaa9f8-750a-47be-8163-f7e80d1ee2be%26cache%3Dv2&w=1920&q=75)
Convolution for images
convolution layer는 image를 위한 neural network의 중요한 구성요소중 하나입니다. Input image tensor X는 W, H, D_in으로 구성됩니다. 이때 Convolution layer의 i,j pixel의 output은 아래와 같습니다.
- ∆K는 kernel size에 따라서 해당 pixel에 영향을 주는 모든 shift의 조합을 의미합니다.
Positional encoding for images
Transformer기반의 아키텍쳐에는 absolute와 relative로 두가지 방식의 positional encoding이 주로 사용되었습니다.
absolute positional encoding: positional encoding을 위한 vector P가 각 pixel p에 할당됩니다.
relative positional encoding: Transformer-XL에서 도입된 방식으로 key pixel의 절대적인 위치 대신 query pixel과 key pixel의 위치 차이를 고려하자는 것입니다.
- key에 대한 absolute position embedding P_k는 r_δ로 대체됩니다.
- query에 대한 absolute position embedding P_qW_qry는 3번째와 4번째 term에서 각각 u,v로 대체되며 u와 v는 학습가능한 parameter를 사용합니다.
- key를 계산하기위한 W_key가 W_key와 ^W_key로 분리됩니다. W_key는 content 기반의 key_vector를 ^W_key는 location 기반의 key_vector를 만들어내기 위함입니다.
Proposed Method
Theorem 1
Multi-head attention layer가 D_p ≥ 3 인 경우 kernel이 (√N_h x √N_h)이고 output channel이 min(D_h, D_out)인 Convolution layer를 표현할 수 있습니다.
- N_h: Head의 수
- D_h: Head의 dimension
- D_out: output dimension
- D_p: relative positional encoding의 dimension
이 theorem은 multi-head self-attention의 parameter를 선택하여 convolution layer 처럼 동작할 수 있음을 보여 증명됩니다. 각 head의 attention score는 relative shift에 따라서 K x K kernel의 다른 pixel shift에 참여해야합니다. 정확한 조건은 Lemma 1에서 확인할 수 있습니다
Lemma 2는 quadratic encoding이 relative encoding으로 사용될때 위에 언급한 조건을 만족함을 보입니다.
학습되는 parameter인 ∆(h) = (∆_1(h), ∆_2(h))와 α(h)는 각 head의 attention의 넓이와 center를 결정합니다. δ=(δ_1, δ_2)는 query와 key 사이의 상대적인 shift를 표현합니다.
여기서 Lemma1을 만족하는 조건이 위에서 언급한 quadratic encoding만이 아니라는 점이 중요합니다. 실험에 따르면 학습된 neural network에 의한 relative encoding 또한 lemma의 조건을 만족하게 됩니다. 위에서 정의한 encoding의 경우 D_p=3 dimension만으로도 relative position encoding을 하는데 충분하고 좋은 성능을 보입니다.
PROOF OF MAIN THEOREM
Lemma 1
MHSA가 N_h = K^2, D_h ≥ D_out이고 f:[N_h] → ∆K가 shift에 따라 head가 일대일 대응인 bijective mapping이라고 가정합니다.
head의 attention probability은 f(h)=q-k일때만 1이고 나머지는 모두 0이 된다고 가정합니다.
그러면 KxK kernel이고 D_out의 output channel을 가지는 Convolution layer가 있을 때
모든 X에 대해 MHSA(X) = Conv(X)인 W_val^(h)가 존재합니다.
Proof
D_h ≥ D_out이라고 가정하면 학습되는 파라미터들을 W(h)로 replace가능합니다.
따라서 pixel에 따른 MHSA은 다음과 같이 쓸 수 있습니다.
Lemma의 가정으로 k=q-f(h)인 경우만 attention probability는 1이고 나머지는 모두 0이됩니다. 따라서 pixel q에 대한 MHSA의 output은 아래와 같습니다.
Convolution layer의 K=√N_h인 경우 동치입니다. 여기서 각 head의 W(h)는 kernel의 W_(k1,k2)에 일대일 대응됩니다.
Lemma 2
D_p ≥ 3이고 parameter W_qry, W_key, ^W_key, u가 있고 D_p ≤ D_k인 경우 모든 ∆에 대해 k − q = ∆인 경우 attention probability가 1이고 나머지는 0인 v가 존재한다.
Proof
input tensor X에 대해 attention probability는 독립적이므로 W_key와 W_qry를 0으로 설정하여 마지막 term만 남깁니다. 그리고 W_key를 identity matrix로 설정하면 아래의 식을 얻게 됩니다.
이를 2차식으로 아래와 같이 쓸 수 있습니다. δ = ∆인 경우에는 Attention score가 −αc로 max가 되게 됩니다.
δ = ∆인 경우에는 softmax를 통해 1인 결과를 얻을 수 있게되고 나머지는 0이됩니다.
그러면 eq 14의 v와 r_δ가 존재하는를 증명해야합니다. 우항을 확장해보면 아래와 같습니다.
여기서 eq(9)에 따라서 v=-α(1, -2∆_1, -2∆_2), r_δ =(||δ||^2, δ_1, δ_2)로 설정하면
이는 다시 eq(14)와 matching이 되므로 c=-||∆||^2로 증명이 완료됩니다.
Experiment
Implementation details
Cifar-10에 대해 ResNet18과 Bello et al의 Attention Augmented Convolutional Networks(2019)를 비교하였습니다.
모든 실험에서 image의 크기를 줄이기 위해서 2x2 invertible down-sampling을 적용하였습니다. 큰 이미지에는 full attention이 적용될 수 없기 때문에 attention coefficient tensor의 크기는 input image의 크기에 따라서 2차적으로 scale됩니다. 최종 feature는 average pooling을 거쳐 마지막 linear classifier를 거치게 됩니다.
Remark on accuracy
self-attention이 잘돌아가는 것을 확인하기 위해서 300 epoch이상 학습을 한 뒤 성능을 비교하고 figure6를 그렸습니다.
Content에 영향을 받지 않는 attention map을 만들기 위해서 100장의 attention map을 average하였습니다. 구현을 gaussian attention prob으로 최적화를 하여 FLOPS를 크게 줄여 parameter의 수 차이로 인한 gap을 빼꿀 수 있겠으나 본 실험의 목적이 아니므로 시도하지 않았습니다.
Quadratic encoding
첫번째로 eq(9) relative position encoding을 도입한 attention layer에서 convolution layer와 같이 행동하도록 학습하는지 확인합니다. 각 레이어는 9개의 attention head를 갖고 ResNet은 3x3 kernel을 가지도록 하고 학습했습니다.각 head의 attention은 ∆(h) ∼ N (0, 2I_2)로 initialize됬습니다.
Figure3에서는 layer4에서 각 head의 attention의 초기 위치가 학습을 통해서 바뀌는 모습을 보여 줍니다. 최적화 이후에는 head의 attention이 이미지의 특정 위치에 관여하도록 query pixel(중앙의 검은색 픽셀) 주변으로 grid를 형성하는 모습을 볼 수 있습니다. 따라서 queried pixel 주변으로 self-attention이 convolution filter를 학습할 수 있는 점을 확인할 수 있습니다.
Figure4는 학습이 끝난 각 레이어의 모든 head의 attention을 display합니다. 앞 레이어(layer1 and 2)들은 local pattern에 집중하게 되고 deeper layer들은 더 큰 pattern에 attention이 가게 됩니다. attention head는 겹치지 않고 입력 공간의 범위를 최대화하는 배열을 하는 것처럼 보입니다.
Learned relative positional encoding
2D relative positional encoding은 SASA를 기반으로 구현했습니다.
각 row와 column pixel shift에 대한 D_p/2 position의 encoding vector를 학습했습니다. query pixel q에 대해 key pixel k의 relative positional encoding은 row shift embedding δ1와 column shift embedding δ2의 concatenation과 같습니다. 실험에서는 D_p = D_out = 400를 설정햇습니다.
구현에서는 다음 두가지가 다릅니다.
1) downsampling에 convolution stem과 ResNet bottleneck 대신 2x2 invertible downsampling layer를 input에 사용했습니다.
2) 학습되는 필터의 개수는 min(D_h, D_out)이라는 이론에 따라서 D_h = D_out/N_h 대신 D_h = D_out으로사용했습니다.
처음에는 eq(8)에서 input 데이터를 버리고 마지막 term만으로 attention score를 계산합니다. 각 head의 Attention probability는 Figure 5에서 확인가능합니다. 이 그림에서 첫번째와 두번째 그리고 3번째 레이어의 일부는 가설을 증명해줍니다. 무작위로 초기화된 벡터로 부터 positional encoding 방식을 학습할 때 self-attention head는 각기 다른 pixel을 attention하도록 합니다. 이는 Lemma 1의 조건과 유사하고 theorem 1을 따릅니다.
practical한 case로 넘어가서 positional과 content-based attention을 모두 고려하여 attention score를 계산하도록 해봅니다. 이때의 각 layer의 attention probabilities는 Figure 6에서 확인할 수 있습니다. 2번째와 3번째 레이어에서 가설을 확인할 수 있습니다. content에서 encoding을 배우도록 해도 특정 self-attention head는 position-based attention만을 사용하여 query pixel에 대해 convolution kernel의 receptive field를 reproducing하여 고정된 shift를 보임을 알 수 있습니다. 다른 Head들은 content-based attention을 더 사용하여 self-attention의 장점을 사용하게 됩니다.
query pixel이 이미지를 따라 convolve할 때 convolution과 multi-head attention 사이의 비슷한 점이 돋보입니다. 다른 query pixel을 받은 Figure7과 비교해보면 이러한 특성을 확인할 수 있습니다. layer2와 layer3의 attention pattern을 비교해보면 query pixel로부터 고정되게 shift된 것을 볼 수 있습니다. 이는 convolution kernel이 receptive field를 convolving하는 것과 유사합니다.
Conclusions
새롭게 알게 된 사실은 없지만 self-attention을 이미지에서 사용할 때 어떻게 동작하게 될지 이해할 수 있는 논문이었습니다. 다양한 module에서 self-attention을 적용하기 때문에 동작 원리를 이해하는데 도움이 되었습니다.
relative positional encoding은 Transformex-XL에서 긴 문장을 segment로 분리하여 reccurrent하게 가져가게 될 경우 다른 문장으로 넘어가더라도 동일한 position으로 인식되는 것을 막기 위해서 relative positional encoding을 적용합니다.
Share article