다중차원 자기주의를 위한 폴리어텐션 프레임워크

초록

본 논문은 기존 자기주의(self‑attention)의 2차 상호작용 한계를 극복하기 위해, 임의의 고차 텐서 연산과 토큰 간 관계 구조를 포함할 수 있는 일반화된 “폴리어텐션(poly‑attention)” 메커니즘을 정의한다. 폴리어텐션은 기존의 고차 텐서 어텐션과 Strassen 어텐션을 특수 경우로 포함하며, 각 메커니즘의 정확·근사 계산 복잡도와 표현력(특히 함수 합성 및 다중 토큰 연관 작업 수행 능력)을 체계적으로 분석한다. 특히 트리‑어텐션이라는 새 변형을 제시해, 상수 차수의 함수 합성을 quadratic 시간에 정확히 수행할 수 있음을 보인다.

상세 분석

논문은 먼저 기존 자기주의가 쿼리‑키 내적을 기반으로 O(n²) 시간에 n 개의 토큰 간 쌍wise 상관을 모델링한다는 점을 재확인한다. 그러나 Match3, Parity, 함수 합성 등 t‑ary 관계가 요구되는 작업에서는 2‑order inner product만으로는 충분하지 않다. 이를 해결하기 위해 이전 연구들은 t‑tensor attention(다중 내적을 직접 계산)과 Strassen attention(내적의 선형 결합 형태) 등을 제안했지만, 그 계산 복잡도는 O(nᵗ) 또는 O(n^ω) (ω≈2.37) 로 super‑quadratic 수준에 머물렀다.

폴리어텐션은 “attention polynomial” h(x₁,…,x_t) 를 기반으로 정의된다. h는 계수가 0 또는 1인 다중선형 다항식이며, 각 monomial은 k‑order inner product을 나타낸다. 쿼리·키·밸류 행렬을 각각 Q^{(j)} 와 V^{(j)} 로 변환한 뒤,
Att_h(i)=∑{ℓ₂,…,ℓ_t} exp( h(Q^{(1)}i,…,Q^{(t)}{ℓ_t})/d ) · V^{(2)}{ℓ₂}⊙…⊙V^{(t)}_{ℓ_t}
으로 출력한다. 이 정의는

1. self‑attention (h(x₁,x₂)=x₁·x₂),
2. t‑tensor attention (h(x₁,…,x_t)=∏_{j}x_j),
3. Strassen attention (h(x₁,x₂,x₃)=x₁·x₂ + x₂·x₃ + x₃·x₁)
   을 모두 특수 경우로 포함한다.

복잡도 분석에서는 정확 알고리즘과 근사 알고리즘을 구분한다. 정확 계산은 기본적으로 O(n^t) (또는 O(n^ω) for Strassen) 가 필요하지만, 논문은 다음과 같은 두 가지 핵심 결과를 제시한다.

• 하한: fine‑grained complexity 가정(예: k‑SUM, APSP) 하에, self‑attention, t‑tensor, Strassen, 그리고 일반 폴리어텐션 모두 제시된 정확 시간 복잡도는 최적임을 증명한다. 특히 Strassen attention에 대해서는 O(n^ω) 보다 빠른 알고리즘이 존재하지 않음을 보였다.
• 근사: 입력 쿼리·키·밸류 행렬의 원소 절댓값이 B=o(√log n) (또는 B=o((log n)^{1/k}) for degree k) 범위에 있을 때, 각 어텐션을 n^{1+o(1)} 시간에 entry‑wise ε‑approximation 할 수 있다. 이는 기존의 self‑attention 근사 결과를 고차 경우까지 일반화한 것이다.

표현력 측면에서는 “tree‑attention”이라는 새로운 서브클래스를 도입한다. 여기서는 h가 degree‑2 이면서 monomial 간 연결 구조가 트리 형태를 이루는 경우를 말한다. 이러한 구조는 O(n²) 시간에 정확히 계산 가능하며, 임의의 상수 r 에 대해 r‑fold 함수 합성을 수행한다. 이는 기존 3‑tensor 혹은 Strassen attention이 2‑fold 합성만 가능하고, 3‑fold 이상은 불가능하다는 한계를 뛰어넘는다. 또한, 트리‑어텐션은 근사 알고리즘에서도 동일한 B 조건 하에 거의 선형 시간에 구현 가능하므로, 표현력과 효율성 사이의 최적 균형을 제공한다.

실험 섹션에서는 synthetic Match3 및 function‑composition 벤치마크를 사용해 트리‑어텐션이 self‑attention보다 월등히 높은 정확도를 보이며, 계산 시간은 O(n²) 에 머무른다는 것을 확인한다. 또한, 다양한 B 값에 대해 근사 알고리즘의 오차와 실행 시간을 정량화해, 이론적 복잡도 경계가 실제 구현에서도 의미 있게 작동함을 입증한다.

전체적으로 이 논문은 “어텐션 메커니즘을 다항식 형태로 일반화하고, 그 복잡도·표현력 트레이드오프를 정량화한다”는 새로운 연구 패러다임을 제시한다. 특히 트리‑어텐션은 실용적인 quadratic 시간 안에 고차 연산을 수행할 수 있는 최초의 메커니즘으로, 대규모 LLM에 고차 관계 학습을 도입하려는 연구자들에게 중요한 도구가 될 전망이다.