Attention à Têtes Frappantes

Knocking-Heads Attention

October 27, 2025

papers.authors: Zhanchao Zhou, Xiaodong Chen, Haoxing Chen, Zhenzhong Lan, Jianguo Li

cs.AI

papers.abstract

L'attention multi-têtes (MTA) est devenue la pierre angulaire des grands modèles de langage modernes, améliorant la capacité de représentation grâce à des têtes d'attention parallèles. Cependant, l'augmentation du nombre de têtes affaiblit intrinsèquement la capacité individuelle de chaque tête, et les mécanismes d'attention existants - qu'il s'agisse de la MTA standard ou de ses variantes comme l'attention à requêtes groupées (GQA) et l'attention liée groupée (GTA - se contentent de concaténer les sorties de têtes isolées sans interaction forte. Pour remédier à cette limitation, nous proposons l'attention par "cognement de têtes" (KHA), qui permet aux têtes d'attention de "cogner" les unes contre les autres - facilitant des interactions au niveau des caractéristiques entre les têtes avant l'attention par produit scalaire mis à l'échelle. Ceci est réalisé en appliquant une matrice de projection partagée, initialisée de manière diagonale, à toutes les têtes. L'initialisation diagonale préserve la spécialisation spécifique à chaque tête au début de l'entraînement tout en permettant au modèle d'apprendre progressivement des représentations intégrées entre les têtes. KHA n'ajoute que des paramètres et des FLOPs minimaux et peut être intégré de manière transparente dans la MTA, la GQA, la GTA et d'autres variantes d'attention. Nous validons KHA en entraînant un modèle MoE de 6,1 milliards de paramètres (1,01 milliard activé) sur 1 000 milliards de tokens de haute qualité. Par rapport aux mécanismes d'attention de référence, KHA offre une dynamique d'entraînement supérieure et plus stable, obtenant de meilleures performances sur diverses tâches en aval.

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.