Knocking-Heads Attention

📄 arXiv: 2510.23052v1 📥 PDF

作者: Zhanchao Zhou, Xiaodong Chen, Haoxing Chen, Zhenzhong Lan, Jianguo Li

分类: cs.CL

发布日期: 2025-10-27

💡 一句话要点

提出Knocking-Heads Attention，通过头间交互提升大型语言模型表征能力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多头注意力 跨头交互 大型语言模型 模型优化 MoE模型

📋 核心要点

  1. 现有MHA及其变体缺乏有效的头间交互，限制了模型表征能力。
  2. KHA通过共享的对角初始化投影矩阵，使注意力头在特征层面进行跨头交互。
  3. 实验表明，KHA在MoE模型上实现了更优越和稳定的训练动态，并在下游任务中表现更好。

📝 摘要（中文）

多头注意力（MHA）已成为现代大型语言模型的基石，通过并行注意力头增强表征能力。然而，增加头的数量会削弱单个头的能力，并且现有的注意力机制（无论是标准MHA还是其变体，如分组查询注意力（GQA）和分组绑定注意力（GTA））只是简单地连接来自孤立头的输出，而没有强大的交互。为了解决这个限制，我们提出了Knocking-Heads Attention（KHA），它使注意力头能够“敲击”彼此——在缩放点积注意力之前促进跨头特征级交互。这是通过在所有头上应用共享的、对角初始化的投影矩阵来实现的。对角初始化在训练开始时保持特定于头的专业化，同时允许模型逐步学习集成的跨头表示。KHA仅添加最少的参数和FLOPs，并且可以无缝集成到MHA、GQA、GTA和其他注意力变体中。我们通过在1T高质量tokens上训练一个61亿参数的MoE模型（激活10.1亿）来验证KHA。与基线注意力机制相比，KHA带来了卓越且更稳定的训练动态，在下游任务中实现了更好的性能。

🔬 方法详解

问题定义：现有的大型语言模型依赖于多头注意力机制（MHA）来提升模型的表征能力。然而，简单地增加注意力头的数量会降低每个头的容量。此外，现有的注意力机制，如GQA和GTA，只是简单地连接各个头的输出，缺乏有效的头间交互，限制了模型学习更复杂表示的能力。因此，需要一种能够促进头间交互，同时保持计算效率的注意力机制。

核心思路：KHA的核心思路是允许不同的注意力头在进行缩放点积注意力计算之前，能够相互“敲击”，即进行特征层面的交互。通过这种方式，模型可以学习到各个头之间的关联性，从而更好地利用不同头学习到的特征。这种设计旨在增强模型的表征能力，同时避免引入过多的计算开销。

技术框架：KHA可以无缝集成到现有的多头注意力机制中，如MHA、GQA和GTA。其主要步骤如下：首先，对输入进行线性变换，得到Query (Q), Key (K), Value (V)。然后，将Q, K, V分别输入到不同的注意力头中。在每个注意力头内部，KHA引入了一个共享的、对角初始化的投影矩阵，用于在不同的头之间进行特征交互。最后，将交互后的特征进行缩放点积注意力计算，并输出最终结果。

关键创新：KHA最重要的创新点在于引入了跨头交互机制，允许不同的注意力头在特征层面进行信息交换。与现有的注意力机制相比，KHA不再是简单地将各个头的输出进行拼接，而是通过共享的投影矩阵，使得各个头之间能够相互影响，从而学习到更丰富的特征表示。对角初始化保证了训练初期各个头的独立性，避免了训练不稳定。

关键设计：KHA的关键设计在于共享的、对角初始化的投影矩阵。该矩阵的对角初始化保证了在训练初期，各个注意力头能够独立地学习不同的特征。随着训练的进行，模型可以逐渐学习到各个头之间的关联性，从而实现跨头交互。此外，KHA还采用了低秩分解等技术，以减少参数量和计算复杂度。损失函数与标准注意力机制相同，没有引入额外的损失项。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，KHA在61亿参数的MoE模型上，相比于基线注意力机制，实现了更优越和稳定的训练动态。具体而言，KHA在下游任务上取得了显著的性能提升，验证了其有效性。该模型在1T高质量tokens上进行训练，证明了KHA在大规模数据上的可扩展性。

🎯 应用场景

Knocking-Heads Attention (KHA) 可广泛应用于自然语言处理领域，尤其是在需要强大表征能力的大型语言模型中。例如，可以用于提升机器翻译、文本生成、问答系统等任务的性能。此外，KHA的设计思想也可以推广到其他领域，如计算机视觉和语音识别，以增强模型对复杂数据的理解能力。其轻量级的设计使其易于部署在资源受限的设备上。

📄 摘要（原文）

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.