HydraHead: De la heterogeneidad funcional a nivel de cabezas a la hibridación de atención especializada

Resumen

La complejidad cuadrática de la atención supone un cuello de botella crítico para el procesamiento de contextos largos, lo que ha despertado interés en diseños de atención híbrida. La mayoría de los modelos híbridos de código abierto adoptan una estrategia por capas. Sin embargo, trabajos previos han señalado la dificultad inherente de integrar la Atención Lineal (LA) con la Atención Completa (FA), sugiriendo que el espacio de diseño de la hibridación de la atención sigue siendo poco explorado. Para explorar este espacio, realizamos un análisis de interpretabilidad y observamos que las capas exhiben una similitud funcional a nivel de bloques, mientras que los cabezales individuales dentro de una misma capa muestran una especialización funcional distinta a pesar de compartir características de entrada. Esta heterogeneidad a nivel de cabezales sugiere que la dimensión de los cabezales proporciona una granularidad natural y fundamentada para fusionar señales de atención heterogéneas. Partiendo de esta idea, presentamos HydraHead, una arquitectura novedosa que hibridiza FA y LA a lo largo del eje de los cabezales. HydraHead cuenta con dos innovaciones clave: (1) una estrategia de selección basada en interpretabilidad que identifica los cabezales críticos para la recuperación de información y preserva FA solo para ellos, y (2) un módulo de fusión con normalización de escala que reconcilia la brecha distribucional entre las salidas de los cabezales FA y LA. Mediante el aprovechamiento de un pipeline de transferencia de tres etapas con reutilización de parámetros y destilación, logramos modelos híbridos de alto rendimiento con una sobrecarga de entrenamiento mínima. Bajo una configuración de entrenamiento unificada, HydraHead supera a otros diseños híbridos en tareas de contexto largo, manteniendo a la vez un sólido razonamiento general. Con una selección de cabezales guiada por interpretabilidad, iguala el rendimiento en contexto largo de un híbrido por capas en proporción 3:1 con una relación LA:FA de 7:1. De manera crucial, entrenado con solo 15B de tokens, HydraHead logra una mejora superior al 69% sobre la línea base en una longitud de contexto de 512K, acercándose a Qwen3.5, un modelo líder de tamaño comparable con una longitud de contexto nativa de 256K. Esto destaca el significativo potencial de escalado de la hibridación a nivel de cabezales.

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.