KV sharing is usually framed as a memory-saving shortcut. The less obvious idea here is that every reuse edge also sends deep-layer training signals back to the layers that produced the shared keys and values. That makes cache reuse a way to shape representations, not just store fewer tensors.
Key Findings
- Deep layers reconstruct their KV states from stride-aligned shallow and middle sources. Learnable, per-dimension coefficients combine the two representations, giving deep layers access to both lower-level patterns and more abstract features.
- The hierarchical topology spreads gradient flow across many layers rather than concentrating it on adjacent layers or a few global anchors. The authors’ analysis links this structure to sharper key-retrieval directions and a richer value subspace.
- Across dense and MoE models from 0.5B to 16B parameters, results improve over matched baselines with both GQA and MLA. In the reported 16B MoE comparison, MMLU-Pro rises by 7.00 points, while CMMLU and C3 rise by 6.61 and 6.41 points.
- The deep stage does not materialize independent KV tensors, cutting cache use by roughly one third. At 128K context, reported time-to-first-token falls by 26.8% for a 9B MLA model and 28.8% for a 16B GQA model; decoding time per output token is nearly unchanged.
Method and Trade-offs
The source layers are arranged in shallow, middle, and deep stages. Each deep layer combines KV from corresponding positions in the first two stages, with learned channel-wise weights. This preserves an early-completion path for prefill while allowing deep-layer losses to update the source KV projections. The gains therefore depend on training the model with this sharing topology; this is an architectural change, not a drop-in inference optimization for an unchanged model.
Who It’s For
A good fit for researchers exploring Transformer depth, KV-cache reduction, and prefill latency, especially when training or adapting model architectures is feasible. Look elsewhere if you need a plug-and-play speedup for an existing model or a substantial reduction in token-by-token decoding latency: the reported benefit is strongest in cache use and prefill, while TPOT remains nearly flat.