PaperlayerView exact revision on arXiv

Source-linked AI summary

Knocking-Heads Attention

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

arXiv:2510.23052v1cs.CL

TL;DR

Attention heads often operate independently, while prior attention-matrix mixing can be computationally costly. KHA adds shared, diagonally initialized feature-level projections across heads. In large-scale experiments, it improves downstream scores and training stability over baseline attention.

  • Problem

    Attention heads operate independently, and talking-heads projections on attention matrices incur dramatically increased computational complexity, especially with many heads.

    Source: §1 · p. 2

  • Method

    KHA applies shared, diagonally initialized transformations across heads at the feature level while retaining head-specific projections.

    Source: §1 · p. 2

  • Results

    Overall average scores improved by 1.26 points, with gains of 4.32 in language understanding, 3.9 in code, and 1.62 in math; training loss spikes also fell.

    Source: §4.3 · p. 8, §4.5.1 · p. 8

  • Takeaways & Limitations

    KHA is presented as a practical enhancement to transformer attention, adding less than 1% computational overhead and serving as a drop-in replacement for standard MHA.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

KHA adds shared, diagonally initialized projections so attention heads can interact at the feature level while retaining head-specific projections. In 1T-token experiments, it improves downstream scores and training stability.

1 Introduction

  • KHA applies shared transformations across heads at the feature level, with diagonal initialization preserving head specialization as cross-head interaction develops.Query and key projections are optional because they provide smaller improvements than value projections.Source: §1 · p. 2
  • The method adds less than 1% in training FLOPs and parameters, and linear projections can be integrated into the original projections at inference.The shared projections supplement the original head-specific projections.Source: §1 · p. 2
  • Overall scores improved by 1.26 points across evaluated tasks, including +4.32 for language understanding, +3.9 for code, and +1.62 for math.General-knowledge and professional-knowledge performance remained comparable; KHA also reduced training-loss spikes.Source: Figure 1 · p. 1, §1 · p. 2
  • Experiments identify MLP-based value projections as the strongest configuration, with benefits evident at four value/key heads.The selected GQA setup uses 32 query heads and 4 KV groups; the 1T-token evaluation uses 6.1B total and 1.01B active parameters.Source: §1 · p. 2
  • The contribution summary reports KHA improvements in training stability and performance across attention mechanisms, MoE and dense models, and model scales.The large-scale experiment trains on 1T tokens.Source: §1 · p. 2

2 Related Work

Prior work uses parameter sharing to reduce redundancy, while attention-head interaction methods seek communication among otherwise independent heads. KHA adds shared feature-level projections while retaining head-specific projections and diagonal initialization.

2.1 Parameter Sharing in Architecture Design

  • Parameter sharing appears in CNN kernels, ALBERT’s cross-layer weights, and DeepSeek-MoE’s shared experts; KHA adds shared projections across attention heads.The paper presents these examples as parameter-sharing approaches serving different architectural purposes.Source: §2.1 · pp. 2–3

2.2 Attention Head Interaction

  • Talking-heads attention mixes attention matrices at substantial computational cost; collaborated attention limits specialization, while mixture-of-head attention offers limited communication and complex routing.KHA instead retains head-specific projections and uses diagonal initialization alongside shared projections.Source: §2.2 · p. 3

2.3 Loss Spikes in Pre-training

  • Prior work addresses loss spikes through scaled embedding initialization, activation gating, or QK-clip; KHA uses shared transformations across heads to introduce regularization effects.The cited approaches associate spikes with parameter gradients, massive activations, or attention logits, respectively.Source: §2.3 · p. 3

3 Method

The method extends classical attention with shared transformations that let heads interact before per-head attention, while retaining head-specific projections. It uses staged initialization for specialization and reports low training-computation overhead.

3.1 Classical Attention

  • MHA computes scaled dot-product attention in n parallel heads, concatenates their outputs, and applies a projection; GQA shares key-value pairs within query-head groups.Source: §3.1 · pp. 3–4

3.2 Knocking-Heads Attention

  • KHA adds shared head-level transformations after Q/K/V projection and before per-head attention, enabling interaction alongside original head-specific projections; value transformations are most critical, while query and key transformations are optional.Source: §3.2 · p. 4
  • Linear KHA transforms Q/K/V and can be absorbed into original projections at inference; MLP KHA focuses on values for nonlinear expressiveness with the linear variant’s parameter count.Source: §3.2 · p. 4

3.3 Initialization Strategy

  • Diagonal/zero initialization makes shared transforms approximate identity initially, preserving specialization before off-diagonal interactions develop; random initialization converged to much higher loss.Source: §3.3 · p. 4

3.4 Complexity Analysis

  • 0.55% of total layer computation and 1.17% of original MHA computation: KHA’s training overhead at L = 2048, d = 1024, and n = 32.Source: §3.4 · p. 5

4 Experiment

The experiments examine KHA’s head-count sensitivity, projection designs, and compatibility with attention variants. They identify value-focused MLP projections as a strong configuration and test KHA across attention mechanisms.

4.1 Training Setup

  • Main experiments use a 6.1B-parameter MoE model with 1.01B active parameters; exploratory experiments test varied model sizes and attention mechanisms.The main and exploratory experiments train on 1T and 100B tokens, respectively.Source: §4.1 · p. 5

4.2.1 KHA with Varying Head Number

  • Both KHA variants improve significantly as GQA KV heads increase from 1 to 4, while KHA-MLP is more stable across configurations than linear KHA.The comparison fixes query heads and uses a 0.8B-activated-parameter model trained on 75B tokens.Source: §4.2.1 · p. 5, Table 1 · p. 6

4.2.2 Ablation Studies on Different Variants of Knocking-heads

  • Shared transformations improve performance at query, key, and value positions; value projections benefit most, and MLPs outperform linear layers for values.Standalone gating mechanisms as shared transformations are detrimental to model performance.Source: §4.2.2 · p. 6, Table 2 · p. 6

4.2.3 Compatibility with Other Attention Variants

  • MLA is compatible only with KHA-Linear, while the other tested attention variants are compatible with KHA-Linear and KHA-MLP.The experiments use KHA-Linear for MLA and KHA-MLP for the other variants.Source: §4.2 · p. 5, §4.2.3 · p. 6
  • A 0.02 loss reduction shrinks the GQA4(32)–MHA16(16) gap to 0.002 when both use knocking-heads projections.KHA also improves the tested GQA, MHA, GTA, and MQA variants.Source: §4.2.3 · pp. 6–7, Table 3 · p. 7

4.3 Large-scale Pretraining On 1T Tokens

In 1T-token pretraining, KHA reduced loss spikes and achieved lower stabilized training loss. The final model also improved overall downstream-task performance, with gains varying by task group.

4.3 Large-scale Pretraining On 1T Tokens

  • The controlled pretraining comparison uses a 6.1B-parameter MoE model with 1.01B activated parameters, GQA with group size 4, and 1T tokens.The baseline configuration has 32 query heads.Source: §4.3 · p. 7
  • After loss stabilized, KHA achieved 0.015 lower loss at equivalent training steps and significantly reduced training spikes.The baseline showed numerous spikes during the first half of training.Source: §4.3 · pp. 7–8
  • Overall average improved 1.26 points; language understanding, code, and math average scores rose 4.32, 3.9, and 1.62 points, while knowledge tasks remained comparable.The comparison evaluates final models trained on 1T tokens with 1.01B active and 6.1B total parameters.Source: §4.3 · p. 8

4.4 Scalability Compared with Transformer

Across MoE and dense model scales, KHA’s reported benefits vary by architecture and scale.

4.4 Scalability Compared with Transformer

  • KHA yields a 0.02 loss reduction for MoE models with negligible training overhead; scaling baseline MoE from 6.6B to 14.6B parameters reduces loss by 0.045.The models are trained on 100B tokens and use GQA with four KV heads and 32 query heads.Source: §4.4 · p. 8
  • KHA’s benefits become more pronounced at larger MoE scales across the reported 2.3B–14.6B-parameter range.The comparison also includes dense models ranging from 0.61B to 3.94B parameters.Source: §4.4 · p. 8
  • KHA-Linear provides greater improvements on dense models than the other reported KHA configuration.The tested dense models range from 0.61B to 3.94B parameters.Source: §4.4 · p. 8

4.5 Visualization

The visualizations show KHA reducing training loss spikes across model scales and revealing learned, layer-dependent patterns of cross-head sharing. Value projections show especially strong sharing, while query and key projections retain a mix of shared and head-specific features.

4.5.1 Training Stablity

  • Across model scales, KHA reduces loss spikes and improves training stability; the 1.6B activated MoE model has more frequent spikes than the 0.4B model.Both KHA-Linear and KHA-MLP reduce spikes in the 0.8B activated model, which the authors hypothesize reflects implicit regularization from head sharing.Source: §4.5.1 · p. 8, Figure 2 · p. 9

4.5.2 Learnt Weight of Shared Matrix

  • Learned T_V projections have consistently low diagonal values across layers, indicating aggressive head sharing that may explain the greater improvements from value projections.T_Q and T_K show mixed patterns: some dimensions retain high diagonal values while others show inter-head interaction; some layers show little interaction.Source: §4.5.2 · p. 8, Figure 3 · p. 9
  • The learned T_Q and T_K patterns vary by feature dimension and layer, while MLP gate matrices also exhibit clear structure.Figure 3 visualizes projection weights across layers and types; all listed projection weights except W gate are clipped to 0–1 for comparison.Source: §4.5.2 · p. 8, Figure 3 · p. 9

5 Conclusion

KHA adds shared, diagonally initialized transformations to enable cross-head communication while keeping computational overhead below 1%. Experiments report superior performance and improved training stability, with diagonal initialization balancing head specialization and collaboration.

5 Conclusion

  • KHA enables cross-head communication through shared, diagonally initialized transformation matrices, with less than 1% computational overhead.The conclusion presents it as a simple enhancement to MHA.Source: §5 · p. 9
  • On 1T tokens with 6.1B-parameter MoE models, KHA achieves superior performance and significantly improves training stability versus standard MHA.Source: §5 · p. 9
  • The authors conclude that diagonal initialization balances head specialization with cross-head collaboration and describe KHA as a drop-in replacement for standard MHA.Source: §5 · p. 9

A.1 Comparison of Interactive-heads Attention Mechanisms

The comparison distinguishes interactive-head mechanisms by how they combine heads, their computational and compatibility costs, and their effects on head specialization and training stability. KHA uses lightweight feature-sharing modules within existing attention variants.

A.1 Comparison of Interactive-heads Attention Mechanisms

A.2 Evaluation Benchmark

The benchmark spans five task categories intended to measure different aspects of model competence, from general knowledge and language understanding to professional knowledge, math, and code.

A.2 Evaluation Benchmark

  • General Knowledge tasks include ARC, AGIEval, PIQA, and HellaSwag.Source: §A.2 · p. 13
  • Language Understanding is represented by RACE.Source: §A.2 · p. 13
  • Professional Knowledge tasks include MMLU, CMMLU, MMLU-Pro, GPQA, C-Eval, and CommonsenseQA.Source: §A.2 · p. 13
  • Math tasks include GSM8K, MATH, and CMATH.Source: §A.2 · p. 13
  • Code tasks include HumanEval-plus, MBPP, and MBPP-Plus.Source: §A.2 · p. 13
Loading 2510.23052v1…