GQA의 3가지 딜레마를 동시에 푸는 Virtual Head Attention
원제 Virtual Head Attention
추천 10댓글 2
Key Point
KV 캐시 메모리와 학습 효율, 모델 품질의 균형을 동시에 달성해야 하는 LLM 개발자들이 GQA의 근본적 한계를 극복할 수 있는 방법을 제시합니다.
핵심 요약
- 긴 문맥 처리 시 KV 캐시 크기, 학습·프리필 연산량, 모델 품질을 동시에 개선하기는 어렵다는 '불가능의 삼각형'에 직면한 어텐션 메커니즘 설계에 대응하기 위해 Virtual Head Attention(VHA)을 제안한다.
- VHA는 두 개의 경량 선형 연산으로 구성되는데, Q Premix가 쿼리 헤드 개수를 절반으로 줄인 상태에서 KV 그룹 내 가상 헤드 다양성을 복구하고, Linear Postmix가 낮은 순위의 I + AB^top 잔여 구조를 통해 헤드 간 특징을 융합한다.
- VHA는 비선형 연산을 도입하지 않으므로 Premix와 Postmix를 추론 시간에 각각 쿼리 프로젝션과 출력 프로젝션으로 폴드할 수 있어, 표준 GQA-2 아키텍처로 축약되면서 기존의 모든 GQA 추론 최적화를 직접 재사용할 수 있다.
- 4K 문맥과 50B 토큰 사전학습에서 4개 스케일에 걸쳐 VHA는 표준 GQA 베이스라인을 모든 규모에서 +0.05%~+0.81% 상회하면서, 밀집 모델에서는 GQA-8 대비 KV 캐시를 4배, MoE 모델에서는 GQA-4 대비 2배 감축했다.
- 학습 및 프리필 FLOPs도 1.7B 모델에서 7.9% 감소했고, 500B 토큰으로 8K 문맥에서 학습한 산업용 30B-A3B MoE 모델에서 VHA는 MLA와 GQA-4를 12개 벤치마크 평가에서 각각 +0.70%, +2.60% 초과했다.
- 별도로 진행한 10B-A1B MoE 비교에서 8K~128K 구간에서 VHA는 평균 RULER 점수 84.4를 달성해 GQA의 73.0보다 앞섰으며, 전체 어텐션 레이어에서 GQA의 4개 KV 헤드 대신 2개만 사용했다.