Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings

기술 요약: 극좌표 위치 임베딩을 통한 "무엇(What)"과 "어디(Where)"의 분리

문제 정의
트랜스포머 아키텍처에서 어텐션 메커니즘은 두 가지 요소, 즉 콘텐츠("무엇")와 시퀀스 위치("어디")를 기반으로 키(key)를 쿼리(query)에 매칭합니다. 저자들은 Llama 3 및 DeepSeek과 같은 프런티어 모델에서 널리 사용되는 회전 위치 임베딩(Rotary Position Embedding, RoPE)이 이 두 요소를 얽히게 만든다고 주장합니다. RoPE에서는 키와 쿼리 벡터에 적용되는 회전이 두 요소의 콘텐츠 크기와 상대적 위치 모두에 의존합니다. 이러한 얽힘은 "무엇-어디" 혼재(confound)를 야기하며, 콘텐츠 또는 위치 단독으로 독립적인 매칭이 필요한 작업에서 모델 성능을 저해합니다. 구체적으로, RoPE의 상호작용 항은 회전의 위상(phase)이 콘텐츠 특징에 의해 영향을 받도록 만들어, 모델이 엄격하게 위치 인덱싱이나 콘텐츠 매칭에만 의존하는 규칙을 학습하는 것을 어렵게 만듭니다.

방법론: 극좌표 위치 임베딩 (Polar Coordinate Positional Embedding, PoPE)
저자들은 키와 쿼리 표현을 극좌표 형식으로 재구성함으로써 콘텐츠와 위치를 분리하는 RoPE의 수정안인 PoPE를 제안합니다.

• 재구성: RoPE에서는 컴포넌트가 위치에 따라 회전하지만, 초기 위상은 콘텐츠에 의해 결정됩니다. PoPE에서 저자들은 실수 값인 키와 쿼리 벡터를 다음과 같은 복소수 벡터로 변환합니다:
  • 크기($\mu$): 비음수성을 보장하기 위해 소프트플러스 활성화 함수($\sigma(x) = \ln(1 + e^x)$)를 사용하여 오직 콘텐츠 특징으로부터만 유도됩니다.
  • 위상($\phi$): 엄격하게 시퀀스 위치($s$ 또는 $t$)와 컴포넌트별 주파수($\theta_c$)에 의해 결정됩니다.
• 어텐션 점수: 어텐션 점수는 쿼리의 켤레 전치(conjugate transpose)와 키의 내적의 실수 부분으로 계산됩니다. 수학적으로 이는 다음과 같습니다:
  $$a^{PoPE}_{ts} = \sum_{c=1}^{d} \mu_{\tilde{q}_{tc}} \mu_{\tilde{k}_{sc}} \cos((s - t)\theta_c + \delta_c)$$
  여기서 $\delta_c$는 학습 가능한 편향(bias) 항입니다.
• RoPE와의 주요 차이점:
  1. 분리(Disentanglement): RoPE에서 콘텐츠와 위치를 혼합하는 상호작용 항($\phi_{k_{sc}} - \phi_{q_{tc}}$)이 제거되었습니다. 크기(콘텐츠)와 위상(위치)이 엄격하게 분리됩니다.
  2. 주파수 해상도: PoPE는 $d/2$ 쌍이 아닌 $d$개의 컴포넌트(개별 요소) 단위로 작동하여, 효과적으로 주파수 채널의 수를 두 배로 늘립니다.
  3. 학습 가능한 편향: 고정되어 있지만 학습 가능한 편향 $\delta_c$를 도입하여, 콘텐츠-위치 혼재를 다시 유발하지 않고도 각 주파수에 대한 최적의 상대적 오프셋을 조정할 수 있습니다.

주요 기여

1. 진단 작업 (간접 인덱싱): 저자들은 모델이 소스 캐릭터로부터 상대적 오프셋을 기준으로 타겟 캐릭터를 식별해야 하는 작업(예: "'z'의 왼쪽 3번째 위치에 있는 문자를 찾아라")을 도입했습니다. 이 작업은 콘텐츠와 위치의 독립적인 조작을 요구합니다.
2. 이론적 통찰: 논문은 RoPE가 어떻게 콘텐츠와 위치를 혼동시키는지에 대한 대수적 유도를 제공하며, 반면 PoPE는 어텐션 메커니즘이 독립적인 "무엇"과 "어디" 매칭의 결합으로서 규칙을 구현할 수 있게 함을 보여줍니다.
3. 효율적인 구현: 전체 복소수 행렬을 실체화하지 않고 복소수 내적의 실수 부분만을 계산하는 Triton 및 Flash Attention 2를 사용한 효율적인 구현이 제공되며, 이는 표준 어텐션에 비해 단 한 번의 곱셈 오버헤드만을 추가합니다.

결과

• 간접 인덱싱: 진단 작업에서 RoPE는 평균 정확도가 11.16%에 불과하여 독립적인 매칭 학습에 어려움을 겪었습니다. 반면, PoPE는 94.82%의 정확도를 달 기록하며 이러한 요인들을 분리하는 능력을 입증했습니다.
• 시퀀스 모델링 (음악 및 유전체학):
  • 음악 (JSB 및 MAESTRO): PoPE는 두 데이터셋 모두에서 음의 로그 가능도(NLL) 측면에서 RoBLE을 능가하며 더 낮은 퍼플렉시티(perplexity)를 달성했습니다.
  • 유전체학 (인간 참조 게놈): PoPE는 RoPE 베이스라인에 비해 NLL의 유의미한 감소를 보여주었습니다.
• 언어 모델링 (OpenWebText): 다양한 모델 규모(124M, 253M, 774M 파라미터)에 걸쳐 PoPE는 RoPE보다 일관되게 낮은 퍼플렉시티를 달성했습니다. 성능 격차는 모델 크기가 커짐에 따라 안정적으로 유지되거나 증가했습니다.
• 다운스트림 작업: 6가지 다운스트림 작업(LAMBADA, BLiMP, CBT, HellaSwag, PIQA, ARC-E)에 대한 제로샷 평가 결과, PoPE 기반 모델은 모든 모델 크기에서 RoPE 기반 모델보다 높은 평균 정확도를 달성했습니다.
• 길이 외삽 (Length Extrapolation):
  • 제로샷: PoPE는 별도의 미세 조정이나 주파수 보간 없이도 훈련 컨텍스트(1024 토큰)보다 최대 10배 긴 시퀀스(10240 토큰)까지 강력한 제로샷 길이 외삽 능력을 보여주었습니다.
  • YaRN과의 비교: 외삽을 위해 설계된 방법인 YaRN은 미세 조정 한계(4096 토큰)까지는 RoPE의 성능 저하를 완화하지만, 그 이상의 구간에서는 급격히 저하됩니다. PoPE는 YaRN이 필요로 하는 보간 기술이나 광범위한 미세 조정을 요구하지 않고도 더 긴 시퀀스에서 성능을 유지했습니다.
  • 주파수 사용 분석: 특징 노름(feature norms) 분석 결과, RoPE 모델은 노이즈를 피하기 위해 고주파 채널을 억제하는 경향이 있는 반면, PoPE 모델은 레이어 전반에 걸쳐 고주파를 포함한 더 넓게 분포된 범위의 주파수를 활용합니다.

의의 및 주장
본 논문은 PoPE가 RoPE에 내재된 "무엇-어디" 얽힘에 대한 원칙적인 해결책을 제공한다고 주장합니다. 콘텐츠와 위치를 분리함으로써 PoPE는 데이터 효율성, 점근적 정확도, 그리고 특히 컨텍스트 길이 일반화 능력을 향상시킵니다. 저자들은 이러한 이점이 기존 RoPE 아키텍처에 복잡한 구조적 변경 없이도 미세한 수정을 통해 달성된다는 점을 강조합니다.

이 연구는 시퀀스 모델링에 관한 근본적인 연구로 제시됩니다. 저자들은 이 방법이 음악, 유전체학, 자연어와 같은 도메인에서의 능력을 향상시켜 (단백질 구조 예측이나 신약 개발과 같은 응용 분야에 도움을 줄 수 있지만), 유해한 콘텐츠 생성이나 정교한 딥페이크와 같이 파운데이션 모델에서 흔히 발생하는 이중 용도 위험(dual-use risks)을 수반한다는 점을 언급합니다. 핵심적인 기여는 콘텐츠와 위치 정보를 분리하는 것이 더 견고하고 일반화 가능한 트랜스포머 모델을 만든다는 것을 입증한 데 있습니다.