[논문 요약] On the relationship between self-attention and convolutional layer (1)

1. Introduction

Computer Vision에 Attention을 적용하는 트렌드

• 모델의 블록을 만들 때 과연 Convolutional layer가 우수한 지 고민하게 되었다
• CNN으로 하여금 long-range dependency를 학습하도록 도와주는 것을 넘어서, Attention이 CNN을 완전히 대체할 수 있음이 입증되었다
• 이러한 추세는 학습된 Attention이 CNN과 유사한 역할을 하는 지 의문을 갖도록 했다

본 논문은 Attention layer가 Convolution처럼 동작할 수 있는 것에 대한 근거를 제시하며,
충분한 개수의 Head를 갖는 MHSA는 적어도 Convolution layer보다 expressive함을 증명한다.

2. Background

Position Encoding을 사용하지 않을 때

$Input\,\,\,\,\, X$
$Attention\,\, Score\,\,\,\,\, A\,\,:=\,XW_{qry}W^T_{key}X^T$

$Self\,Attention(X)_{t,:}\,:=\,softmax(A_{t,:})XW_{val}$

이 경우, 문장을 구성하는 토큰의 순서와 상관없이 Self Attention이 항상 같은 결과를 출력한다.

Position Encoding을 사용할 때

$Input\,\,\,\,\, X$
$Attention\,\, Score\,\,\,\,\, A\,\,:=\,(X+P)W_{qry}W^T_{key}(X+P)^T$

$Self\,Attention(X)_{t,:}\,:=\,softmax(A_{t,:})XW_{val}$

Attention Score를 계산할 때 인풋에 포지션 인코딩을 더한 값을 넣으면 위와 같은 문제를 해결할 수 있다.

Multi-head self attention

$MHSA(X) := \underset{h \in [N_h]}{concat}[\,Self-Attention_h(X)\,]W_{out}+b_{out}$

• 여러 개의 head를 갖는 self-attention으로, 각각의 Query, Key, Value의 weight가 다르다
• 각 head에 따라 계산한 self-attention 결과 행렬을 옆으로 concat하고, $W_{out}$ 행렬을 곱해 Head의 개수 x $D_{in}$ 차원을 $D_{out}$ 차원으로 project한다

Convolutional layer

   Conv 필터 수식

                                  $Conv(X)_{i, j} := \sum_{{(\delta_1, \delta_2)\in \bigtriangleup_K}} X_{i+\delta_1, j+\delta_2, :} W_{\delta_1, \delta_2, :, :} + b$

이미지의 $i$, $j$번째 픽셀을 기준으로 $(\delta_1, \delta_2)$만큼 이동한 픽셀 값과 $(\delta_1, \delta_2)$에 해당하는 필터 값을 곱하는 과정을 모든 shift에 대해 반복하고 총합해서 마지막에 bias를 더한다.

Absolute positional encoding

$A^{abs}_{q, k} = (X_{q,:}+P_{q,:})W_{qry}W^T_{key}(X_{k,:}+P_{k,:})^T$

$\,\,\,\,\,\,\,\,\,\,\,\,=X_{q,:}W_{qry}W^T_{key}X^T_{k,:} + X_{q,:}W_{qry}W^T_{key}P^T_{k,:}+P_{q,:}W_{qry}W^T_{key}X^T_{k,:}+P_{q,:}W_{qry}W^T_{key}P^T_{k,:}$

이 방식은 학습과 평가에서 사용할 수 있는 문장의 최대 길이가 정해져 있다는 문제가 있다.

Relative positional encoding

일련의 과정을 지나 최종 수식을 도출해낸다.

(1) 토큰 사이의 거리 $k-q$에 대한 수식으로 바꾼다

$A^{rel}_{q, k} = (X_{q,:}+\textcolor{red}{\sout{P_{q,:}}})^TW_{qry}^TW_{key}(X_{k,:}+P_{\textcolor{red}{k-q}})$

$\,\,\,\,\,\,\,\,\,\,\,\,=X_{q,:}^TW_{qry}^TW_{key}X_{k,:} + X_{q,:}^TW_{qry}^TW_{key}P_{k,:}+\textcolor{red}{\sout{P_{q,:}^TW_{qry}^T}}W_{key}X_{k,:}+\textcolor{red}{\sout{P_{q,:}^TW_{qry}^T}}W_{key}P_{k,:}$

(2) $P_{q,:}^TW_{qry}^T$를 학습 가능한 파라미터 $u, v$로 대체한다

$A^{rel}_{q, k} = X_{q,:}^TW_{qry}^TW_{key}(X_{k,:}+P_{\textcolor{red}{k-q}})$

$\,\,\,\,\,\,\,\,\,\,\,\,=X_{q,:}^TW_{qry}^TW_{key}X_{k,:} + X_{q,:}^TW_{qry}^TW_{key}P_{\textcolor{red}{k-q}}+\textcolor{red}{u^T}W_{key}X_{k,:}+\textcolor{red}{v^T}W_{key}P_{\textcolor{red}{k-q}}$

(3) $k-q$를 간단하게 $\delta$로 치환하고, $P_{k-q}$를 $r_{\delta}$로 표기한다.

$A^{rel}_{q, k} = X_{q,:}^TW_{qry}^TW_{key}(X_{k,:}+\textcolor{red}{r_{\delta}})$

$\,\,\,\,\,\,\,\,\,\,\,\,=X_{q,:}^TW_{qry}^TW_{key}X_{k,:} + X_{q,:}^TW_{qry}^TW_{key}\textcolor{red}{r_{\delta}}+\textcolor{red}{u^T}W_{key}X_{k,:}+\textcolor{red}{v^T}W_{key}\textcolor{red}{r_{\delta}}$

(4) 토큰 임베딩을 이용하는 $W_{key}$와 토큰 사이의 거리를 이용하는 $\widehat{W}_{key}$를 구분한다

$A^{rel}_{q, k} = X_{q,:}^TW_{qry}^TW_{key}(X_{k,:}+\textcolor{red}{r_{\delta}})$

$\,\,\,\,\,\,\,\,\,\,\,\,=X_{q,:}^TW_{qry}^TW_{key}X_{k,:} + X_{q,:}^TW_{qry}^T\textcolor{blue}{\widehat{W}_{key}}\textcolor{red}{r_{\delta}}+\textcolor{red}{u^T}W_{key}X_{k,:}+\textcolor{red}{v^T}\textcolor{blue}{\widehat{W}_{key}}\textcolor{red}{r_{\delta}}$

3. Implementation

Theorem 1

$N_h$개의 Head를 갖는 multi-head self attention layer($N_hD_h\,\times\,D_{out}$)와 $D_p\,\geq\,3$ 이상의 차원을 갖는 relative positional encoding은 $\sqrt{N_h}\,\times\,\sqrt{N_h}$ kernel과 $min(D_h, D_{out})$ 차원의 output channel을 갖는 어떤 Convolutional layer도 표현할 수 있다

Theorem 1의 보조 정리 Lemma 1은 다음과 같다.

Lemma 1

$N_h=K^2$개의 $head$를 갖는 $Multi-head self-attention$의 각각의 $head$가 $Conv$ 필터의 $shift$와 일대일 대응을 이루고 모든 $Head$가
$softmax(A_{q,:}^{(h)})_k=\left\{\begin{matrix} 1\,\,if\,\,f(h)=q-k\\ 0\,\,otherwise\,\,\,\,\,\,\,\,\,\,\,\,\,\, \end{matrix}\right.$를 만족하면,

어떤 $K\times K\,\,kernel$을 갖는 Conv Layer에 대해서도 $MHSA(X)=Conv(X)$를 만족하는 $W_{val}^{(h)}$가 존재한다

Lemma 1의 중요한 가정은 Conv 필터의 shift와 Head의 Query, Key 사이의 거리가 같을 때만 스코어가 1이고 나머지는 0이라는 것이다.

실제로 이것이 가능함은 Lemma 2에서 증명하기 때문에 우선은 그렇다고 가정하고 Lemma 1을 증명하자.

Proof

$Self-Attention(X)_{t,:}\,\,:=\,\,softmax(A_{t,:})XW_{val}$

$MHSA(X)\,\,:=\,\,concat[Self-Attention_h(X)]W_{out}\,+\,b_{out}$
$\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\scriptsize{h\in[N_h]}$

위의 두 식을 조합하면 $MHSA(X)$를 아래와 같이 한번에 표현할 수 있다.

$MHSA(X)\,\,:=\,\,b_{out}\,+\,\sum_{h\in[N_h]}\,\,softmax(A^{(h)})XW_{val}^{(h)}W_{out}[(h-1)D_h\,+\,1\,:\,hD_h\,+\,1]$

각 $h$에 따른 $W_{val}^{(h)}W_{out}[(h-1)D_h\,+\,1\,:\,hD_h\,+\,1]$들을 한 번에 묶어서 $W^{(h)}$로 만들어 다음과 같이 식을 간소화한다.

$MHSA(X)\,\,:=\,\,b_{out}\,+\,\sum_{h\in[N_h]}\,\,softmax(A^{(h)})XW^{(h)}$

지금은 인풋 $X$의 모든 토큰에 대해 병렬적으로 한번에 행렬 연산을 진행하는 식인데,

각각의 토큰 $q$에 대해 연산하는 식으로 바꾸면 다음과 같다.

$MHSA(X)_q\,\,:=\,\,\sum_{h\in[N_h]}\,(\sum_k\,softmax(A_{q,:}^{(h)})_kX_{k,:})W^{(h)}\,+\,b_{out}$

위 식에 Lemma 1을 반영하면 모든 $k$ 중에서 $k=q-f(h)$일 때만 살아남게 된다.

$MHSA(X)_q\,:=\,\sum_{h\in[N_h]}\,X_{q-f(h),:}W^{(h)}\,+\,b_{out}$

이 최종 결과를 다음과 같이 해석해보면,

• $h\in[N_h]$ : 필터의 원소
• $X_{q-f(h)}$ : 이미지의 픽셀 값
• $W^{(h)}$ : 필터의 가중치

$Convolution \,\,layer$의 $q$번째 출력값의 계산식과 동일하다는 것을 알 수 있다.

이로써, 특정 상황에서 $multi-head\,\, self\,\, attention$이 마치 $Convolution$ 처럼 작동한다는 $Lemma \,\,1$을 증명했다.

Lemma 2

모든 $\bigtriangleup \in \bigtriangleup_K$에 대해 $k-q=\bigtriangleup$일 때만 $softmax(A_{q,:}^{(h)})_k=1$로 만드는 어떤 vector $v$가 존재하도록 하는

Relative positional encoding $r_\delta(D_p\geq 3), W_{qry}, W_{key}, \widehat{W}_{key}, u$가 존재한다.

Proof

$A^{rel}_{q,k}=X^T_{q,:}W^T_{qry}W_{key}X_{k,:}+X^T_{q,:}W^T_{qry}\widehat{W}_{key}r_\delta+u^TW_{key}X_{k,:}+v^T\widehat{W}_{key}r_\delta$

$W_{qry}=W_{key}=0,\,\,\widehat{W}_{key}=I\,\,(identity\,\,matrix)$ 로 설정한다

그러면, $A^{rel}_{q,k}=v^Tr_\delta\,\,where\,\,\delta:=k-1$로 간소화된다

$A^{rel}_{q,k}=-\alpha(||\delta-\bigtriangleup||^2+c)$ 의 꼴을 갖는다고 가정하자  ($\bigtriangleup$ : $Conv$ 필터의 $shift$)

그러면 softmax attention 식은 다음과 같다.
$softmax(A_{q,:})_k={e^{-\alpha(||\delta-\bigtriangleup||^2+c)}\over \sum_{k'}e^{-\alpha(||(k'-q)-\bigtriangleup||^2+c)}}$

$\alpha$를무한대로보내면 ,
$\underset{\alpha \to \infty}{\lim}\,softmax(A_{q,:})_k=\underset{\alpha \to \infty}{\lim}\,{e^{-\alpha(||\delta-\bigtriangleup||^2+c)}\over \sum_{k'}e^{-\alpha(||(k'-q)-\bigtriangleup||^2+c)}}$

이 경우 분모와 분자에 모두 있는 상수 $c$가 의미 없어지므로,
$\underset{\alpha \to \infty}{\lim}\,softmax(A_{q,:})_k=\underset{\alpha \to \infty}{\lim}\,{e^{-\alpha(||\delta-\bigtriangleup||^2)}\over \sum_{k'}e^{-\alpha(||(k'-q)-\bigtriangleup||^2)}}$

$\delta$가 $\bigtriangleup$가 같을 때와 같지 않을 때를 계산하면 다음과 같다.

$For\,\,\delta=\bigtriangleup,$

$\underset{\alpha \to \infty}{\lim}\,softmax(A_{q,:})_k=\underset{\alpha \to \infty}{\lim}\,{1\over 1+\sum_{k'\neq k}e^{-\alpha(||(k'-q)-\bigtriangleup||^2)}}$=1

분자는 $e^{-\alpha\times 0}$이므로 1이 되고,
분모는 $k=k'$일 때만 1이고 나머지는 $e^{-\inf\times \cdot\cdot\cdot}=0$이므로 합하면 1이 된다.
따라서 결과값은 ${1 \over 1}=1$이다.

$For\,\,\delta\neq\bigtriangleup,$

$\underset{\alpha \to \infty}{\lim}\,softmax(A_{q,:})_k=\underset{\alpha \to \infty}{\lim}\,{e^{-\alpha(||\delta-\bigtriangleup||^2)}\over \sum_{k'}e^{-\alpha(||(k'-q)-\bigtriangleup||^2)}}$=0

분자는 $e^{-\inf\times \cdot\cdot\cdot}$이므로 0이 되고,
분모는 $k=k'$일 때만 1이고 나머지는 $e^{-\inf\times \cdot\cdot\cdot}=0$이므로 합하면 1이 된다.
따라서 결과값은 ${0 \over 1}=0$이다.

이제, attention score의 값이 $\alpha(||\delta-\bigtriangleup||^2+c)$의 형태를 갖도록 해주는 파라미터 벡터 $v$와 positional encoding $r_\delta$가 존재하는지만 증명하면 $Lemma\,\,2$가 입증된다.

$v\,=\,-\alpha(1, -2\bigtriangleup_1, -2\bigtriangleup_2),$
$r_\delta\,=\,(||\delta||^2,\delta_1,\delta_2)$ 로 설정해보자

그러면,
$A_{q,k}=v^Tr_\delta$
$\,\,\,\,\,\,\,\,\,\,\,\,=-\alpha(||\delta||^2-2\bigtriangleup_1\delta_1-2\bigtriangleup_2\delta_2)$
$\,\,\,\,\,\,\,\,\,\,\,\,=-\alpha(||\delta||^2-2<\delta, \bigtriangleup>)$
$\,\,\,\,\,\,\,\,\,\,\,\,=-\alpha(||\delta-\bigtriangleup||^2-||\bigtriangleup||^2)$
$\,\,\,\,\,\,\,\,\,\,\,\,=-\alpha(||\delta||^2+c)\,\,when\,\,c=-||\bigtriangleup||^2$