AI/ Rotary Positional Embedding (RoPE)

Rotary Positional Embedding (RoPE)

원래의 Transformer는 sin과 cos을 이용하여 절대적 위치 정보를 반영하는 sinusoidal position encoding을 사용했으나, 이후 이런 방식이 long-context에서 성능이 떨어진다는 것이 발견되고 이를 개선하기 위해 Rotary Positional Embedding을 이용한 상대적 위치 인코딩 방식이 등장했다.

RoPE는 지수 형식을 사용하여 다음과 같은 수식으로 정의된다.

\text{RoPE}(\mathbf{x}) = \mathbf{x}e^{i\theta}

원래의 sinusoidal positional encoding이 입력 전체에 대해 position 정보를 추가했다면 RoPE는 Query와 Key에 대해서만 position 정보를 추가하기 때문에(value는 query와 key의 유사도 비교를 통해 선택되는 값이기 때문에 position 정보가 불필요하다) 아래와 같이 적용된다. 여기서

W_Q, W_K

는 일반적인 self-attention에서 입력을 query와 key로 변환하는 행렬이고,

\mathbf{h}_t

는

t

시점의 입력(이전 step의 출력 또는 이전 layer의 출력)에 해당한다.

\mathbf{q}_t = \text{RoPE}(W_Q\mathbf{h}_t) \\\mathbf{k}_t = \text{RoPE}(W_K\mathbf{h}_t)

또한

\mathbf{x}e^{i\theta}

에서

\theta

는 기존 sinusoidal positional encoding의 주파수 특성을 활용하여 아래처럼 정의된다.

\theta_d = 10000^{-2(d-1)/d_{\text{model}}}

이 값을 이용하여 각 토큰 위치

t

(0 이상의 정수)에서

\theta

는 아래처럼 계산된다.

\theta_{t,d} = t\cdot \theta_d

실제로 구현할 때 허수 부분

i

를 사용할 수는 없기 때문에 삼각함수를 이용하여 처리한다. 우선 아래가 성립하고

e^{i\theta} = \cos \theta + i\sin \theta

\mathbf{x}

를 실수 부분과 허수 부분으로 나눠서

\mathbf{x} = \mathbf{x}_r + i\mathbf{x}_i

로 표현한 다음 위의 식을 적용하면 아래와 같이 나타낼 수 있다.

\mathbf{x}e^{i\theta} = (\mathbf{x}_r + i\mathbf{x}_i)(\cos\theta + i\sin\theta) = \mathbf{x}_r \cos\theta -\mathbf{x}_i\sin\theta + i(\mathbf{x}_r\sin\theta + \mathbf{x}_i\cos\theta)

이것을 다시 행렬 곱셈 형식으로 변환하면 아래와 같이

i

를 제거한 형태로 만들 수 있다.

\begin{bmatrix}\mathbf{x}_r' \\ \mathbf{x}_i' \end{bmatrix} = \begin{bmatrix}\cos \theta & -\sin\theta \\ \sin \theta & \cos \theta \end{bmatrix}\begin{bmatrix}\mathbf{x}_r \\ \mathbf{x}_i \end{bmatrix}

이를 활용하여 실제에서는 벡터의 짝수 인덱스를 실수, 홀수 인덱스를 허수 부분으로 간주하여

\mathbf{x}

를 분해하고 위의 행렬을 곱하고 그것을 다시 합쳐 최종적으로 RoPE가 적용된

\mathbf{x}

를 구한다. 아래 샘플 코드 참조.

Sample

실제 코드 구현 예시는 아래 참조

def apply_rope(x, dim):
    """
    x: (batch_size, seq_len, d_model)
    dim: 임베딩 차원의 절반 크기
    """
    batch_size, seq_len, d_model = x.shape
    theta = 10000 ** (-2 * (torch.arange(0, dim, 2).float()) / d_model)  # RoPE 각도 값
    theta = theta.to(x.device)  # GPU에 맞게 이동
    t = torch.arange(seq_len, device=x.device).float().unsqueeze(1)  # 위치 인덱스

    # 회전 변환을 위한 사인 및 코사인 값 계산
    sin_t = torch.sin(t * theta)
    cos_t = torch.cos(t * theta)

    # x를 (batch, seq_len, dim//2, 2) 형태로 변환하여 적용
    x_reshaped = x.view(batch_size, seq_len, -1, 2)  # [batch, seq_len, dim//2, 2]
    x_rotated = torch.cat([
        x_reshaped[..., 0] * cos_t - x_reshaped[..., 1] * sin_t,
        x_reshaped[..., 0] * sin_t + x_reshaped[..., 1] * cos_t
    ], dim=-1)  # 다시 원래 차원으로 변환

    return x_rotated.view(batch_size, seq_len, d_model)  # 원래 차원으로 변환
Python
복사