Attenzione a Testate Contrapposte

Knocking-Heads Attention

October 27, 2025

Autori: Zhanchao Zhou, Xiaodong Chen, Haoxing Chen, Zhenzhong Lan, Jianguo Li

cs.AI

Abstract

L'attenzione multi-testa (MHA) è diventata la pietra angolare dei moderni modelli linguistici di grandi dimensioni, potenziando la capacità rappresentativa attraverso teste di attenzione parallele. Tuttavia, l'aumento del numero di teste indebolisce intrinsecamente la capacità delle singole teste, e i meccanismi di attenzione esistenti - che si tratti della MHA standard o delle sue varianti come l'attenzione a query raggruppate (GQA) e l'attenzione a gruppi vincolati (GTA) - si limitano a concatenare gli output di teste isolate senza una forte interazione. Per affrontare questa limitazione, proponiamo l'attenzione a teste interagenti (KHA), che consente alle teste di attenzione di "bussarsi" reciprocamente - facilitando interazioni a livello di feature tra le teste prima dell'attenzione prodotto-scalare. Ciò è ottenuto applicando una matrice di proiezione condivisa, inizializzata diagonalmente, su tutte le teste. L'inizializzazione diagonale preserva la specializzazione specifica di ciascuna testa all'inizio dell'addestramento, consentendo al modello di apprendere progressivamente rappresentazioni integrate tra le teste. KHA aggiunge solo parametri e FLOPs minimi e può essere integrato senza soluzione di continuità in MHA, GQA, GTA e altre varianti di attenzione. Convalidiamo KHA addestrando un modello MoE da 6.1B di parametri (1.01B attivati) su 1T di token di alta qualità. Rispetto ai meccanismi di attenzione baseline, KHA offre dinamiche di addestramento superiori e più stabili, raggiungendo prestazioni migliori in tutti i task downstream.

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.