Gruppierte Query-Experten: Mixture-of-Experts auf GQA-Self-Attention

Zusammenfassung

Self-Attention ist zentral für die Leistungsfähigkeit von Transformatoren und stellt bei langen Kontextlängen oft den rechenintensivsten Teil dar, da ihre paarweisen Token-Interaktionen quadratisch mit der Sequenzlänge skalieren. Standardmäßige dichte Aufmerksamkeit wendet zudem unabhängig von der Schwierigkeit oder dem Informationsgehalt eines Tokens denselben Satz von Aufmerksamkeitsköpfen auf jedes Token an. Diese gleichförmige Aktivierung kann Rechenleistung verschwenden, insbesondere wenn Sequenzen länger werden und die Kosten für die Aufmerksamkeit rapide steigen. Wir schlagen Grouped Query Experts (GQE) vor – eine Mixture-of-Experts-Schicht oberhalb der gruppierten Query-Aufmerksamkeit (GQA). Innerhalb jeder GQA-Gruppe wählt ein Router pro Token k Query-Head-Experten aus, während alle Key-Value-(KV)-Köpfe dicht und unverändert bleiben. Somit bewahrt GQE die Vorteile des KV-Cache von GQA und reduziert nur die aktive Berechnung der Query-Köpfe. Bei einem festen Budget von 30 Milliarden Token auf der 250-Millionen-Parameter-Skala erreicht GQE die gleiche Downstream-Genauigkeit wie die vollaktive GQA-Baseline, während pro Token nur die Hälfte der Query-Köpfe aktiviert wird.

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.