PaTH 어텐션: 누적 하우스홀더 변환으로 구현하는 데이터‑의존 위치 인코딩

초록

PaTH는 입력에 따라 동적으로 변하는 하우스홀더‑유사 행렬을 누적 곱으로 결합해 위치 정보를 인코딩한다. 기존 RoPE가 상대 위치만을 반영하는 반면, PaTH는 입력 내용까지 고려해 표현력을 크게 확장한다. UT 변환을 이용한 압축 표현과 FlashAttention‑스타일 블록 연산을 통해 학습·추론 비용을 기존와 동등하게 유지하면서, 합성 추론 과제와 760M 규모 언어 모델에서 RoPE·FoX를 능가하는 성능을 보인다. 또한 사전 학습된 RoPE 모델을 PaTH로 전이시켜 추가 학습만으로도 개선이 가능하다.

상세 분석

PaTH Attention은 “데이터‑의존적 위치 인코딩”이라는 새로운 패러다임을 제시한다. 핵심 아이디어는 각 토큰 t에 대해 (H_t = I - \beta_t w_t w_t^\top) 형태의 하우스홀더‑유사 변환을 정의하고, 쿼리‑키 쌍 ((i,j)) 사이의 위치 변환을 (H_{j+1} H_{j+2}\dots H_i) 의 누적 곱으로 표현한다는 점이다. 여기서 (w_t)와 (\beta_t)는 현재 입력 (x_t) 에 의해 비선형(시그모이드)으로 결정되므로, 위치 변환이 입력 내용에 따라 실시간으로 변한다. 이는 RoPE가 고정된 회전 행렬 (R^{i-j}) 만을 사용해 상대 위치를 인코딩하는 것과 근본적으로 다르다.

이러한 동적 변환을 효율적으로 구현하기 위해 논문은 두 가지 기술적 공헌을 제시한다. 첫째, 하우스홀더 행렬들의 곱을 (I - W^\top T^{-1} W) 형태의 UT 변환으로 압축한다. 여기서 (W) 는 모든 (w_t) 를 행으로 쌓은 행렬, (T^{-1}) 는 (L\times L) 삼각 행렬의 역으로, 삼각 연산과 행렬 곱만으로 전체 누적 변환을 한 번에 구할 수 있다. 둘째, 이 UT 표현을 블록 단위로 마스킹해 (H_{j+1}\dots H_i) 를 부분 구간마다 재사용한다. 이를 FlashAttention‑style의 블록‑와이즈 스캔과 결합하면, 각 블록 쌍에 대해 (O(B^2 d + B d^2)) 의 연산량으로 전체 (O(L^2 d + L d^2 / B)) 복잡도를 유지하면서도 메모리 이동을 최소화한다.

이론적으로 PaTH는 한 층, 두 헤드, (\log n) 정밀도 설정에서 NC¹‑완전 문제를 AC⁰‑감소를 통해 해결할 수 있음을 증명한다. 이는 RoPE 기반 트랜스포머가 TC⁰에 머무르는 한계를 넘어서는 것으로, 데이터‑의존적 변환이 순차적 상태 추적 능력을 크게 강화함을 의미한다. 실험에서는 “flip‑flop” 언어 모델링, 상태 추적 벤치마크 등 TC⁰‑한계에 걸리는 합성 작업에서 PaTH가 RoPE와 Forget‑Transformer(FoX)를 모두 앞선다. 760M 파라미터 규모의 실제 언어 모델에서도 Perplexity 감소와 길이 외삽 능력 향상을 입증하였다.

또한 PaTH는 FoX와 자연스럽게 결합될 수 있다. FoX는 로그‑스케일의 데이터‑의존적 “forget” 게이트 (f_s) 를 추가해 (A_{ij} \propto \exp(k_j^\top q_i) \cdot \prod_{s=j+1}^{i} f_s) 형태로 변형한다. PaTH‑FoX는 두 변환을 곱해 (A_{ij} \propto \bigl(\prod_{s=j+1}^{i} f_s\bigr) \exp!\bigl(k_j^\top H_{j+1}\dots H_i q_i\bigr)) 로 만들며, 실험에서 특히 긴 시퀀스에 대한 일반화가 크게 개선된다.

마지막으로, 사전 학습된 RoPE 모델을 PaTH로 변환하는 방법을 제시한다. 기존 파라미터를 그대로 유지하면서 (H_t) 를 새롭게 학습하고, 짧은 추가 학습 단계만 거치면 RoPE 기반 모델보다 일관된 성능 향상을 얻을 수 있다. 이는 PaTH가 기존 인프라와 호환 가능함을 보여주며, 실제 서비스에 적용하기 위한 진입 장벽을 낮춘다.

요약하면, PaTH는 하우스홀더‑유사 행렬을 데이터‑의존적으로 누적해 위치 인코딩을 수행하고, UT 변환과 블록‑와이즈 알고리즘을 통해 효율성을 확보한다. 이 설계는 이론적 복잡도 확장(NC¹)과 실험적 성능 향상(합성·실제 언어 모델) 모두를 동시에 달성한다는 점에서 트랜스포머 연구에 중요한 전진을 의미한다.