HydraHead: Von funktionaler Heterogenität auf Kopfebene zur spezialisierten Aufmerksamkeitshybridisierung

Zusammenfassung

Die quadratische Komplexität von Attention stellt einen kritischen Engpass für die Verarbeitung langer Kontexte dar und weckt Interesse an hybriden Attention-Designs. Die meisten Open-Source-Hybridmodelle verwenden eine schichtweise Strategie. Doch frühere Arbeiten haben auf die inhärente Schwierigkeit der Integration von Linear Attention (LA) mit Full Attention (FA) hingewiesen, was darauf hindeutet, dass der Entwurfsraum der Attention-Hybridisierung noch wenig erforscht ist. Um diesen Raum zu untersuchen, führen wir eine Interpretierbarkeitsanalyse durch und beobachten, dass Schichten blockweise funktionale Ähnlichkeit aufweisen, während einzelne Köpfe innerhalb derselben Schicht trotz gemeinsamer Eingabemerkmale eine unterschiedliche funktionale Spezialisierung zeigen. Diese Heterogenität auf Kopfebene legt nahe, dass die Kopfdimension eine natürliche und prinzipientreue Granularität für die Fusion heterogener Attention-Signale bietet. Aufbauend auf dieser Erkenntnis stellen wir HydraHead vor, eine neuartige Architektur, die FA und LA entlang der Kopfachse hybridisiert. HydraHead zeichnet sich durch zwei zentrale Innovationen aus: (1) eine interpretierbarkeitsgesteuerte Auswahlstrategie, die abrufkritische Köpfe identifiziert und FA nur für diese beibehält, und (2) ein skalennormalisiertes Fusionsmodul, das den Verteilungsunterschied zwischen FA- und LA-Kopfausgaben ausgleicht. Durch die Nutzung einer dreistufigen Transfer-Pipeline mit Parameterwiederverwendung und Distillation erreichen wir leistungsstarke Hybridmodelle mit minimalem Trainingsaufwand. Unter einem einheitlichen Trainingssetup übertrifft HydraHead andere hybride Designs bei langen Kontextaufgaben und behält gleichzeitig eine starke allgemeine Argumentationsfähigkeit bei. Mit interpretierbarkeitsgesteuerter Kopfauswahl erreicht es die Langkontextleistung eines 3:1-schichtweisen Hybriden bei einem LA-zu-FA-Verhältnis von 7:1. Entscheidend ist, dass HydraHead, trainiert mit nur 15B Tokens, eine Verbesserung von über 69% gegenüber der Baseline bei einer Kontextlänge von 512K erzielt und damit an Qwen3.5 heranreicht, ein führendes Modell vergleichbarer Größe mit einer nativen Kontextlänge von 256K. Dies unterstreicht das erhebliche Skalierungspotenzial der Hybridisierung auf Kopfebene.

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.