Knocking-Heads Attention

TL;DR

Knocking-Heads Attention (KHA) addresses the lack of inter-head interaction in standard multi-head attention by introducing shared, diagonally-initialized projection matrices across all heads, enabling cross-head feature-level communication while preserving head specialization. The approach adds minimal overhead and is compatible with MHA, GQA, GTA, and other attention variants, with both linear and MLP-based implementations for the shared projections; initialization is designed to start with near-isolated heads and gradually encourage collaboration. Extensive experiments on 1T-token, 6.1B total-parameter MoE models demonstrate improved training stability (fewer loss spikes) and consistent improvements across downstream tasks, including Language Understanding, Code, and Math, yielding an average gain of approximately $1.26$ points. The results suggest that KHA is a practical, scalable augmentation for transformer architectures, particularly beneficial in large-scale pretraining where cross-head interaction can be leveraged without sacrificing head specialization or adding significant computational costs.

Abstract

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.

Figures (3)

• Figure 1: (Left) The knocking-heads attention architecture. Purple represents the original multi-head attention, while pink represents the added knocking-heads projections. $T^Q$ and $T^K$ within the dashed box are optional projections due to their lower importance compared to $T^V$. (Right) Training loss over 1T tokens for 6.1B MoE models (1.01B activated parameters): baseline vs. knocking-heads attention. KHA reduces loss spikes and maintains consistently lower training loss.
• Figure 2: Training loss curves before and after applying knocking-heads across different model sizes, and the loss curves in (c) and (d) are smoothed for better visualization.
• Figure 3: Visualization of learned knocking-heads projection weights across different layers and types. We apply 0-1 clipping to all knocking-heads projection weights except $W^{gate}$, including $T^K$, $T^Q$, $T^V$, $W^{up}$, and $W^{down}$, for comparative analysis.