One-sentence Summary

HydraHead, a head-level attention hybridization architecture, employs interpretability-driven selection to preserve full attention only for retrieval-critical heads and a scale-normalized fusion module to reconcile output distributions, achieving over 69% improvement over the baseline at 512K context length and matching a 3:1 layer-wise hybrid's long-context performance at a 7:1 linear-to-full attention ratio after training on only 15B tokens.

Key Contributions

• Interpretability analysis reveals that transformer layers exhibit block-wise functional similarity while attention heads within a layer specialize in distinct functions, establishing the head dimension as a principled granularity for fusing heterogeneous attention mechanisms.
• HydraHead hybridizes full and linear attention along the head axis, using an interpretability-driven selection strategy to retain full attention only on retrieval-critical heads and a scale-normalized fusion module to reconcile the distributional mismatch between the two attention types.
• A three-stage transfer pipeline with parameter reuse and distillation trains HydraHead, achieving over 69% improvement on NIAH benchmarks at 512K context with 15B training tokens, matching a 3:1 layer-wise hybrid’s long-context performance at a 7:1 LA-to-FA ratio, and approaching Qwen3.5 without degrading general reasoning.

Introduction

The shift from static LLMs to autonomous agents that plan and reason over long horizons has made extending context windows essential, but standard full attention scales quadratically with sequence length. Linear attention mechanisms offer linear complexity yet often suffer from expressivity collapse, struggling with precise retrieval. Hybrid architectures that interleave different attention types across layers have emerged as a compromise, but layer-wise designs face a challenge: layer outputs vary smoothly, making it hard to decide where to place each mechanism, and training such hybrids remains difficult, leaving the design space underexplored. The authors leverage mechanistic interpretability to show that individual attention heads within a layer exhibit sharp functional heterogeneity, with only a sparse subset critical for retrieval, while layers lack clear boundaries. They propose HydraHead, a fine-grained head-level hybridization architecture that uses a lightweight, one-shot interpretability-based selection to assign full attention only to retrieval-critical heads and linear attention to the rest, combined with a head-wise scale-normalized fusion and a three-stage transfer pipeline to effectively train the hybrid model.

Method

The authors introduce HydraHead, a novel architecture that hybridizes Full Attention (FA) and Linear Attention (LA) along the head axis to balance computational efficiency and representational expressivity. The method comprises three key components: an interpretability-driven head selection strategy, the head-wise hybrid attention module, and a three-stage transfer learning pipeline.

Head Importance Estimation via Causal Intervention

To determine which heads require the precision of FA and which can utilize the efficiency of LA, the authors employ causal intervention techniques to estimate head necessity. They use activation patching to measure the direct causal effect of each head on target behaviors (receivers) and path patching to trace upstream contributions (senders). By fusing these scores across multiple target capabilities, they compute a unified importance score $S_h$Sh​ for each head. Heads are then ranked, and FA is retained only for the top-K causally indispensable heads, while the remainder are assigned to LA.

Head-wise Hybridization

The authors explore various hybridization granularities to identify the optimal design space. As illustrated in the figure below, they compare layer-wise, token-wise, and head-wise paradigms, ultimately adopting a head-wise selection strategy where specific heads are exclusively assigned to either the FA or LA branch based on functional importance.

The core architecture partitions the set of query heads $\mathcal{H}$H into two disjoint subsets: $\mathcal{H}_F$HF​ for FA and $\mathcal{H}_L$HL​ for the LA branch (specifically Gated DeltaNet, GDN). The framework diagram below details this process. Input projections are routed to their respective branches in parallel.

A fundamental challenge in hybridizing FA and LA is the distributional gap between their outputs. FA produces sharp, low-entropy distributions modulated by query norm, while LA yields smoother, higher-entropy representations. To reconcile this, the authors propose a Head-wise Scale-normalized Fusion module. They apply RMSNorm independently to each head's output $\mathbf{O}_h$Oh​ to unify feature scales:

$$\hat{\mathbf{O}}_h = \mathrm{Norm}(\mathbf{O}_h)$$O^h​=Norm(Oh​)

These normalized outputs are concatenated along the head dimension to preserve functional identity. A learnable head-wise scaling vector $\boldsymbol{\gamma} \in \mathbb{R}^H$γ∈RH is then introduced to adaptively recalibrate the contribution of each head:

$$\tilde{\mathbf{O}}_{:,h:} = \boldsymbol{\gamma}_h \cdot \hat{\mathbf{O}}_{:,h:}, \quad \forall h \in [1, H]$$O~:,h:​=γh​⋅O^:,h:​,∀h∈[1,H]

The modulated tensor is reshaped and projected to produce the final attention output.

Branch-specific Refinements

The internal structural designs of the two attention types are also optimized to enhance performance under head-wise hybridization. The figure below illustrates the specific configurations for the FA and LA branches.

For the FA branch, the authors remove Rotary Position Embedding (RoPE) and instead apply a log-scale coefficient to query features to stabilize attention distributions in long-context scenarios. An auxiliary gate branch is also introduced to boost representational capacity and alleviate the attention sink phenomenon. For the GDN branch, RoPE is explicitly integrated into query and key projections to compensate for the limited positional sensitivity of linear recurrence. Additionally, the number of key-value heads is expanded to match the query heads, transitioning from a Grouped-Query Attention configuration to a Multi-Head Attention-like setup to enhance representational capacity.

Efficient Hybrid Transfer Learning

To convert a pre-trained standard FA model into the hybrid architecture with minimal training overhead, the authors leverage a three-stage transfer pipeline.

Stage 1: Parameter Migration and Layer-wise Output Alignment The FA heads are initialized with pre-trained weights, and a lightweight gate branch is added with near-zero weights to approximate an identity function initially. The LA heads reuse the $Q, K, V$Q,K,V projection weights from the original FA layers, using channel-wise repetition to handle dimension mismatches. The model is trained to align the hidden states of each hybrid layer $\mathbf{H}_{Hybrid}^{(l)}$HHybrid(l)​ with the original FA layer $\mathbf{H}_{FA}^{(l)}$HFA(l)​ using Mean Squared Error loss:

$$\mathcal{L}_{align} = \sum_{l=1}^{L} || \mathbf{H}_{FA}^{(l)}(x) - \mathbf{H}_{Hybrid}^{(l)}(x) ||_2^2$$Lalign​=l=1∑L​∣∣HFA(l)​(x)−HHybrid(l)​(x)∣∣22​

Stage 2: Global Logits Distillation The entire model is unfrozen, and global knowledge distillation is performed to align the final output distribution of the student model with the teacher model. This stage combines KL divergence loss with cross-entropy loss to ensure global semantic coherence:

$$\mathcal{L}_{KD} = D_{KL}(P_{teacher}(\cdot|x) \parallel P_{student}(\cdot|x))$$LKD​=DKL​(Pteacher​(⋅∣x)∥Pstudent​(⋅∣x))

Stage 3: Long-Context Fine-tuning Finally, the model undergoes standard supervised fine-tuning using the Next Token Prediction objective with extended context lengths to consolidate long-context capabilities:

$$\mathcal{L}_{NTP} = - \sum_{t} \log P_{student}(x_{t+1} | x_{1:t})$$LNTP​=−t∑​logPstudent​(xt+1​∣x1:t​)

Experiment

The experiments evaluate HydraHead, a head-wise hybrid attention architecture that converts a pretrained Transformer by assigning full attention to a small subset of heads selected through causal interpretability and linear attention to the rest, using a three-stage transfer learning pipeline on Qwen3-1.7B. Controlled comparisons show that head-wise hybridization consistently outperforms layer-wise and token-wise alternatives, achieving a superior balance between long-context retrieval and general reasoning, while ablation studies confirm that feature normalization, query decomposition, and head-wise scale modulation are critical for stable fusion of heterogeneous attention features. Interpretability-guided selection proves essential because retrieval-critical heads are sparse, scattered across layers, and causally verified, making per-head allocation far more effective than uniform or random strategies. Scaling to over 15 billion tokens, HydraHead surpasses existing hybrid models and matches standard Transformers in reasoning while maintaining robust extrapolation to 256K context, demonstrating that fine-grained head-wise mixing is a promising design principle for efficient long-context models.

The authors investigate feature fusion strategies for their head-wise hybrid architecture, comparing direct concatenation without normalization, head-wise scale modulation, and head-wise gated competition. Results show that feature normalization is critical, as removing it causes substantial performance drops in long-context retrieval tasks. Head-wise scale modulation emerges as the most effective fusion method, significantly outperforming gated competition in extended context scenarios while maintaining robust general reasoning capabilities. Removing feature normalization leads to significant degradation in single-key retrieval performance across both native and extended context lengths. Head-wise scale modulation provides the most stable representations for long-range dependencies, achieving substantial improvements over gated competition in extended context tasks. The head-wise scale modulation variant consistently outperforms other fusion strategies across nearly all evaluation metrics, establishing it as the optimal default model.

The authors optimize the three-stage training pipeline for their head-wise hybrid architecture by scaling up data volume, increasing batch sizes, and extending context lengths in the final stage. This optimized configuration yields consistent performance improvements across evaluation benchmarks, particularly enhancing long-context multi-key retrieval and complex general reasoning capabilities. In contrast, baseline layer-wise hybrid models show limited or divergent scaling behavior under the same settings, highlighting the proposed architecture's superior capacity to absorb additional training data. The optimized training setup significantly boosts the hybrid model's accuracy on intricate long-context dependencies and challenging reasoning tasks. Increasing the context length and batch size in the later training stages stabilizes optimization dynamics across the heterogeneous attention branches. The head-wise hybrid design demonstrates a higher capacity for leveraging scaled training data compared to layer-wise alternatives, which suffer from performance degradation in certain retrieval tasks.

The authors evaluate their proposed head-wise hybrid architecture against standard full attention models and other hybrid baselines on long-context retrieval benchmarks. Results show that the proposed method maintains high retrieval accuracy across extended context lengths up to 256K tokens, whereas competing hybrid models and standard attention variants suffer severe performance degradation. This demonstrates the model's superior ability to balance efficient long-context extrapolation with robust retrieval capabilities. The proposed head-wise hybrid architecture achieves state-of-the-art performance on long-context retrieval tasks, significantly outperforming other hybrid designs at extended sequence lengths. While standard full attention models fail to retrieve information beyond their native context window, the proposed method sustains high accuracy even at extreme context lengths. The head-wise hybrid design effectively prevents the severe performance collapse observed in competing hybrid models when handling sequences far beyond the training context length.

The authors compare a layer-wise hybrid architecture with a standard ratio against their head-wise Global-Interp-C model using a more aggressive ratio of linear attention. Results indicate that the head-wise approach maintains long-context retrieval performance broadly comparable to the layer-wise baseline. Crucially, the head-wise model demonstrates a substantial advantage in general reasoning capabilities, particularly on challenging benchmarks, despite the higher proportion of linear attention. The head-wise model with a higher linear attention ratio achieves long-context performance comparable to the layer-wise baseline with a lower ratio. The head-wise architecture significantly outperforms the layer-wise hybrid in general reasoning tasks, showing large gains on hard benchmarks. Interpretability-guided head selection enables aggressive compression while preserving strong general-domain capabilities.

The authors optimize a three-stage training pipeline for their head-wise hybrid architecture by progressively increasing the batch size and context length. Scaling up the training data and adjusting hyperparameters across these stages significantly enhances knowledge transfer and stabilizes optimization dynamics for heterogeneous attention branches. This optimized configuration leads to consistent performance improvements across long-context retrieval and complex reasoning benchmarks, demonstrating that head-wise hybridization absorbs additional training data more effectively than layer-wise alternatives. The training configuration progressively increases batch sizes and dramatically extends the context length in the final stage to accommodate long-context dependencies. The learning rate schedule utilizes cosine decay during the initial training phases before switching to a constant rate for the final long-context adaptation. Scaling up the token count and batch size across the three stages enables the head-wise architecture to effectively internalize heterogeneous attention patterns.

The experiments evaluate a head-wise hybrid attention architecture through fusion strategy comparisons, training pipeline optimization, long-context retrieval benchmarks up to 256K tokens, and trade-offs with aggressive linear attention ratios. Feature normalization proves essential, and head-wise scale modulation emerges as the most effective fusion method, yielding stable long-range representations. Scaling data, batch sizes, and context lengths in a three-stage pipeline consistently improves long-context retrieval and complex reasoning, with the head-wise design absorbing additional data far better than layer-wise alternatives. The architecture maintains high retrieval accuracy at extreme lengths where other models collapse, and interpretability-guided head selection allows a higher proportion of linear attention while preserving strong general reasoning, establishing the head-wise hybrid as a robust and scalable approach.