HydraHead: От функциональной гетерогенности на уровне голов к гибридизации специализированного внимания

Аннотация

Квадратичная сложность внимания создает критическое узкое место для обработки длинных контекстов, стимулируя интерес к гибридным архитектурам внимания. Большинство открытых гибридных моделей используют послойную стратегию. Однако предыдущие работы отмечали принципиальную сложность интеграции линейного внимания (LA) с полным вниманием (FA), что указывает на недостаточную изученность пространства проектирования гибридного внимания. Для исследования этого пространства мы проводим анализ интерпретируемости и обнаруживаем, что слои демонстрируют функциональное сходство на уровне блоков, в то время как отдельные головы внутри одного слоя, несмотря на общие входные признаки, проявляют отчетливую функциональную специализацию. Эта гетерогенность на уровне голов указывает на то, что размерность голов обеспечивает естественную и обоснованную степень детализации для слияния гетерогенных сигналов внимания. Основываясь на этом понимании, мы представляем HydraHead — новую архитектуру, которая гибридизирует FA и LA вдоль оси голов. HydraHead отличается двумя ключевыми инновациями: (1) стратегией отбора на основе интерпретируемости, которая идентифицирует головы, критически важные для извлечения информации, и сохраняет FA только для них, и (2) модулем слияния с нормализацией масштаба, устраняющим расхождение в распределениях выходов голов FA и LA. Используя трехэтапный конвейер переноса с повторным использованием параметров и дистилляцией, мы достигаем высокопроизводительных гибридных моделей с минимальными вычислительными затратами на обучение. В единой схеме обучения HydraHead превосходит другие гибридные конструкции в задачах с длинным контекстом, сохраняя при этом высокое качество общих рассуждений. Благодаря отбору голов на основе интерпретируемости, она достигает производительности на длинных контекстах, сопоставимой с послойным гибридом 3:1, при соотношении LA к FA 7:1. Важно, что, обученная всего на 15 миллиардах токенов, HydraHead демонстрирует улучшение более чем на 69% по сравнению с базовой моделью при длине контекста 512 000, приближаясь к Qwen3.5 — ведущей модели сопоставимого размера с родной длиной контекста 256 000. Это подчеркивает значительный потенциал масштабирования гибридизации на уровне голов.

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.