기술 요약: RoPE 어텐션은 거의 선형 시간에 학습될 수 있다

문제 정의

회전 위치 임베딩(Rotary Position Embedding, RoPE) 메커니즘은 Llama, Claude, Apple의 모델들과 같은 최첨단 대규모 언어 모델(LLM)의 표준 구성 요소가 되었으며, 전통적인 위치 인코딩에 비해 토큰 관계를 포착하는 데 있어 탁월한 표현력을 제공한다. 그러나 RoPE에 내재된 위치 의존적 회전은 어텐션 메커니즘의 계산을 복잡하게 만든다.

최근 연구([AS24a])는 "유계 엔트리(bounded entry)" 체제(행렬 엔트리가 파라미터 $B$에 의해 제한되는 경우) 하에서 RoPE 어텐션의 순방향(forward) 계산에 대해 거의 선형 시간($n^{1+o(1)}$) 알고리즘을 확립했지만, 역방향(backward) 계산(학습을 위한 그래디언트 계산)은 다뤄지지 않았다. 역방향 계산은 어텐션 행렬과 위치 임베딩의 비선형 변환을 포함하기 때문에 본질적으로 더 복-잡하다. 본 연구가 다루는 핵심 질문은 유계 엔트리 조건 하에서 RoPE 어텐션의 역방만 그래디언트 계산이 순방향 계산과 동일한 거의 선형 시간 효율성을 달칠 수 있는지 여부이다.

방법론

저자들은 거의 선형 시간으로 실행되는 역방향 RoPE 어텐션 계산을 위한 최초의 알고리즘을 개발한다. 이 접근 방식은 폐쇄형 그래디언트 유도, 저계수 근사(low-rank approximation), 다항식 방법 및 고속 푸리에 변환(FFT)의 결합에 의존한다.

1. 폐쇄형 그래디언트 재정식화

논문은 먼저 가중치 행렬에 대한 RoPE 어텐션 손실 함수의 그래디언트에 대한 폐쇄형 표현식을 유도한다. "텐서 트릭"(Kronecker 곱)을 활용하고 어텐션 행렬 $A(X)$를 재정식화함으로써, 그래디언트는 다음과 같이 표현된다:
$$\frac{d\text{Loss}(x)}{dx} = \tilde{A}^\top \text{vec}(\gamma(x))$$
여기서 $\gamma(x)$는 다음을 포함하는 복소 행렬 함수이다:

• $s(x)$: 정규화된 Softmax 벡터.
• $\ell(x)$: 어텐션 출력과 타겟 사이의 차이에서 유도된 오차 항.
• $\beta(x)$: 오차와 가치(value) 행렬을 결합한 항.
• $\gamma(x)$: $s(x)$의 대각 성분과 $s(x)s(x)^\top$가 $\beta(x)$에 작용하는 항을 포함하는 항.

2. 저계수 근사 전략

거의 선형 시간 복잡도를 달성하기 위해, 저자들은 $\gamma(x)$의 구성 요소들을 저계수 행렬을 사용하여 근사한다. 이 전략은 $\gamma(x)$를 두 부분인 $\gamma_1(x)$와 $\gamma_2(x)$로 분해하고 각각을 별도로 근사하는 것을 포함한다:

• $s(x)$ 및 $\ell(x)$ 근사: [AS24a]의 순방향 알고리즘을 바탕으로, 저자들은 정규화된 Softmax $s(x)$가 $n^{1+o(1)}$ 시간 내에 저계수 행렬 $U_1 V_1^\top$로 근사될 수 있음을 보여준다. 오차 항 $\ell(x)$ 또한 이 결과를 사용하여 근사된다.
• $\beta(x)$ 근사: $\beta(x)$는 가치 행렬과 오차 항의 곱이므로, 구성 요소들의 근사를 기반으로 저계수 인자들을 구축하여 근사된다.
• $\gamma(x)$ 근사:
  • $\gamma_1(x) = \text{diag}(s(x))\beta(x)$는 행 단위 Kronecker 곱을 사용하여 $s(x)$와 $\beta(x)$의 저계수 인자들을 결합함으로써 근사된다.
  • $\gamma_2(x) = s(x)s(x)^\top\beta(x)$는 중간 항들을 사전 계산하고 $s(x)$와 $\beta(x)$의 저계수 구조를 활용하여 근사된다.

3. 난해도 분석 (Hardness Analysis)

유계 엔트리 조건의 필요성을 확립하기 위해, 저자들은 강한 지수 시간 가설(SETH)에 기반한 하한(lower bounds)을 도출한다. 저자들은 만약 엔트리 경계 $B$가 특정 임계값(구체적으로 $B = \omega(\sqrt{\log n})$)을 초과하면, SETH를 가정할 때 어떤 알고리즘도 이차 미만의 시간($O(n^{2-q})$) 내에 그래디언트를 계산할 수 없음을 증명한다. 이는 유계 엔트리 가정이 단순히 기술적인 편의가 아니라, 이차 미만 성능을 위한 근본적인 요구 사항임을 확인시켜 준다.

주요 기여

1. 폐쇄형 그래디언트: 본 논문은 RoPE 어텐션의 그래디언트에 대한 최초의 폐쇄형 정식화(Lemma 4.1)를 제공하고 그 정확한 시간 복잡도를 분석하여, 나이브한 계산에서의 이차 복잡도 병목 현상을 식별한다.
2. 거의 선형 시간 알고리즘: 저자들은 유계 엔트리 조건 하에서 $n^{1+o(1)}$ 시간 내에 RoPE 어텐션의 역방향 그래디언트를 근사하는 최초의 알고리즘을 제시한다(Theorem 5.7). 이는 순방향 패스의 효율성과 일치한다.
3. 이론적 하한: 본 연구는 이차 미만 성능을 위해 유계 엔트리 조건이 필수적임을 입증하며, SETH로부터 도출된 난해도 결과(Theorem 6.1)를 제공한다.
4. 알고리즘 기법: 이 접근 방식은 RoPE의 구조적 제약에 특화된 저계수 근사 기술과 다항식 근사 방법 및 FFT를 통합한다.

결과

주요 결과(Theorem 5.7)는 $d = O(\log n)$ 및 $B = o(\sqrt{\log n})$인 경우, $n^{1+o(1)}$ 시간 내에 $1/\text{poly}(n)$으로 제한된 가산 오차(additive error)를 가지며 RoPE 어텐션 그래디언트 계산 문제를 해결하는 알고리즘이 존재함을 보여준다.

반대로, 난해도 결과(Theorem 6.1)는 만약 $B = \omega(\sqrt{\log n})$라면, SETH 가정 하에 $O(n^{2-q})$ 시간 내에 그래디언트를 계산하는 것은 불가능함을 보여준다.

의의

본 연구는 RoPE 기반 트랜스포머의 이론적 이해에 있어 중요한 간극을 메운다. 유계 엔트리 조건 하에서 역방향 계산이 순방향 계산만큼 효율적일 수 있음을 증명함으로써, 이 논문은 RoPE를 사용하는 대규모 모델을 학습시키는 데 있어 중요한 계산적 장벽을 제거한다. 이러한 결과는 유계 엔트리 체제가 유지되는 한, RoPE 기반 모델의 학습 효율성이 표준 어텐션을 사용하는 모델과 이론적으로 대등함을 시사한다.

본 논문은 RoPE 역방향 계산의 미세한 복잡도(fine-grained complexity)를 규명하며, 순방향 계산에 대한 이전 연구를 확장한다. 또한 알고리즘 설계와 계산 복잡도 이론 사이의 상호작용을 강조하며, 다른 고급 어텐션 변형 및 위치 인코딩 메커니즘에 대한 서브 그래디언트(sub-gradient) 계산 연구를 위한 토대를 제공한다. 저자들은 향후 연구에서 엔트리가 제한되지 않은 경우와 이러한 이론적 경계가 실제 LLM 학습에 미치는 실질적인 영향에 대해 탐구할 수 있다고 언급한다.