Knock-Head-Attention

papers.abstract

Multi-Head-Attention (MHA) hat sich zum Eckpfeiler moderner großer Sprachmodelle entwickelt und erhöht die Repräsentationskapazität durch parallele Aufmerksamkeitsköpfe. Die Erhöhung der Anzahl der Köpfe schwächt jedoch inhärent die Kapazität einzelner Köpfe, und bestehende Aufmerksamkeitsmechanismen – ob Standard-MHA oder Varianten wie Grouped-Query-Attention (GQA) und Grouped-Tied-Attention (GTA) – verketteten lediglich die Ausgaben isolierter Köpfe ohne starke Interaktion. Um diese Einschränkung zu adressieren, schlagen wir Knocking-Heads-Attention (KHA) vor, das es Aufmerksamkeitsköpfen ermöglicht, aneinander „anzuklopfen“ und so feature-basierte Interaktionen zwischen Köpfen vor der skalierten dot-product-Attention zu ermöglichen. Dies wird durch Anwendung einer gemeinsamen, diagonal initialisierten Projektionsmatrix über alle Köpfe erreicht. Die diagonale Initialisierung bewahrt zu Beginn des Trainings die kopf-spezifische Spezialisierung, während das Modell schrittweise integrierte, kopfübergreifende Repräsentationen erlernt. KHA fügt nur minimale Parameter und FLOPs hinzu und kann nahtlos in MHA, GQA, GTA und andere Aufmerksamkeitsvarianten integriert werden. Wir validieren KHA durch das Training eines 6,1B-Parameter-MoE-Modells (1,01B aktiviert) mit 1T hochwertigen Tokens. Im Vergleich zu Baseline-Aufmerksamkeitsmechanismen erzielt KHA überlegene und stabilere Trainingsdynamiken und erreicht eine bessere Leistung in nachgelagerten Aufgaben.

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.

Knock-Head-Attention

Knocking-Heads Attention

papers.abstract

Support