Attention 구조의 논리에서 출발하는 Positional Embedding 종류 정리

Transformer 구조에서 등장한 Positional Embedding은 토큰의 순서 정보를 어떻게 모델에 전달할 것인가로부터 시작된 방법론이다.

Transformer의 연산 방식의 핵심은 “Self-Attention”인데, 이 연산 구조 특성상 연산에 포함되는 컨텍스트 (이 글에서 일괄 계산에 포함되는 토큰 정보/임베딩 정보를 컨텍스트라고 통일하겠다)의 요소만 반영하고 각자의 위치를 반영하지 못하는 것이다.

예컨대 다음 두 문장이 있다고 가정했을때, 두 문장을 구분하지 못한다:

1. I love AI studying.
2. AI love I studying.

1, 2번 문장은 엄연히 다른 의미를 가진다. 심지어 2번 문장의 경우 문법적으로 틀렸다. 그럼에도 불구하고 두 문장에 포함된 토큰 집합이 동일하다면 Attention 연산은 같은 입력으로 인식하게 된다.

Transformer Encoder의 구조를 가장 간단하게 표현하면 다음과 같다: 입력이 $X$라면 $Q, K, V$를 추출하고, $Attention$ 연산을 수행한 뒤 $MLP$를 거친다.

예컨대 입력이 여러 토큰(e.g. 단어)로 구성된 일련의 시퀀스라고 가정하고, 다음과 같이 표현된다고 생각해보자.

$$X = [x_1, x_2, \cdots, x_n]$$

해당 시퀀스에 포함된 토큰에 대해 어텐션 연산을 수행하기 위해 Query, Key, Value로 변환한다.

$$Q = XW_Q,~K = XW_K,~V=XW_V$$

$Attention$은 Query($Q$)와 Key($K$) 사이의 내적에 대한 상대적 확률값에, 각 토큰별 Value를 가중합하는 논리가 된다. 이때, 임베딩 dimension에 따라 score값 자체가 스케일링되는 문제가 생길 수 있기 때문에 (기본 내적은 vector norm 크기를 고려하지 않기 때문), 많이 사용되는 scaled dot-product formulation은 다음 형태를 가지게 된다.

$$Attention (Q, K, V) = softmax\left(\frac{QK^\top}{\sqrt{d}} \right)V$$

$i$번째 토큰과 $j$번째 토큰 사이의 attention score는 $Q_i K_j^\top$이다. 이 연산에는 토큰 임베딩의 내적만 들어가고 토큰의 위치는 들어가지 않는다.

결국 앞서 봤던 예시를 다시 가져와보면

1. I love AI studying. $\rightarrow$ [I, love, AI, studying]
2. AI love I studying. $\rightarrow$ [AI, love, I, studying]

1번 문장의 I, AI는 2번 문장의 I, AI와 위치가 달라지지만, 동일한 임베딩으로 나오므로 attention score는 같은 값을 내보내게 된다. 즉, 위치에 따른 구분이 불가능하다.

수식으로 다시보는 Permutation Invariant

예컨대, Permutation matrix $P$를 생각해보자. Permutation matrix란 앞에 어떤 요소를 곱하면 해당 요소의 순서를 바꾸는 연산을 수행한다. 보통 permutation matrix는 Identity Matrix의 요소를 행별로 랜덤 배치해서 얻어낼 수 있는데:

$$X = [1, 2, 3, 4],~P = \begin{bmatrix} 0 & 1& 0 & 0 \newline 1 & 0 & 0& 0 \newline 0 & 0 & 0 & 1 \newline 0 & 0 & 1 & 0 \end{bmatrix}$$

예를 들어 위와 같이 구성된 $X$라는 벡터와 $P$라는 permutation matrix가 있다고 가정하면, 두 행렬 연산을 통해 우리는 permutation된 $X^\prime$를 획득할 수 있다.

$$X^\prime = XP = [2, 1, 4, 3]$$

다시 원래 목적으로 돌아와서, $X$가 길이 $n$짜리의 토큰 시퀀스라고 가정해보자. 그렇게 되면 $n\times n$ 크기를 가지는 permutation $P$에 대해서 랜덤하게 요소별로 위치를 바꾼 $X^\prime$을 얻을 수 있을 것이다. 그러면 원래 시퀀스가 아닌 랜덤하게 위치가 바뀐 시퀀스에 대해서 Query, Key, Value를 구하면,

$$\begin{aligned} Q^\prime =& P(XW_Q) = PQ \newline K^\prime =& P(XW_K) = PK \newline V^\prime =& P(XW_V) = PV \end{aligned}$$

위와 같이 원래의 Query, Key, Value에 permutation matrix를 곱한 형태가 된다. 이 상태에서 Attention 연산을 동일하게 적용한다면,

$$Attention (PQ, PK, PV) = softmax\left(\frac{PQ(PK)^\top}{\sqrt{d}} \right)PV$$

Permutation 연산은 결괏값의 크기에는 무관하고 위치만 바꾸기 때문에 softmax 연산과는 무관하다. 그리고 permutation matrix $P$는 orthogonal matrix이므로 $P^\top = P^{-1}$이고, 결국 $PP^\top = P^\top P = I$라는 결론에 다다른다. 해당 정보를 가지고 위의 식을 단순화하면 결국 다음과 같이 정리된다:

$$Attention (PQ, PK, PV) = P \cdot Attention(Q,K,V)$$

‘토큰’의 위치를 바꾸면 Attention의 위치도 함께 바뀐다. 입력의 위치 변화에 맞춰 출력의 위치도 바뀌는 연산을 “Permutation Equivariant”하다고 한다.

여기서 Equivariant하다는 것은 입력의 변화가 출력에 ‘그대로’ 반영된다는 것인데, 이를 Transformer 모델 출력 전체를 기준으로 말해보면, 어텐션 연산의 변화가 토큰의 위치에 따라 구분력을 가지지 못하기 때문에 결국 동일한 결괏값을 내게 되어 Invariant하다는 특징을 가진다.

조금 복잡하기 때문에 다시 풀어서 설명하면, 연산 자체는 Equivariant하지만 모델 출력 관점에서는 의미적으로 Invariant하므로, Position에 대한 정보가 없는 어텐션 기반의 Transformer는 “Permutation Invariant하다”라고 볼 수 있는 것이다.

그래서, 우리에겐 위치 정보가 필요하다.

Attention은 결국 토큰 간의 유사도 / 연관도만 계산한다. 즉, 토큰 간의 위치가 달라져도 어텐션 값에는 아무런 변화가 없다. 하지만 언어는 순서가 문법적으로도, 의미적으로도 매우 중요하다.

1. Dog bites man
2. Man bites dog

1번은 [bite]의 주체가 [Dog]이며 [Man]이 그 대상이지만, 2번은 주체와 대상이 서로 반대가 된다. 그렇기 때문에 Transformer는 순서를 알려주는 함수를 추가해야 한다.

1. Absolute Positional Embedding

Absolute PE(Positional Embedding)은 입력에 절대적 위치에 따른 유니크한 값을 더해준다.

$$X + APE = [x_1+p_1,x_2+p_2,\cdots,x_n+p_n]$$

가장 초창기에 제시된 방법이며, 이를 구현하는 가장 간단하고 대표적인 방법이 Learned Positional Embedding / Sinusoidal Positional Embedding이다.

위치마다 하나의 embedding을 학습한다고 가정해보자. 0번째 인덱스에서는 $[0.42, -0.31, \cdots]$, $k$번째 인덱스에서는 $[-0.14, 0.82]$ 이런 식으로 말이다. 마치 Positional Embedding 자체가 하나의 Key-Value 해시 테이블처럼 더해지는 구조이다.

그러나 이러한 방식은 학습된 토큰의 길이를 넘어서는 순간 out-of-distribution이 된다. 예컨대 2048개의 토큰이 최대인 학습 모델은 4096 토큰이 들어오면 2049번째 임베딩부터는 더해줄 임베딩을 룩업 테이블에서 찾을 수 없게 된다.

이에 Transformer 원래 논문에서는 Embedding을 따로 학습하지 않고 사인/코사인을 사용하였다.

$$\begin{aligned} PE(pos, 2i) =& \sin\left(\frac{pos}{10000^{2i/d}}\right) \newline PE(pos, 2i+1) =& \cos\left(\frac{pos}{10000^{2i/d}}\right) \end{aligned}$$

이렇게 되면 임베딩의 각 디멘션에 대해 다른 주파수의 주기를 사용하면서 그와 동시에 모든 인덱스별로 유니크한 임베딩을 획득할 수 있게 된다.

그러나 이러한 APE에는 근본적인 문제점이 있다. 이는 바로 Semantic과 Position의 정보가 Projection 이전에 완전히 섞여버린다는 것이다.

$$\begin{aligned} Q_iK_j^\top =& (x_i+p_i)^\top W_Q^\top W_K (x_j + p_j) \newline =& x_i^\top A x_j + x_i^\top A p_j + p_i^\top A x_j + p_i^\top A p_j \end{aligned}$$

수식을 보게 되면, Query, Key 계산 이전에 positional embedding을 더해버린 결과, 실제 score 결과에 토큰 간의 관계, positional embedding 간의 관계, 그리고 토큰과 positional embedding 간의 관계 전체가 섞여서 더해지는 것을 볼 수 있다. 즉, 이러한 entanglement를 온전히 이후의 non-linearity 연산에 의존해서 해결해야하는데, 사실상 추가적인 모델링 없이 이를 학습만으로 해결하기엔 어려웠던 것이다. 또한 길이가 매우 길어지는 Context에서는 APE를 활용했을때 연산 안정성이 떨어지거나 주기 함수를 활용할 경우 표현력이 감소하는 이슈가 발생할 수 있다.

2. Relative Positional Embedding

Absolute PE(APE)는 절대적인 위치에 따라서 Unique한 값이 더해지므로 동일한 문장인데 시작점과 위치가 달라지는 경우 이에 대한 정보를 반영할 수 없다.

예컨대,

1. The cat sits on the box
2. I know the box where the cat sits usually

두 문장이 있을때, ‘the box’라는 문구와 ‘cat sits’라는 두 연속된 단어 간의 거리는 항상 1이다. 이를 토큰 단위로 해석했을때, 만약 Absolute PE가 더해졌다고 가정하면 위치에 따라 더해지는 임베딩 값 자체가 달라지기 때문에 동일한 토큰으로 인식되지 못한다. 개별 토큰은 그렇다고 하더라도, 연속된 토큰 간의 관계성은 Attention map에 동일하게 반영되어야 학습 시에 토큰 간 거리에 대한 일관성이 유지되는데 그렇지 못한 것이다.

Relative PE는 attention score에 토큰 간 상대적 거리를 반영한다.

$$Q_iK_j^\top + Q_i R_{i-j}$$

$i$번째 토큰과 $j$번째 토큰의 거리와 순서에 따라 Bias가 정해진다. $i$가 먼저 오면 음수이고 나중에 오면 양수이므로, attention은 거리와 앞뒤 순서를 함께 반영할 수 있다.

그러나 해당 방식은 메모리와 연산량을 폭증시키는 문제가 있다. 앞서 설명했던 APE의 경우 token 시퀀스에 더해지는 방식이었다면, 이번에는 attention score에 더해지는 방식이므로 각 원소마다 distance lookup (거리에 따라 더해지는 임베딩이 달라지므로)를 조회해야하므로, 기존에 비해 저장하는 메모리 양과 연산량이 증가하게 된다.

3. Rotary Positional Embedding

Absolute PE는 표현력 학습의 문제와 상대적 거리 부재의 이슈, Relative PE는 연산량과 메모리 폭증의 문제가 있었다. 그렇다면 Absolute PE의 장점과 Relative PE의 장점만 살릴 수 있는 Positional Embedding 방식은 없을까?라는 생각에서 등장한 방식이 RoPE이다. 최근 LLM이 대거 Attention 구조를 차용하면서, 점차 처리할 수 있는 Context 길이도 증가하고 그에 따른 리소스 오버로드가 커지기 시작했다. 기존 Transformer에서 단순히 접근했던 Attention 방식을 떠나서, 학습 시에 활용할 수 있는 Positional Embedding의 효율과 효과에 대해 집중하게 된 계기라고 볼 수 있다.

RoPE는 Positional Embedding을 더하는 방식에서 문제가 있다고 지적한다. 그 이유는 실제 Attention Score가 가지고 있는, 그리고 시퀀스 토큰 각각의 임베딩이 가지는 표현자 feature가 있는데, 이에 positional embedding을 더하게 되면 위치 관계가 틀어지는 문제가 발생할 수 있기 때문이다. 그렇기 때문에 RoPE(Rotary Positional Embedding)은 임베딩을 회전변환하는 관점으로 접근한다.

핵심 아이디어는 Embedding 공간에서의 다차원 벡터를 토큰 간의 거리 차이에 대한 각도의 배수로 회전변환을 수행한다. 논문의 Ideation은 결국 Euclidean 공간에서 토큰 간의 내적(유사도) 관계를 기반으로 특징점을 구성해 나가는데, 이때 결국 중요한 것은 벡터 간의 거리가 아니라 각도 차이라는 관점이다. 또한 장점이 명확하다. 회전변환에 대해서는 굳이 토큰 간의 거리를 조회하고, 이를 기반으로 임베딩을 설계할 필요가 없다. 회전 변환 자체가 다음 특징이 있기 때문이다:

$$R_i^\top R_j = R_{j-i}$$

직관적으로 이해하면 쉬운데, 어떤 벡터를 $\theta_i$만큼 반시계 방향으로 회전시키는 행렬($R_i$)이 있고 $\theta_j$만큼 회전시키는 행렬($R_j$)이 있다고 가정해보자. 회전행렬 또한 각각의 역행렬이 Transpose와 같은($R^\top = R^{-1}$) orthogonal matrix에 해당된다.

$R_i^\top$는 어떤 벡터를 $\theta_i$만큼 시계 방향, 즉 원래 방향과 반대로 회전시키는 행렬이다. 회전행렬을 곱하면 각 변환이 벡터에 순서대로 적용되므로, 최종 결과는 원래 벡터를 $\theta_j - \theta_i$만큼 반시계 방향으로 회전시키는 행렬($R_{j-i}$)이 된다.

각 벡터를 위치별 고유 각도만큼 회전시키면 Absolute PE처럼 위치 정보가 들어간다. 동시에 다른 토큰과의 인덱스 차이에 따른 상대적 위치도 회전 각도의 차이로 남는다.

RoPE는 이와 같이 사전에 정의된 절댓값을 회전 변환의 기준으로 사용하기 때문에 추가 파라미터를 가지거나 Lookup Table을 가질 필요가 없다.

다만 각도에 대한 표현력 $e^{i\theta}$ 또한 장문에 대해 주파수 aliasing이 발생하고 주파수가 올라갈수록 표현력의 한계점을 가지는 문제가 있으나, 그럼에도 불구하고 APE의 장점과 RPE의 장점을 혼합한 좋은 방식이라는 점에서 인정받았다. 현재 RoPE 방식은 LLaMA, Qwen, Gemma 등 다양한 LLM에서 활용하고 있다.