HydraHead：從頭級功能異質性到專業化注意力混合

摘要

注意力機制的二次複雜度成為長上下文處理的關鍵瓶頸，促使學界轉而關注混合注意力設計。目前多數開源混合模型採用逐層策略，然而先前研究已指出線性注意力（Linear Attention, LA）與全注意力（Full Attention, FA）整合上的固有困難，顯示注意力混合的設計空間仍有待深入探索。為此，我們透過可解釋性分析觀察到：層與層之間展現區塊式功能相似性，而同一層內不同注意力頭雖共享輸入特徵，卻呈現明確的功能特化。這種頭層級異質性表明，頭維度是融合異質注意力訊號的自然且具原則性的粒度。基於此發現，我們提出HydraHead——一種沿頭軸混合FA與LA的新型架構。HydraHead包含兩項關鍵創新：（1）基於可解釋性驅動的篩選策略，辨識出對檢索至關重要的注意力頭，並僅為其保留FA；（2）尺度歸一化融合模組，用以調和FA與LA頭輸出之間的分佈差異。透過結合參數重複利用與知識蒸餾的三階段遷移流程，我們以最少訓練開銷實現高效能混合模型。在統一訓練設定下，HydraHead在長上下文任務中優於其他混合設計，同時維持強大的通用推理能力。採用可解釋性驅動的頭篩選時，僅以7:1的LA對FA比例，即可達到3:1逐層混合模型的長上下文效能。關鍵的是，HydraHead僅使用150億個詞元的訓練資料，便在512K上下文長度下較基線模型提升超過69%，逼近規模相當且原生上下文長度為256K的領先模型Qwen3.5。這凸顯了頭層級混合的顯著擴展潛力。

English

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.