헤드(head) 단위에서 풀 어텐션(full-attention)과 선형 어텐션(linear-attention)을 병합하면, 다운스트림 품질을 크게 해치지 않으면서도 Transformer의 FLOPs(연산량)를 대폭 줄일 수 있습니다. 핵심 비결은 비용이 많이 드는 이차식 경로(quadratic path)를 정말로 중요한 곳에만 유지하고, 나머지 부분은 저렴한 선형 경로(linear path)가 처리하도록 하는 것입니다.

HydraHead 이전의 대부분의 하이브리드 설계는 레이어(layer) 수준에서 작동했습니다. 즉, 전체 레이어의 어텐션 메커니즘을 선형 변형으로 교체하거나, 풀 어텐션 레이어와 선형 레이어의 비율을 고정된 비율로 배치하는 방식이었습니다. 이러한 방식은 두 어텐션 계열 간의 분포 불일치(distributional mismatch)를 해결하는 데 어려움을 겪었으며, 연구자들은 주로 레이어 축(layer axis)만이 유일하게 실용적인 혼합 지점이라고 가정했습니다.

HydraHead는 단 4분의 1의 헤드에 대해서만 풀 어텐션 연산을 보존하고, 나머지 4분의 3은 (논문에 설명된 대로) GDN 선형 모듈에 위임합니다. “기본적으로 우리는 헤드의 25%에 대해 FA(full-attention) 연산을 유지하는 반면, 나머지 75%는 GDN 구조를 사용합니다.” [1]

이러한 공격적인 프루닝(pruning)에도 불구하고, 선형 대 풀 어텐션 헤드 비율이 7:1까지 올라가더라도 이 모델은 전통적인 3:1 레이어 단위 하이브리드와 동일한 벤치마크 점수에 도달합니다. “해석 가능성(Interpretability)에 기반한 헤드 선택은 성능 저하를 최소화하면서 공격적인 FA 압축을 가능하게 하며... 상당히 높은 LA(linear-attention) 대 FA 혼합 비율(예: 7:1)에서도 3:1 레이어 단위 하이브리드의 전반적인 성능과 일치합니다.” [1]

저자들은 주로 긴 문맥 읽기 및 추론 세트에서 HydraHead를 평가했으며, 증류(distillation) 전 15B 토큰으로 학습을 진행했습니다. 이는 헤드 선택 파이프라인이 더 작은 사전 학습(pre-training) 예산으로 확장 가능한지, 또는 토큰 수준의 분류나 엄격한 충실도 제약이 있는 생성과 같이 미세한 토큰 수준의 상호작용에 크게 의존하는 작업에서 어떻게 작동하는지에 대한 과제를 남겨두고 있습니다.

보고된 비율이 일반적인 워크로드(workloads) 전반에 걸쳐 유지된다면, 일반적인 트랜스포머 블록(vanilla transformer block)을 HydraHead 블록으로 교체함으로써 정확도를 유지하면서도 어텐션 관련 연산량(FLOPs)을 최대 약 40%(약 1/3)까지 줄일 수 있을 것이며, 이는 잠재적으로 더 큰 컨텍스트 윈도우(context windows)를 가능하게 하거나 엣지급(edge-class) GPU에 더 큰 모델을 탑재할 수 있게 해줄 것입니다.

참고 문헌

1. HydraHead: From Head-Level Functional Heterogeneity to Specialized Attention Hybridization