提出显式建模注意力头间交互的新方法，提升大模型性能与效率
Explicit Multi-head Attention for Inter-head Interaction in Large Language Models
- 通过头级线性组合与归一化显式建模头间交互
- 支持更大学习率，收敛更快，验证损失更低
- 可压缩50%键值缓存，适合资源受限部署
在基于Transformer架构的大语言模型中，近期研究发现头间交互能提升注意力表现。受此启发，我们提出多头显式注意力（MEA），一种简单而有效的注意力变体，能够显式建模跨头交互。MEA包含两个关键组件：头级线性组合（HLC）模块，对各头的键和值向量分别应用可学习的线性组合，实现丰富的头间通信；以及头级分组归一化层，对重组后的头进行统计特性对齐。MEA在预训练中表现出强鲁棒性，允许使用更大的学习率，加速收敛，最终在多种任务上取得更低的验证损失和更优性能。此外，我们探索了MEA的参数效率：通过减少注意力头数并利用低秩‘虚拟头’与HLC重构，实现一种实用的键值缓存压缩策略，在知识密集型与科学推理任务上将KV缓存内存使用降低50%，仅导致奥数级数学基准3.59%准确率下降。
原文摘要 · Abstract (English)
In large language models built upon the Transformer architecture, recent studies have shown that inter-head interaction can enhance attention performance. Motivated by this, we propose Multi-head Explicit Attention (MEA), a simple yet effective attention variant that explicitly models cross-head interaction. MEA consists of two key components: a Head-level Linear Composition (HLC) module that separately applies learnable linear combinations to the key and value vectors across heads, thereby enabling rich inter-head communication; and a head-level Group Normalization layer that aligns the statistical properties of the recombined heads. MEA shows strong robustness in pretraining, which allows the use of larger learning rates that lead to faster convergence, ultimately resulting in lower validation loss and improved performance across a range of tasks. Furthermore, we explore the parameter efficiency of MEA by reducing the number of attention heads and leveraging HLC to reconstruct them using low-rank "virtual heads". This enables a practical key-value cache compression strategy that reduces KV-cache memory usage by 50% with negligible performance loss on knowledge-intensive and scientific reasoning tasks, and only a 3.59% accuracy drop for Olympiad-level mathematical benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。