HydraHead: From Head-Level Functional Heterogeneity to Specialized Attention Hybridization

TL;DR

HydraHead hybridizes Full Attention (FA) and Linear Attention (LA) along the head axis within a single layer, rather than the now-standard layer-wise interleaving used by Kimi Linear, Jet-Nemotron, MiniMax-M1, and Qwen3-Next. The argument is interpretability-driven: layers exhibit block-wise functional similarity, but individual heads in the same layer specialize despite sharing input features — so the head, not the layer, is the natural granularity for fusing heterogeneous attention signals. A retrieval-critical head selector preserves FA only on heads that need it, a scale-normalized fusion module reconciles the FA/LA distributional gap, and a three-stage transfer pipeline (parameter reuse + distillation) keeps training cost low. Trained on only 15B tokens, the resulting hybrid matches a 3:1 layer-wise hybrid’s long-context score at a 7:1 LA-to-FA ratio and approaches Qwen3.5 at 512K context with >69% improvement over baseline.

Key claims

Layers in a Transformer exhibit block-wise functional similarity, while individual heads in the same layer display distinct functional specialization despite sharing input features — making the head dimension a principled granularity for fusing heterogeneous attention signals [Abstract, §Interpretability Analysis].
An interpretability-driven head selector identifies retrieval-critical heads and preserves FA only on those, with the remaining heads converted to LA [Abstract].
A scale-normalized fusion module is required to reconcile the distributional gap between FA and LA head outputs before concatenation [Abstract].
A three-stage transfer pipeline with parameter reuse and distillation converts a pretrained FA model to HydraHead with minimal training overhead [Abstract].
Under a unified training setup, HydraHead outperforms other hybrid designs (layer-wise hybrids) on long-context tasks while maintaining strong general reasoning [Abstract].
With interpretability-driven head selection, HydraHead matches a 3:1 layer-wise hybrid’s long-context performance at a 7:1 LA-to-FA ratio — i.e. retains less FA for the same long-context quality [Abstract].
Trained on only 15B tokens, HydraHead achieves >69% improvement over the baseline at 512K context length and approaches Qwen3.5 (a comparable-size model with native 256K context) [Abstract].

Method

HydraHead is positioned as an intra-layer, head-dimension hybrid. The construction begins with an interpretability analysis of a pretrained FA model: across layers, behavior is grouped block-wise, but within a layer, individual attention heads route to distinct functional roles. From this observation, the selector picks a small subset of heads judged retrieval-critical (these stay FA) and converts the remainder to a linear-attention variant; the resulting per-head outputs are then concatenated through a scale-normalized fusion module that corrects for the very different output magnitudes FA and LA produce on the same input. Conversion is done via a three-stage transfer pipeline that reuses the original FA parameters (KV projection weights for the LA-side heads, fusion-module initialization) and distills the original model on a 15B-token budget. The headline configuration is a 7:1 LA-to-FA ratio at the head level.

Results

Outperforms other hybrid designs in long-context tasks under a unified training setup while maintaining strong general reasoning [Abstract].
Matches a 3:1 layer-wise hybrid’s long-context performance at a 7:1 LA-to-FA ratio [Abstract].
69% improvement over the FA baseline at 512K context length, on only 15B training tokens [Abstract].
Approaches Qwen3.5 (a comparable-size model with native 256K context) at 512K [Abstract].

Why it’s interesting

This is the first filed datapoint that argues the head, not the layer, is the right hybridization granularity — directly orthogonal to every layer-wise recipe currently in Hybrid Linear Attention (Kimi Linear’s 3:1 KDA:MLA, MiniMax-M1’s 7:1 lightning:softmax, Qwen3-Next, Jet-Nemotron, Nemotron-3-Super). If the 7:1 head-axis result actually matches 3:1 layer-axis quality, it implies the layer-wise recipes are over-allocating FA capacity to heads that don’t need it, which would reset the open question on “right linear:softmax ratio” that the concept tracks. The interpretability-first construction also makes this a clean instance of “mechanistic story → architectural intervention”, complementing Why Can't Transformers Learn Multiplication? Reverse-Engineering Reveals Long-Range Dependency Pitfalls‘s mechanistic-story-→-training-intervention result and reinforcing the within-layer head-specialization line started by Interleaved Head Attention (which broke the one-pattern-per-head bottleneck in a different way) and the Sparsity is Cool: Reverse-Engineering MoBA and NSA reverse-engineering of MoBA/NSA head behavior.