기술 요약: HydraHead

문제 정의

표준 풀 어텐션(Full Attention, FA) 메커니즘의 이차 복잡도(quadratic computational complexity)는 대규모 언어 모델(LLM)의 긴 문맥 처리에 있어 결정적인 병목 현상을 초래합니다. 선형 어텐션(Linear Attention, LA) 패러다임은 선형 시간 복잡도를 제공하지만, 정교한 추론에 필요한 고정밀 검색 능력을 유지하는 데 어려움을 겪는 "표현력 붕괴(expressivity collapse)" 현상이 빈번하게 발생합니다. 기존의 하이브리드 아키텍처는 주로 FA와 LA 레이어를 교차 배치하는 레이어 단위(layer-wise) 전략을 채택합니다. 그러나 선행 연구들은 LA와 FA를 통합하는 것이 본질적으로 어렵다는 점을 시사하며, 어텐션 하이브리드화의 설계 공간은 여전히 미개척 상태로 남아 있습니다. 또한, 레이어 단위 하이브리드는 문맥 길이를 확장하면 일반적인 추론 능력이 저하되거나, 반대로 추론 능력을 유지하면 문맥 확장이 어려운 트레이드오프 문제에 직면하곤 합니다.

방법론

본 논문은 레이어 수준이 아닌 **헤드 수준(head level)**에서 어텐션 하이브리드화를 수행하는 새로운 아키텍처인 HydraHead를 제안합니다. 이 접근 방식은 메커니즘적 해석 가능성(mechanistic interpretability) 분석에 근거합니다. 이 분석에 따르면 레이어들은 블록 단위의 기능적 유사성을 보이지만, 동일한 레이어 내의 개별 어텐션 헤드들은 뚜렷한 기능적 특수성을 나타냅니다.

1. 해석 가능성 기반 헤드 선택

HydraHead는 임의로 어텐션 메커니즘을 할당하는 대신, 인과적 개입(causal intervention) 기술을 사용하여 특정 작업에 결정적인 역할을 하는 헤드를 식별합니다:

• 인과적 패칭(Causal Patching): 이 방법은 액티베이션 패칭(activation patching)(직접적인 효과 측정)과 패스 패칭(path patching)(상류 기여도 추적)을 사용하여 타겟 능력(예: 긴 문맥 검색 및 일반 추론)에 대한 각 헤드의 인과적 필요성을 정량화합니다.
• 선택 전략: 헤드들은 이러한 인과적 측정값으로부터 도출된 복합 중요도 점수에 따라 순위가 매겨집니다. 상위 순위의 "검색 결정적(retrieval-critical)" 헤드들만 풀 어텐션(FA)으로 유지되며, 나머지 헤드들은 선형 어텐션(구체적으로 Gated DeltaNet 또는 GDN)으로 전환됩니다.
• 효율성: 이 선택 과정은 가볍게 설계되어, 정밀한 토큰 수준 검색에 필수적인 희소 부분 집합을 식별하기 위해 작은 캘리브레이션 데이터셋에 대해 단 몇 번의 포워드 패스(forward pass)만을 요구합니다.

2. 헤드 단위 하이브리드 아키텍처

이 아키텍처는 각 레이어 내에서 헤드 축을 따라 FA와 LA를 융합합니다:

• 병렬 브랜치 계산: 쿼리 헤드는 두 개의 서로 소인 집합인 $H_F$(FA 할당)와 $H_L$(GDN 할당)로 분할됩니다.
• 스케일 정규화 융합(Scale-Normalized Fusion): FA와 LA 사이에는 근본적인 비대칭성이 존재합니다. FA는 쿼리 크기에 의해 조절되는 날카롭고 낮은 엔트로피의 분포를 생성하는 반면, LA는 더 부드럽고 높은 엔트로피의 분포를 생성합니다. 이러한 분포 차이를 해결하기 위해 HydraHead는 헤드 단위 스케일 정규화 융합 모듈을 도입합니다. 여기에는 다음이 포함됩니다:
  • 각 헤드의 출력을 통일하기 위해 적용되는 독립적인 RMSNorm.
  • 구조적 무결성을 보존하면서 특징 차이를 메우기 위해 각 헤드의 기여도를 적응적으로 재조정하는 학습 가능한 헤드 단위 스케일링 벡터($\gamma$).
• 브랜치별 세부 개선:
  • FA 브랜치: 긴 문맥 어텐션을 안정화하기 위해 표준 회전 위치 임베딩(RoPE) 대신 로그 스케일 계수를 사용하며, "어텐션 싱크(attention sink)" 현상을 완화하기 위한 보조 게이트 브랜치를 도입합니다.
  • LA (GDN) 브랜치: 위치 인지 능력을 향 향상시키기 위해 쿼리와 키 투영(projection)에 RoPE를 통합하고, 표현 능력을 높이기 위해 키-값 헤드를 쿼리 헤드와 일치하도록 확장합니다(MHA 방식).
• 3단계 전이 학습 파이프라인: 사전 학습된 FA 모델을 최소한의 오버헤드로 하이브리드 아키텍처로 변환하기 위해, 저자들은 다음과 같은 3단계 파이프라인을 제안합니다:

1. 파라미터 이주 및 레이어별 정렬: 사전 학습된 가중치를 재사용하여(GDN을 위한 채널 반복 포함) 하이브리드 모델을 초기화하고, MSE 손실을 사용하여 하이브리드 레이어의 은닉 상태를 원래의 FA 레이어와 정렬합니다.
2. 글로벌 로짓 증류(Global Logits Distillation): 모델의 가중치를 해제(unfreeze)하고 지식 증류(KL 발산)를 수행하여 최종 출력 분포를 사전 학습된 교사 모델과 일치시킴으로써 성능 저하를 회복합니다.
3. 긴 문맥 미세 조정(Long-Context Fine-tuning): 긴 문맥 시퀀스에 대해 다음 토큰 예측(NTP)을 포함한 표준 지도 학습 미세 조정을 적용하여 긴 문맥 능력을 공고히 합니다.

주요 기여

• 헤드 수준의 세밀함: 본 논문은 어텐션 헤드를 하이브리드화의 자연스럽고 원칙적인 단위로 식별하며, 레이어 단위 경계는 기능적 이질성을 포착하기에 너무 거칠다고 주장합니다.
• 해석 가능성 가이드 할당: 인과적 해석 가능성 분석이 검색에 중요한 희소 헤드 집합을 효과적으로 식별할 수 있음을 입증하여, 상당한 성능 손실 없이 공격적인 FA 압축(예: 7:1 LA-to-FA 비율)이 가능함을 보여줍니다.
• 스케일 정규화 융합: 정규화와 학습 가능한 스케일링을 통해 분포 차이를 해결함으로써 이질적인 어텐션 신호(FA와 LA)를 효과적으로 융합하는 메커니즘을 도입합니다.
• 학습 효율성: 제안된 3단계 전이 파이프라인을 통해 최소한의 학습 오버헤드(15B 토큰만으로 학습)로 고성능 하이브리드 모델을 구축할 수 있습니다.

실험 결과

저자들은 Qwen3-1.7B를 기반으로 한 HydraHead를 통일된 훈련 설정 하에서 레이어 단위, 토큰 단위 및 기타 헤드 단위 베이스라인들과 비교 평가했습니다.

• 긴 문맥 성능: HydraHead는 긴 문맥 작업에서 다른 하이브리드 설계들을 크게 앞섭니다. 512K 문맥 길이에서, Needle In A Haystack(NIAH) 작업에 대해 베이스라인 대비 69% 이상의 향상을 달-성했습니다. 이는 훨씬 적은 토큰으로 학습되었음에도 불구하고 네이티브 256K를 지원하는 선도적인 모델인 Qwen3.5의 성능에 근접합니다.
• 일반 추론: 일반적인 추론 작업에서 성능이 저하되는 레이어 단위 하이브리드와 달리, HydraHead는 강력한 일반 추론 능력을 유지하며 어려운 추론 벤치마크에서 레이어 단위 베이스라인보다 10% 이상 우수한 성능을 보입니다.
• 효율성 대 성능: 7:1 LA-to-FA 비율을 사용하는 HydraHead는 3:1 레이어 단위 하이브리드와 대등한 긴 문맥 성능을 보이면서도 더 우수한 추론 능력을 유지합니다.
• 확장성: 이 모델은 견고한 외삽(extrapolation) 능력을 입증하여, 다른 하이브리드 모델들이 해당 길이에서 성능이 거의 제로로 떨어지는 것과 달리 256K 문맥에서도 높은 정확도(Single-Key 94.53%, Multi-Key 52.70%)를 유지합니다.

의의

본 논문은 HydraHead가 헤드 단위 하이브리드화의 상당한 스케일링 잠재력을 검증한다고 주장합니다. 메커니즘적 해석 가능성을 활용하여 아키텍처 설계를 유도함으로써, 이 방법은 계산 효율성과 표현력 사이의 긴장을 성공적으로 해결합니다. 이는 긴 문맥 능력이 단순히 단일 어텐션 메커니즘을 선택하는 문제가 아니라, 가장 미세한 기능적 단위에서 이질적인 신호들을 원칙적으로 융합하는 것이 필요함을 보여줍니다. 결과적으로 헤드 단위 하이브리드화는 현재의 최첨단 하이브리드 모델들보다 더 균형 잡힌 솔루션을 제공하며, 매우 긴 문맥을 효율적으로 처리하면서도 고정밀 검색과 복잡한 추론 능력을 보존할 수 있음을 시사합니다.