[Daily morning study] Transformer 위치 인코딩의 발전: RoPE와 ALiBi
#daily morning study
위치 인코딩이 필요한 이유
Transformer의 Self-Attention은 본질적으로 순서 정보를 갖지 않는다. 입력 토큰을 순열해도 Attention 결과가 동일하다는 뜻이다. 언어는 순서가 의미를 결정하기 때문에, 위치 정보를 어떻게 주입하느냐는 모델 성능과 직결된다.
초기 Transformer(Vaswani et al. 2017)는 Sinusoidal Positional Encoding을 사용했다.
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
각 위치마다 고정된 사인/코사인 벡터를 임베딩에 더하는 방식이다. 단순하고 학습 파라미터가 없지만, 학습 시보다 긴 시퀀스에 일반화하기 어렵다는 한계가 있다.
이후 학습 가능한 위치 임베딩(Learned Positional Embedding) 방식이 GPT-2, BERT 등에서 채택됐다. 위치별로 임베딩 벡터를 별도로 학습하는 방법이다. 직관적이지만 학습 시의 최대 시퀀스 길이를 벗어나면 성능이 급락한다.
RoPE (Rotary Position Embedding)
핵심 아이디어
RoPE는 Su et al. (2021)이 제안했고, LLaMA, Gemma, Mistral, Qwen 등 최신 LLM 대부분이 채택하고 있다.
절대 위치를 임베딩에 더하는 대신, Query와 Key 벡터를 위치에 따라 회전(rotation)시킨다는 아이디어다. 회전 행렬을 이용해 위치 정보를 인코딩한다.
수식으로는 2차원 벡터에서 다음과 같이 표현된다.
RoPE(q, pos) = q * e^(i * pos * θ)
실제로는 d차원 벡터를 d/2개의 2D 블록으로 나눠 각각 회전을 적용한다.
| 방식 | 위치 정보 주입 시점 | 특징 |
|---|---|---|
| Sinusoidal PE | 임베딩에 덧셈 | 파라미터 없음, 일반화 제한 |
| Learned PE | 임베딩에 덧셈 | 학습 길이 초과 시 성능 저하 |
| RoPE | Q, K 벡터 회전 | 상대 위치 자연스럽게 반영 |
상대 위치 자동 반영
RoPE의 핵심 장점은 Q와 K의 내적에서 상대 위치 차이 (pos_q - pos_k)만 남는다는 점이다.
<RoPE(q, m), RoPE(k, n)> = f(q, k, m-n)
절대 위치가 아닌 토큰 간 상대 거리가 자동으로 Attention Score에 반영된다.
장거리 외삽 (Extrapolation)
학습 시보다 긴 컨텍스트를 처리하려면 위치 인덱스가 학습 범위를 벗어난다. RoPE는 주파수 스케일링으로 이 문제를 어느 정도 해결할 수 있다.
- Position Interpolation: 기존 위치 범위 안으로 새 위치를 압축해서 매핑
- YaRN (Yet another RoPE extensioN): 주파수 대역별로 다른 스케일링 적용, LLaMA 3.1의 128k 컨텍스트에 활용
# RoPE 스케일링 예시 (개념적)
original_seq_len = 4096
target_seq_len = 32768
scale = target_seq_len / original_seq_len # 8배
# 각 주파수(theta)를 scale배로 늘려 위치 범위 확장
theta_scaled = theta * scale
ALiBi (Attention with Linear Biases)
핵심 아이디어
Press et al. (2021)이 제안한 방법으로, 학습 길이보다 긴 시퀀스에서도 성능을 유지하는 것을 목표로 한다.
위치 정보를 임베딩에 더하거나 Q/K 벡터를 변형하는 대신, Attention Score에 거리에 비례하는 음의 편향(penalty)을 직접 더하는 방식이다.
Attention Score = Q * K^T / sqrt(d) - m * |i - j|
i, j: 쿼리와 키의 위치 인덱스m: 헤드마다 다른 기울기 (slope), 학습하지 않는 고정 하이퍼파라미터
헤드별 기울기 설정
헤드 개수 n에 따라 기울기 집합을 미리 정의한다.
m_k = 2^(-8k/n) for k = 1, ..., n
n=8이라면 기울기는 {1/2, 1/4, 1/8, …, 1/256}이 된다. 기울기가 작은 헤드는 먼 토큰도 어느 정도 참조하고, 기울기가 큰 헤드는 가까운 토큰에 집중한다.
| 헤드 인덱스 | 기울기 m (n=8 기준) | 특성 |
|---|---|---|
| 1 | 1/2 | 근거리 집중 |
| 4 | 1/16 | 중거리 |
| 8 | 1/256 | 원거리도 참조 |
외삽 강점
ALiBi는 학습 범위 밖의 위치가 나와도 동일한 선형 패널티 공식이 유지된다. 학습 길이의 2~4배까지 큰 성능 저하 없이 동작한다고 알려져 있다. MPT, BLOOM 등에서 채택됐다.
RoPE vs ALiBi 비교
| 항목 | RoPE | ALiBi |
|---|---|---|
| 위치 정보 주입 위치 | Q, K 벡터 회전 | Attention Score에 bias 추가 |
| 파라미터 | 없음 (회전 주파수는 고정) | 없음 (기울기는 공식으로 결정) |
| 학습 범위 외 처리 | 스케일링 기법 필요 (YaRN 등) | 선형 외삽으로 자연스럽게 처리 |
| 채택 모델 | LLaMA 시리즈, Gemma, Mistral | MPT, BLOOM |
| 상대 위치 반영 | 수학적으로 자연스럽게 내적에 반영 | 거리 기반 penalty로 강제 |
Sinusoidal PE → RoPE 전환이 의미하는 것
BERT/GPT-2 시대에는 Learned PE나 Sinusoidal PE가 주류였다. 이 방식은 임베딩 레이어에서 한 번 더해지므로, 이후 레이어 전체에 위치 정보가 섞여 흐른다.
RoPE는 Attention이 계산될 때마다 Q/K에 위치 정보를 주입하므로, 레이어마다 신선하게 위치 정보가 반영된다. 이 차이가 긴 컨텍스트에서 성능 차이를 만든다.
실습: RoPE 간단 구현
import torch
import math
def precompute_freqs_cis(dim: int, seq_len: int, theta: float = 10000.0):
# 주파수 계산
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
t = torch.arange(seq_len)
freqs = torch.outer(t, freqs) # (seq_len, dim/2)
# 복소수로 표현
freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
return freqs_cis
def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
# xq, xk: (batch, seq_len, n_heads, head_dim)
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
freqs_cis = freqs_cis[:, None, :] # 브로드캐스팅용
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
return xq_out.type_as(xq), xk_out.type_as(xk)
정리
- Sinusoidal PE: 첫 Transformer에서 사용, 단순하지만 외삽 한계 있음
- Learned PE: 학습 가능하지만 최대 시퀀스 길이에 묶임
- RoPE: Q/K 회전으로 상대 위치를 수학적으로 반영, 현재 LLM 표준
- ALiBi: Attention Score에 선형 패널티 추가, 외삽 성능 우수
- 컨텍스트 길이 확장(YaRN, Position Interpolation)은 RoPE의 주파수 스케일링을 활용
