HydraHead: Da Heterogeneidade Funcional de Nível de Cabeça à Hibridização de Atenção Especializada

Resumo

A complexidade quadrática da atenção impõe um gargalo crítico para o processamento de contextos longos, despertando interesse em designs híbridos de atenção. A maioria dos modelos híbridos de código aberto adota uma estratégia por camadas. No entanto, trabalhos anteriores notaram a dificuldade inerente de integrar Atenção Linear (LA) com Atenção Total (FA), sugerindo que o espaço de projeto da hibridização da atenção permanece pouco explorado. Para investigar esse espaço, realizamos análises de interpretabilidade e observamos que as camadas exibem similaridade funcional em blocos, enquanto cabeças individuais dentro da mesma camada apresentam especialização funcional distinta, apesar de compartilharem características de entrada. Essa heterogeneidade ao nível das cabeças sugere que a dimensão das cabeças fornece uma granularidade natural e fundamentada para fundir sinais de atenção heterogêneos. Com base nessa percepção, apresentamos o HydraHead, uma nova arquitetura que hibridiza FA e LA ao longo do eixo das cabeças. O HydraHead apresenta duas inovações principais: (1) uma estratégia de seleção orientada por interpretabilidade que identifica cabeças críticas para recuperação e preserva FA apenas para elas, e (2) um módulo de fusão com normalização de escala que reconcilia a lacuna distribucional entre as saídas das cabeças FA e LA. Ao alavancar um pipeline de transferência de três estágios com reutilização de parâmetros e destilação, alcançamos modelos híbridos de alto desempenho com mínimo custo de treinamento. Sob uma configuração de treinamento unificada, o HydraHead supera outros designs híbridos em tarefas de contexto longo, mantendo um raciocínio geral forte. Com a seleção de cabeças orientada por interpretabilidade, ele iguala o desempenho em contexto longo de um híbrido por camadas na proporção 3:1 com uma proporção LA para FA de 7:1. Crucialmente, treinado em apenas 15B tokens, o HydraHead alcança mais de 69% de melhoria em relação à linha de base em comprimento de contexto de 512K, aproximando-se do Qwen3.5, um modelo líder de tamanho comparável com comprimento de contexto nativo de 256K. Isso destaca o significativo potencial de escalabilidade da hibridização ao nível das cabeças.

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.