[논문리뷰] Enhanced Transformer with Rotary Position Embedding(RoFormer)

Enhanced Transformer with Rotary Position Embedding
Paper | GitHub | Guide 

 

💡 Introduction

단어의 sequence order은 자연어 분야에서 큰 가치가 있습니다. RNN 기반의 모델은 각 시퀸스를 재귀적으로 처리하면서 이전 시퀸스에서 계산한 hidden state를 현재 시퀸스로 전달하여 연산을 진행하며, 이를 통해 토큰의 순서를 인코딩할 수 있게 됩니다. 또한, CNN 기반 모델 (Gehring et al.[2017])은 position-agonistic으로 고려되었지만, Islam et al.[2020]의 논문에 의해 padding operation이 암묵적으로 위치 정보를 학습할 수 있다고 주장하였습니다. 

 

Vaswani et al.(2017)을 통해 소개 되었던 Transformer는 pre-trained LLM에서 다양한 downstream task에 적용이 되면서 SOTA을 달성하였습니다. 위의 소개되었던 RNN 및 CNN과 다르게 Pre-trained LLM에서는 self-attention를 사용하여 특정 말뭉치에 대한 문맥적 표현을 파악할 수 있게 되었습니다. 결과적으로 사전 훈련된 LLM은 병렬화 측면에서 큰 개선을 이루었으며, longer intra-token 모델링 가능한 ability가 대폭 상승하였습니다. 

 

Self-attention 구조가 적용된 pre-trained LLM은 각각의 토큰의 유사성을 기준으로 중요도가 결정되기 때문에 토큰의 위치가 중요하지 않습니다. 이러한 주장에 따라 학습과정에서 토큰의 위치 정보를 인코딩하는 다양한 방법론이 소개되었습니다만, 단순히 context representation에 위치정보를 더하는 것은 linear-self attention 구조에 부적합하다고 합니다. 

 

이미지 출처 : DeepLearning Hero(YouTube)

 

이에 본 논문의 저자들은 이러한 문제점을 해결하기 위해 RoPE(Rotary Positional Embedding)를 제안하였습니다. 이는 

회전행렬(Rotation Matrix)를 통해 각 토큰의 고유한 각도를 부여하여 absolute position 인코딩이 가능하며, relative position의 dependency를 통합할 수 있게 되었습니다. 

---

💡 Background and Related Work

Position Embedding에 대한 선행 연구가 다양하게 이루어졌으며, 입력 시퀸스의 순서 정보가 중요한 기계 번역, 텍스트 요약와 같은 분야에서 활용되었습니다.

✔️  Vaswani et al.[2017]

사전에 정의된 함수를 사용하여 특정 위치에 대한 패턴을 생성하여, context-representation에 위치정보를 추가하는 self-attention 개념을 도입하였습니다. Transformer 모델은 본질적으로 순서에 무관한 구조이기에, position embedding을 통해 각 토큰의 위치 정보를 모델에 직접적으로 통함하므로써 입력된 시퀸스의 토큰에 대한 순서 정보를 보존할 수 있게 되었습니다. 

 

✔️  Shaw et al.[2018] (Self-attention with relative position representations)

이미지 출처 : 딥러닝논문읽기모임

 

본 연구에서는 변형된 Transformer 모델의 토큰 위치 간의 relative distance가 Self-Attention Score에 어떻게 영향을 미치는지 연구하였습니다. 하지만 이러한 가중치 값들이 종종 너무 크거나 작아서 학습이 원할하게 진행이 되지 않는 경우가 있습니다. 이러한 문제점을 해결하기 위해 입력 벡터 간의 관계를 제한하여 가중치를 Clipping 하는 방법론을 고안하였습니다. 즉, Clipping은 일정한 임계값 이상의 가중치를 제한하는 것입니다. 이러한 연구를 기반으로 Transformer 모델이 문맥을 더 잘 이해하고, 특히 위치 정보가 중요한 다양한 언어 처리 작업에서 더욱 정확한 결과를 낼 수 있도록 돕는데 큰 기여를 했습니다. 상대 위치 임베딩의 도입은 이후의 많은 연구에 영감을 주었으며, 모델의 해석력과 성능을 높이는 중요한 단계로 평가받고 있습니다.

 

✔️  Dai et al.[2019] (Transformer-XL: Attentive language models beyond a fixed-length context) 

본 연구는 기존 Transformer의 구조를 변형하여, 각 요소들이(Query,Key,Value) 위치 정보를 다루는 방식을 개선하였습니다. 이는 토큰의 위치 정보를 더 정교하게 처리하여 모델의 문맥적 의미를 더욱 더 효과적으로 파악하고, 긴 문서나 복잡한 문장 구조를 가진 텍스트에서 더 좋은 결과를 도출할 수 있게 되었습니다. 이 논문의 Contribution은 아래와 같습니다. 

• Sinusoid-Encoded Relative Counterpart : relative position embedding 대신 sinusoidal-encoded를 사용함으로써, 토큰 간의 상대적 거리를 더 정교하게 인코딩하였습니다. 또한, 위치 정보 처리에 유연성올 높이기 위해 trainable vector u와 v를 사용하였습니다. 이에 결과적으로 다양한 길이의 텍스트나 다른 구조의 문장에서 일관된 성능을 극대화하였습니다. 
  
• Context-Based Key Vector와 Location-Based Key Vector :  Key Vector가 context 혹은 location으로 분리되면서, 모델의 문맥 정보와 위치 정보를 더욱 정확하게 분석할 수 있게 되었습니다. 
  

✔️  Raffel et al., He et al., Huang et al. [2020] 

상대적 위치 임베딩을 사용하여 Self-Attention의 메커니즘을 개선하는 다양한 방법이 제안되었습니다. 이러한 연구들은 Vaswani et al. [2017]의  연구를 기반으로 하였으며, trainable bias 및 projection matrix를 추가하여 방정식을 개선하는 연구가 진행 되었습니다. 

---

💡 Proposed Approach 

위의 방식들은 context representation에 position information를 더하는 것이 공통적인 부분입니다. 하지만 Rotary position embedding의 저자들은 선행연구와는 다르게 아래의 수식에서 일부 제약 조건하에 relative position embedding을 유도하고자 하였습니다. 

 

본 논문의 저자들은 relative position embedding를 context representation의 회전행렬과 결합하는 방식이 조금 더 직관적이라고 주장하였으며, 조금 더 자세히 알아보도록 하겠습니다.  

📌 Formulation 

대체적으로 Transformer 기반의 언어 모델들은 self-attention mechanism를 통해 토큰의 위치정보를 추출하게 됩니다. 아래의 왼쪽 수식에서 \(q_{T}^{m}\) \(k_{n}\)는 시퀸스 내 서로 다른 위치에 있는 토큰 간의 Knowledge Conveyance(지식 전달)을 가능하게 합니다. 

 

일반적으로 relative position information를 구체화하기 위해서는 query (\(q_{m}\))와 key (\(k_{n}\))의 내적이 함수 g로 정의되어야 하며, 입력 변수로는 단어 임베딩 (\(x_{m}\), \(x_{n}\))과 relative position(m-n)로 사용이 되어야 합니다. 하지만 선행연구들은 대부분 additive 기반으로 g함수를 구성하였지만, 본 논문의 저자들은 multiplicative 기반으로 위의 수식 같이 g함수를 정의하였습니다.

---

📌 Rotary Position Embedding (2D Case) 

✔️ Position-Encoded Query and Key Vectors 

 

위치 인코딩을 적용하는 구체적인 방법과 이 인코딩이 쿼리와 키 벡터에 어떻게 적용되는지를 설명하는 부분입니다. 복소 평면상에서의 회전을 통해 각 토큰의 위치정보를 벡터에 통합하는 방식을 설명합니다. 위의 수식에서 쿼리와 키 벡터가 각각의 위치 m과 n에 따라 어떻게 변형되는지를 표현하고 있습니다. 

 

• 쿼리 벡터  \(f_{q}\left(x_{m},m\right)\) : 위치 m에 있는 토큰에 대한 쿼리 벡터를 뜻합니다. 쿼리 벡터는 입력 토큰의 특성을 반영하며, Self-Attention 메커니즘에서 다른 토큰과의 관계를 평가하는데 사용이 됩니다. 
• 변환 가중치 행렬 \(W_{q}\) : 가중치 행렬은 입력 벡터 \(x_{m}\)을 쿼리 벡터의 차원으로 변환합니다. 이 변환은 입력 데이터의 특성을 새로운 특징 공간으로 매핑하게 됩니다. 
• 임베딩 벡터 \(x_{m}\) : 이는 NLP에서 사용하는 가장 기본적인 데이터 표현 방식입니다. 임베딩 벡터는 단어,문장 등의 텍스트 데이터를 신경망이 처리할 수 있는 고정된 크기의 실수 벡터로 변환하는 역할을 수행합니다. 
• 회전 \(e^{im\theta }\) : 이는 복소수 평면에서 \(m\theta \)만큼 토큰을 회전시키는 연산입니다. \(\theta \)는 회전 각도를 정의하는 상수이며, i는 허수 단위를 뜻합니다. 이러한 회전은 각 토큰의 절대 위치 정보를 쿼리 벡터에 통합하게 됩니다. 토큰이 시퀸스 내에서 서로 다른 위치에 있을 때, 이 회전은 각 토큰에 고유한 변화를 부여하여, 모델이 토큰의 순서와 위치를 인식할 수 있게 됩니다. 이러한 회전은 특히 문맥에 따라 달라지는 의미를 가진 단어(나라를 의미하는 "인도"와 길을 걷는 "인도")를 구별하는데 유용하게 사용됩니다. 
  
  

✔️ Define Function g

 

함수 g는 쿼리 벡터 \(x_{m}\)와 키 벡터 \(x_{n}\) 간의 내적을  복소수 곱셉을 통해 계산하며, 위치 m과 n 사이의 상대적 거리에 따라 회전을 적용합니다. Query의 복소수 표현과 Key의 켤레 복소수 표현을 곱한 뒤, 이들의 실수 부분을 취합니다. 이 과정은 쿼리와 키 벡터 사이의 상대적 위치에 따라 변화하는 내적값을 제공하며, 이 값은 최종적으로 attention score을 계산하는데 기여합니다.

 

* 실수 값만 취하는 이유 :  transformer 모델의 attention score는 쿼리와 벡터 간의 유사도를 평가하며, 이는 softmax 함수를 통해 확률로 변환되어 각 토큰의 중요도를 결정합니다. 이 과정에서 실수 값이 필요하기 때문입니다. 

---

📌 General Form 

2차원에서의 결과를 더 높은 차원 d로 일반화할 수 있다고 주장합니다. 여기에서 d는 짝수이며, 모든 d 차원 공간에서 d/2개의 2차원 부공간으로 나누고, 이들을 내적의 선형성을 이용해 결합하는 방법을 아래와 같이 정의할 수 있습니다. 

 

여기에서 \(R_{\theta ,m}^{d}\)는 아래와 같이 정의할 수 있습니다. \(R_{\theta ,m}^{d}\)는 사전 정의된 매개변수 (\Theta = {\theta_i = 10000^{-2(i-1)/d}, i \in [1, 2, \ldots, d/2]}}을 가지는 회전행렬 입니다. 

 

이러한 수식을 Self-Attention 매커니즘에 적용하면 아래와 같은 수식을 얻을 수 있습니다. 

✔️ Derivation

 

(1) 회전 행렬을 통한 변환:

쿼리 벡터 \(q_{m}\)과 키 벡터 \(k_{n}\)는 회전 행렬 \(R_{\theta ,m}^{d}\) 및 \(R_{\theta ,n}^{d}\)과 가중치 행렬 \(W_{q}\) 및 \(W_{K}\)를 사용하여 변환합니다. 

(2) 내적 계산: 

변환된 쿼리 벡터와 키 벡터의 내적을 계산합니다.

 

(3) 전치 행렬 사용 

내적 계산에서 전치 행렬의 성질을 사용하여 아래와 같이 전개합니다. 

(4) 회전 행렬의 성질  

회전 행렬 \(R_{\theta ,m}^{d}\) 및 \(R_{\theta ,n}^{d}\)는 직교 행렬이므로, 이들의 전치 행렬은 역행렬과 같습니다. 

 

* \(R_{\theta ,m}^{d}\) 및 \(R_{\theta ,n}^{d}\) 직교 행렬인 이유 

 

행렬의 열 벡터(또는 행 벡터)들이 서로 직교하고, 모두 단위 벡터이기 떄문입니다. 즉,
\(R_{\theta ,m}^{d}\) 및 \(R_{\theta ,n}^{d}\)의 열 벡터 혹은 행 벡터들 간의 내전이 0이고, 각 벡터의 크기가 1입니다. 

 

직교 행렬의 중요한 성질 중 하나는, 직교 행렬의 전치 행렬이 그 행렬의 역행렬이라는 점이며, \(R^{T}R\ =\ RR^{T}\ =\ I\) 가 성립될 수 있습니다. 이러한 성질은 회전행렬에도 적용이 됩니다.

 

회전 행렬의 구조를 보면, 각 2차원 부분 공간에 대해 아래와 같이 표현할 수 있으며, 내적의 값이 0이고 크기가 1이므로 직교 행렬이라고 할 수 있습니다. 

 

(5) 상대적 회전 

n과 m의 차이에 따른 회전 행렬은 아래와 같이 표현됩니다. 

 

(6) 최종 결과 

전체 내적 계산은 다음과 같이 됩니다. 

이는 회전 행렬의 성질과 전치 행렬의 성질을 이용하여 쿼리와 키의 내적을 계산하는 방법이 제안되었습니다. 

 

위의 내용을 기반으로 RoPE가 어떻게 작동하는지 다시 한번 그려보면 좋을거 같습니다. 

---

📌 Properties of RoPE

✔️ Long-term Decay 

Long-term decay 속성은 Transformer 기반의 모델에서 위치 인코딩이 멀리 떨어진 토큰들 간의 관계를 적절히 모델링하기 위한 중요한 개념입니다. 이는 상대적 거리가 큰 토큰 쌍의 연결 강도가 감소해야 한다는 직관을 반영합니다. 쉽게 말하자면, 두 토큰 사이의 거리가 멀어질수록 이들 간의 내적 값이 작아져야 한다는 것을 의미합니다. 

 

앞서, Vaswani et al.[2017]에서 제안한 위치 인코딩 방식에서는 각도 매개변수 \(\theta _{i}\)가 아래와 같이 설정이 됩니다. 

 

여기에서 d는 모델의 차원 수, i는 차원의 인덱스를 뜻합니다. 이러한 설정은 인코딩 벡터의 주기성을 조절하여, 상대적인 위치 정보가 내적 계산에 미치는 영향을 점진적으로 감소하게 만듭니다. 이를 조금 더 직관적으로 설명한 것은 아래와 같습니다. 

 

• 위치 인코딩 백터의 주기성 :  \(\theta _{i}\)는 차원이 커질수록 작은 값이 됩니다. 이는 고차원에서 위치 인코딩 벡터가 더 느리게 변하게 함으로써, 멀리 떨어진 토큰들 간의 유사도를 자연스럽게 낮추는 역할을 수행합니다. 
  
  
• 상대적 거리 증가에 따른 내적 값 감소 : 상대적 위치 m과 n이 증가할수록, 두 벡터간의 내적 값이 감소하여 트랜스포머 모델에서 멀리 떨이진 토큰들 간의 연결 강도가 약해지는 것을 의미합니다. 이에 모델의 성능 향상에 기여할 수 있게 됩니다. 

✔️ RoPE with Linear Attention 

선형 어텐션(Linear Attention)은 기본적으로 어텐션 메커니즘의 계산 복잡도를 줄이기 위해 제안된 방법입니다. 트랜스포머 모델에서의 어텐션 메커니즘은 아래와 같은 형태를 가집니다. 

 

기존의 Attention Mechanism은 모든 쿼리와 키 간의 유사도를 계산해야 하기 떄문에, 계산 복잡도가 \(O\left(N^{2}d_{k}\right)\)으로 급격하게 상승하게 됩니다. 이러한 계산 복잡도를 줄이기 위해 제안된 방법이 선형 어텐션(Linear Attention)입니다. 

 

쿼리와 키 간의 내적 계산을 선형화하여 계산 복잡도를 줄인다는 것이 핵심입니다. 이를 위해서 유사도 함수(similarity function) sim(Q,K)를 분해 가능한 형태로 아래와 같이 정의 하였습니다.

 

여기서 \(\phi \left(⋅\right)\)는 비음수 함수(non-negative function)를 쿼리와 키 벡터에 적용하여 새로운 벡터 \(\phi \left(Q\right)\)와 \(\phi \left(K\right)\)를 생성합니다. 대부분 \(\phi \left(x\right)\ =\ elu\left(x\right)+1\) 사용하는데, 이는 모든 입력값에 대해 비음수 출력을 생성하여 학습과정에서 더 나은 수렴 특성을 가질 수 있습니다. 결과적으로 행렬 곱셉의 결합 법칙을 사용하여 계산 복잡도를 \(O\left(Nd_{k}\right)\)로 줄일 수 있습니다.