让注意力头相互协作,提升大模型表现。
Knocking-Heads Attention
- 通过共享对角初始化矩阵实现头间特征交互
- 6.1B参数模型训练中稳定收敛,性能优于基线
- 零成本适配主流注意力结构,适合大模型优化
多头注意力(MHA)是现代大语言模型的核心组件,通过并行多个注意力头提升表征能力。然而,增加头数会削弱单个头的表达能力,现有机制如标准MHA、分组查询注意力(GQA)和分组绑定注意力(GTA)仅简单拼接独立头的输出,缺乏强交互。为此,我们提出敲头注意力(KHA),使注意力头之间“互相敲击”——在缩放点积注意力前引入跨头特征级交互。通过所有头共享一个对角初始化的投影矩阵实现,该初始化在训练初期保留头的特异性,同时逐步学习整合的跨头表示。KHA仅增加极少参数与浮点运算量,可无缝集成至MHA、GQA、GTA等注意力变体。我们在1T高质量语料上训练了一个6.1B参数的MoE模型(激活参数1.01B)。相比基线注意力机制,KHA展现出更优且更稳定的训练动态,在下游任务中取得更好效果。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。