Grouped Value Attention: Efficient KV Caching via On-Demand Key Reconstruction

AI summary

Overview

• Research area: Efficient Transformer inference — specifically KV-cache compression and attention architecture design in the hardware/systems-oriented area (cs.AR), with language modeling experiments.
• Technical level: Advanced. The paper assumes familiarity with multi-head attention, grouped-query attention, RoPE, softmax scaling, and decode-time memory-bandwidth limits.
• Scope: The paper proposes Grouped Value Attention (GVA), an attention variant that caches only grouped value streams plus a small shared positional key and reconstructs per-head content keys with a learned linear map, and it reports 350M-parameter pretraining and zero-shot benchmark results against matched GQA and MLA baselines.

What This Paper Is About

Transformer decoding keeps a KV cache holding the keys and values of all preceding tokens, and that cache grows with sequence length, becoming a dominant memory-capacity and bandwidth cost. Existing fixes such as grouped-query attention (GQA) share key–value heads across query heads, but still write both a key stream and a value stream at every step. GVA asks whether the content key can instead be derived from the stored value — keeping the value itself as the persistent content state — so that the content-key stream disappears from the cache while positional information is preserved by a small, separately cached shared RoPE key.

Key Contributions

1. Grouped Value Attention (GVA): an attention design that caches grouped values and reconstructs a per-head content key with a learned linear map, K_h = V_{g(h)} M_h. Because M_h is fixed at inference, it can be absorbed into the query (q̃_h = q_h M_h^T), so the intended decode path never writes or materializes a content-key stream.
2. A key/value grouping asymmetry: GVA caches only G grouped value streams but reconstructs H distinct content keys — one per query head — whereas GQA gives every query head in a group the same key. Head-specific keys are retained while the persistent cache shrinks.
3. A decoupled RoPE compatibility scheme: a small shared positional channel following DeepSeek MLA, where an unrotated content slice is reconstructed from the stored value and a short rotated slice is shared across heads, costing d_r scalars per token rather than G·d_r.
4. An initialization rule for the reconstruction map: each M_h is initialized so that reconstructed keys and queries start at the same RMS, σ_M = σ_Q / (σ_V √d_in), addressing the "projection of a projection" scale mismatch that otherwise flattens attention early in training.

Main Findings

• Cache reduction: For the configurations studied, the decoupled-RoPE GVA representation reduces persistent cache scalars by approximately 45–47% relative to matched GQA. The reported ratio is 1/2 + d_r/(2Gd_h), corresponding to about 47% saved at d_r = 16 and 45% at d_r = 24 for the paper's grouping. Without a separate positional slice, shared KV and GVA store exactly half of GQA (T·G·d_h versus 2·T·G·d_h). These are representation-level counts, not measured serving-memory reductions.
• Accuracy parity band: At the 350M-parameter scale trained on 30B FineWeb-Edu tokens, the 16-dimensional positional variant (GVA + DRoPE, d_r = 16) reaches 44.35 average accuracy across five tasks, compared with 44.36 for GQA — a gap of 0.01 points, which the paper describes as within seed variation — and 43.88 for MLA.
• Best GVA row by average: Scale-matched GVA with query RMSNorm scores 44.41, described as the strongest GVA row, sitting close to GQA's 44.36.
• Default variant: GVA variance-fixed without query RMSNorm scores 43.77, slightly lower than 43.77... i.e. close to MLA's 43.88.
• GVA baseline alone is competitive: With standard initialization of M, GVA baseline reaches 43.91, indicating that reconstruction alone already lands in this band; scale matching mainly improves early training loss rather than consistently improving the final average.
• d_r = 24 is behind GQA: GVA + DRoPE with d_r = 24 averages 44.29, and the paper states neither DRoPE row beats GQA on the average.
• Shared KV fails: The first design, setting K = V, halves the GQA cache but the training loss never recovers to the GQA baseline over the whole displayed run — one vector cannot serve both to score and to be retrieved.
• Training trajectories: After the initial transient, scale-matched GVA tracks GQA and MLA in loss; the baseline GVA is worse early and narrows the gap later, consistent with the scale analysis. No second collapse is observed once M is scale-matched.
• Full task-level table (three seeds per configuration, averaged):
  • GQA: HellaSwag 43.41, WinoGrande 52.48, OBQA 33.40, ARC-E 63.42, ARC-C 29.09, Avg 44.36
  • MLA: 43.20, 51.61, 34.60, 62.87, 27.13, Avg 43.88
  • GVA baseline: 42.12, 52.96, 34.80, 61.95, 27.73, Avg 43.91
  • GVA, scale-matched + Q-norm: 42.05, 53.51, 34.40, 62.75, 29.35, Avg 44.41
  • GVA, variance-fixed, no Q-norm: 42.71, 53.35, 32.20, 62.33, 28.24, Avg 43.77
  • GVA + DRoPE d_r = 24: 42.94, 53.12, 33.20, 63.69, 28.50, Avg 44.29
  • GVA + DRoPE d_r = 16: 42.69, 53.35, 33.60, 63.81, 28.32, Avg 44.35
• No systems gains claimed yet: The paper states that neither latency nor decode-throughput gains are established by the present experiments. Fused decode throughput, peak serving memory, and batch capacity measurements are not reported. Prefill is compute-bound rather than cache-bound, so GVA seeks no advantage there; prefill cost stays close to GQA provided the combined content and positional width does not exceed the head-dimension tile the baseline already occupies.
• An earlier released artifact: The initial shared-KV design with query normalization was trained as Lumma-0.6B and open-sourced on Hugging Face as Lumma-0.6B-Base.

Methodology in Plain English

The authors keep GQA's grouping structure for the value side: H query heads share G value heads. They then change only the key. Instead of storing a grouped key, they learn one small matrix per query head that turns the stored grouped value into that head's content key. Because that matrix does not change at inference, the matrix multiplication can be pushed onto the query instead: instead of building a key and comparing the query to it, they transform the query once and take a single inner product against the stored value. This is an exact algebraic identity, not an approximation.

The obstacle is rotary position encoding, which normally rotates keys and queries so that the score depends on the relative distance between positions. If the key is reconstructed from the value and then rotated, the rotation sits between the query and the reconstruction matrix and varies with every cached position, so it cannot be folded into one transformed query. Their fix is to split each head into two parts: an unrotated content slice reconstructed from the value (which absorbs), and a short rotated positional slice whose key is shared across all heads and cached directly. Nothing learned sits between the two rotations, so relative position is preserved, and because the positional key is shared it costs d_r scalars per token instead of G·d_r.

They also had to fix scale. Since the reconstructed key is a projection of the value, which is itself a projection of the hidden state, naive initialization left keys far smaller than queries, producing nearly uniform attention and wasting early training. They derive an initialization for the reconstruction map so that keys and queries start at the same RMS.

Experiments train decoder-only Transformers from scratch at approximately 350M parameters on a 30B-token sample of FineWeb-Edu, using the same data order, token budget, optimizer, and context length across configurations. ZClip was used during pretraining to mitigate gradient spikes. Each configuration is run three times with different random seeds, and zero-shot accuracy is reported on HellaSwag, WinoGrande, OpenBookQA, and the Easy and Challenge splits of ARC; the "Average" is the unweighted mean of the five task-level means. Baselines are GQA and MLA trained with the same recipe. Appendices provide the scale derivation, attention heatmaps illustrating flat versus peaked softmax under different scale regimes, and per-task accuracy curves over training.

Why This Matters

The KV cache is described as a primary bottleneck for Transformer decoding, with footprint and read traffic growing with sequence length. If a cache can be made roughly half the size at essentially unchanged benchmark accuracy, that is a meaningful lever for long-context serving, where capacity and memory bandwidth often bind before compute does. The paper is also notable for what it does not claim: it separates the representation-level cache argument from systems-level measurements, and explicitly says throughput, latency, peak memory, and batch-capacity results are not reported here, with kernels developed and evaluation ongoing.

Real-world applications:

• Long-context serving and chat assistants: smaller per-request cache state directly affects how many concurrent sessions fit on a given accelerator.
• Memory-constrained deployment: edge, on-device, or single-GPU inference where cache footprint, not compute, limits usable context length.
• Retrieval-augmented and document-processing pipelines: workloads that feed very long inputs and therefore produce long caches.
• Batched high-throughput inference services: reducing cache bytes per sequence is relevant to how many sequences can be batched, though the paper does not report batch-capacity measurements.

Industry relevance: GQA is already a standard ingredient in open models (the paper cites Llama 2 70B as an example), and MLA is a leading alternative compression scheme, so any method that claims near-GQA quality at roughly half the intended persistent cache sits directly in the design space that inference-serving teams care about. The paper also frames its contribution relative to MLA honestly — per-head key up-projections over a shared latent are already used by MLA, so head-specific keys are not a GVA-only advantage; GVA's distinction is using the value itself as the persistent content state rather than a separate joint latent, keeping a direct value path.

Future Directions

1. Complete the systems evaluation: the custom decoding kernels have been developed and are being tested; fused decoding throughput, end-to-end inference performance, peak serving memory, and batch capacity against MLA and GQA are all unreported and planned for a release.
2. Sweep RoPE width and allocation: the paper tested only d_r ∈ {16, 24} and did not systematically sweep RoPE width, additive versus carved content/positional allocation, or longer contexts. It also notes the runs do not isolate the effects of positional width versus content-width allocation.
3. Scale and seed coverage: all comparisons use one scale (approximately 350M parameters), one data mix (30B FineWeb-Edu tokens), and three random seeds; broader scale and seed sweeps are named as further work.
4. Clarify the value of scale matching: scale matching primarily improved early loss rather than consistently improving final averages, so the conditions under which the initialization rule pays off at the final benchmark level remain an open question.

Target Audience

Researchers and engineers working on efficient Transformer inference, KV-cache compression, and attention architecture design — particularly those comparing GQA-style head sharing against latent-compression approaches such as MLA. It is also useful for systems-oriented readers in cs.AR interested in how a cache-representation change maps onto decode-path kernels, provided they note that the systems measurements are not yet included. Readers looking for a ready-to-deploy speedup, or for results beyond a single approximately 350M-parameter scale on 30B FineWeb-Edu tokens, will not find them in this version.

Authors’ abstract

The KV cache is a primary bottleneck for Transformer decoding: its memory footprint and cache-read traffic grow with sequence length. Grouped-query attention (GQA) reduces this cost by sharing key-value heads, but still stores both a key and a value at every step. We introduce Grouped Value Attention (GVA), which stores grouped values and reconstructs content keys with a learned linear map. At inference, the map can be absorbed into the query, eliminating the need to materialize content keys in the intended decode path. A small shared decoupled RoPE channel retains positional information through a separately cached positional key. For the configurations studied, this representation reduces persistent cache scalars by approximately 45-47% relative to matched GQA. At the 350M-parameter scale with 30B FineWeb-Edu tokens, the 16-dimensional positional variant reaches 44.18 average accuracy across five tasks, compared with 44.36 for GQA and 43.88 for MLA. These results demonstrate near-GQA benchmark accuracy with a more compact cache representation. To translate this compact representation into faster autoregressive inference, we have developed custom decoding kernels and are currently evaluating their end-to-end inference performance with an open-source release planned soon.