Search papers, labs, and topics across Lattice.
This paper introduces Grouped Query Experts (GQE), a novel mixture-of-experts layer that optimizes the self-attention mechanism in Transformers by selectively activating query-head experts based on token difficulty. By leveraging grouped-query attention (GQA) and maintaining a dense key-value structure, GQE significantly reduces computational overhead while preserving accuracy. The key finding reveals that GQE can achieve the same downstream performance as a fully active GQA model while activating only half of the query heads per token, demonstrating a more efficient use of resources in long-context scenarios.
Activating only half the query heads in Transformers can match the performance of fully active models, slashing computational costs while maintaining accuracy.
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.
