기술 요약: PaTH Attention

1. 문제 정의

어텐션 메커니즘은 현대 대규모 언어 모델(LLM)의 근간이지만, 본질적으로 순열 불변성(permutation-invariant)을 가집니다. 언어와 같은 구조화된 도메인을 모델링하기 위해서는 위치 인코딩이 필요합니다. Rotary Position Embedding(RoPE)이 사실상의 표준(de facto standard)이 되었으나, 이는 결정적인 한계를 지니고 있습니다. 즉, 키(key)와 쿼리(query)에 적용되는 변환이 **데이터 독립적(data-independent)**이라는 점입니다. RoPE에서 두 토큰 사이의 상호작용은 토큰 자체의 내용이 아니라 오직 상대적 거리에 의해서만 결정됩니다.

이러한 데이터 독립성은 RoPE 기반 트랜스포머의 표현력을 제한합니다. 이론적 분석에 따르면, RoPE 기반 모델은 $TC^0$ 복잡도 클래스(절대적 위치 임베딩을 사용하는 일반적인 트랜스포머와 동일)에 머물러 있으며, 이는 순차적 추론이나 상태 추적(예: 플립플롭 언어 모델링, 상태 추적 문제)이 필요한 과제를 해결하는 능력을 제한합니다. 결과적으로, 이러한 모델들은 입력 시퀀스에 대한 동적 적응을 요구하는 합성 태스크(synthetic tasks)를 수행하는 데 어려움을 겪습니다.

2. 방법론: PaTH Attention

저자들은 유연하고 데이터 의존적인 위치 인코딩 방식인 **PaTH (Position Encoding via Accumulating Householder Transformations)**를 제안합니다.

핵심 메커니즘

고정된 회전 행렬을 적용하는 RoPE와 달리, PaTH는 쿼리와 키 위치 사이의 경로를 따라 **데이터 의존적인 Householder 유사 변환(Householder-like transformations)**의 누적 곱을 사용하여 어텐션 로짓(attention logits)을 계산합니다.

• 어텐션 로짓 공식: 쿼리 $q_i$와 키 $k_j$ 사이의 어텐션 점수는 다음과 같이 정의됩니다:
  $$A_{ij} \propto \exp\left(k_j^\top \left(\prod_{s=j+1}^{i} H_s\right) q_i\right)$$
• 데이터 의여적 전이 행렬 ($H_s$): 시퀀스의 각 행렬 $H_s$는 Identity 행렬에 계수(rank-one)를 더한 형태의 Householder 유사 행렬입니다:
  $$H_s = I - \beta_s w_s w_s^\top$$
  여기서 $w_s$와 $\beta_s$는 입력 토큰 $x_s$의 함수입니다. 구체적으로 $\beta_s = 2 \times \text{sigmoid}(u^\top x_s + b)$이며, 이를 통해 전이 행렬이 음의 고윳값(eigenvalues)을 가질 수 있게 하여 상태 추적 성능에 핵심적인 역할을 수행합니다.

이론적 표현력

본 논문은 PaTH가 트랜스포머의 이론적 역량을 확장함을 입증합니다.

• 복잡도 클래스: 두 개의 어텐션 헤드를 가진 1계층 PaTH 트랜스포머는 $AC^0$ 리덕션 하에서 $NC^1$-완전(complete) 문제를 해결할 수 있습니다. 이는 PaTH가 RoPE 기반 모델을 제약하는 $TC^0$ 장벽을 깰 수 있음을 보여줍니다.
• 선형 RNN과의 관계: PaTH는 표현력이 풍부한 선형 RNN(예: DeltaNet)과 밀접하게 연관되어 있으며, 소프트맥스 어텐션의 연상 회상(associative recall) 이점을 유지하면서 상태를 추적하는 능력을 상속받습니다.

확장: PaTH-FoX

저자들은 PaTH가 다른 어텐션 수정 사항들과 호환됨을 보여줍니다. 이들은 PaTH의 곱셈적 위치 인코딩과 **Forgetting Transformer (FoX)**의 가산적(additive) 데이터 의존적 망각 게이트를 결합한 PaTH-FoX를 제안합니다. 이 하이브리드 접근 방식은 특히 길이 외삽(length extrapolation) 성능을 더욱 향상시킵니다.

3. 효율적인 학습 및 추론

PaTH의 주요 과제는 전체 $L \times L$ 어텐션 행렬을 실체화하지 않고도 Householder 행렬의 누적 곱을 효율적으로 계산하는 것입니다.

• UT 변환 표현: 저자들은 Householder 행렬의 곱을 압축하여 표현하기 위해 UT 변환을 활용합니다. $L$개의 변환의 곱은 $P = I - W^\top T^{-1} W$로 표현될 수 있으며, 여기서 $W$는 벡터 $w_s$를 포함하고 $T$는 $\beta$ 값들로부터 유도된 삼각 행렬입니다.
• FlashAttention 스타일 알고리즘: 하드웨어 효율적인 학습을 위해 저자들은 블록 단위 알고리즘을 개발했습니다:
  • 각 블록에 대해 "경계 조정된(boundary-adjusted)" 쿼리 및 키 행 matrices를 사전 계산합니다.
  • 키/값 블록에 대해 우측에서 좌측으로 스캔(right-to-left scan)을 수행하여 스트리밍 방식으로 누적 Householder 곱을 축적합니다.
  • 이 방식은 동적 위치 인코딩을 포함하면서도 (SRAM 내에서 작동하는) FlashAttention의 I/O 효율성을 유지합니다.
• 추론: 추론 시에는 현재 타임스텝의 전이 행렬을 사용하여 기존 키들을 인플레이스(in-place)로 업데이트하므로, 별도의 변환 벡터 캐시가 필요하지 않습니다. 이를 통해 기존의 FlashDecoding과 같은 추론 커널과 호환성을 유지할 수 있습니다

4. 실험 결과

합성 벤치마크

PaTH는 순차적 추론 및 상태 추적을 조사하기 위해 설계된 태스크에서 평가되었습니다:

• 플립플롭 언어 모델링 (FFLM): PaTH 기반 모델(1계층, 2헤드)은 이 태스크에서 완벽에 가까운 정확도를 달성한 반면, RoPE 기반 모델은 크게 고전했습니다.
• 수리 문제 ($A_5$ 군): PaTH는 베이스라인(RoPE, SBA, FoX)보다 적은 계층으로 $NC^1$-완전 수리 문제를 해결했습니다.
• Multi-query Repeated Associative Recall (MQRAR-N): PaTH는 $N$-back 회상($N < 4$)을 위해 변수 값을 성공적으로 추적하여, 정렬된 이력을 추적하는 데 실패한 최근 베이스라인들을 능가했습니다.

언어 모델링 (760M 파라미터)

Fineweb-Edu 코퍼스(50B 토큰)를 대상으로 실험이 진행되었습니다:

• 표준 벤치마크: PaTH는 제로샷 상식 추론 태스크(LAMBADA, PIQA, HellaSwag 등)에서 RoPE 및 FoX를 일관되게 앞질렀으며, Wikitext에서 더 낮은 퍼플렉시티(perplexity)를 달성했습니다.
• 길이 외삽 (Length Extrapolation): 4K 컨텍스트로 학습된 모델은 더 긴 시퀀스에서도 잘 일반화되었습니다.
  • RoPE: 4K를 넘어서면 급격히 실패했습니다.
  • PaTH: 32K까지 안정적인 성능을 유지했습니다.
  • PaTH-FoX: 상태 추적이 중요한 코드 생성(CodeParrot)에서 특히 강점을 보이며 64K 토큰까지 견고하게 일반화되었습니다.
• 롱 컨텍스트 벤치마크: RULER, BABILONG, PhoneBook에서 PaTH-FoX는 가장 높은 검색 및 상태 추적 점수를 기록하며, Multi-Needle-In-A-Haystack 및 변수 추적 태스크에서 RoPE와 FoX를 크게 압도했습니다.

모델 변환

저자들은 사전 학습된 RoPE 모델을 PaTH로 변환하는 실험을 탐구했습니다:

• SmolLM2-1.7B: 고품질 수학/코드 데이터에 노출되기 전에 모델을 변환한 후 지속 학습을 진행했을 때, PaTH가 원래의 RoPE 모델보다 수학(GSM-8K) 및 코딩(HumanEval, MBMBP) 벤치마크에서 더 우수한 성능을 보였습니다.
• Qwen2.5-7B: 이미 과학습된(overtrained) RoPE 모델로부터 증류(distillation)한 결과는 혼재되어 나타났습니다. 일부 성능 향상은 있었으나 특정 벤치마크에서의 저하도 관찰되었는데, 이는 변환이 모델이 "경직(ossify)"되기 전에 수행되는 것이 가장 효과적임을 시사합니다.

5. 의의 및 주장

본 논문은 데이터 의존적 곱셈 변환을 도입함으로써 현재의 위치 인코딩 방식이 가진 이론적, 경험적 한계를 해결한다고 주장합니다.

• 표현력: 트랜스포머가 $NC^1$-완전 문제를 해결할 수 있도록 하여, $TC^0$ 복잡도 클래스를 넘어설 수 있는 원칙적인 방법을 제공합니다.
• 실용성: UT 변환과 블록 단위 알고리즘을 통해 PaTH는 표준 어텐션 메커니즘과 대등한 효율성을 달성하여 대규모 학습이 가능하도록 합니다.
• 상태 추적: Householder 변환의 데이터 의존적 특성 덕분에 모델은 입력 시퀀스에 동적으로 적응할 수 있으며, RoPE 기반 모델로는 다룰 수 없는 상태 추적 태스크를 해결합니다.

6. 한계점

저자들은 몇 가지 주의 사항을 인정합니다:

• 수치적 안정성: 학습 시 정밀도에 민다할 수 있습니다. 특히 BF16 환경에서는 누적 곱 과정에서 불안정성이 발생할 수 있어, 스케일링 인자 $\beta$에 대한 클리핑(clipping)이 필요합니다.
• 연산 오버헤드: 효율적이긴 하지만, Householder 곱 계산은 표준 어텐션에 비해 더 큰 헤드 차원(head dimension)을 가질 경우 계산 및 메모리 오버헤드를 증가시킵니다.
• 기하학적 편향: PaTH는 (Householder 행렬이 반사(reflection)를 수행하므로) 회전을 네이티브하게 모델링하지는 않습니다. 이는 RoPE가 가진 상대적 위치에 대한 구조적 의존성과 같은 특정 기하학적 귀납 편향(inductive bias)을 제한할 수 있습니다. 다만 저자들은 향후 연구를 통해 반사의 합성을 통해 회전을 근사할 수 있을 것이라고 제안합니다.