PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Vishesh Tripathi, Abhay Kumar

arXiv:2606.20945v2cs.LG

TL;DR

Self-attention applies every query head to every token, leaving open whether mixture-of-experts routing can reduce attention computation without sacrificing quality. GQE routes tokens to sparse query-head experts within GQA groups while keeping the KV path dense, matching all-active GQA accuracy within 0.2 average points while activating half the routed query heads.

  • Problem

    It remains unclear whether mixture-of-experts routing can reduce GQA’s uniformly activated query-head computation without sacrificing model quality.

  • Method

    GQE routes each token to k query-head experts within each GQA group while keeping all key-value heads dense and unchanged.

  • Results

    Within 0.2 average points of all-active GQA accuracy, GQE activates half of the routed query heads, or 9 of 16 total query-attention computations.

  • Takeaways & Limitations

    GQE reduces active query-attention computation while preserving downstream accuracy relative to the corresponding all-active GQA baseline under the reported setting.

  • Takeaways & Limitations

    Results are limited to the 250M-parameter scale and 30B-token budget, and require confirmation across multiple seeds and larger scales.

Abstract

from arXiv · show

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.

1 Introduction

Self-attention becomes a long-context compute bottleneck because token interactions scale quadratically and every token activates every head. Grouped Query Experts (GQE) adds conditional query-head routing to GQA, preserving dense KV computation while matching GQA accuracy with fewer active query heads.

  • Self-attention is costly for long contexts because token–token interactions grow quadratically, while standard attention activates every head for every token.This uniform allocation can be inefficient for heterogeneous token contexts.
  • GQE routes each token to k query-head experts within each GQA group while keeping all key-value heads dense and unchanged.This preserves GQA’s KV-cache memory and bandwidth benefits while making query-head computation conditional.
  • 1.7–1.8× long-context speedup: GQE prefill is faster than the GQA baseline as avoided query-side computation becomes increasingly valuable.The speedup is calculated as baseline prefill time divided by GQE prefill time.
  • Sparse query-head routing matches dense GQA only with a proper router learning signal and an always-on shared head.The shared head stabilizes routing and is required for matching the dense baseline.
  • 30B tokens at 250M parameters: GQE matches the all-active GQA baseline in downstream accuracy while activating half of the routed query-head experts per token.The main setting computes 9 of 16 query-attention heads when the shared head is included.

2 Background and Related Work

This section introduces mixture-of-experts sparsity and grouped-query attention as the foundations for GQE. GQE routes query experts within fixed GQA groups while retaining dense KV sharing, distinguishing it from prior attention-routing approaches.

  • Mixture-of-Experts: MoE layers route each token to a small subset of experts, creating conditional capacity by activating only part of a larger parameter set.The selected experts’ outputs are combined using router probabilities.
  • Grouped-Query Attention: GQA shares one KV head within each query group, positioning it between MHA’s separate KV heads and MQA’s single shared KV head.A 16-query-head model with 8 KV heads has two query heads per KV group, reducing KV-cache storage relative to MHA.
  • GQE Design: GQA still evaluates all query heads per token, whereas GQE preserves GQA-style KV sharing and selectively activates query experts within each group.The reduced active slots require WO to be resized from N × d to (kG + 2) × d.
  • Related Work: Prior attention-routing methods include post-training head pruning, per-token head selection, and MoE-style modules that gate head outputs while retaining full per-head KV caches.These approaches motivate conditional computation inside attention but differ in routing granularity and KV-cache treatment.
  • GQE Design: GQE selects top-k query experts inside every fixed GQA group, keeps all GQA KV heads dense, and lets tokens skip only some query experts.This design is simpler and more constrained than routing over heads from the full model or changing the KV cache.

3 Grouped Query Experts

GQE turns query heads within each GQA group into routed experts while keeping shared KV heads dense and unchanged. It combines hard top-k expert selection with shared and router-supervised output slots, reducing active query computation while preserving the GQA KV cache.

  • Expert construction: GQE treats each GQA group’s query heads as experts that share one KV head, while KV projections remain always computed and the KV cache stays identical to GQA.Each expert has its own query projection and attention-head output.
  • Routing: For each token and group, the router applies softmax over that group’s experts and selects the top-k highest-scoring experts.Selection is performed independently within each GQA group.
  • Output composition: Selected experts are concatenated as unscaled hard-routed slots rather than averaged, alongside a renormalized weighted-sum slot and an always-on shared-head slot.The weighted-sum slot supplies a differentiable path for router learning because top-k decisions are discrete.
  • Output shape: With 16 query heads, 8 KV heads, and k = 1, GQE produces 8 hard-concatenated outputs plus weighted-sum and shared-head slots, yielding 10 slots before WO.The output projection therefore receives 10 slots instead of the dense baseline’s 16.
  • Compute sparsity: 9/16 ≈56% is the total active query-attention fraction in the 16-query-head / 8-KV-head setting with k = 1, while “50%” counts routed experts only.The weighted-sum slot reuses selected expert outputs and adds no attention computation; increasing M lowers routed activity for fixed k.

4 Experiments and Results

Experiments compare GQE with 8-KV-head GQA baselines under a fixed 30B-token, 250M-parameter budget, evaluating downstream quality and long-context throughput. Corrected GQE matches the all-active GQA baseline within 0.2 average points while activating half of routed query-head experts, with routing design determining whether quality is preserved.

  • Experimental setup: 30B tokens and 250M parameters are fixed across ablations to isolate architectural and routing choices.The training data is a 30B-token sample from FineWeb-Edu.
  • Evaluation: Downstream quality is reported on HellaSwag, PIQA, and ARC-Easy, while throughput is measured as context length increases.Long contexts are emphasized because attention compute dominates there and sparse query-head activation should provide the clearest benefit.
  • Headline result: Within 0.2 average points, corrected GQE matches the all-active GQA baseline while activating half of routed query-head experts, or 9 of 16 total query-attention computations.The comparison is made at the same training budget.
  • Routing ablations: 55.18 average is 0.68 below the dense baseline for weighted concatenation without a renormalized router-supervised slot, while hard concatenation reaches 55.43, still 0.43 below baseline.The ablation ladder uses 16 query heads, 8 KV heads, 30B tokens, and k = 1 active expert per group.
  • Routing ablations: 56.04 recovers and marginally exceeds the all-active GQA baseline at 55.86 while activating 8 of 16 routed query heads per token and 9 of 16 query-attention computations in total.This configuration combines a renormalized weighted-sum slot with an always-on shared head.
  • Throughput: At 2k tokens, GQE achieves a measured prefill speedup of 1.15× over GQA.The speedup is defined as GQA baseline latency divided by GQE latency; routing and dispatch overheads are relatively large at this length.

5 Limitations & Future Work

At the 250M-parameter scale and 30B-token budget, GQE is reported as matching rather than improving over the baseline, but this result requires broader validation. Future work will test alternative long-context architectures and larger Transformers.

  • Current limitations: 250M parameters and 30B tokens define the current evaluation scope, limiting conclusions about GQE’s broader robustness.The reported accuracy margin over the baseline should be confirmed with multiple seeds and at larger scales.
  • Current limitations: GQE matches rather than improves over the baseline, so the small accuracy margin should not yet be treated as robust.The paper explicitly reports the result as a match and calls for confirmation with multiple seeds and larger models.
  • Future work: Future work will compare GQE with Mamba and evaluate whether its routing benefits extend to larger Transformer architectures.These experiments target alternative long-context designs and larger-scale Transformer settings.

6 Conclusion

The paper presents a mixture-of-experts layer on GQA self-attention that reduces computation by conditionally activating query-head experts while keeping the KV path dense. It preserves many possible attention patterns, but only with a proper router learning signal.

  • MoE on GQA reduces computation by applying conditional computation to query heads while leaving the KV path dense.The method routes tokens to query-head experts instead of evaluating every GQA query head.
  • Each token is routed to a sparse subset of query-head experts within its GQA groups rather than activating all query heads.
  • The approach preserves many possible attention patterns while reducing active computation only when the router receives a proper learning signal.

Appendix · A Loss Graphs

Figure 4 presents the training-loss curves for the four variants compared in Table 2.

  • A Loss Graphs: The plotted variants are the same four variants compared in Table 2.The figure is explicitly tied to the Table 2 comparison.
  • A Loss Graphs: Figure 4 shows training-loss curves for four variants.The figure summarizes the training-loss behavior of the variants evaluated in Table 2.

B Downstream Task Accuracy Graphs

Figures 5–7 trace downstream accuracy over the full 30B-token training budget on HellaSwag, ARC-Easy, and PIQA, comparing GQA, routing ablations, and final GQE. Each graph shows accuracy over training tokens for one task.

  • Overview: Figures 5–7 compare the GQA baseline, intermediate routing ablations, and final GQE configuration across the full 30B-token training budget.The comparisons cover downstream accuracy trends on HellaSwag, ARC-Easy, and PIQA.
  • HellaSwag: Figure 5 shows HellaSwag accuracy over training tokens for the GQA baseline, routing ablations, and final GQE model.
  • ARC-Easy: Figure 6 shows ARC-Easy accuracy over training tokens for the GQA baseline, routing ablations, and final GQE model.
  • PIQA: Figure 7 shows PIQA accuracy over training tokens for the GQA baseline, routing ablations, and final GQE model.
Loading 2606.20945v2…