Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention

Self-attention is central to Transformer performance and is often the most expensive part of the Transformer at long context lengths because its pairwise token interactions scale quadratically with sequence length. Standard dense attention also applies the same set of attention heads to every token regardless of token difficulty or information content. This uniform activation can waste compute, especially as sequences grow longer and attention cost increases rapidly. We propose Grouped Query Experts (GQE), a mixture-of-experts layer on top of grouped-query attention (GQA). Within each GQA group, a router selects k query-head experts per token while all key-value (KV) heads remain dense and unchanged. Thus, GQE keeps the KV cache benefits of GQA and reduces only the active query-head computation. On a fixed 30B token budget at the 250M parameter scale, GQE matches the all-active GQA baseline in downstream accuracy while activating half the query heads per token.

한 줄 요약

GQE는 GQA 기반의 self-attention에서 query head 계산을 반으로 줄이며 정확도를 유지하는 MoE 기법이다.

핵심 기여도

GQE는 GQA 구조 위에 MoE 레이어를 적용하여 query head 계산을 반으로 줄임.
30B 토큰 예산, 250M 파라미터 규모에서 GQA 베이스라인과 동일한 downstream 정확도 달성.
긴 문맥에서 prefill 속도가 1.7–1.8× 향상됨.
모든 KV head는 유지되며, query head만 조건부로 활성화됨.

핵심 아이디어

기존의 self-attention은 모든 토큰에 대해 동일한 attention head를 사용하여 계산 효율성이 떨어진다. GQE는 이 문제를 해결하기 위해 GQA 구조를 기반으로, 각 GQA 그룹 내에서 토큰마다 `k`개의 query-head expert만 선택적으로 활성화하는 방식을 제안한다. 이는 MoE 아이디어를 attention 블록에 적용한 것으로, query head의 조건부 활성화를 통해 계산량을 줄인다. 특히, `W_O` 행렬의 크기를 `(kG+2) × d`로 조정하여 모델 구조에 맞게 확장한다. 학습 시에는 `renormalized scoring`과 `always-on shared head`를 통해 라우터의 안정성을 보장한다.

기술적 접근법

**GQE 구조**: GQA 그룹 내에서 `k`개의 query-head expert만 선택적으로 활성화.
**KV head**: 모든 KV head는 유지되며, GQA의 KV-cache 이점을 활용.
**W_O 행렬**: `(kG+2) × d` 크기로 재설계.
**라우터 학습**: `renormalized scoring`과 `always-on shared head`를 통해 안정성 확보.
**하이퍼파라미터**: `k = 9`, `G = 16` (main setting).
**모델 규모**: 250M 파라미터, 30B 토큰 예산.

주요 결과

**정확도**: GQE는 GQA 베이스라인과 동일한 downstream 정확도를 달성.
**계산량**: 토큰당 query head 활성화 수를 50% 감소 (9/16).
**속도**: 긴 문맥에서 prefill 속도가 1.7–1.8× 향상됨.
**모델 파라미터**: 250M 규모에서 30B 토큰 예산으로 실험.

의의 및 한계

GQE는 self-attention의 계산 효율성을 향상시키는 새로운 접근법으로, 특히 긴 문맥 처리에서 유용하다. 기존 GQA의 KV-cache 이점을 유지하면서 query head 계산만 줄임으로써 전체 attention 계산량을 반으로 줄인다. 그러나 라우터 학습이 불안정하면 성능이 저하될 수 있으며, 이에 대한 해결책으로 `renormalized scoring`과 `shared head`가 필요하다는 점이 한계이다. 또한, 토큰 별로 필요한 attention pattern이 다를 경우에만 효과적이며, 모든 토큰이 동일한 패턴을 필요로 하는 경우 이 방법의 이점이 줄어든다.

실용적 활용

GQE는 대규모 언어 모델에서 긴 문맥 처리를 효율적으로 수행해야 하는 상황에 적용 가능하다. 특히, 서버 리소스가 제한된 환경이나 실시간 추론이 필요한 애플리케이션에서 유용하다. 또한, 다양한 attention head가 필요한 복잡한 텍스트 분석 작업에도 활용 가능하다.