기술 요약: 커널화된 선형 어텐션(Kernelized Linear Attention): 대칭 원뿔을 통한 용량의 벽 돌파

1. 문제 정의

선형 어텐션 메커니즘은 표준 소프트맥스(softmax) 어텐션의 $O(T^2)$ 학습 비용과 $O(T)$ KV-캐시 메모리 문제를 해결하며, 상수 시간의 순환 추론(constant-time recurrent inference)과 고정 크기의 상태 표현을 제공한다. 그러나 기존의 선형 어텐션 모델들은 시퀀스 길이가 증가함에 따라 연상 회상(associative recall) 능력이 급격히 저하되는 문제를 겪는다. 이러한 실패는 "메모리 충돌(memory collision)"에 기인하는데, 이는 많은 키-값 바인딩(key-value bindings)이 하나의 고정된 상태를 공유해야 함으로써 검색 성능이 붕괴하기 때문이다.

현재의 게이트형 선형 RNN(예: DeltaNet, Mamba2, Gated DeltaNet)은 학습된 망각 게이트(forget gates)와 콘텐츠 주소 지정 삭제(content-addressed erase) 연산을 통해 이를 완화하려고 시도한다. 하지만 이들의 특징 기하학(feature geometries)이 갖는 이론적 한계는 여전히 불투명하다. 구체적으로 다음과 같은 사항들에 대한 원칙적인 이해가 부족하다:

• 고정된 상태 크기에서 회상 용량을 제한하는 요소는 무엇인가?
• 어떤 유한 특징 맵(finite feature maps)이 소프트맥스의 등방성(isotropy)을 보존하는가?
• 특징 맵은 반드시 비음수(nonnegative)여야 하는가? 만약 그렇다면, 어떤 기하학이 허용되는가?

2. 방법론

2.1 이론적 프레임워크: 대칭 원뿔(Symmetric Cones) 및 KATA

저자들은 어텐션 회상을 **구형 패킹 문제(spherical-packing problem)**로 정식화한다. 이들은 **커널화된 선형 어텐션 활성화(Kernelized Linear Attention Activations, KATA)**를 도입하며, 이는 **자기 쌍대 동차 원뿔(self-dual homogeneous cones)**을 통해 비음수 어텐션 가중치를 인증함으로써 제1원리로부터 도출된 프레임워크이다.

• 기하학적 가정: 이 프레임워크는 (1) 자기 쌍대 원뿔($K = K^*$)을 통한 비음수 어텐션 가중치, (2) 특징 맵의 등거리 불변성(isometric invariance), (3) 잠재 기하학의 동차성(homogeneity)을 가정한다.
• Koecher–Vinberg 분류: 이러한 가정 하에서 허용 가능한 특징 기하학은 Koecher–Vinberg 정리의 기약 인자(irreducible factors)로 제한된다: 양의 직교 원뿔(positive orthant), 로렌츠 원뿔(Lorentz cone), 그리고 대칭 양의 준정부호(PSD) 원뿔이다.
• 출력 게이팅: 저자들은 정규화된 선형 어텐션의 분모가 **매개변수가 없는, 토큰 조건부 볼록 출력 게이트(parameter-free, token-conditioned convex output gate)**로 자연스럽게 도출됨을 증명한다. 이 게이트는 이전 리드아웃(readout)과 현재 값 사이의 볼록 보간(convex interpolation)을 수행하며, 기본 공식에서 명시적인 학습된 게이팅 메커니즘의 필요성을 제거한다.

2.2 용량 이론(Capacity Theory)

본 논문은 각 원뿔 기하학에 대한 "연상 용량(associative capacity, 구별 가능한 키 주소의 수)"을 분석한다:

• 직교 원뿔 및 로렌츠 원뿔: 이 기하학들은 **웰치 경계(Welch bound)**와 **랭킨 벽(Rankin wall)**의 영향을 받는다. 이들은 선형 용량($O(d)$)을 제공하거나 간섭 임계값(interference thresholds)에 의해 엄격히 제한된다 (예: 로렌츠 간섭은 상호 둔각인 키를 강제하지 않고서는 $1/2$ 미만으로 떨어질 수 없다).
• PSD 원뿔 (Rank-One): 저자들은 랭크-원 양의 준정부호(rank-one PSD) 특징 $\psi(u) = uu^\top$를 사용하는 것을 제안한다. 이 맵은 가공되지 않은 내적($\langle u_i, u_j \rangle$)을 제곱하여 간섭 기하학을 변환한다.
  • 이론적 결과: 가공되지 않은 벡터들은 웰치 하한(Welch floor)의 영향을 받지만, 제곱된 간섭은 임의의 고정된 허용 오차 $\epsilon > 0$에 대해 투영 차원 $p$에 대한 **지수적 용량(exponential capacity)**을 허용한다. 구체적으로, 용량은 $T \geq \frac{1}{2}(1-\epsilon)^{-(p-1)/2}$로 스케일링된다.
  • 신호 대 잡음비: 저자들은 pSNR(power signal-to-noise ratio) 분석을 통해, PSD 맵이 이상적인 검색 임계값을 $\Theta(p)$(선형)에서 $\Theta(p^2)$로 확장하여, 파라미터를 추가하지 않고도 효과적으로 상태를 확장함을 보여준다.

2.3 하드웨어 정렬 구현

이러한 이론적 이득을 실현하기 위해, 저자들은 NVIDIA GPU 상에서 융합된 Triton 커널을 사용하여 KATA를 구현한다.

• 청크 병렬성(Chunked Parallelism): 시퀀스를 청크로 나눈다. 청크 내부(intra-chunk) 어텐션은 텐서 코어를 사용하며, 청크 간(inter-chunk) 상태 전파는 **결합적 스캔(associative scan, parallel prefix sum)**을 통해 처리된다.
• 복잡도 감소: 이를 통해 청크 간 순환 깊이를 $O(N_C)$(순차적)에서 $O(\log N_C)$로 줄인다 (여기서 $N_C$는 청크의 수).
• 두 가지 동작 지점:
  1. 이차 시간($O(T^2)$) 순방향: 전체 상태를 생성하지 않는 FlashAttention 스타일의 커널로, FlashAttention-2보다 최대 1.6배 높은 처리량을 달성한다.
  2. 선형 시간($O(T)$) 상태: 정확한 청크 상태 형태를 사용하여 131k 토큰에서 FlashAttention-2 대비 $\sim$11배의 처리량에 도달한다.
• 변형 모델: 프레임워크에는 상태 크기와 패킹 효율성을 조절하기 위한 KATA-Mg(다중 PSD 인자) 및 KKATA-$\Sigma$g(합산된 PSD 인자)가 포함된다.

3. 주요 기여

1. 선형 어텐션을 위한 원뿔 분류: 본 논문은 비음수 가중치와 등거리 불변성 제약 조건 하에서 허용 가능한 비음수 특징 기하학(직교 원뿔, 로렌츠 원뿔, PSD 원뿔)을 규명하여, 선형 어텐션을 Koecher–Vinberg 분류에 근거하여 정립하였다.
2. 연상 회상 용량 이론: 회상을 이산 구형 패킹(discrete spherical packing)으로 간주함으로써, 저자들은 로렌츠 랭킨 벽과 rank-one PSD 광선(rays)의 차원 의존적 패킹 경계를 증명하였다. 이들은 rank-one PSD 특징이 웰치 하한 아래에서 지수적 용량 이점을 제공함을 입증하였다.
3. 하드웨어 최적화 커널: 저자들은 KATA 순환을 위한 융합된 Triton 커널을 공개한다. 이 커널은 결합적 스캔을 활용하여 청크 간 깊이를 낮춤으로써, 순차적 선형 어텐션 베이스라인 및 긴 컨텍스트에서의 FlashAttention-2보다 훨씬 높은 처리량을 달성한다.
4. 실험적 검증: 합성 태스크(MQAR, 반복 키 덮어쓰기, 유도) 및 340M 파라미터 언어 모델 실험을 통해, KATA 변형들(특히 PSD 원뿔 기반)이 Gated DeltaNet 및 기타 선형 베이스라인보다 연상 회상 능력이 뛰어나며, 16배의 분포 외(out-of-distribution) 길이 외삽에서도 거의 완벽한 회상 성능을 유지함을 보여준다.

4. 결과

• 합성 벤치마크:
  • MQAR (Multi-Query Associative Recall): KATA-M1은 훈련 컨텍스트 길이의 16배에서도 0.985의 정확도를 유지하며, KV-캐시 항목을 약 4분의 1만 사용하면서도 소프트맥스 성능에 근접한다.
  • 반복 키 덮어쓰기 (Repeated-Key Overwrite): 델타 규칙(delta rule)을 적용한 변형(DeltaKATA)은 반복되는 주소에서의 오래된 바인딩을 수정하며 거의 완벽한 덮어쓰기 정확도를 달성한다.
  • Needle-in-a-Haystack (NIAH): 고엔트로피 태스크(UUID)에서 KATA-M2 및 DeltaKATA-M1은 회상 성능이 거의 0(0.004)으로 떨어지는 Gated DeltaNet보다 현저히 우수한 성능을 보인다.
• 언어 모델 사전 학습:
  • 사전 학습된 340M 파라미터 모델은 PSD 변형들이 표준 벤치마크(ARC, PIQA 등)에서 경쟁력 있는 제로샷 정확도를 유지하면서도, "유창성 트레이드오프"(소프트맥스보다 퍼플렉시티는 높지만 다른 선형 모델보다 회상 능력은 좋음)를 보임을 보여준다.
  • 유도 (Induction): KATA는 선형 어텐션이 흔히 실패하는 작업인 유도 능력에서 거의 완벽한 성능을 보존한다.
• 성능:
  • 선형 상태 KATA-M2 커널은 131k 토큰에서 FlashAttention-2보다 $\sim$11배 빠른 순방향 처리량을 달성한다.
  • 이차 시간 KATA-M1 순방향은 FlashAttention-2 대비 $\sim$1.6배의 처리량에 도달한다.
  • 결합적 스캔 구현은 매칭된 순차적 선형 어텐션 청크보다 평균 $\sim$2.4배 높은 처리량을 기록한다.

5. 의의 및 주장

본 논문은 비음수 어텐션을 위한 올바른 기하학적 기질(대칭 원뿔, 특히 PSD 원뿔)을 식별함으로써 선형 어텐션의 "용량 벽"을 해결한다고 주장한다.

• 이론적 통찰: 저자들은 선형 어텐션의 상수 상태(constant state)에 따르는 대가는 선형성 자체의 문제가 아니라 특징 기하학의 선택에 달려 있음을 밝혀냈다. rank-one PSD 특징으로 이동함으로써, 파라미터 수를 늘리지 않고도(단, 상태 차원만 증가) 용량을 지수적으로 증가시킬 수 있다.
• 실질적 영향: 이 연구는 하드웨어와 정렬된 특징 맵(융합 커널을 통해)과 이론적으로 건전한 기하학(대칭 원뿔을 통해)을 제공한다면, 선형 시간 추론 및 고정 크기 상태를 사용하면서도 높은 용량의 연상 회상이 가능하다는 것을 보여준다.
• 한계점: 저자들은 기하학적 용량은 높지만, **학습 가능성(learnability)**이 변수라고 언급한다. 즉, 모델이 고차원 구형 패킹을 학습하는 데 어려움을 겪을 수 있다(예: KATA-M1은 이론적 용량은 더 크지만, 훈련 역학 때문에 문맥적 태스크에서 KATA-M2보다 성능이 떨어질 때가 있음). 또한, 현재 구현은 실수 원뿔(real-valued cones)에 집중되어 있어 향 이상의 복소수 확장을 향후 과제로 남겨두었다.

요약하자면, KATA는 연상 메모리의 이론적 한계와 하드웨어 효율적인 추론의 실제적 제약 사이를 원칙적으로 연결하며, 장거리 회상을 희생하지 않는 선형 시간 트랜스포머로 가는 길을 제시한다.