HydraHead : De l'hétérogénéité fonctionnelle au niveau des têtes à l'hybridation d'attention spécialisée

Résumé

La complexité quadratique de l’attention constitue un goulot d’étranglement critique pour le traitement de longs contextes, suscitant un intérêt accru pour les architectures d’attention hybrides. La plupart des modèles hybrides open source adoptent une stratégie par couche. Cependant, des travaux antérieurs ont souligné la difficulté intrinsèque d’intégrer l’attention linéaire (LA) avec l’attention complète (FA), suggérant que l’espace de conception de l’hybridation de l’attention reste sous-exploré. Pour sonder cet espace, nous menons une analyse d’interprétabilité et observons que les couches présentent une similarité fonctionnelle par blocs, tandis que les têtes individuelles au sein d’une même couche affichent une spécialisation fonctionnelle distincte malgré le partage des caractéristiques d’entrée. Cette hétérogénéité au niveau des têtes suggère que la dimension des têtes constitue une granularité naturelle et fondée pour fusionner des signaux d’attention hétérogènes. Sur la base de cette observation, nous introduisons HydraHead, une nouvelle architecture qui hybride FA et LA selon l’axe des têtes. HydraHead présente deux innovations clés : (1) une stratégie de sélection guidée par l’interprétabilité qui identifie les têtes critiques pour la récupération d’information et ne conserve FA que pour celles-ci, et (2) un module de fusion normalisé par échelle qui concilie l’écart distributionnel entre les sorties des têtes FA et LA. En exploitant un pipeline de transfert en trois étapes avec réutilisation des paramètres et distillation, nous obtenons des modèles hybrides performants avec un surcoût d’entraînement minimal. Dans un cadre d’entraînement unifié, HydraHead surpasse les autres architectures hybrides dans les tâches à long contexte tout en maintenant de solides capacités de raisonnement général. Grâce à une sélection des têtes guidée par l’interprétabilité, elle atteint les performances d’un hybride par couche avec un ratio 3:1 pour un ratio LA/FA de 7:1. Plus important encore, entraînée sur seulement 15 milliards de jetons, HydraHead améliore de plus de 69 % les performances de la ligne de base pour une longueur de contexte de 512 000, se rapprochant de Qwen3.5, un modèle leader de taille comparable avec une longueur de contexte native de 256 000. Cela souligne le potentiel de passage à l’échelle significatif de l’hybridation au niveau des têtes.

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.