让注意力头之间能交流,提升大模型推理能力。
Interleaved Head Attention
- 每个注意力头生成多个伪头,通过线性组合实现跨头信息融合。
- 在真实任务中,检索准确率提升10%-20%,数学推理提升5.8%。
- 适合需要多步逻辑推理的场景,如数学题、复杂问答。
多头注意力(MHA)是现代大语言模型的核心计算单元,但存在根本性的线性扩展瓶颈:H个注意力头产生完全独立的H个注意力矩阵,计算过程中头间无信息交互。这在多步推理任务中尤为不利,因正确答案需整合上下文多部分证据并构建链式中间推理的隐含关系。为此,我们提出交错头注意力(IHA),通过为每个头构造P个伪头(通常P=H),使每个伪查询/键/值成为原始所有H个查询/键/值的可学习线性组合,从而在伪查询与伪键之间诱导最多每头P²个注意力模式,仅引入适度参数开销𝒪(H²P)。理论分析表明,在合成多项式任务中,IHA参数量仅为MHA的Θ(√k n²) vs. Θ(k n²);在合成有序敏感的CPM-3任务中,所需头数从N_max降至⌈√N_max⌉。在真实基准测试中,IHA在RULER多键检索任务上提升10%-20%(4k-16k上下文),在针对推理微调后的OpenThoughts数据集上,相比全注意力,GSM8K提升5.8%,MATH-500(多数投票)提升2.8%。
原文摘要 · Abstract (English)
Multi-Head Attention (MHA) is the core computational primitive underlying modern Large Language Models (LLMs). However, MHA suffers from a fundamental linear scaling limitation: $H$ attention heads produce exactly $H$ independent attention matrices, with no communication between heads during attention computation. This becomes problematic for multi-step reasoning, where correct answers depend on aggregating evidence from multiple parts of the context and composing latent token-to-token relations over a chain of intermediate inferences. To address this, we propose Interleaved Head Attention (IHA), which enables cross-head mixing by constructing $P$ pseudo-heads per head (typically $P=H$), where each pseudo query/key/value is a learned linear combination of all $H$ original queries, keys and values respectively. Interactions between pseudo-query and pseudo-key heads induce up to $P^2$ attention patterns per head with modest parameter overhead $\mathcal{O}(H^2P)$. We provide theory showing improved efficiency in terms of number of parameters on the synthetic Polynomial task (IHA uses $Θ(\sqrt{k}n^2)$ parameters vs. $Θ(kn^2)$ for MHA) and on the synthetic order-sensitive CPM-3 task (IHA uses $\lceil\sqrt{N_{\max}}\rceil$ heads vs. $N_{\max}$ for MHA). On real-world benchmarks, IHA improves Multi-Key retrieval on RULER by 10-20% (4k-16k) and, after fine-tuning for reasoning on OpenThoughts, improves GSM8K by 5.8% and MATH-500 by 2.8% (Majority Vote) over full attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。