DepthGraft: Structural Regularization through Hierarchical Cross-Layer KV Reconstruction

DepthGraft. Structural Regularization through Hierarchical Cross-Layer KV Reconstruction. By Xiaohan Qin and six others, from Shanghai Jiao Tong University, the Shanghai Innovation Institute, and one other institution. Published on the second of November.

Making a large language model deeper is a key route to making it stronger. Pre-LN architectures, though, often suffer from depth-wise information dilution, which limits how effectively the deep layers get used. Cross-layer connectivity has proven effective against that dilution. The catch is that existing designs for it usually bring extra computation or extra memory access, so the fix for depth charges its own price.

The authors return to cross-layer KV sharing and look at it as an efficient form of that connectivity. Usually KV sharing is thought of as cache compression. Their observation goes further. Each edge along which a KV is reused creates a direct gradient path from a deep layer back to the source KV projections. Sharing therefore acts as an implicit structural regularizer on the geometry of the KV representations.

A systematic analysis of this mechanism shows that the topology matters. Proper sharing topologies sharpen the directions along which keys are retrieved and enrich the subspace that the values carry content in. Existing adjacent designs, and global-anchor designs, do not fully exploit that effect.

DepthGraft is the architecture built on that finding. It is hierarchical: deep-layer KV is reconstructed from pairs of source layers in the shallow and middle part of the network, chosen to be stride-aligned. Spreading the sources this way distributes non-local gradient flow across several source layers rather than concentrating it. Two practical properties are kept. Prefill early completion is preserved, and the KV cache shrinks by roughly one third.

Experiments cover dense models and MoE models, ranging from zero point five billion to sixteen billion parameters, and the downstream improvements are consistent under both GQA and MLA. The largest model gives the headline numbers. On the sixteen-billion-parameter MoE model, DepthGraft gains seven point zero percent absolute on MMLU-Pro. It also gains six point six percent on CMMLU, and six point four percent on C3. The authors read these results as validating DepthGraft as a scalable design for large-scale pretraining.

What changes here is the reason for sharing a cache. Read as compression, KV sharing is simply a cost saving. Read through its gradients, each reuse edge is a path that ties deep layers back to shallow projections, and the shape of the sharing pattern decides how well that regularizes the representations. DepthGraft picks the pattern for the second reason and still collects the first, with a cache about a third smaller.