12 분 소요

0. Introduction

Paper link

Grouped Query Experts, 이하 GQE는 “attention head를 sparse하게 쓰자”라는 익숙한 아이디어를 꽤 조심스럽게 다시 설계한 논문이다. 핵심은 MoE를 FFN이 아니라 attention block 안으로 가져오되, GQA의 KV-cache 구조는 건드리지 않고 query-head computation만 conditional하게 줄이는 것이다.

GQA는 이미 modern LLM serving에서 중요한 default가 되었다. 여러 query head가 더 적은 수의 key/value head를 공유하므로 KV cache memory and bandwidth를 줄일 수 있다. 하지만 GQA에서도 모든 query head는 모든 token에서 항상 계산된다. Long context에서는 attention compute가 sequence length에 따라 커지므로, “모든 token이 모든 query head를 정말 필요로 하는가”라는 질문이 자연스럽다.

GQE의 답은 per-token sparse query expert routing이다. 각 GQA group 안에 여러 query-head expert를 두고, router가 token마다 top-k expert만 선택한다. KV head는 group마다 dense하게 유지한다. 즉 KV cache layout은 GQA와 같고, 줄이는 것은 query-side attention computation이다.

한 줄 요약: GQE는 GQA의 dense KV path는 유지하면서 각 GQA group 내부의 query head를 MoE expert처럼 routing해, main 16-query-head / 8-KV-head setting에서 8 routed query experts plus 1 shared head만 계산하고, 250M scale 30B-token training에서 GQA baseline과 유사한 downstream accuracy 및 long-context prefill speedup을 보이는 attention-side conditional computation 논문이다.

이 논문을 지금 볼 가치가 있는 이유는 다음과 같음.

  • MoE를 MLP block이 아니라 self-attention query path에 적용한다.
  • KV cache를 sparse하게 만들지 않고, GQA의 serving-friendly memory profile을 유지한다.
  • Sparse routing이 그냥 잘 되는 것이 아니라, router learning signal and shared head가 필요하다는 ablation을 보여준다.
  • Long-context prefill에서 query-side attention compute를 줄이는 방향을 제시한다.
  • 250M scale의 small experiment지만, attention MoE design에서 무엇을 조심해야 하는지 명확한 engineering lesson을 준다.

이 글에서는 GQE를 “새 attention SOTA”보다, GQA를 보존한 상태에서 query-head conditional compute를 넣으려면 어떤 output construction과 router signal이 필요한가를 보여주는 design note로 읽는다.

1. Problem Setting

1-1. Problem definition

Transformer attention은 sequence length가 길어질수록 compute cost가 커진다. Standard multi-head attention에서는 각 token이 모든 head를 계산한다.

\[\mathrm{Attention}(Q,K,V) = \mathrm{softmax} \left( \frac{QK^\top}{\sqrt{d}} \right)V\]

Head 수가 $H_q$이고 sequence length가 $T$라면, query-side attention computation은 대략 $H_q T^2$에 비례한다. GQA는 key/value head 수를 줄여 KV cache and bandwidth를 줄이지만, query head computation은 그대로 남는다.

GQA를 간단히 쓰면 다음과 같다.

  • Query head 수: $H_q$
  • KV head 수: $H_{kv}$
  • Group 수: $G=H_{kv}$
  • Group당 query head 수: $E=H_q/G$

각 group은 하나의 KV head를 공유하지만, group 안의 모든 query head는 여전히 active다. 따라서 GQA의 장점은 KV cache에 있고, query-side compute를 conditional하게 줄이지는 않는다.

GQE가 해결하려는 문제는 다음이다.

GQA의 KV-cache benefit은 유지하면서, token마다 필요한 query head만 계산할 수 있는가?

1-2. Why previous approaches are insufficient

1) Head pruning

Head pruning은 학습 후 중요하지 않은 head를 제거한다. 이 방식은 static하다. 어떤 token에서는 유용한 head가 다른 token에서는 불필요할 수 있는데, pruning은 token-dependent choice를 허용하지 않는다.

2) Attention-head MoE

Mixture-of-attention or head routing 계열은 token별로 attention head를 선택한다. 하지만 많은 접근은 전체 head pool에서 routing하거나, full per-head KV cache를 유지하거나, GQA의 group structure와 직접 맞지 않는다.

GQE는 더 좁은 design을 택한다.

  • Routing은 전체 head pool이 아니라 각 GQA group 안에서 수행한다.
  • 각 group은 여전히 dense KV head를 가진다.
  • KV cache shape and serving benefit은 GQA와 동일하게 유지한다.

3) KV sparsification

Token-wise KV selection이나 KV pruning은 KV cache memory를 줄일 수 있지만, cache correctness and retrieval behavior를 바꾼다. GQE는 이 길을 가지 않는다. KV path는 dense and unchanged다. 따라서 GQE의 saving은 KV memory가 아니라 active query attention operations다.

이 distinction이 중요하다. GQE를 KV-cache reduction method로 이해하면 오해다. GQE는 GQA 위에서 query-head compute를 줄이는 방법이다.

2. Core Idea

2-1. Main contribution

GQE의 핵심 contribution은 세 가지다.

  1. Within-group query expert routing
    • 각 GQA group 안의 query heads를 expert pool로 본다.
    • Router가 token마다 group별 top-k expert를 선택한다.
    • KV head는 dense and always computed다.
  2. Router-supervised output construction
    • Hard-selected expert output은 ordinary head slot으로 concatenate한다.
    • Router probability를 사용하는 renormalized weighted-sum slot을 추가한다.
    • Always-on shared head를 추가해 stable attention path를 제공한다.
  3. Small-scale controlled validation
    • 250M parameter scale
    • 30B token budget
    • FineWeb-Edu sample
    • GQA baseline과 routing ablation을 비교
    • Long context prefill speedup을 측정

2-2. Design intuition

GQE의 직관은 token마다 필요한 attention pattern이 다르다는 것이다.

  • Content-bearing token은 specialized long-range head를 필요로 할 수 있다.
  • Punctuation or stop word는 모든 head를 쓸 필요가 없을 수 있다.
  • Code identifier and rare term은 특정 expert가 더 유용할 수 있다.
  • Long context에서는 불필요한 head compute가 더 크게 누적된다.

그러나 routing을 단순히 넣으면 잘 되지 않는다. Top-k selection은 discrete operation이므로 router가 language modeling loss에서 충분한 signal을 받기 어렵다. 논문은 이 문제를 weighted-sum slot으로 해결한다.

Router가 expert를 선택하더라도, selected expert output을 그냥 hard concat하면 router probability가 output value에 직접 들어가지 않는다. 이 경우 router가 어떤 expert를 골라야 하는지 학습 signal이 약해질 수 있다. GQE는 selected expert outputs의 probability-weighted sum을 별도 slot으로 넣어 router probability가 output and loss에 직접 영향을 주게 한다.

그리고 always-on shared head는 routing이 아직 불안정한 training early stage에서 token-independent fallback path를 제공한다.

3. Architecture / Method

3-1. Overview

Item Description
Goal GQA의 KV path는 유지하고 query-head compute만 sparse하게 줄임
Base attention Grouped-Query Attention
Routing unit Token과 GQA group
Expert GQA group 내부에 있는 query-head expert
KV path Dense하게 유지되며 항상 계산됨
Main active setting 16 query head와 8 KV head를 두고, group마다 top-1을 선택
Computed query attentions 선택된 query expert 8개와 shared head 1개
Output slots hard expert slot 8개, weighted-sum slot 1개, shared slot 1개
Auxiliary loss Load-balancing router loss
Main scale 250M parameters, 30B tokens

3-2. Module breakdown

1) GQA group structure

Query head를 $G$개의 group으로 나눈다고 하자. 각 group은 하나의 key/value head를 공유한다. Group $g$가 $E$개의 query expert를 포함하면, routed query expert pool의 전체 크기는 다음과 같다.

\[N = G E\]

Dense GQA는 모든 token에 대해 모든 query head를 계산한다. GQE는 대신 group마다 선택된 $k$개 expert만 계산해 query-side compute를 줄인다.

2) Per-group router

Token representation $x_t$와 group $g$에 대해, router는 해당 group 안의 expert들에 대한 score를 만든다.

\[s_{t,g,e} = r_g(x_t)_e\]

Softmax는 group 내부에서 적용된다.

\[p_{t,g,e} = \frac{ \exp(s_{t,g,e}) }{ \sum_{e'=1}^{E} \exp(s_{t,g,e'}) }\]

그다음 top-k expert를 선택한다.

\[\mathcal{S}_{t,g} = \mathrm{TopK} \left( p_{t,g,1:E}, k \right)\]

선택된 expert들은 query projection을 계산하고, group $g$의 동일한 dense KV head를 상대로 attention을 수행한다.

3) Hard-concatenated expert slots

선택된 expert output은 일반 attention head slot처럼 concatenate된다.

\[h_{t}^{\mathrm{hard}} = \mathrm{Concat} \left( \{o_{t,g,e}: e \in \mathcal{S}_{t,g}, g=1,\ldots,G\} \right)\]

중요한 점은 선택된 expert output들을 평균하지 않는다는 것이다. 이들은 별도의 head slot으로 남는다.

4) Router-supervised weighted-sum slot

GQE는 선택된 expert들에 대해 renormalized router probability를 사용하는 weighted-sum slot을 추가한다.

\[h_t^{\mathrm{route}} = \sum_{g=1}^{G} \sum_{e \in \mathcal{S}_{t,g}} \tilde{p}_{t,g,e} o_{t,g,e}\]

여기서 $\tilde{p}$는 선택된 expert output들 위에서 normalize된 값이다. 이 slot 덕분에 router는 language modeling loss로부터 differentiable learning path를 얻는다.

이 slot이 없으면 hard routing은 expert output을 선택할 수는 있지만, router probability가 output magnitude를 직접 제어하지 않는다. Ablation은 이 경우 quality가 떨어진다는 것을 보여준다.

5) Always-on shared head

GQE는 routing과 무관하게 모든 token에서 계산되는 shared head를 추가한다.

\[h_t^{\mathrm{shared}} = \mathrm{Head}_{\mathrm{shared}}(x_t)\]

이 path는 training을 안정화하고 모든 token에 안정적인 attention component를 제공한다. 최종 attention output은 hard expert slot, weighted-sum slot, shared head를 concatenate해 만든다.

\[h_t = \mathrm{Concat} \left( h_t^{\mathrm{hard}}, h_t^{\mathrm{route}}, h_t^{\mathrm{shared}} \right)\]

Main setting에서 $G=8$, $k=1$이면 hard selected slot은 8개다. Weighted-sum slot이 output slot 1개를 추가하고, shared head가 output slot 1개를 추가한다. 따라서 output projection은 16개가 아니라 10개 slot을 받는다.

6) Compute profile

Routed expert active fraction은 다음과 같다.

\[\frac{kG}{N} = \frac{k}{E}\]

하지만 전체 query-attention computation에는 shared head도 포함된다.

\[H_{\mathrm{active}} = kG + 1\]

Main 16-query-head / 8-KV-head setting에서는 $G=8$, $E=2$, $k=1$이다.

\[H_{\mathrm{active}} = 8+1=9\]

따라서 “half of routed experts”라는 표현은 routed query head 16개 중 8개를 뜻한다. Shared head까지 포함하면 실제 query-attention computation은 16개 중 9개다. Weighted-sum slot은 선택된 expert output을 재사용하므로 attention computation을 추가로 만들지 않는다.

7) Output projection caveat

최종 output이 $kG+2$개의 slot을 가지기 때문에, output projection input width는 dense GQA와 다르다. Main setting에서는 16 slot이 아니라 10 slot이다. 논문은 이 caveat을 인정한다. 비교는 data, training budget, per-head dimension, KV layout, token budget을 고정하지만, attention output projection까지 완전히 parameter-matched된 비교는 아니다.

이 점은 accuracy equality를 해석할 때 중요하다. 일부 saving은 active query computation 감소에서 나오고, output projection도 더 좁아진다.

4. Training / Data / Recipe

4-1. Data

모든 모델은 FineWeb2의 FineWeb-Edu sample을 사용해 고정된 30B-token budget으로 학습된다. 목표는 data exposure가 아니라 architecture와 routing effect를 분리하는 것이다.

Item Value
Training data 30B-token sample from FineWeb-Edu
Scale 250M parameters
Context length 2048
Baseline 8 KV head를 쓰는 GQA
Main GQE setting 16 query head / 8 KV head, group마다 top-1

4-2. Training strategy

논문이 보고한 training setup은 다음과 같다.

Hyperparameter Value
Optimizer Fused AdamW
LR schedule WSD
Warmup 3B tokens
Weight decay 0.1
Global batch size 1.05M tokens
Sequence length 2048
Precision Mixed precision BF16
Stabilization ZClip for loss spikes

논문은 GQE를 post-training conversion으로 설명하지 않는다. Router와 query expert는 모델 안에서 처음부터 함께 학습된다.

4-3. Routing auxiliary loss

GQE는 expert collapse를 막기 위해 load-balancing auxiliary loss를 사용한다. Router가 각 group 안에서 항상 같은 expert만 고르지 않아야 한다. 이는 MoE load balancing과 유사하지만, query expert group 내부에 적용된다.

전체 objective는 개념적으로 다음처럼 쓸 수 있다.

\[\mathcal{L} = \mathcal{L}_{\mathrm{LM}} + \lambda \mathcal{L}_{\mathrm{balance}}\]

정확한 weight와 implementation detail은 최종 논문이나 코드를 확인해야 한다. 중요한 점은 language modeling loss와 load balance가 함께 routing을 만든다는 것이다.

4-4. Engineering notes

  1. KV path를 바꾸지 않는다
    • GQE는 KV cache가 GQA와 동일하게 유지될 때 가장 이해하기 쉽다.
  2. Router에는 differentiable output path가 필요하다
    • Hard top-k selection만으로는 충분하지 않다.
    • Weighted-sum slot은 단순한 detail이 아니라 router가 학습 신호를 받는 통로다.
  3. Shared head가 early training을 안정화한다
    • Sparse expert path는 specialization이 느릴 수 있다.
    • Always-on head는 안전한 default attention channel을 제공한다.
  4. Compute saving은 query-side에서 나온다
    • GQE를 GQA 위의 추가 KV-cache compression으로 설명하면 안 된다.
  5. Benefit은 long context에서 커진다
    • Short context에서는 routing overhead가 중요하다.
    • $T$가 커질수록 회피한 attention head computation의 이득이 커진다.
  6. Kernel과 dispatch가 중요하다
    • Sparse head routing은 dispatch overhead가 지배하면 더 느려질 수 있다.
    • 측정된 speedup은 implementation, batch, sequence length, hardware에 의존한다.

5. Evaluation

5-1. Accuracy results

Main downstream evaluation은 HellaSwag, ARC-Easy, PIQA를 사용한다.

Variant HellaSwag ARC-E PIQA Average
GQA baseline, all 16 heads active 41.31 61.36 64.90 55.86
Weighted concat, no renormalized slot 40.16 60.52 64.85 55.18
Hard concat only 40.66 60.56 65.07 55.43
GQE, renorm scoring plus shared head 41.01 62.41 64.69 56.04

핵심 결과는 GQE가 baseline을 결정적으로 이겼다는 것이 아니다. Margin은 작다. 더 강한 claim은 corrected GQE가 더 적은 routed query expert를 활성화하면서 all-active GQA baseline과 quality를 맞췄다는 것이다.

Ablation ladder가 최종 average보다 더 중요하다.

  • Renormalized slot 없는 weighted concat은 성능이 낮다.
  • Hard concat은 나아지지만 여전히 낮다.
  • Renormalized scoring과 shared head를 함께 쓰면 baseline quality를 회복한다.

이는 sparse query routing에 router learning path와 stable shared attention path가 모두 필요하다는 주장을 지지한다.

5-2. Throughput

논문은 2K부터 1024K token까지 sequence length를 바꾸며 GQA baseline 대비 prefill speedup을 측정한다. Short context에서는 routing overhead가 상대적으로 크기 때문에 speedup이 제한적이다. 4K 이후부터 논문은 long-context speedup이 약 1.7-1.8x 수준이라고 보고한다.

이 결과는 조심해서 읽어야 한다.

  • Speedup은 prefill에 대해 측정되었고, decode speedup을 의미한다고 볼 수는 없다.
  • KV path는 dense하게 유지된다.
  • 이득은 inactive query expert를 건너뛰는 데서 나온다.
  • 정확한 speedup은 sparse dispatch와 hardware utilization에 의존한다.
  • Very long context setting에서는 다른 memory와 kernel bottleneck이 드러날 수 있다.

5-3. What really matters in the experiments

1) The ablation is the main result

최종 average 56.04 vs 55.86은 명확한 improvement라고 부르기에는 충분히 robust하지 않다. 의미 있는 결과는 naive sparse routing은 실패하고, corrected output construction을 넣어야 quality가 회복된다는 점이다.

2) GQE is not KV compression

모든 KV head가 dense하게 유지되므로 GQE는 GQA cache layout을 보존한다. 이는 실용적인 장점이지만 동시에 boundary이기도 하다. GQE는 GQA가 이미 제공하는 것 이상으로 KV cache memory를 줄이지 않는다.

3) Output projection caveat matters

GQE에서는 output projection의 input slot 수가 더 적다. Parameter count, FLOPs, model quality를 비교할 때 이 점을 반영해야 한다.

4) Scale is small

250M parameters와 30B tokens는 design exploration에는 유용하지만, large-scale LLM에서도 통한다는 점을 증명하기에는 부족하다.

5) Expert pool is limited

Main setting은 사실상 group마다 2개 expert 중 1개를 고르는 구조다. 더 큰 expert pool은 specialization을 더 보여줄 수 있지만, routing과 dispatch complexity도 더 키운다.

6. Limitations

  1. Scale이 작다
    • Main experiments는 250M parameter scale과 30B tokens에서 수행된다.
    • Large model behavior는 아직 검증되지 않았다.
  2. Accuracy margin이 작다
    • GQE average는 56.04이고 GQA는 55.86이다.
    • 이는 robust improvement라기보다 quality match로 읽는 편이 맞다.
  3. Single-seed concern이 있다
    • 논문은 multiple seeds로 확인할 필요가 있음을 시사한다.
    • 작은 margin은 seed에 민감할 수 있다.
  4. 완전히 parameter-matched 비교는 아니다
    • Main setting에서 output projection shape이 dense 16 slot에서 10 slot으로 바뀐다.
    • Parameter와 compute comparison은 조심해서 해석해야 한다.
  5. 추가 KV-cache saving은 없다
    • GQE는 dense GQA KV path를 그대로 둔다.
    • 줄이는 것은 KV memory가 아니라 query-head compute다.
  6. Routing overhead
    • Router computation과 dispatch overhead는 short context에서 이득을 지배할 수 있다.
    • Speedup은 long-context prefill에서 가장 설득력 있다.
  7. Benchmark coverage가 제한적이다
    • HellaSwag, ARC-Easy, PIQA는 좁은 quality check만 제공한다.
    • Perplexity, long-context retrieval, coding, reasoning evaluation이 추가되면 좋다.
  8. Large expert pool sweep이 없다
    • Group당 expert 수가 더 많으면 specialization이 좋아질 수 있지만, 넓게 탐색되지는 않았다.
  9. 다른 attention MoE와의 large-scale 비교가 부족하다
    • SwitchHead, MoH, MoA, MoMHA-like design은 논의되지만, 동일 조건에서 exhaustive하게 비교되지는 않는다.
  10. Serving integration이 쉽지 않다
    • Sparse query expert dispatch에는 실제 kernel support가 필요하다.
    • 논문이 보고한 speedup은 target inference stack에서 다시 측정해야 한다.

7. My Take

7-1. Why this matters for my work

GQE의 가장 중요한 메시지는 “attention head도 MoE로 만들 수 있다”보다, attention MoE를 GQA serving contract 안에서 설계해야 한다는 점이다.

Modern LLM serving에서 KV cache shape는 사소한 detail이 아니다. Cache layout, memory bandwidth, batching, prefix sharing, quantization, paged attention은 모두 여기에 의존한다. Attention-side conditional compute가 KV path를 흔들면 구현 비용이 커진다.

GQE는 아주 좁은 길을 택한다.

Keep KV dense.
Keep GQA groups fixed.
Route query heads only.
Add router learning slot.
Add shared safety path.

이 설계는 보수적이지만 실용적이다. 특히 long-context prefill bottleneck을 줄이려면 KV cache compression과 별도로 query-side compute sparsity도 봐야 한다.

7-2. Reuse potential

Long-context prefill

RAG, codebase reading, agent context assembly처럼 long prompt prefill이 큰 service에서는 query-side attention sparsity가 도움이 될 수 있다. 특히 KV cache는 유지하되 prefill compute만 줄이고 싶은 경우 GQE-style routing이 맞다.

Layer-wise specialization

모든 layer에 GQE를 넣는 것이 최선일지는 모른다. Early layer는 shared lexical pattern이 많고, later layer는 token-specific specialization이 더 클 수 있다. Layer별 routing 필요성을 측정할 가치가 있다.

Token-type adaptive attention

Router가 어떤 token에서 어떤 expert를 고르는지 분석하면 punctuation, identifier, rare term, delimiter, long-range dependency에 따른 head usage pattern을 볼 수 있다. 이는 token-level architecture diagnostic과 연결된다.

Mixture with GQA variants

GQE는 GQA 위에 얹는 방법이므로 MQA, MLA, grouped KV compression, query/key/value sharing과 조합 가능성을 검토할 수 있다. 다만 KV path를 건드리면 serving contract가 달라진다.

7-3. Production considerations

  • Short context에서는 routing overhead가 saving을 상쇄할 수 있다.
  • Prefill과 decode speedup을 분리해서 측정해야 한다.
  • Expert routing은 batching에 맞게 구현되어야 한다.
  • Router collapse를 monitoring해야 한다.
  • Load-balance loss weight가 quality와 speed에 미치는 영향을 봐야 한다.
  • Output projection shape difference를 FLOP accounting에 반영해야 한다.
  • Long-context retrieval benchmark에서 quality degradation을 확인해야 한다.

7-4. Follow-up papers

  • GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
  • Fast Transformer Decoding: One Write-Head is All You Need
  • Mixture of Attention Heads
  • SwitchHead
  • MoH: Multi-Head Attention as Mixture-of-Head Attention
  • LLaMA-MoE v2
  • ZClip
  • Grouped Query Attention and KV cache optimization studies

8. Summary

  • GQE는 고정된 GQA group 내부의 query head에 MoE-style routing을 적용한다.
  • KV head는 dense하고 변경되지 않은 상태로 유지되므로, GQA KV-cache profile이 보존된다.
  • Router-supervised weighted-sum slot과 always-on shared head는 quality recovery에 필수적이다.
  • Main result는 250M scale에서의 quality matching이지, 결정적인 quality improvement는 아니다.
  • 이 아이디어는 long-context prefill에 유망하지만, large-scale과 serving-stack validation은 아직 open problem으로 남아 있다.

댓글남기기

참고