Interleaved Head Attention¶

会议: NeurIPS 2026
arXiv: 2602.21371
OpenReview: https://openreview.net/forum?id=AoSjvxQp5K
领域: LLM效率
关键词: 注意力机制, 多头注意力, 伪头交互, 多跳推理, 长上下文检索

一句话总结¶

针对标准多头注意力因头间隔离导致的线性表达瓶颈，本文提出交错头注意力（IHA），通过在注意算子前将各头线性组合生成伪头并交错排布，使单头诱导多达二次方的注意力模式，在多跳推理与长程检索任务中显著降低参数复杂度并提升准确率。

研究背景与动机¶

标准多头注意力（Multi-Head Attention, MHA）是当今大语言模型的核心计算基石，其基本设定是让 \(H\) 个注意力头彼此独立计算，每个查询头仅与对应的键值投影交互。这种隔离性导致了一个固有的线性扩展瓶颈：\(H\) 个头在单层内只能产生 \(H\) 个独立的注意力矩阵。然而，真实世界中的多跳组合推理（例如“《霍比特人》作者的出生地是哪儿”，需要先推理出《霍比特人》→托尔金，再推理出托尔金→南非）要求单个 token 能够跨序列聚合分散证据并串联中间潜在关系。以经典的多项式图滤波器（Polynomial Graph Filter）作为可控代理分析可知，\(r\) 步依赖对应于关系算子的 \(r\) 次迭代；MHA 的单个头一次只能表示一个固定的依赖模式，若要在单层中同时捕获 \(k\) 种不同长度的关系链，MHA 必须线性增加至 \(\Theta(k)\) 个头或依赖更深的网络层级堆叠。

这一瓶颈的核心矛盾在于：多步推理对中间关系复合能力的需求随跳数快速增长，而传统 MHA 单层内的关系组合容量却严格受限于头数的线性约束。以往试图打破头间隔离的方法通常在注意力计算之后混合 logits 或注意力图（如 Talking-Heads、Knocking-Heads），但这破坏了标准自注意力算子内部的软最大化结构，不仅无法兼容如 FlashAttention 这类现代 GPU 硬件加速算子，且难以赋予每个位置不同的空间相位表示；另一些高阶注意力机制（如 2-Simplicial Attention）则面临显著的参数膨胀与高昂的计算代价。

本文的切入角度是：保留标准缩放点积注意力算子的底层结构不变，将跨头交互提前至注意力算子计算之前，为每个头构造原头投影的线性组合伪头，并通过虚拟交错扩展序列维度。核心 idea：为每个注意力头构建多个由原始全部投影线性混合而成的伪查询、伪键与伪值，将其交错展开为虚拟序列并运行标准因果注意力，使每头能诱导出多达二次方的注意力组合模式，在保持与 FlashAttention 算子兼容的同时以更少头数实现指数级跳数表征。

方法详解¶

交错头注意力（Interleaved Head Attention, IHA）旨在打破 MHA 头与头之间“一对一”刚性绑定的局限，用极小的参数代价引入头间混合。其端到端执行流程包括三个阶段：投影与跨头伪头线性组合、虚拟位置交错与因果注意力计算、以及伪头折叠与输出映射。

整体框架¶

在输入序列进入 IHA 层后，模型首先计算标准的查询、键、值投影矩阵。随后，IHA 引入学习得来的跨头混合权重张量，在注意力头维度上执行线性组合，为每个注意力头生成 \(P\) 个伪查询（pseudo-queries）、伪键（pseudo-keys）和伪值（pseudo-values）。为使伪头能够无缝适配现有的硬件加速算子与旋转位置编码（RoPE），IHA 将生成的 \(P\) 个伪头沿序列维度交错排布，把长度为 \(N\) 的原始序列展开为长度为 \(N \times P\) 的虚拟 token 序列。随后直接调用标准的因果自注意力算子计算注意力得分与加权聚合。最后，通过一个投影/折叠矩阵（Collapse Map）将多余的伪头通道降采样或对齐回标准头维度，并经由输出线性层送入后续的前馈网络。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 X<br/>维度 N × D"] --> B["标准 Q, K, V 投影<br/>维度 H × N × d"]
    B --> C["跨头线性混合生成伪头<br/>通过 αQ, αK, αV 生成 P 个伪头"]
    C --> D["虚拟标记交错展开<br/>序列长度扩展为 N·P"]
    D --> E["标准因果自注意力计算<br/>诱导高达 P² 种注意力模式"]
    E --> F["伪头折叠与输出映射<br/>通过矩阵 R 折叠并投影"]
    F --> G["多头特征拼接与输出"]

关键设计¶

1. 跨头伪头线性组合：打破单头独立性以诱导二次方关系模式

针对 MHA 中各头完全解耦、无法直接在单层内表征多步复合关系的痛点，IHA 在注意力打分之前利用学习到的三阶混合张量对原始投影进行通道线性重组。对每个基础头 \(h \in \{1,\dots,H\}\)，IHA 构造 \(P\) 个伪查询 \(\widetilde{\bm{Q}}_{h,j}\)、伪键 \(\widetilde{\bm{K}}_{h,j}\) 及伪值 \(\widetilde{\bm{V}}_{h,j}\)（通常取 \(P = H\)）。其形式为原始 \(H\) 个头相应投影的加权求和：

\[\widetilde{\bm{Q}}_{h,j} = \sum_{m=1}^{H} \alpha^{Q}_{m,h,j} \bm{Q}_{m}, \quad \widetilde{\bm{K}}_{h,j} = \sum_{m=1}^{H} \alpha^{K}_{m,h,j} \bm{K}_{m}, \quad \widetilde{\bm{V}}_{h,j} = \sum_{m=1}^{H} \alpha^{V}_{m,h,j} \bm{V}_{m}\]

其中 \(\alpha^{Q}, \alpha^{K}, \alpha^{V} \in \mathbb{R}^{H \times H \times P}\) 为可学习的混合权重张量。通过让 \(P\) 个伪查询与 \(P\) 个伪键在一个注意力头内部相互作用，单头即可诱导出多达 \(P^2\) 种相异的注意力模式，整层总计可形成多达 \(H P^2\) 种交互模式。相比于模型主维度 \(d_{\text{model}}\)，头数 \(H\) 和伪头数 \(P\) 通常极小（如 \(H=20\)），因此新增的参数量仅为 \(\mathcal{O}(H^2 P)\)，几乎不增加整体模型参数量。

2. 虚拟序列交错排布：保留标准注意力算子与 RoPE 相位兼容性

针对多数跨头混合算法需要自定义复杂 CUDA kernel 且破坏相对位置编码兼容性的缺陷，IHA 巧妙地通过序列维度的交错操作实现伪头交互。在每个原始位置 \(n \in \{1,\dots,N\}\) 上，原本由该位置派生出的 \(P\) 个伪头 token 被排列为连续的虚拟 token：\((1,1), (1,2), \dots, (1,P), (2,1), \dots, (N,P)\)，从而构成长度为 \(N \cdot P\) 的扩展序列。这种交错设计具有两大独特优势：其一，由于底层计算形态完全等价于一个加长序列的标准点积注意力，因此可以无缝复用 FlashAttention 等高效加速库；其二，当与旋转位置编码（RoPE）配合时，每个虚拟 token \((n, p)\) 自然获得独立的旋转相位，使不同伪头能够在微观位置上学习到平滑的相对偏移区分度，从而敏锐捕捉细粒度的次序关系。

3. 局部滑窗与全局混合调度：严格控制 FLOPs 与显存复杂度

交错展开将有效序列长度放大为原来的 \(P\) 倍，若在所有层全部采用全局注意力，单个头的复杂度将达到 \(\mathcal{O}((NP)^2 d) = \mathcal{O}(P^2 N^2 d)\)，带来 \(P^2\) 倍的计算开销。为了在严格匹配总算力预算（FLOP-matched）的前提下发挥表达力优势，IHA 采用了混合局部-全局调度架构。具体而言，模型以 4:1 的比例交替堆叠 4 层局部滑动窗口 IHA 层与 1 层全局 IHA 层。局部 IHA 层的滑动窗口大小设置为：

\[W \coloneqq \frac{N}{2 P^2}\]

在滑动窗口层中，注意力的计算范围被严格限制在 \(W \cdot P\) 个虚拟 token 内，使得该层的计算复杂度严格压缩为 \(\mathcal{O}(H \cdot (NP) \cdot (WP) \cdot d) = \mathcal{O}(H N^2 d / 2)\)。经过 4 层局部层与 1 层全局层的加权平摊，平均单层复杂度仅为 \(\frac{4 \times (1/2) + 1}{5} \mathcal{O}(H N^2 d) = 0.6 \times \mathcal{O}(H N^2 d)\)，甚至低于普通全局 MHA 层的计算开销，从而在完全公平的算力约束下验证纯架构创新的增益。

一个完整示例¶

以合成的阶数为 \(k=4\) 的多项式图滤波器计算为例，输入为 \(\widehat{\bm{X}} = [\bm{X}, \bm{I}]\)，目标是单层并行输出 \([\bm{X}, \bm{A}\bm{X}, \bm{A}^2\bm{X}, \bm{A}^3\bm{X}]\)： 1. 参数配置与投影：传统 MHA 必须分配 \(H=4\) 个独立头；而 IHA 仅需设置 \(H = \lceil\sqrt{k}\rceil = 2\) 个头，每个头配置 \(P = 2\) 个伪头。查询和键投影矩阵将图邻接算子分块分解为步长块：\(\bm{W}_{Q}^{(1)} = [\bm{0}; \bm{I}]\)，\(\bm{W}_{Q}^{(2)} = [\bm{0}; \bm{A}^2]\)，\(\bm{W}_{K}^{(1)} = [\bm{0}; \bm{I}]\)，\(\bm{W}_{K}^{(2)} = [\bm{0}; \bm{A}^\top]\)。 2. 伪头组合与注意力：在头 1 中，伪查询 1 分别与两组伪键交互，得到注意力矩阵块 \([\bm{I}, \bm{A}]\)，加权聚合值投影后直接生成前两项 \([\bm{X}, \bm{A}\bm{X}]\)；在头 2 中，伪查询 2 同样与两组伪键交互，诱导出偏移后的矩阵块 \([\bm{A}^2, \bm{A}^3]\)，生成后两项 \([\bm{A}^2\bm{X}, \bm{A}^3\bm{X}]\)。 3. 折叠输出：两头拼接后一次性得到全部 4 阶图滤波器输出，在参数量由 \(\Theta(k N^2)\) 降低至 \(\Theta(\sqrt{k} N^2)\) 的同时，头数从 4 缩减为 2。

实验关键数据¶

主实验¶

论文基于 2.4B 参数量的 Decoder-only Transformer 架构（26 层，隐层 2560，20 头，头维度 128），在 240B tokens 上统一预训练，并在 OpenThoughts 数据集上进行了 8B tokens 的长链思维监督微调（SFT）。所有对比模型均在完全相同的硬件环境（128 张 H200 GPU）与 FLOP 预算下完成训练与评测。

在 OpenThoughts 监督微调后的主评测结果如下表所示：

模型	GSM8K (Pass@1)	GSM8K (Maj@16)	MATH-500 (Pass@1)	MATH-500 (Maj@16)	MBPP (Pass@1)	MBPP (Pass@10)	平均排名 (↓)
IHA (本文)	34.3% (+4.8)	54.2% (+5.8)	10.0% (+1.2)	18.4% (+2.8)	15.5% (+0.8)	41.6% (+0.4)	1.5
Global Attention	29.5% (–)	48.4% (–)	8.8% (–)	15.6% (–)	14.7% (–)	41.2% (–)	3.8
Global+Local	26.5% (-3.0)	46.9% (-1.5)	7.6% (-1.2)	15.0% (-0.6)	15.0% (+0.3)	41.9% (+0.7)	4.3
Talking Heads	29.3% (-0.2)	49.4% (+1.0)	7.8% (-1.0)	18.2% (+2.6)	15.9% (+1.2)	43.1% (+1.9)	2.5
Diff Transformer	31.6% (+2.1)	53.5% (+5.1)	9.0% (+0.2)	18.0% (+2.4)	15.3% (+0.6)	39.2% (-2.0)	2.8

在长上下文 RULER 基准测试（微调至 64k 上下文）中，各模型 Multi-Key 检索与全套平均准确率如下：

模型	RULER Multi-Key (4k)	RULER Multi-Key (8k)	RULER Multi-Key (16k)	RULER 平均 Exact Match (全套)
IHA (本文)	显著领先 (+27%)	显著领先 (+32%)	大幅领先 (+112%)	44.0%
Global+Local	基线水平	基线水平	严重衰减	40.6%
Diff Transformer	次优水平	次优水平	中幅衰减	37.2%
Global Attention	基线参考	基线参考	基线参考	35.0%

消融实验¶

论文在预训练阶段（5-shot 无微调状态）对注意力机制的基础表征能力进行了对比，验证无 SFT 条件下各机制的固有推理增益：

配置	GSM8K (EM)	GSM8K (Maj@5)	MATH-500 (EM)	MBPP (Pass@1)	HumanEval (Pass@1)	平均排名 (↓)
IHA (本文完整模型)	8.34% (+2.73)	8.42% (+2.81)	3.54% (+0.66)	24.5% (+1.1)	17.1% (-0.1)	1.4
Global Attention	5.61% (–)	5.61% (–)	2.88% (–)	23.4% (–)	17.2% (–)	2.9
Global+Local	6.82% (+1.21)	6.90% (+1.29)	2.26% (-0.62)	23.6% (+0.2)	16.0% (-1.2)	2.9
Talking Heads	5.46% (-0.15)	5.38% (-0.23)	— (未收敛/退化)	23.8% (+0.4)	16.0% (-1.2)	4.0
Diff Transformer	5.46% (-0.15)	5.61% (0.00)	— (未收敛/退化)	25.0% (+1.6)	15.4% (-1.8)	3.5

在合成推理任务（Binary 2-hop 与 Ternary 3-hop 关系组合）中，IHA 与高阶 Simplicial Attention 及 MHA 进行了单层 8 头设置下的受控对比： - 2-hop 二元关系组合：IHA 相比于最强 baseline（Simplicial Attention）准确率提升高达 +4.7%。 - 3-hop 三元关系组合：IHA 相比于 Simplicial Attention 准确率提升达 +3.3%，且收敛速度明显更快。

关键发现¶

数学与逻辑推理收益最为显著：在 GSM8K 和 MATH-500 上，无论预训练还是 SFT 阶段，IHA 均取得最大幅度领先（SFT 后 GSM8K Maj@16 提升 5.8 个百分点），这印证了其在多步潜在依赖建模与长链证据关联上的理论优势。
长上下文检索能力质的跃升：在长达 16k 的 RULER Multi-Key 跨位置检索测试中，IHA 相对全量注意力实现了超 100% 的准确率翻倍，证明跨头注意力的虚拟序列能够大幅增强多键匹配与抗干扰能力。
代码生成呈现差异化特性：Talking Heads 在 MBPP 上略占优势，表明代码生成更多依赖函数级语法特征的头间权重混合，而 IHA 的离散状态追踪和链式组合更契合符号与数学逻辑。

亮点与洞察¶

将多头交互转化为单序列虚拟交错：巧妙地将通道维度的线性混合转化为序列维度的虚拟 token 交错，既规避了非标算子的编译与显存瓶颈，又完美复用了成熟的 FlashAttention 内核与 RoPE 相位结构。
严格证明了对传统 MHA 的真包含关系：通过重复 token 空间上的非线性响应构造，在理论上严谨证明了 IHA 是 MHA 的真超集（Strict Generalization），打破了“注意力头越多才能表示越多模式”的传统认知。
极佳的算力-参数扩展性价比：仅以 \(\mathcal{O}(H^2 P)\) 的极微小参数开销（远小于 \(d_{\text{model}}\)），便在图滤波器任务中将参数量从 \(\Theta(k N^2)\) 降低到 \(\Theta(\sqrt{k} N^2)\)，实现了多跳推理容量的二次方缩放。

局限与展望¶

展开导致显存峰值增加：虽然局部滑窗调度拉平了总 FLOPs，但将序列长度在物理上展开 \(P\) 倍在局部层中仍可能增加 KV 缓存和部分中间激活显存峰值。
伪头数量静态固定：当前所有层和所有头均静态分配 \(P=H\) 个伪头，未来可探索动态、自适应的伪头数量分配机制（如浅层稀疏、深层密集）。
模型模态与任务拓展：目前实验主要集中于自回归因果语言模型，尚未在编码器-解码器架构（如机器翻译）及视觉 Transformer（ViT）中广泛检验其通用性。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 伪头跨头线性组合与虚拟序列交错构想极为巧妙，兼顾高维表达与硬件加速。
实验充分度: ⭐⭐⭐⭐⭐ 包含严谨的理论证明、受控合成基准、2.4B 级别全量预训练及 64k 长文本与推理微调。
写作质量: ⭐⭐⭐⭐⭐ 逻辑演进严密，从图多项式滤波器理论推演直达工业级语言模型实证。
价值: ⭐⭐⭐⭐⭐ 对克服 Transformer 线性头表达瓶颈、增强深层数学与多跳推理具有重要学术与应用价值。