Experts de requêtes groupées : mélange d’experts sur l’auto-attention GQA

Résumé

L'auto-attention est au cœur des performances du Transformer et constitue souvent la partie la plus coûteuse du modèle pour des longueurs de contexte élevées, car ses interactions par paires de tokens augmentent de manière quadratique avec la longueur de la séquence. L'attention dense standard applique également le même ensemble de têtes d'attention à chaque token, indépendamment de sa difficulté ou de son contenu informationnel. Cette activation uniforme peut entraîner un gaspillage de calcul, en particulier lorsque les séquences s'allongent et que le coût de l'attention augmente rapidement. Nous proposons Grouped Query Experts (GQE), une couche de mélange d'experts placée au-dessus de l'attention par requêtes groupées (GQA). Au sein de chaque groupe GQA, un routeur sélectionne k experts de têtes de requête par token, tandis que toutes les têtes clé-valeur (KV) restent denses et inchangées. Ainsi, GQE préserve les avantages du cache KV de GQA et réduit uniquement le calcul actif des têtes de requête. Sur un budget fixe de 30 milliards de tokens à l'échelle de 250 millions de paramètres, GQE atteint la même précision en aval que la ligne de base GQA à activation complète, tout en activant la moitié des têtes de requête par token.

English

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.