토큰 임베딩이 어텐션의 기하학적 기반임을 증명하며 No-Q 어텐션 제안

한줄 요약

토큰 임베딩이 단순한 조회 테이블이 아니라 어텐션의 질의(Query) 역할을 한다는 가설을 제시하고, 이를 바탕으로 No-Q 어텐션을 도입하여 성능 향상과 파라미터 감소를 달성했다.

풀어야 하는 문제

트랜스포머에서 임베딩 레이어는 보통 단순한 입력 변환으로 간주되지만, 저자들은 임베딩이 어텐션의 기하학적 구조에 결정적 영향을 미친다고 주장한다. 기존 연구는 임베딩의 역할을 간과하거나, 푸리에 도메인에서의 조작이 언어 모델에 안전하지 않다는 문제가 있었다.

접근 방법

네 단계로 접근한다. (1) Prescribed Fourier Frequency Training (PFFT): 임베딩 그래디언트에 나이퀴스트 주파수 모드를 부과하여 grokking 가속화. (2) Sounding Hammer 진단: BPE 어휘의 그래디언트가 스펙트럼 평탄하여 푸리에 조종이 안전하지 않음을 발견하고 Natural Ordering Conditions (NOC) 제안. (3) Fourier Gradient Projection (FGP): 훈련 중 중요한 주파수 모드를 동적으로 추적하는 일반적 그래디언트 도구. (4) 행동 가중치 궤적 분석: 모든 가중치 행렬(Q, K, V, MLP)이 임베딩으로부터 동일한 two-arm 궤적을 상속받음을 발견하고, 이를 바탕으로 No-Q 어텐션(Query 투영 제거)을 제안.

결과·기여

PFFT는 모듈러 산술에서 grokking 에폭을 92.7% 감소시켰다. No-Q 어텐션은 TinyStories에서 BPC 3.18%, FineWeb에서 2.24% 개선하고 파라미터 8% 감소, grokking 51% 가속화했다. 이 연구는 임베딩이 어텐션의 질의 역할을 한다는 새로운 관점을 제시하며, 향후 트랜스포머 설계에 중요한 통찰을 제공한다.