Mixture of Attention Schemes (MoAS): Learning to Route Between MHA, GQA, and MQA

TL;DR

The paper tackles KV-cache-induced memory bottlenecks in autoregressive Transformers by introducing Mixture of Attention Schemes (MoAS), a per-token routing mechanism that dynamically selects among MHA, GQA, and MQA. A lightweight router assigns tokens to attention schemes, with an auxiliary load-balancing loss to avoid trivial solutions, and all three attention variants are computed in parallel to enable dynamic selection. Experiments on WikiText-2 show that dynamic routing outperforms a static mixture (2.3074 vs. 2.3093 val loss), approaching MHA performance while offering avenues for conditional compute efficiency. This work suggests a practical path to balance modeling quality and memory efficiency in Transformer inference, with potential benefits amplified at larger scales where KV-cache constraints are more pronounced.

Abstract

The choice of attention mechanism in Transformer models involves a critical trade-off between modeling quality and inference efficiency. Multi-Head Attention (MHA) offers the best quality but suffers from large Key-Value (KV) cache memory requirements during inference. Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) reduce memory usage but often at the cost of model performance. In this work, we propose Mixture of Attention Schemes (MoAS), a novel architecture that dynamically selects the optimal attention scheme (MHA, GQA, or MQA) for each token via a learned router. We demonstrate that dynamic routing performs better than static averaging of schemes and achieves performance competitive with the MHA baseline while offering potential for conditional compute efficiency. Experimental results on WikiText-2 show that dynamic routing (val loss 2.3074) outperforms a static mixture (2.3093), validating the effectiveness of the proposed method. Our code is available at https://github.com/Esmail-ibraheem/Mixture-of-Attention-Schemes-MoAS.