PJ-RoPE: A Fourier-Jet-Affine Position Space for Relative Attention

기술 요약: PJ-RoPE: 상대적 어텐션을 위한 푸리에-제트-아핀 위치 공간

문제 정의

상대적 위치 표현은 트랜스포머의 긴 문맥(long-context) 어텐션에서 핵심적인 역할을 하지만, 기존 방법들은 통합된 구조의 구성 요소라기보다 서로 경쟁하는 설계들로 취급되는 경우가 많습니다.

• RoPE는 푸리에 특성(회전)을 통해 상대적 지연(lag)을 인코딩합니다.
• ALiBi는 어텐션 로짓에 단조적인 아핀 최근성(recency) 항을 직접 추가합니다.
• Jordan-RoPE는 반단순(semisimple) 회전 주파수를 유한 조르단 블록(finite Jordan blocks)으로 대체하여, 미분과 유사한 "유한 제트(finite-jet)" 특징을 생성합니다.

본 논문은 이러한 방법들이 상수 계수 차분 모듈(constant-coefficient difference modules)에 뿌리를 둔 공통된 대수적 구조를 공유한다고 주장합니다. 구체적으로, 단순 근(simple roots)은 푸리에/RoPE 특성에 해당하고, 중복된 0이 아닌 근은 조르단/푸리에 제트를 생성하며, 중복된 단위 근(unit roots)은 ALiBi와 같은 아핀 최근성을 생성합니다. 핵심 문제는 이러한 방법들을 하나의 학습 가능한 공간 내에서 선택 가능한 섹터로 취급하는 통합 프레임워크가 부족하다는 점과, 특정 작업이 실제로 어떤 섹터를 활용하는지 결정하는 진단 도구가 없다는 점입니다. 또한, 고차 제트 특징(거리의 거듭제곱을 포함하는)은 긴 문맥에서 안정성 문제를 야기하며, 이로 인해 변환된 특징과 로짓이 급격히 증가합니다.

방법론

저자들은 이러한 프리미티브들을 학습 가능한 푸리에-제트-아핀(Fourier–Jet–Affine) 상대 위치 공간으로 통합하는 프레임워크인 PJ-RoPE(Poincaré-type RoPE)를 제안합니다.

1. 이론적 정식화

이 프레임워크는 상수 계수 차분 모듈의 대수에 기초합니다. 저자들은 위치 공간 $P_{PJ}$를 세 가지 섹터의 직합(direct sum)으로 정의합니다:

• 동차 푸리에-제트 섹터 ($H_{FJ}$): RoPE(0차)와 Jordan-RoPE(r차 > 0)를 포함합니다. 대수적으로 이들은 중복된 0이 아닌 근을 갖는 방정식의 해이며, $(d/L)^r e^{-cd/L} e^{i\omega d}$와 같은 항을 생성합니다.
• 아핀 최근성 섹터 ($A_{rec}$): 상수 및 선형 항(예: $-sd/L$)을 포함하며, 이는 중복된 단위 근에 대응합니다. 이는 로렌츠 군의 아핀 확장과 유사한 "병진(translation-like)" 방향을 제공합니다 (따라서 "Poincaré-type"이라 명명됨).
• 라이트 콘(Light-Cone, LC) 섹터 ($H_{LC}$): 안정성을 해결하기 위한 고차 동작의 압축된 좌표 차트입니다. 원시 거리 $d$를 래피디티(rapidity)와 유사한 위상 $\phi_L(d) = L \text{asinh}(d/L)$와 포화 진폭 $\beta_L(d) = d/\sqrt{d^2 + L^2}$로 대체합니다.

프레임워크는 두 가지 구현 체계를 구분합니다:

• PJ-bias: 푸리에 특성, 감쇠된 제트, 아핀 항 및 LC 변형을 포함하는 스칼라 가산 어텐션 커널입니다. 이는 학습 가능하지만 정확한 특징 변환(feature transform)은 아닙니다.
• PJ-rotary: 쿼리 및 키 특징에 상대적 작용을 적용하는 정확한 특징 변환 구현으로, RoPE와 Jordan-RoPE의 폐쇄성(closure) 속성을 검증하는 데 사용됩니다.

2. 적응형 진단 (Adaptive Diagnostics)

작업이 공간의 어느 영역을 선택하는지 결정하기 위해, 저자들은 적응형 섹터 게이트($g_{FJ}, g_{aff}, g_{LC}$)와 세 가지 진단 도구를 도입합니다:

• 유효 질량 (Effective Mass): 게이트된 파라미터 크기가 각 섹터와 제트 차원에 어떻게 할당되는지 측정합니다.
• 함수 에너지 (Functional Energy): 감쇠 및 위상 상쇄를 고려하여 평가 윈도우 동안 실현된 커널 구성 요소의 규모를 측정합니다.
• Leave-One-Order-Out (LOO): 특정 제트 차원을 제거했을 때 적합된 함수가 저하되는지 평가하여, 파라미터 선택과 실제 기능적 기여를 분리합니다.

3. 라이트 콘 안정화 (Light-Cone Stabilization)

긴 문맥에서 고차 제트의 지수적 성장을 방지하기 위해, LC 브랜치는 좌표 변환을 적용합니다. 이는 변환된 특징과 로짓의 크기를 제한하지만, 긴 범위에서의 위상 해상도를 압축하며, 이를 통해 안정성과 해상도 사이의 트레이드오프를 명시적으로 설정합니다.

주요 기여

1. 위치 공간 정식화: 상대적 어텐션을 학습 가능한 푸리에-제트-아핀 공간으로 정식화하고, 스칼라 바이어스 커널과 정확한 회전 특징 변환을 구분했습니다.
2. 섹터 포함 관계: RoPE, Jordan-RoPE(유한 제트), ALiBi(아핀 최근성), LC 좌표를 단일 수학적 구조 내의 구별되는 섹터로 식별했습니다.
3. 적응형 진단: 섹터 게이트, 유효 질량, 함수 에너지 및 LOO 절제(ablation)를 도입하여 경험적으로 작업 수준의 섹터 선택을 측정하는 방법을 제공합니다.
4. 증거 체인: 통제된 커널, 합성 작업, 바이트 수준 언어 모델, 상징적 음악 토큰 스트림을 통해 섹터 복구, 학습 가능한 사용성 및 자연스러운 작업 할당을 평가했습니다.

실험 결과

실험은 짧은 문맥에서 학습되고 최대 32,768 토큰의 긴 문맥에서 평가되는 소규모 인과적 트랜스포머(약 0.15M–0.20M 파라미터)를 사용합니다.

• 고정 커널 복구: PJ-bias 스칼라는 순수 위상, 유한 제트, 아핀, LC 타겟을 성공적으로 복구하여, 해당 공간이 설계된 프리미티브 패밀리를 포괄함을 확인했습니다.
• 적응형 섹터 복구: "교사(teacher)" 신호(예: 순수 아핀 또는 특정 제트 차수)와 함께 학습될 때, 학습된 게이트와 진단 도구는 해당 섹터로 질량을 올바르게 이동시킵로 합니다.
• 합성 시퀀스 작업: 학습 가능한 어텐션 모델은 위치 브랜치가 신호를 전달할 때 이러한 섹터들을 활용할 수 있습니다. 아핀 교사는 최근성 브랜치를 활성화하고, 부호가 있는 제트 교사는 FJ 변형을 활성화합니다.
• 자연어 (바이트 수준): 언어 작업(Tiny Shakespeare, WikiText-2, War and Peace)에서 모델은 주로 아핀/최근성 동작을 선택합니다. 고차 FJ/LC 질량은 미미한 진단적 보정으로만 나타납니다. NTK 스타일의 RoPE 스케일링을 아핀 최근성과 결합하는 것이 가장 좋은 긴 문맥 성능을 보였습니다.
• 음악 토큰 스트림: 상징적 음악 작업(MAESTRO, MusicNet)에서는 할당 방식이 다릅니다. LC-아핀 변형이 여전히 강력하지만, 고차 보정에 대한 일관되고 측정 가능한 질량 할당이 관찰됩니다. 이는 모티프의 귀환과 리듬 엔벨로프를 모델링해야 하는 필요성과 일치합니다.
• LC 안정성 vs. 해상도: LC 변형은 32,768 토큰에서 Q/K 노름과 캐시/로짓 규모를 성공적으로 제한합니다. 그러나 이는 비용을 수반합니다: "원거리 위상 범위(far-range phase span)"가 크게 감소하여, 원시 좌표에 비해 하드 네거티브 프로브(hard-negative probe)에서의 검색 해상도가 낮아집니다.

의의 및 주장

본 논문은 PJ-RoPE를 단일 새로운 아키텍처가 아니라, 위치 공간 진단 프레임워크로 규정합니다. 주요 의의는 상대적 위치 메커니즘을 서로 다른 작업이 서로 다른 영역을 선택하는 학습 가능한 통합 공간으로 재구성했다는 점에 있습니다:

• 작업 의존적 선택: 경험적 증거는 언어 작업이 아핀/최근성 구조를 선호하는 반면, 구조화된 상징적 스트림(음악 등)은 아핀과 고차 제트 보정의 혼합을 활용한다는 것을 보여줍니다.
• 안정성-해상도 트레이드오프: 라이트 콘 좌표가 도입하는 트레이드오프를 명시적으로 정량화했습니다: 즉, 긴 문맥 확장을 위해 성장을 제한하여 안정화하지만 위상 해상도를 압축합니다.
• 관심사의 분리: 논문은 "표현의 정확성"(정확한 PJ-rotary 폐쇄성)과 "긴 문맥 안정성"(스칼라 LC 경로)을 구분하며, 정확한 특징 변환이 안정화 없이는 긴 범위에서 급격히 저하될 수 있음을 지적합니다.

저자들은 안정적인 긴 문맥 사용이 아핀 최근성, 푸리에-제트 구조, 그리고 LC 압축 사이의 작업 의존적 할당에 달려 있다고 결론짓습니다. 이들은 이러한 발견이 소규모 모델과 특정 토큰화 방식에 기반한 것임을 언급하며, 이러한 할당을 배포 규칙으로 간주하기 위해서는 더 큰 규모의 검증이 필요하다는 점에서 신중한 태도를 유지합니다.