Virtual Head Attention cuts cache while preserving model quality
In plain words
When a language model reads a long passage, its attention mechanism keeps a key-value (KV) cache, which takes memory. Sharing fewer cached heads can save memory but may lower answer quality.
The paper introduces Virtual Head Attention. It learns query and output mixing while training, then folds those operations into ordinary attention for use at inference.
Across four small models, its average accuracy on eight multiple-choice tests was slightly above the matching grouped-query attention models; that average does not mean every test improved.
The research
As language models handle longer contexts, their attention cache uses more memory, while reducing shared key-value heads can damage model quality. Can an attention design reduce this memory burden and training computation without losing accuracy?
The paper tests whether learned transformations can recover useful attention behavior with fewer physical heads.
Virtual Head Attention learns query diversity and cross-head mixing during training, then folds both into standard grouped-query attention for inference.
The authors pretrained four Qwen3-architecture model scales on 50B tokens and compared VHA with grouped-query attention using the standard per-scale key-value head counts. On eight zero-shot multiple-choice tests, average accuracy at each scale was 0.05 to 0.81 percentage points higher with VHA.
At 1.7B, VHA also used 1.00 KB of BF16 key-value cache per layer per token versus 4.00 KB for GQA-8; its training computation was 7.9% lower at sequence length 4096.
Together, these measurements support a quality and efficiency balance in these tests, while the accuracy averages do not imply gains on every test.
The small-model accuracy figures are averages across eight tests, so they do not show that VHA improved every test.
The comparisons also do not establish that all training settings were identical; the larger-model and long-context results come from separate studies with their own model and training setups.
Analysis
Deep
- Original paper authors
- Xibo Ding, Guoxia Wang, Jinle Zeng, Jiabin Yang, Sirui Tao
- Submitted to arXiv
- License
- arXiv non-exclusive