Expertos en Consultas Agrupadas: Mezcla de Expertos en Autoatención GQA

Resumen

La auto-atención es central para el rendimiento del Transformer y suele ser la parte más costosa del mismo en contextos largos, debido a que sus interacciones entre pares de tokens escalan cuadráticamente con la longitud de la secuencia. La atención densa estándar también aplica el mismo conjunto de cabezas de atención a cada token, independientemente de su dificultad o contenido informativo. Esta activación uniforme puede desperdiciar cómputo, especialmente a medida que las secuencias se alargan y el costo de la atención aumenta rápidamente. Proponemos Grouped Query Experts (GQE), una capa de mezcla de expertos (mixture-of-experts) sobre la atención de consultas agrupadas (GQA). Dentro de cada grupo de GQA, un enrutador selecciona k expertos de cabezas de consulta por token, mientras que todas las cabezas clave-valor (KV) permanecen densas e inalteradas. Así, GQE conserva las ventajas de la caché KV de GQA y reduce únicamente el cómputo activo de las cabezas de consulta. Con un presupuesto fijo de 30 mil millones de tokens a la escala de 250 millones de parámetros, GQE iguala al modelo base GQA con todas las cabezas activas en precisión descendente, mientras activa la mitad de las cabezas de consulta por 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.