击脑袋注意力机制
计算与语言
2025-10-28 v1
摘要
多头注意力（MHA）已成为现代大型语言模型的基石，通过并行注意力头提升表征能力。然而，增加注意力头的数量会固化地削弱单个头的容量，而现有注意力机制——无论是标准MHA还是其变体如分组查询注意力（GQA）和分组绑定注意力（GTA）——仅对孤立注意力头的输出进行简单拼接，缺乏强大的相互作用。为此，我们提出击脑袋注意力（KHA），使注意力头之间“敲击”——在缩放点积注意力前实现跨头特征级交互。这通过在所有注意力头上应用共享的、对角初始化的投影矩阵实现。对角初始化在训练初期保留头的特定化专业性，同时允许模型逐渐学习集成的跨头表征。KHA仅增加极少的参数和FLOPs，可无缝集成到MHA、GQA、GTA及其他注意力变体中。我们通过在1T高质量token上训练一个61亿参数的Mixture of Experts模型（10亿激活参数）来验证KHA。与基线注意力机制相比，KHA实现更优更稳定的训练动力学，在多个下游任务上取得更好性能。
引用
@article{arxiv.2510.23052,
title = {Knocking-Heads Attention},
author = {Zhanchao Zhou and Xiaodong Chen and Haoxing Chen and Zhenzhong Lan and Jianguo Li},
journal= {arXiv preprint arXiv:2510.23052},
year = {2025}
}