HydraHead: Van functionele heterogeniteit op head-niveau naar gespecialiseerde aandachtshybridisatie

Samenvatting

De kwadratische complexiteit van aandacht vormt een kritieke bottleneck voor het verwerken van lange contexten, wat de interesse in hybride aandachtontwerpen aanwakkert. De meeste open-source hybride modellen hanteren een laagsgewijze strategie. Eerder werk heeft echter de inherente moeilijkheid opgemerkt van het integreren van Lineaire Aandacht (LA) met Volledige Aandacht (FA), wat suggereert dat de ontwerpruimte van aandachthybridisatie nog onderbelicht is. Om deze ruimte te verkennen, voeren we een interpreteerbaarheidsanalyse uit en observeren we dat lagen bloksgewijze functionele gelijkenis vertonen, terwijl individuele heads binnen dezelfde laag een duidelijke functionele specialisatie vertonen, ondanks het delen van invoerkenmerken. Deze heterogeniteit op head-niveau suggereert dat de head-dimensie een natuurlijke en principiële granulariteit biedt voor het fuseren van heterogene aandachtssignalen. Voortbouwend op dit inzicht introduceren we HydraHead, een nieuwe architectuur die FA en LA hybridiseert langs de head-as. HydraHead kent twee belangrijke innovaties: (1) een op interpreteerbaarheid gebaseerde selectiestrategie die retrieval-kritische heads identificeert en alleen voor hen FA behoudt, en (2) een schaal-genormaliseerde fusiemodule die de distributionele kloof tussen FA- en LA-head-uitgangen overbrugt. Door gebruik te maken van een drietraps overdrachtspijplijn met parameterhergebruik en distillatie, bereiken we hoogwaardige hybride modellen met minimale trainingsoverhead. Onder een uniforme trainingsopzet presteert HydraHead beter dan andere hybride ontwerpen in taken met een lange context, terwijl het een sterke algemene redeneervaardigheid behoudt. Met op interpreteerbaarheid gebaseerde head-selectie evenaart het de prestaties bij lange context van een 3:1 laagsgewijze hybride bij een LA-tot-FA verhouding van 7:1. Cruciaal is dat HydraHead, getraind op slechts 15B tokens, een verbetering van meer dan 69% behaalt ten opzichte van de baseline bij een contextlengte van 512K, en daarmee Qwen3.5 benadert, een toonaangevend model van vergelijkbare grootte met een native contextlengte van 256K. Dit benadrukt het significante schaalpotentieel van hybridisatie op head-niveau.

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.