뒤로가기 (back)
2runo's blog

논문 리뷰 - Self-Attention with Relative Position Representations

Relative Positional Embedding 논문을 리뷰합니다.

논문 정리

methodology

기존에 사용되던 positional encoding 방식은 크게 두 가지였다. 이들은 모두 특정 값을 input 벡터에 더하여 위치 정보를 표현한다.

그러나 이 방법은 모두 큰 문제를 가지고 있었다. 바로 ‘새로운 길이의 입력에 제한적이라는 점’이다. learned positional embedding을 사용한 모델은 입력 sequence의 최대 길이가 고정되어 있다. 이는 학습된 위치 벡터가 특정 길이에 맞춰져 있기 때문이다. 반면 sinusoidal positional encoding의 경우 이론적으로 토큰 길이의 한계는 없다. 그러나 학습 데이터보다 긴 sequence의 경우 모델 입장에서 처음 보는 positional representation 값이 사용되므로 성능이 저하될 수밖에 없다.

따라서 저자들은 효율적이고, 길이의 제약이 없는 relative positional representation 기법을 제안한다. 이는 상대적인 거리 정보를 attention 계산 시에 주입하는 방식으로 작동한다. 설명에 앞서, vanilla transformer의 attention 계산 방법을 살펴보자:

$$ z_i = \sum_{j=1}^n{\alpha}_{ij}(x_jW^V) $$

$$ \alpha_{ij} = \frac{\text{exp}\ e_{ij}}{\sum_{k=1}^{n}{\text{exp}\ e_{ik}}} $$

$$ e_{ij} = \frac{(x_i W^Q)(x_j W^K)^T}{\sqrt{d_z}} $$

본 논문에서는 위 수식을 아래처럼 바꾼다.

$$ z_i = \sum_{j=1}^n{\alpha}_{ij}(x_jW^V + a^V_{ij}) $$

$$ e_{ij} = \frac{x_i W^Q(x_j W^K + a^K_{ij})^T}{\sqrt{d_z}} $$

\(a^V_{ij}\)와 \(a^K_{ij}\)는 \(i\)번째 토큰과 \(j\)번째 토큰 간의 edge(거리)를 표현하는 벡터이다. 논문에서는 일반화를 위해 edge라는 표현을 사용하지만 ‘거리’로 이해하면 된다. 이 벡터를 attention 연산 시에 더해주어 상대적인 위치 정보를 주입한다.

$$ a^K_{ij} = w^K_{clip(j-i,k)} \\ a^V_{ij} = w^V_{clip(j-i,k)} $$

\(a^V_{ij}\)와 \(a^K_{ij}\)의 값은 key, query 토큰 간의 상대적인 위치에 따라 결정된다. 이때의 \(w^{K,V}\)는 trainable parameters이다.

하지만 \(e_{ij}\)를 계산하는 수식은 효율적이지 않다. vanilla transformer에서는 query matrix(\(xW^Q\))와 key matrix(\(xW^K\))를 미리 구한 뒤 행렬곱을 하면 한번에 모든 \(i,j\)에 대한 \(e_{ij}\) 값을 구할 수 있었다. 그러나 \(a^K_{ij}\)를 더하는 항이 추가되면서 key matrix(\(xW^K + a^K_{ij}\))를 미리 구할 수 없게 되었다. 때문에 일일이 for 문을 돌리며 key matrix(\(xW^K + a^K_{ij}\))를 구하고 query matrix와 행렬곱하는 과정을 n_tokens번 반복해야 한다.

이 문제를 해결하기 위해 저자들은 아래처럼 수식을 변형했다.

$$ e_{ij} = \frac{x_i W^Q(x_j W^K)^T + x_iW^Q(a^K_{ij})^T}{\sqrt{d_z}} $$

수학적으로 같은 식이지만, 이제 연산을 좀 더 효율적으로 할 수 있게 되었다. 왼쪽 항은 vanilla transformer와 동일한 attention score 계산 수식이다. query&key matrix를 미리 구하고 한번의 행렬곱으로 모든 \(i, j\)에 대한 값을 구할 수 있다. 오른쪽 항도 edge matrix(\(\begin{bmatrix} a^K_{11} \cdots a^K_{1j} \\ \vdots\ \ \ddots\ \vdots \\ a^K_{i1} \cdots a^K_{ij} \end{bmatrix}\))를 미리 구하고 query matrix와 한 번만 행렬곱하면 모든 \(i, j\)에 대한 값을 구할 수 있다. 다시 말해 query, key, edge matrix를 미리 구하고 두 번의 행렬곱으로 모든 \(e_{ij}\)를 구할 수 있는 것이다.

또한 위 수식은 제안 기법의 작동 방식을 해석하는 데도 유용하다. 왼쪽 항은 vanilla transformer와 동일한 attention score 계산 식이라고 언급했었다. 그렇다면 오른쪽 항은 상대적인 위치 정보를 바탕으로 attention score 값에 관여하는 term이라고 볼 수 있다. 즉, vanilla transformer의 수식대로 attention score를 계산하고 여기에 상대적인 위치에 따른 bias를 더해주는 방식으로 작동한다.

직관적으로 생각해보면, 멀리 떨어진 토큰들 사이에서는 bias가 작게 설정되어 attention score가 낮아지는 것이 자연스러워 보인다. 가까이 있는 토큰들은 의미적으로 긴밀하게 연결되어 있지만, 거리가 멀어질수록 이러한 연결성은 점차 희미해지기 때문이다. 해당 내용은 아래 ‘가설 검증’ 단계에서 살펴볼 것이다.

experiments

Figure 1: 실험 결과. 논문 발췌

실험 결과는 제안 기법을 적용한 것이 기존 sinusoidal positional encoding 방식보다 성능이 좋았다. 다만 learned positional embedding과 비교한 실험은 없어서 아쉬웠다. 또한 다른 여러 태스크로도 성능을 평가해봤으면 좋았을 것 같다.

제안 기법은 상대적인 거리 정보를 이용하기 때문에 입력 sequence 길이 제한이 없다. 이 점이 제안 기법의 가장 큰 장점이라고 생각한다.

탐구하기

여기부터는 오로지 내 생각으로 쓴 글이다. 논문을 읽으며 들었던 생각을 정리하고 확인해본다.

멀리 떨어진 토큰일수록 attention score가 작아질까?

위에서 언급했듯 멀리 떨어진 토큰들 사이에서는 bias가 작게 설정되어 attention score가 낮아지는 것이 자연스러워 보인다. 글에서 단어는 서로 멀리 떨어져 있을수록 대개 의미 연관성이 적어지기 때문이다. 또한 단어를 tokenizing하여 모델 입력에 넣기에 하나의 단어가 여러 토큰으로 잘리는 경우가 있다. 때문에 이웃하는 토큰끼리의 연관성은 매우 커진다.

실제 attention score가 작아지는지 확인해보기 위해 본 논문의 제안 기법을 적용한 모델을 살펴봤다. 허깅페이스에 공개된 twhin-bert-base라는 모델이다.

Figure 2: 상대적 거리에 따른 attention score bias 시각화

Figure 2는 상대적인 위치에 따른 attention score bias를 시각화한 것이다. 이때 attention score bias는 \(e_{ij}\)를 구하는 수식의 오른쪽 term이다. 이 값이 크다면 그만큼 attention score가 커진다. 예를 들어 Figure 2-(a)는 모델의 첫 번째 레이어에서, ‘quick’이라는 토큰이 query일 때 임의의 key 토큰과의 상대적 위치에 따른 attention score bias를 나타낸다. x축에서 빨간 선이 위치한 0은 자기 자신 토큰을, +2는 key 토큰이 ‘quick’ 토큰보다 2토큰 오른쪽에 있을 경우를 나타낸다.

Figure 2를 봤을 때, 본 가설은 틀렸음을 알 수 있다. Figure 2-(a)처럼 멀어질수록 attention score가 줄어드는 경향을 보이기도 하지만, (b)나 (d)처럼 다른 variation을 보이는 경우도 많기 때문이다.

특이한 것은 바로 왼쪽의 토큰(x=-1)에 대해서 bias 값이 갑자기 작아진다는 것이다. 바로 왼쪽의 토큰에 대해 이런 현상이 나타나는 이유는 모르지만, absolute한 위치 정보를 얻기 위한 수단이거나 문맥을 이해할 때 바로 왼쪽 토큰이 특별한 역할을 한다는 것 정도가 당장 떠오르는 가설이다. 이는 입력 sequence를 reverse한 채로 모델을 pre-training해보면 가설을 확인해볼 수 있을 것 같다. 만약 정말 문맥 이해를 위한 것이라면 reverse 모델은 오른쪽(x=+1) bias 값이 갑자기 작아질 것이기 때문이다.

왜 k=2 이후로는 성능 개선이 없는가?

Figure 3: k에 따른 성능. 논문 발췌

여기서 k는 모델이 표현할 상대적인 위치의 최대 길이를 의미한다. 만약 k=16이라면 16토큰을 넘어가는 거리일 경우 거리가 16일 경우의 위치 벡터를 사용하여 퉁 친다는 뜻이다. 상대적 거리에 따라 무한하게 위치 벡터를 사용할 수는 없기 때문이다.

실험에 따르면 k=2 이후로는 k가 증가해도 성능에 큰 차이가 없었다. 그 이유가 무엇일까?

Figure 4: 위치 0 벡터와 각 위치 벡터 간의 코사인 유사도

Figure 4는 위치 0(토큰 자기 자신)일 때의 위치벡터와, 다른 위치 벡터 간의 코사인 유사도를 시각화한 것이다. x축은 상대적인 위치를 나타낸다. 예를 들어 x=2는 기준 토큰에서 오른쪽으로 2토큰 만큼 떨어졌을 때의 위치 벡터를 의미한다.

그래프에서 코사인 유사도는 x=0을 기준으로 -2 ≤ x ≤ 2까지만 값의 변화가 크다. 그 외의 벡터는 값의 변화가 미미해진다. 이것이 k=2 이후로는 성능에 차이가 없는 이유다. 상대적인 거리가 2를 넘어갔을 때의 위치벡터는 값의 차이가 미미해진다. 따라서 이들을 모두 하나의 벡터(상대적 거리가 2일 때의 위치벡터)로 퉁쳐도 괜찮은 것이다.

방향이 달라지면 위치 벡터에 마이너스(-)를 곱하여 사용하면 안될까?

제안 논문에서는 왼쪽과 오른쪽 방향을 구분하여 위치 벡터를 마련했다. 예를 들어 \(w^V_{2}\)와 \(w^V_{-2}\)는 구분되는 서로 다른 벡터이다. 그런데 직관적으로 생각했을 때 정확히 반대되는 위치 벡터(ex. \(w^V_{2}\)와 \(w^V_{-2}\))는 마이너스(-) 관계로 표현할 수 있지 않을까? 다시 말해, 상대 위치가 오른쪽인 위치 벡터(ex. \(w^V_{2}\))만 남기고, 왼쪽인 위치 벡터(ex. \(w^V_{-2}\))는 학습하지 않는다. 대신 왼쪽인 위치 벡터가 필요하면 대응되는 오른쪽 위치 벡터에 마이너스를 곱하여 사용한다.

$$ a^K_{ij} = \begin{cases}w^K_{clip(j-i,k)} \ \ \ \ \ \text{if} \ \ j-i \ge 0 \\-w^K_{clip(i-j,k)} \ \ \text{otherwise} \end{cases} $$

그러나 학습된 모델을 조금만 살펴보면 해당 아이디어는 현실성이 없음을 알 수 있다. 음수를 곱하면 내적의 결과는 정 반대가 되기 때문에 attention score bias 그래프는 점 대칭에 가까워져야 한다. 그러나 Figure 2의 그래프는 선 대칭에 가깝다면 가까웠지, 전혀 점 대칭이 아니다.