Knocking-Heads 어텐션

Knocking-Heads Attention

October 27, 2025

저자: Zhanchao Zhou, Xiaodong Chen, Haoxing Chen, Zhenzhong Lan, Jianguo Li

cs.AI

초록

다중 헤드 어텐션(MHA)은 병렬 어텐션 헤드를 통해 표현 능력을 향상시키며 현대 대규모 언어 모델의 핵심 요소로 자리 잡았다. 그러나 헤드 수를 증가시키는 것은 본질적으로 개별 헤드의 능력을 약화시키며, 기존의 어텐션 메커니즘(표준 MHA 또는 GQA, GTA와 같은 변형 포함)은 강력한 상호작용 없이 분리된 헤드들의 출력을 단순히 연결한다. 이러한 한계를 해결하기 위해 우리는 노킹 헤드 어텐션(KHA)을 제안한다. KHA는 어텐션 헤드들이 서로 "노크"하여, 스케일드 닷-프로덕트 어텐션 이전에 헤드 간 특징 수준 상호작용을 촉진한다. 이는 모든 헤드에 걸쳐 공유되고 대각선으로 초기화된 프로젝션 행렬을 적용하여 달성된다. 대각선 초기화는 훈련 시작 시 헤드별 특화 능력을 보존하면서 모델이 점진적으로 통합된 헤드 간 표현을 학습할 수 있게 한다. KHA는 최소한의 매개변수와 FLOPs만 추가하며 MHA, GQA, GTA 및 기타 어텐션 변형에 원활하게 통합될 수 있다. 우리는 1조 개의 고품질 토큰으로 61억 개의 매개변수(10.1억 개 활성화)를 가진 MoE 모델을 학습하여 KHA를 검증했다. 기준 어텐션 메커니즘과 비교했을 때, KHA는 더 우수하고 안정적인 훈련 동역학을 보여주며 다운스트림 작업 전반에서 더 나은 성능을 달성했다.

English

Multi-head attention (MHA) has become the cornerstone of modern large language models, enhancing representational capacity through parallel attention heads. However, increasing the number of heads inherently weakens individual head capacity, and existing attention mechanisms - whether standard MHA or its variants like grouped-query attention (GQA) and grouped-tied attention (GTA) - simply concatenate outputs from isolated heads without strong interaction. To address this limitation, we propose knocking-heads attention (KHA), which enables attention heads to "knock" on each other - facilitating cross-head feature-level interactions before the scaled dot-product attention. This is achieved by applying a shared, diagonally-initialized projection matrix across all heads. The diagonal initialization preserves head-specific specialization at the start of training while allowing the model to progressively learn integrated cross-head representations. KHA adds only minimal parameters and FLOPs and can be seamlessly integrated into MHA, GQA, GTA, and other attention variants. We validate KHA by training a 6.1B parameter MoE model (1.01B activated) on 1T high-quality tokens. Compared to baseline attention mechanisms, KHA brings superior and more stable training dynamics, achieving better performance across downstream tasks.