[FAIR/Meta] Interleaved Head Attention：打破 MHA 线性瓶颈，开启二次方表达力新范式

TL;DR

在 Transformer 架构中，多头注意力（MHA）的“头独立性”一直是限制其推理效率的隐形枷锁。Meta AI、MIT 及 UT Austin 的研究者联合提出的 Interleaved Head Attention (IHA)，通过在注意力计算前通过线性投影构造“伪头（Pseudo-heads）”，让 $H$ 个头能够产生 $H^2$ 级别的注意力模式。在不增加模型深度的前提下，IHA 在长文本检索上提升了 20% 以上，推理任务（GSM8K）提升了 5.8%。

1. 痛点：为什么 MHA 处理不好复杂推理？

在标准 MHA 中，每个头都是一个“孤岛”：头 $h$ 的 Query 只与头 $h$ 的 Key 计算。这意味着 $H$ 个头最多只能捕捉 $H$ 种不同的关系模式。

然而，现实中的逻辑推理往往是组合性的。例如回答“《霍比特人》作者的出生地在哪？”，模型必须：

1. 建立关系 A：《霍比特人》 $\to$ J.R.R. Tolkien；
2. 建立关系 B：Tolkien $\to$ 南非。

论文通过 Polynomial Filter (多项式滤波器) 理论分析证明，要同时表达这种 $k$ 步依赖关系，MHA 需要 $\Theta (k)$ 个头或更深的层数。这种线性扩展性使得模型在面对极长上下文或深层嵌套逻辑时，参数效率利用极低。

2. 核心机制：伪头交织 (Interleaved Pseudo-heads)

IHA 的直觉非常优雅：如果能在进入 Softmax 之前，让 Q、K、V 就先进行一次“头际”大融合呢？

2.1 算法流程

1. 生成伪头：利用可学习参数 ${\alpha }_Q,{\alpha }_K,{\alpha }_V$，将 $H$ 个原始投影线性混合成 $P$ 个伪投影。
2. 交织合并：将伪头维度合并到序列维度 $N$。这意味着序列长度逻辑上变成了 $NimesP$。
3. 标准注意力：执行标准的高效注意力计算（兼容 FlashAttention）。
4. 折叠输出：通过减少矩阵 $R$ 将多出的维度折叠回原始尺寸。

2.2 物理直觉：二次方表达力的飞跃

通过这种构造，每个物理头实际上承载了 $P$ 个伪头的交互。理论证明，IHA 在处理 $k$ 阶多项式任务时，所需的参数量仅为 MHA 的 $\sqrt{k}$。这种“开方级”的效率提升，让模型能以极小的代价建模极其复杂的关系图谱。

3. 实验战绩：推理与长文本的双重重击

研究团队在 2.4B 参数规模的模型上进行了严谨的 FLOP-matched 对比实验。

3.1 长文本检索 (RULER Benchmark)

在 Multi-Key 检索任务（需要在万级 Token 中精准定位多个关键信息）中，IHA 展现了压倒性优势。

• 16k 文本长度：相比全局注意力提升了 112%。
• RULER 平均准确率：达到 44.0%，远超 Diff Transformer (37.2%) 和滑动窗口方法。

3.2 逻辑推理 (GSM8K & MATH)

在经过 OpenThoughts 强化指令微调后：

• GSM8K (数学应用题)：准确率 54.2%（+5.8%）。
• MATH-500：准确率 18.4%（+2.8%）。

实验表明，IHA 在 逻辑状态跟踪 方面尤其出色，这种跨头混合机制似乎充当了某种“层内工作记忆”，加速了信息的合成。

4. 深度洞察：IHA vs. Talking Heads

很多人会联想到 Noam Shazeer 提出的 Talking-Heads Attention。两者的本质区别在于：

• Talking-Heads：是在计算出注意力分数（Logits）后再进行混合，本质上是“结果的再分配”。
• IHA：是在计算分数前对特征（Q/K/V）进行混合。 从理论上讲，IHA 提供的非线性潜力更大，且因为它依然遵循标准注意力公式，在算子优化（Kernel Fusion）上更具天然优势。

5. 总结与局限

IHA 的价值在于它在不增加深度、不改变注意力算子的前提下，通过数学上的精妙设计打破了 MHA 的表达力天花板。

局限性： 由于序列长度逻辑上变成了 $NP$，虽然论文采用了混合滑动窗口（Hybrid Layout）来匹配 FLOPs，但在实现复杂度上略高于标准模型。未来的方向可能在于如何动态地分配伪头资源，或将其引入端到端的视觉感知任务中。

Takeaway: IHA 告诉我们，与其追求更深的模型，不如通过精细化注意力头的交互逻辑，挖掘每一层参数的极限。