슬라이스 ReLU 어텐션: 정렬 기반 준선형 컨텍스트 표현력

초록

**
슬라이스 ReLU 어텐션은 키와 쿼리 차이를 1차원으로 투사한 뒤 ReLU와 정렬을 이용해 O(n log n) 시간에 전역 어텐션을 계산한다. 비대칭이지만 조건부 양정정성을 갖는 커널을 사용해 소프트맥스와 동일한 이론적 표현력을 유지하며, 정렬 기반 구현으로 긴 시퀀스에서도 효율적이다. 실험에서는 Long‑Range Arena, 포인트‑클라우드 분류, 소규모 언어 모델 사전학습 등에서 기존 방법과 경쟁력 있는 성능을 보였다.

**

상세 분석

**
본 논문은 기존 트랜스포머에서 사용되는 소프트맥스 어텐션이 갖는 O(n²) 복잡도를 근본적으로 탈피하고자 한다. 핵심 아이디어는 (i) 키‑쿼리 차이를 고정된 1차원 방향으로 투사하고, (ii) 투사값들의 차이에 ReLU 함수를 적용한 비대칭 커널을 정의하며, (iii) 이 커널을 정렬(sorting)과 누적합(cumulative sum) 연산만으로 O(n log n) 시간에 정확히 계산한다는 점이다.

1. 커널 정의와 수학적 성질

   • 투사 연산 Π: ℝᵈ → ℝ 은 선형 또는 작은 MLP 형태로 구현된다.
   • 핵심 커널: K(x, y) = ReLU(ΠQx − ΠKy). 이는 비대칭이지만, 절대값 정규화 |ΠQx − ΠKy| 로 나누어 주면 전체 어텐션 스코어가 양수이며, 조건부 양정정성(conditional positive definiteness) 특성을 갖는다. 이는 Energy Distance 커널과 동일한 1차원 형태이며, γᵀKγ ≥ 0 가 γ의 합이 0일 때 성립한다.

2. 정렬 기반 O(n log n) 알고리즘

   • 모든 투사값 z₁,…,z₂ₙ을 정렬하고, 각 위치 i에 대해 누적합 aᵢ = Σ_{j≤i} γⱼ, bᵢ = Σ_{j≤i} γⱼ zⱼ 를 선형 시간에 구한다.
   • ReLU 차이 합은 aᵢ·zᵢ − bᵢ 로 표현되므로, 정렬 후 한 번의 스캔만으로 전체 어텐션 행렬을 얻을 수 있다.
   • 쿼리와 키를 하나의 배열에 합쳐 정렬하면, 쿼리‑키 쌍에 대한 모든 상호작용을 한 번에 계산한다.

3. 값(Value) 중심화

   • 실험에서 값 벡터를 평균 0 으로 중심화하면 학습이 안정된다. 이는 ReLU 커널이 γᵀKγ ≥ 0 를 보장하려면 γ가 평균 0이어야 하는 조건과 일치한다.

4. ReLU‑Bump 변형

   • 보다 국소적인 상호작용을 원할 경우, 세 개의 시프트된 ReLU 커널을 선형 결합해 삼각형(‘hat’) 형태의 bump 커널을 만든다.
   • 이 변형은 여전히 정렬‑스캔 방식으로 O(n log n) 에 계산 가능하고, 1차원에서는 양정정성을 유지한다.

5. 표현력 이론

   • 논문은 두 가지 주요 정리를 제시한다. 첫째, 소프트맥스 어텐션이 증명된 “시퀀스‑투‑시퀀스 디스엔탱글링” 능력을 슬라이스 ReLU 어텐션도 동일하게 수행한다. 둘째, 컨텍스트 유니버설 어프로시메이션(Universal Approximation) 정리를 확장해, 충분히 깊은 슬라이스 ReLU 트랜스포머가 임의의 연속적인 시퀀스‑투‑시퀀스 함수들을 근사할 수 있음을 보인다. 깊이만 증가시키면 차원 d 와 헤드 수 H 에 대한 제한이 사라진다.

6. 실험 결과

   • Long‑Range Arena: 1k~4k 길이 토큰에서 기존 선형 어텐션보다 높은 정확도와 비슷한 메모리 사용량을 기록.
   • PointNet++ 기반 포인트 클라우드 분류: ModelNet40에서 ReLU‑Bump 변형이 특히 국소 구조를 잘 포착해 기존 슬라이스 어텐션보다 1~2% 향상.
   • 소규모 언어 모델: BERT‑like 인코더를 슬라이스 ReLU 어텐션으로 교체하고 2‑epoch 사전학습 후 GLUE 벤치마크에서 평균 1.3% 상승.

7. 한계와 향후 연구

   • 현재는 투사 방향을 학습된 MLP 하나에 의존하고 있어, 다중 방향(다중 슬라이스) 확장이 필요하다.
   • 정규화 방식이 절대값 기반이지만, 확률적 정규화(softmax와 유사)와의 성능 차이를 체계적으로 분석할 여지가 있다.
   • 대규모 사전학습에서의 수렴 속도와 최적화 역학을 이해하기 위한 이론적 분석이 아직 부족하다.

**