Deep layers stop building their own KV: cross-layer reconstruction lifts a 16B MoE by 7.0 points on MMLU-Pro
Synopsis
DepthGraft turns cross-layer KV sharing from a cache-saving inference trick into a structural regularizer during training: deep layers no longer compute their own KV but rebuild it from a stride-aligned shallow–middle source pair using learnable channel-wise coefficients. Gradients from deep layers then flow back along each sharing edge into the source KV projections, sharpening key retrieval directions while enriching the value content subspace. Across dense and MoE models from 0.5B to 16B parameters, downstream averages improve consistently, with a 7.0-point absolute gain on MMLU-Pro for the 16B MoE and roughly one third less KV cache.
Interpretation
Cross-layer KV sharing is treated as a form of cross-layer connectivity rather than only a way to compress the inference cache. Deep layers read shallow and middle memory representations directly, bypassing the information channel that the residual stream dilutes with depth. Earlier adjacent-layer and global-anchor sharing designs were framed around cache compression and inference acceleration, without a systematic account of how the cross-layer gradients created by sharing reshape KV representations. On the 16B A3B MoE under the same three-stage layout, stopping gradients through the reused KV gives an average of 42.86, close to the 42.38 GQA baseline; enabling the backward path raises it to 45.30.
Deep-layer KV is reconstructed from a stride-aligned pair of shallow and middle sources, so each deep layer receives memory bases at two abstraction levels. Single-source copying connects only the middle or only the shallow stage to the sharing graph, while global-anchor designs concentrate all deep-layer objectives on one or two source layers. In the same 16B A3B MoE setting, DepthGraft reaches an average of 47.56, above 45.30 for shallow single-source copy and 43.58 for the two-stage single-source variant.
The deep stage no longer materializes independent KV tensors, so KV cache and prefill cost fall while decoding overhead stays essentially flat. A standard Transformer caches KV for every layer, and prefill must push the whole prompt through all layers before generation can start. With a fixed 2K generation length on 9B MLA and 16B GQA, time to first token drops by 26.8% and 28.8% at 128K context, while per-token latency changes by less than 1% across all tested lengths.
Perspective
The scope is dense and DeepSeekMoE-style models from 0.5B to 16B parameters, all pretrained from scratch with 8192-token context under one OpenCompass few-shot protocol; long-context results come from RULER on the 16B A3B model after continued training at 32K context. It is aimed at researchers and engineers working on large-model architecture and inference efficiency.
The theoretical part is explicitly a conditional spectral analysis and does not claim that effective rank must move monotonically under arbitrary training dynamics; the relative training-loss reduction is described as moderate, with downstream gains larger than the loss difference alone would suggest. The authors list a more rigorous theoretical account, further sharing topologies and larger pretraining scales as future work.