论文
HydraHead:从头部功能异质性到专业注意力杂交
HydraHead: From Head-Level Functional Heterogeneity to Specialized Attention Hybridization
摘要
注意力的二次复杂度构成了长上下文处理的关键瓶颈,激发了人们对混合注意力设计的兴趣。大多数开源混合模型采用分层策略。然而,先前的工作已经注意到将线性注意力(LA)与完全注意力(FA)整合的固有困难,这表明注意力混合的设计空间仍未得到充分探索。为了探索这个空间,我们进行了可解释性分析,并观察到各层表现出块级功能相似性,而同一层内的各个头尽管共享输入特征,但显示出不同的功能专业化。这种头部层面的异质性表明头部维度为融合异质注意力信号提供了自然且有原则的粒度。基于这一见解,我们引入了 HydraHead,这是一种沿头轴混合 FA 和 LA 的新颖架构。 HydraHead 具有两项关键创新:(1) 可解释性驱动的选择策略,可识别检索关键头并仅为其保留 FA;(2) 尺度归一化融合模块,可协调 FA 和 LA 头输出之间的分布差距。通过利用具有参数重用和 蒸馏 的三级传输管道,我们以最小的训练开销实现了高性能混合模型。在统一的训练设置下,HydraHead 在长上下文任务中优于其他混合设计,同时保持强大的一般推理能力。通过可解释性驱动的头部选择,它可以以 7:1 LA 与 FA 的比率匹配 3:1 分层混合的长上下文性能。至关重要的是,仅在 15B 个 token 上进行训练,HydraHead 在 512K 上下文长度下比基线提高了 69% 以上,接近 Qwen3.5,这是一个具有 256K 原生上下文长度的同等大小的领先模型。这凸显了头级杂交的巨大扩展潜力。