一句话总结

HydraHead 是一种头级别的注意力混合架构，采用可解释性驱动的选择策略，仅在检索关键头上保留全注意力，并通过尺度归一化融合模块调和输出分布，在 512K 上下文长度下相较基线提升超过 69%，在仅用 15B tokens 训练后，以 7:1 的线性注意力与全注意力比例匹配 3:1 层级混合模型的长上下文性能。

核心贡献

• 可解释性分析表明，Transformer 层呈现块状功能相似性，而层内的注意力头则专注于不同功能，确立了头维度作为融合异构注意力机制的合理粒度。
• HydraHead 沿头轴混合全注意力和线性注意力，采用可解释性驱动的选择策略，仅在检索关键头上保留全注意力，并利用尺度归一化融合模块调和两种注意力类型之间的分布失配。
• 结合参数复用和蒸馏的三阶段迁移流程训练 HydraHead，在 512K 上下文、15B 训练 tokens 下，NIAH 基准提升超过 69%，以 7:1 的 LA-to-FA 比例匹配 3:1 层级混合模型的长上下文性能，并在不损害通用推理能力的前提下接近 Qwen3.5。

引言

从静态大语言模型向能够长期规划和推理的自主 agent 的转变，使得扩展上下文窗口变得至关重要，但标准的全注意力机制随序列长度呈二次方增长。线性注意力机制提供线性复杂度，但常常遭遇表达性坍缩，难以进行精确检索。在不同层间交错不同注意力类型的混合架构作为折中方案出现，但层级设计面临挑战：层输出变化平滑，难以决定每种机制的放置位置，且训练此类混合模型仍然困难，导致设计空间未被充分探索。作者利用机制可解释性表明，单层内的各个注意力头表现出鲜明的功能异质性，仅稀疏子集对检索至关重要，而层间缺乏清晰边界。他们提出 HydraHead，一种细粒度的头级别混合架构，使用轻量级、一次性基于可解释性的选择，仅将全注意力分配给检索关键头，其余分配给线性注意力，并结合头级尺度归一化融合和三阶段迁移流程，有效训练混合模型。

方法

作者提出 HydraHead，一种沿头轴混合全注意力（FA）和线性注意力（LA）的新颖架构，以平衡计算效率与表示表达性。该方法包含三个关键组件：可解释性驱动的头选择策略、头级混合注意力模块以及三阶段迁移学习流程。

基于因果干预的头重要性估计

为确定哪些头需要 FA 的精度，哪些可以利用 LA 的效率，作者采用因果干预技术估计头的必要性。他们使用激活修补测量每个头对目标行为（接收者）的直接因果效应，并使用路径修补追踪上游贡献（发送者）。通过融合多个目标能力上的这些分数，为每个头计算统一的重要性分数 $S_h$Sh​。然后对头进行排序，仅对因果上不可或缺的前 K 个头保留 FA，其余分配给 LA。

头级混合

作者探索了各种混合粒度以确定最佳设计空间。如下图所示，他们比较了层级、token 级和头级范式，最终采用头级选择策略，根据功能重要性将特定头专门分配给 FA 或 LA 分支。

核心架构将查询头集合 $\mathcal{H}$H 划分为两个不相交的子集：用于 FA 的 $\mathcal{H}_F$HF​ 和用于 LA 分支（具体为门控 DeltaNet，GDN）的 $\mathcal{H}_L$HL​。下方的框架图详细展示了这一过程。输入投影被并行路由到各自的分支。

混合 FA 和 LA 的一个根本挑战在于它们输出之间的分布差距。FA 产生由查询范数调制的尖锐、低熵分布，而 LA 产生更平滑、更高熵的表示。为调和这一差异，作者提出头级尺度归一化融合模块。他们对每个头的输出 $\mathbf{O}_h$Oh​ 独立应用 RMSNorm，以统一特征尺度：

$$\hat{\mathbf{O}}_h = \mathrm{Norm}(\mathbf{O}_h)$$O^h​=Norm(Oh​)

这些归一化后的输出沿头维度拼接，以保留功能身份。然后引入可学习的头级缩放向量 $\boldsymbol{\gamma} \in \mathbb{R}^H$γ∈RH，自适应地重新校准每个头的贡献：

$$\tilde{\mathbf{O}}_{:,h:} = \boldsymbol{\gamma}_h \cdot \hat{\mathbf{O}}_{:,h:}, \quad \forall h \in [1, H]$$O~:,h:​=γh​⋅O^:,h:​,∀h∈[1,H]

调制后的张量被重塑并投影，产生最终的注意力输出。

分支特定的优化

两种注意力类型的内部结构设计也经过优化，以增强头级混合下的性能。下图展示了 FA 和 LA 分支的具体配置。

对于 FA 分支，作者移除旋转位置嵌入（RoPE），改为对查询特征应用对数尺度系数，以在长上下文场景下稳定注意力分布。同时引入辅助门控分支，以提升表示容量并缓解注意力沉没现象。对于 GDN 分支，RoPE 被显式集成到查询和键投影中，以弥补线性递归有限的位置敏感性。此外，键值头的数量被扩展以匹配查询头，从分组查询注意力配置过渡到类似多头注意力的设置，以增强表示容量。

高效混合迁移学习

为将预训练的标准 FA 模型转换为混合架构，并最小化训练开销，作者利用三阶段迁移流程。

阶段一：参数迁移与层级输出对齐 FA 头用预训练权重初始化，并添加一个轻量级门控分支，其权重接近零，以初始近似恒等函数。LA 头复用原始 FA 层的 $Q, K, V$Q,K,V 投影权重，通过通道级重复处理维度不匹配。模型训练以使用均方误差损失对齐每个混合层 $\mathbf{H}_{Hybrid}^{(l)}$HHybrid(l)​ 与原始 FA 层 $\mathbf{H}_{FA}^{(l)}$HFA(l)​ 的隐藏状态：

$$\mathcal{L}_{align} = \sum_{l=1}^{L} || \mathbf{H}_{FA}^{(l)}(x) - \mathbf{H}_{Hybrid}^{(l)}(x) ||_2^2$$Lalign​=l=1∑L​∣∣HFA(l)​(x)−HHybrid(l)​(x)∣∣22​

阶段二：全局 Logits 蒸馏 解冻整个模型，进行全局知识蒸馏，将学生模型的最终输出分布与教师模型对齐。该阶段结合 KL 散度损失和交叉熵损失，以确保全局语义一致性：

$$\mathcal{L}_{KD} = D_{KL}(P_{teacher}(\cdot|x) \parallel P_{student}(\cdot|x))$$LKD​=DKL​(Pteacher​(⋅∣x)∥Pstudent​(⋅∣x))

阶段三：长上下文微调 最后，模型使用下一 token 预测目标进行标准监督微调，采用扩展的上下文长度，以巩固长上下文能力：

$$\mathcal{L}_{NTP} = - \sum_{t} \log P_{student}(x_{t+1} | x_{1:t})$$LNTP​=−t∑​logPstudent​(xt+1​∣x1:t​)

实验

实验评估了 HydraHead，一种头级混合注意力架构，它通过因果可解释性选择少数头分配全注意力，其余分配线性注意力，并利用三阶段迁移学习流程在 Qwen3-1.7B 上转换预训练 Transformer。受控比较表明，头级混合持续优于层级和 token 级替代方案，在长上下文检索与通用推理之间实现了更优平衡，而消融研究证实特征归一化、查询分解和头级尺度调制对于异构注意力特征的稳定融合至关重要。可解释性引导的选择至关重要，因为检索关键头稀疏、分散在各层且经过因果验证，使得逐头分配远比均匀或随机策略有效。在扩展到超过 150 亿 tokens 后，HydraHead 超越现有混合模型，在推理能力上匹配标准 Transformer，同时保持对 256K 上下文的鲁棒外推，证明细粒度头级混合是高效长上下文模型的一个有前景的设计原则。

作者研究了头级混合架构的特征融合策略，比较了无归一化的直接拼接、头级尺度调制和头级门控竞争。结果表明，特征归一化至关重要，移除它会导致长上下文检索任务性能大幅下降。头级尺度调制成为最有效的融合方法，在扩展上下文场景中显著优于门控竞争，同时保持稳健的通用推理能力。移除特征归一化会导致在原生和扩展上下文长度下的单键检索性能显著退化。头级尺度调制为长程依赖提供了最稳定的表示，在扩展上下文任务中相较门控竞争取得了显著提升。头级尺度调制变体在几乎所有评估指标上持续优于其他融合策略，确立了其最优默认模型的地位。

作者通过扩大数据量、增加批量大小并在最后阶段扩展上下文长度，优化了头级混合架构的三阶段训练流程。这一优化配置在评估基准上带来了持续的性能提升，尤其增强了长上下文多键检索和复杂通用推理能力。相比之下，基线层级混合模型在相同设置下表现出有限甚至发散的扩展行为，凸显了所提架构吸收额外训练数据的优越能力。优化后的训练设置显著提升了混合模型在复杂长程依赖和挑战性推理任务上的准确率。在后期训练阶段增加上下文长度和批量大小稳定了异构注意力分支的优化动态。与在某些检索任务中性能下降的层级替代方案相比，头级混合设计展现出更高的利用规模化训练数据的能力。

作者在长上下文检索基准上评估了所提头级混合架构与标准全注意力模型及其他混合基线的性能。结果表明，所提方法在高达 256K tokens 的扩展上下文长度下保持高检索准确率，而竞争混合模型和标准注意力变体则遭受严重的性能退化。这证明了该模型在平衡高效长上下文外推与鲁棒检索能力方面的优越性。所提头级混合架构在长上下文检索任务上达到了最先进性能，在扩展序列长度下显著优于其他混合设计。标准全注意力模型无法在其原生上下文窗口之外检索信息，而所提方法即使在极端上下文长度下仍保持高准确率。头级混合设计有效防止了竞争混合模型在处理远超训练上下文长度的序列时出现的严重性能崩溃。

作者将具有标准比例的层级混合架构与采用更激进线性注意力比例的头级 Global-Interp-C 模型进行比较。结果表明，头级方法维持了与层级基线大致相当的长上下文检索性能。至关重要的是，尽管线性注意力比例更高，头级模型在通用推理能力上展现出显著优势，尤其在挑战性基准上。具有更高线性注意力比例的头级模型实现了与较低比例层级基线相当的长上下文性能。头级架构在通用推理任务上显著优于层级混合模型，在困难基准上取得了大幅提升。可解释性引导的头选择使得在保持强大通用领域能力的同时实现激进压缩。

作者通过逐步增加批量大小和上下文长度，优化了头级混合架构的三阶段训练流程。在这些阶段中扩展训练数据并调整超参数，显著增强了知识迁移并稳定了异构注意力分支的优化动态。这一优化配置在长上下文检索和复杂推理基准上带来了持续的性能提升，表明头级混合比层级替代方案更有效地吸收额外训练数据。训练配置逐步增加批量大小，并在最后阶段大幅扩展上下文长度以适应长程依赖。学习率调度在初始训练阶段采用余弦衰减，然后切换为恒定速率进行最终的长上下文适应。在三个阶段中扩展 token 数量和批量大小，使头级架构能够有效内化异构注意力模式。

实验通过融合策略比较、训练流程优化、高达 256K tokens 的长上下文检索基准以及激进线性注意力比例下的权衡，评估了头级混合注意力架构。特征归一化被证明至关重要，头级尺度调制成为最有效的融合方法，产生稳定的长程表示。在三阶段流程中扩展数据、批量大小和上下文长度，持续改善了长上下文检索和复杂推理，头级设计吸收额外数据的能力远优于层级替代方案。该架构在其他模型崩溃的极端长度下保持高检索准确率，且可解释性引导的头选择允许更高比例的线性注意力，同时保持强大的通用推理能力，确立了头级混合作为一种鲁棒且可扩展的方法。