멀티헤드 어텐션의 용량 기반 설계 원리

초록

본 논문은 자기‑어텐션의 키‑쿼리 채널이 가질 수 있는 용량을 정량화한다. 전체 키 차원 (D_K = h,d_k) 을 고정했을 때, 하나의 레이어가 얼마나 많은 토큰‑토큰 관계(그래프의 엣지)를 정확히 인코딩할 수 있는지를 “Relational Graph Recognition(RGR)”이라는 프레임워크로 분석한다. 정보‑이론적 하한과 실제 구현을 통한 상한을 제시하며, 특히 헤드 수를 늘려 (D_K) 를 여러 작은 헤드에 분산시키면 임베딩 중첩에 의한 간섭이 감소해 용량이 크게 향상된다는 새로운 다중‑헤드 어텐션의 근거를 제공한다. 실험은 이론적 예측을 검증하고, 소프트맥스, 값‑루팅, 전체 트랜스포머 블록까지 확장했을 때도 동일한 현상이 나타남을 보여준다.

상세 분석

이 논문은 트랜스포머의 핵심 메커니즘인 자기‑어텐션을 “어디에 집중할 것인가”라는 질문으로 추상화하고, 이를 그래프 복원 문제인 Relational Graph Recognition(RGR)으로 공식화한다. RGR에서는 정점 (V) 와 방향성 엣지 (E) 를 가진 그래프 (G=(V,E)) 를 임베딩 (x_v\in\mathbb{R}^{d_{\text{model}}}) 에 매핑한 뒤, 임의의 컨텍스트 (C\subseteq V) 에 대해 각 정점 (v\in C) 의 인‑컨텍스트 이웃 (N_G(v;C)={v’\in C\mid (v,v’)\in E}) 을 정확히 예측하도록 키‑쿼리 채널을 설계한다.

핵심 자원은 전체 키 차원 (D_K = h d_k) 이며, 여기서 (h) 는 헤드 수, (d_k) 는 각 헤드의 키(및 쿼리) 차원이다. 논문은 두 가지 모델 변형을 다룬다. 첫 번째는 “max‑over‑heads” 방식으로, 각 헤드가 독립적인 스코어 (S^{(k)}_{pq}=q^{(k)}_p\cdot k^{(k)}q) 를 계산하고, 최종 스코어를 (\max_k S^{(k)}{pq}) 로 집계한다. 이는 소프트맥스의 경쟁적 라우팅을 최소한의 비선형성으로 모사하면서 분석을 용이하게 만든다. 두 번째는 실제 트랜스포머에서 사용되는 스케일드 소프트맥스이며, 여기서는 각 헤드의 확률 분포를 합산한다.

정보‑이론적 하한은 모든 그래프 (G) 에 대해 관계 수 (m’ = |E|) 를 복원하려면
\