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.
한 줄 요약
HydraHead는 해드 수준의 FA와 LA 혼합을 통해 512K 컨텍스트 길이에서 기존 모델 대비 69% 성능 향상을 달성한 새로운 어텐션 하이브리드 아키텍처이다.
핵심 기여도
- 해드 수준 기반 FA/LA 혼합 설계를 제안, 기존 레이어 수준 접근 대비 7:1의 LA-FA 비율로 동일 성능 달성
- 해드 선택 전략과 스케일 정규화 퓨전 모듈을 도입, 헤테로제네이티 어텐션 신호를 효과적으로 결합
- 15B 토큰 훈련으로 512K 컨텍스트 길이에서 기존 베이스라인 대비 69% 성능 향상
- Qwen3.5(256K) 수준의 성능을 512K 길이에서 달성
핵심 아이디어
기존 어텐션 하이브리드는 레이어 수준에서 FA와 LA를 혼합하지만, 이는 컨텍스트 확장 성능과 일반 추론 능력 사이에 트레이드오프를 유발한다. 본 연구는 해드 수준에서 기능적 이질성을 분석한 결과, 동일 레이어 내에서도 해드 간 기능적 특화가 뚜렷하게 나타남을 발견했다. 이에 따라, 해드 단위로 FA와 LA를 선택적으로 결합하는 새로운 설계 원칙을 제안한다. 구체적으로, 해드 선택 알고리즘을 통해 정밀 토큰 검색에 중요한 해드는 FA로 유지하고, 나머지는 LA로 대체한다. 또한, FA와 LA의 분포 차이를 해소하기 위해 스케일 정규화 퓨전 모듈을 도입하여, 서로 다른 어텐션 패턴을 안정적으로 결합한다.
기술적 접근법
- **HydraHead 아키텍처**: 해드 단위로 FA와 LA를 혼합, 25% 해드는 FA, 75%는 GDN(LA) 구조 사용
- **해드 선택 알고리즘**: 사전 학습된 모델에서 정밀 토큰 검색에 기여하는 해드를 해드별 로짓 기여도 분석을 통해 식별
- **스케일 정규화 퓨전 모듈**: FA와 LA의 출력 분포 차이를 해소, 헤테로제네이티 어텐션 신호 간 간섭 최소화
- **3단계 전이 학습 파이프라인**: 파라미터 재사용, 글로벌 디스틸레이션, 장문 훈련을 통해 훈련 오버헤드 최소화
- **데이터셋**: FineWeb-Edu, 15B 토큰으로 훈련
주요 결과
- **Needle in a Haystack 벤치마크**: 512K 컨텍스트 길이에서 기존 베이스라인 대비 69% 성능 향상
- **Qwen3.5 대비 성능**: 256K 컨텍스트 길이에서 Qwen3.5 수준의 성능을 512K 길이에서도 달성
- **레이어 수준 하이브리드 대비 효율성**: 3:1 레이어 수준 하이브리드의 성능을 7:1 LA-FA 비율로 달성
- **일반 추론 성능 유지**: 장문 추론 성능 향상과 동시에 일반 추론 능력 저하 없이 유지
의의 및 한계
HydraHead는 어텐션 하이브리드 설계의 새로운 차원을 제시하며, 해드 수준의 기능적 이질성을 기반으로 한 설계 원칙을 입증한다. 특히, 해드 선택 알고리즘과 스케일 정규화 퓨전 모듈은 기존 레이어 수준 접근에서 발생하는 성능 저하를 극복하고, 장문 처리와 일반 추론 사이의 트레이드오프를 완화한다. 그러나, 해드 선택 알고리즘은 사전 학습된 모델에 의존하며, 이는 모델 구조 변경 시 재설계가 필요할 수 있다. 또한, 토큰 수준 하이브리드와 비교하면 복잡도가 높아, 실시간 처리 환경에서는 제한이 있을 수 있다.
실용적 활용
HydraHead는 장문 처리가 필수적인 대형 언어 모델, 문서 검색 시스템, 대화형 에이전트 등에 적용 가능하다. 특히, 컨텍스트 길이 확장이 필요한 의료, 법률, 금융 분야의 분석 및 추론 작업에서 유용하게 활용될 수 있다.