HydraHead:从头部级别的功能异质性到专门化注意力的混合化
HydraHead: From Head-Level Functional Heterogeneity to Specialized Attention Hybridization
摘要
注意力机制的二次复杂度在处理长上下文信息时构成了重要的障碍,因此人们开始研究混合式注意力机制。大多数开源混合模型都采用逐层处理的方式。然而,已有研究指出将线性注意力与全注意力机制结合起来存在困难,这表明注意力机制混合的设计空间尚未得到充分探索。为了进一步研究这一领域,我们进行了可解释性分析,发现各层之间存在功能上的相似性,而同一层中的各个注意力头则具有不同的功能特性,尽管它们共享相同的输入特征。这种不同注意力头的特性表明,注意力头的维度为融合不同注意力信号提供了自然的、合理的划分方式。基于这一发现,我们提出了HydraHead这一新型架构,该架构将全注意力与线性注意力结合在注意力头层面。HydraHead具有两个关键创新点:(1)一种以可解释性为导向的选择策略,能够识别出对检索至关重要的注意力头,并仅保留这些注意力头使用全注意力机制;(2)一个比例归一化的融合模块,能够解决全注意力与线性注意力输出之间的分布差异问题。通过采用三阶段传输流程以及参数重用和蒸馏技术,我们能够获得高性能的混合模型,同时减少训练成本。在统一的训练环境下,HydraHead在长上下文任务中表现优于其他混合模型,同时保持了良好的通用推理能力。由于采用了以可解释性为导向的注意力头选择策略,其在7:1的线性注意力与全注意力比例下,其性能与3:1的逐层混合模型在长上下文任务中的性能相当。重要的是,仅在15B个标记上训练后的HydraHead,在512K上下文长度下就能实现超过69%的性能提升,接近Qwen3.5的性能水平——后者是一个规模相当的模型,其原生上下文长度为256K。这凸显了注意力头层面混合机制的巨大扩展潜力。
English Abstract
The quadratic complexity of attention poses a critical bottleneck for long-context processing, spurring interest in hybrid attention designs. Most open-source hybrid models adopt a layer-wise strategy. Yet, prior work has noted the inherent difficulty of integrating Linear Attention (LA) with Full Attention (FA), suggesting that the design space of attention hybridization remains underexplored. To probe this space, we conduct interpretability analysis and observe that layers exhibit block-wise functional similarity, while individual heads within the same layer display distinct functional specialization despite sharing input features. This head-level heterogeneity suggests that the head dimension provides a natural and principled granularity for fusing heterogeneous attention signals. Building on this insight, we introduce HydraHead, a novel architecture that hybridizes FA and LA along the head axis. HydraHead features two key innovations: (1) an interpretability-driven selection strategy that identifies retrieval-critical heads and preserves FA only for them, and (2) a scale-normalized fusion module that reconciles the distributional gap between FA and LA head outputs. By leveraging a three-stage transfer pipeline with parameter reuse and distillation, we achieve high-performance hybrid models with minimal training overhead. Under a unified training setup, HydraHead outperforms other hybrid designs in long-context tasks while maintaining strong general reasoning. With interpretability-driven head selection, it matches a 3:1 layer-wise hybrid's long-context performance at a 7:1 LA-to-FA ratio. Crucially, trained on only 15B tokens, HydraHead achieves over 69% improvement over the baseline at 512K context length, approaching Qwen3.5, a leading model of comparable size with a native context length of 256K. This highlights the significant scaling potential of head-level hybridization.