Virtual Head Attention

Virtual Head Attention. By Xibo Ding and six others, from Baidu and Peking University. Published on the thirteenth of October. Its code repository is worked out: about seven hundred source files, with a readme and a licence.

Designing attention for long-context language models runs into what the authors call an impossible triangle. Its three corners are the size of the KV cache, the FLOPs spent in training and prefill, and the quality of the model. Existing designs each give something up. GQA delivers strong quality, but its KV cache is expensive. MLA and MFA both compress the cache, but they pay in different ways: MLA loses quality, and MFA substantially increases the compute needed for training and prefill. Pick any two corners, the framing goes, and the third suffers.

Virtual Head Attention, or VHA, claims to improve all three corners at once, and it does so with two lightweight linear operations. The starting point is to halve the number of physical query heads. The first operation, Q Premix, then applies a near-identity transformation to the queries inside each KV group, recovering a diversity of virtual heads from the halved physical ones. The second, Linear Postmix, fuses features across heads through a low-rank residual structure: the identity plus a product of two thin matrices.

The choice of linear operations is what makes inference cheap. Since VHA contains nothing nonlinear, Premix can be folded into the query projection and Postmix into the output projection once training is done. What is left at inference time is an ordinary GQA architecture with two KV heads, so every existing GQA inference optimization applies directly. With only two KV heads, the group size in large models easily reaches sixty-four, which fully exploits Tensor Core parallelism during decoding.

The first round of experiments pretrains at four scales, with a context of four thousand tokens and fifty billion training tokens. VHA beats the standard GQA baseline at every scale, by margins from zero point zero five to zero point eight one percent. At the same time it shrinks the KV cache by four times relative to GQA-8 at the dense scales, and by two times relative to GQA-4 at the mixture-of-experts scale. Training and prefill FLOPs drop by seven point nine percent, measured on the one point seven billion model.

At industrial scale, on a thirty billion parameter MoE model trained on five hundred billion tokens at a context of eight thousand, VHA outperforms MLA by zero point seven zero percent and GQA-4 by two point six zero percent, on a base-model evaluation suite of twelve benchmarks.

A separate comparison uses matched ten billion parameter MoE models after long-context mid-training. There VHA reaches an average RULER score of eighty-four point four across contexts from eight thousand to one hundred and twenty-eight thousand tokens. GQA reaches seventy-three. And VHA does it with two KV heads in the full-attention layers, rather than four.

The triangle was framed as a forced choice, and the response here is to add expressiveness only during training, in a form that vanishes at deployment. Premix and Postmix give the model back the head diversity it lost by halving its query heads, and then fold away, leaving an architecture that existing inference tooling already knows how to run fast. Quality, cache and compute all move in the right direction, according to these results, and nothing about serving the model has to change.