[Review] On the Relationship between Self-Attention and Convolution layers

Hwan Heo·2021년 12월 8일

DeepLearning_Paper_Review

목록 보기

5/5

Transformer 의 성공으로 NLP 및 다른 분야에서 점차 그 영향력을 넓혀가고 있다. 이 논문은 Vision Task 에서의 multi-heads self-attention 이 임의의 convolution 을 잘 모사할 수 있기 때문이라고 주장하고, 이에 대한 수학적인 증명을 제시한다.
ICLR2020

1. Introduction

ViT, Stand-Alone Self-Attention in Vision 등으로 인해 이미 Vision 에서도 Self-Attention 이 convolution 을 훌륭하게 대체할 수 있는 것으로 밝혀졌다.

Vision 에서 convolution 과 self-attention 의 가장 주효한 차이점은 receptive filed 가 local region 에 국한되는지 (Convolution), global 인지 (Self-Attention) 에 따른 것이다.

저자들은 적절한 수의 multi-head 를 통해서 multi-perspective self-attention 을 구현하는 것이, convolution kernel (shared)에서 grid 위치마다 다르게 weight 를 축적하는 것을 근사할 수 있다고 보았다.

또한 self-attention 의 translation equivariance 특성을 보정하기 위하야 positional encoding 을 더해주는 작업을 통해서 이러한 MSHA (Multi-Head Self-Attention) 가 convolution 의 translation invarient 특성을 근사할 수 있는 주요한 attribute 라고 주장한다.

2. Method

2.1. Convolution

Convolution 은 shared parameter 를 통해서 image (or feature) 를 weighted aggregation 하여 정보를 축약하는 feature map 을 쌓아 model 을 학습시킨다. 이는 다음과 같은 특징을 가지는데,

Weight (parameter) 가 shared 되기 때문에 translation invarient 한 특징을 가지고 있으며,
downsampling 과 conv layer 의 적축을 통해 large receptive filed 로 정보를 확대할 수 있기 때문에 global dependency 에도 어느 정도 강건하게 학습할 수 있다.

식으로 표현하면 다음과 같다. (bias term 은 생략)

\text{Conv(X)} = \sum _{\delta \in \triangle} \mathbb X_{i+\delta_1 \ j+\delta_2, \ :} \mathbb W_{\delta_1 , \ \delta_2 , \ :}

2.2. Multi-Head Self-Attention (MSHA)

'Attention is all you need' 에서 주장했듯이, self-attention 에서의 multi-head 의 역할은 각각 다른 linear transformation (hyperspace projection) 을 통해 multi-perspective attention 을 구현하여 다양한 관점에서의 feature 정보를 통합하는 것이다.

위 그림에서 이와같은 MHSA 의 역할이 잘 나타나 있는데, 각 head 에 따라서 attention map 이 주로 참조하는 (attention score 가 높은) region 이 달라지고, 이를 전체적으로 병합하여 self-attention 의 결과로써 사용하게 된다.

2.2.1. Relation Between Convolution and Self-Attention

이제 image 에서의 MSHA 를 다음과 같이 기술하자.

\text{MHSA} (X) \ = \ \sum _{h \in N_h} \text{softmax} (A^h ) X W_v^h W_o [ (h-1)D_h +1 : h D_h +1 ]

수식을 살펴보면 attention score map 으로의 projection 이후 multi-head concat 이 이루어짐을 알 수 있다. 이때 다음의 Lemma1 이 성립한다.

💡 Lemma1. Consider a multi-head self-attention layer consisting of $N_h = K^2$ heads, $D_h \ge D_o$ and let $f : N_h \rightarrow \triangle_K$ be a bijective mapping of heads onto shifts. Further, suppose that for every head the following holds:
$\text{softmax} (A^h_q)_k = \begin{cases} 1 \quad f(h) = >q- k \\ 0 \quad o.w \end{cases}$
Then, for any convolution layer with $K \times K$ kernel, there exists $W_v ^h$ such that $\text{ MHSA} (x) = \text{Conv} (x)$ for every X.

이는 self-attention map 이 query 와 key 의 relative position 에 따른 attention map 으로 구성된다면, 이를 convolution 으로 해석할 수 있다는 의미이다.

예를 들어 3x3 Convolution 의 경우는, 다음 그림과 같이 9개의 heads 가 각각 9개의 grid 영역에만 집중해서 학습 된다면, 각 grid 위치에 따른 linear transformation ( $W_k, \ W_q )$ 가 convolution kernel value 처럼 행동하게 될 것이다.

즉 MSHA 을 다음과 같이 기술할 수 있고, 이는 즉 convolution 의 정의와 같음을 알 수 있다. (bias term 은 모두 생략)

\text{MHSA} (X)_q= \sum _h X_{q-f(h), :} W^h

우리는 이미 Self-attention 이 global dependency 를 훌륭하게 학습할 수 있음을 알고 있기 때문에, 위와 같은 이론적 배경을 바탕으로 self-attention 이 shallow layer 에서도 convolution 의 고품질의 feature extracting 을 모사할 수 있음을 알 수 있다.

또한 fully self-attention model 이

Shallow layer 에서는 convolution 처럼 (local attention),
Deeper layer 에서는 local feature 에 대한 global dependency 를 평가하는 'attention' 처럼

행동하여 CNN architecture 를 대체할 수 있게 된다.

2.4. Quadratic Encoding

Self-attention 이 convolution 의 translation invarient 를 효과적으로 모사할 수 있는 relative encoding 의 존재와 그 제약조건들에 대해서 서술하고 있다.

다음과 같이 Self-attention with positional encoding 을 기술할 때,

\begin{aligned} A &= (X_q + P_q ) W_q W_k ^T (X_k + P )^T \\ & = X_q W_q W_k ^T X_k^T + X_q W_q W_k ^T P_k^T + P_q W_q W_k ^T X_k + P_q W_q W_k ^T P_k \end{aligned}

Absolute position encoding 이 아닌 relative position encoding 을 사용하면, 해당 self-attention 수식에서 relative position information 을 가지고 있는 4번째 항에만 dependent 하게됨을 알 수 있다.

이로부터 key 와 query 를 생성하는 weight parameter 를 0으로 간주하거나, softmax 의 temperature 를 contrastive learning 처럼 임의로 낮게 조정하여 그 밀도를 높이는 등의 강대한 제약조건을 도입는데, 이는 보통의 다른 Self-attention 구조에서 차용하는 방식이 아니다.

또한 제시하는 quadratic encoding 외의 다른 positional encoding 방식으로도 확장할 수 있다고 주장하나 리뷰어들에 의해서 그 부분에 대한 강건성을 공격받았기 때문에 이를 비판적으로 수용할 필요가 있다고 생각한다.

3. Experiments

실험은 large dataset 에서 적용되지 않았는데, 이또한 선행연구인 stand-alone self-attention 이 MS-COCO 등에서의 실험 성능을 report 한것에 비하면 아쉽다.

Visualization 을 통해 실제로 shallow layer 에서의 self attention 은 convolution 처럼 local grid 에 더 집중하는 attention score 를, deeper layer 에는 그보다는 더 넓은 영역에서의 attention map 을 구성하는 모양새를 보이고 있다.

하지만 layer1, 2의 9th attention map 은 흡사 global image 를 모두 참조하는 것으로 보이는데, 이는 projection 에 따른 attention perspective pruning 이 redundant 없이 optimal 하게 학습되는 것이 매우 어렵기 때문이다.

이런 문제점 때문에 적절한 multi-head attention perspective pruning 을 보다 효과적으로 진행하는 방법에 대한 연구도 진행되고 있다.

Hwan Heo

기타치는AI Researcher

이전 포스트