Interleaved Head Attention

多头注意力机制（MHA）是当前大语言模型（LLM）的核心计算原语。然而，MHA存在一个根本性的线性扩展瓶颈：$H$ 个注意力头在计算过程中各自独立地生成 $H$ 个互不通信的注意力矩阵，各头之间在注意力计算阶段完全缺乏信息交互。这一局限在多步推理任务中尤为突出——此类任务的正确答案往往依赖于从上下文不同位置聚合证据，并通过一系列中间推理步骤，对隐含的词元到词元关系进行组合与传导。为解决该问题，我们提出交错头注意力机制（IHA）。该机制通过为每个原始注意力头构建 $P$ 个伪头（通常取 $P = H$）来实现跨头信息混合：每个伪查询（pseudo-query）、伪键（pseudo-key）和伪值（pseudo-value）均为全部 $H$ 个原始查询、键和值的可学习线性组合。伪查询头与伪键头之间的交互，可在每个原始头内诱导出最多 $P^2$ 种注意力模式，且仅引入较低的参数开销 $\mathcal{O}(H^2P)$。我们在理论层面证明了IHA在两类合成任务上的参数效率提升：在多项式任务（Polynomial task）中，IHA 的参数量为 $\Theta(\sqrt{k}n^2)$，而 MHA 为 $\Theta(kn^2)$；在对顺序敏感的 CPM-3 任务中，IHA 仅需 $\lceil\sqrt{N_{\max}}\rceil$ 个头，而 MHA 需要 $N_{\max}$ 个头。在真实世界基准测试中，IHA 在 RULER 数据集的多关键词检索任务上将性能提升 10%–20%（上下文长度为 4k–16k）；在 OpenThoughts 数据集上针对推理能力微调后，IHA 在 GSM8K 上较全注意力基线提升 5.8%，在 MATH-500 上（采用多数投票策略）提升 2.8%。

多头注意力（MHA）中各头完全独立，缺乏跨头信息交互，导致其在需要多步推理、证据聚合与关系组合的任务（如长链逻辑推理、多跳检索）中表征能力受限；该问题虽被广泛观察，但尚未被形式化为‘线性头数-表达能力瓶颈’并系统解决。

提出交错头注意力（IHA）：每原始头生成P个可学习的伪头（pseudo-heads），每个伪查询/键/值均为全部H个原始头对应张量的线性组合；通过伪头间交互实现单层内跨头混合，以O(H²P)参数开销激增（而非O(H³)）换取最多P²种注意力模式/头，突破MHA的线性头数-模式数耦合约束。

理论证明IHA在多项式合成任务（参数复杂度Θ(√k n²) vs MHA的Θ(k n²)）和顺序敏感CPM-3任务（所需头数⌈√N_max⌉ vs N_max）上显著提升参数效率；实证验证于RULER（Multi-Key retrieval提升10–20% at 4k–16k context）、OpenThoughts微调后GSM8K+5.8%、MATH-500+2.8%（Majority Vote）；未提及其开源代码状态；值得深入方向包括伪头结构的稀疏化设计、动态P自适应机制、以及IHA在MoE架构中的协同优化。

Lin et al., 'A Survey of Multi-Head Attention Mechanisms' (2023); Bhojanapalli et al., 'Cross-Head Attention via Shared Projections' (ICLR 2022); Liu et al., 'Head-Mixing Transformers' (NeurIPS 2022); Zhang et al., 'Expressivity of Self-Attention in Reasoning Tasks' (ACL 2023); Tay et al., 'Efficient Transformers: A Survey' (JMLR 2022)