DepthGraft: How Hierarchical KV Reconstruction Becomes a Structural Regularization Tool for Deep Models

Depth scaling is a key path to improving large language model capabilities, but Pre-LN architectures often suffer from information dilution along the depth dimension, making it difficult for deeper networks to effectively utilize information. Cross-layer connections have been shown to alleviate this problem, yet existing designs often introduce additional computation or memory access overhead. A recent work titled "DepthGraft: Structural Regularization through Hierarchical Cross-Layer KV Reconstruction" re-examines cross-layer KV sharing, revealing its implicit regularization role beyond cache compression, and proposes a hierarchical sharing architecture called DepthGraft.

Cross-Layer KV Sharing: More Than Just Cache Compression

KV sharing is commonly used in inference acceleration to reduce KV cache, but this study points out that each KV reuse edge establishes a direct gradient path between deep-layer and source-layer KV projections. This means KV sharing actually acts as an implicit structural regularizer for the geometry of KV representations. The authors systematically analyze this mechanism and find that an appropriate sharing topology can sharpen the retrieval direction of keys while enriching the content subspace of values. However, existing adjacent-layer sharing or global anchor designs do not fully exploit this effect.

This perspective elevates KV sharing from a mere efficiency optimization to a structural regularization approach, offering new ideas for the design of deep models.

Pre-LN architecture: Pre-LN (Pre-Layer Normalization) refers to applying LayerNorm before each sublayer (self-attention, feed-forward network) of the Transformer, rather than the original Post-LN which normalizes after residual addition. Pre-LN keeps the residual path as an identity mapping, making training more stable and scalable to deeper networks, and is the mainstream choice for current large models. However, its residual stream in deep layers tends to converge in normalization scale due to repeated accumulation, causing the representational differences between layers to shrink—that is, the "information dilution along the depth dimension" mentioned in the article—making it difficult for deep layers to obtain sufficiently rich and discriminative signals from shallow layers. This is also the structural motivation for introducing cross-layer connections.

DepthGraft Architecture: Hierarchical Reconstruction and Gradient Distribution

Inspired by the above analysis, the authors propose DepthGraft, a hierarchical KV sharing architecture. It reconstructs deep-layer KV from stride-aligned shallow-middle source pairs, thereby distributing non-local gradient flow across multiple source layers. Meanwhile, the design preserves the early prefill completion property and reduces KV cache by about one third.

The core of DepthGraft lies in hierarchical reconstruction, enabling deep networks to receive gradient signals from source layers at different depths, thereby alleviating the information dilution problem. This design introduces no additional computation or memory access, yet achieves the regularization effect of cross-layer connections.

Stride-aligned source pair selection: DepthGraft does not simply make adjacent layers share; instead, it selects one layer each from shallow and middle layers at a fixed stride to form source pairs, from which deep-layer KV is reconstructed. Stride alignment ensures that source layers are evenly distributed in depth, avoiding excessive concentration of gradient paths; shallow sources provide low-order, local features, while middle sources provide more abstract semantics, allowing deep-layer KV representations to absorb gradient signals at different abstraction levels simultaneously. Compared with adjacent sharing, gradients can be back-propagated across multiple layers, resulting in richer paths; compared with global anchors, it avoids the information bottleneck caused by all deep layers relying on a single source layer.

Experimental Validation: Consistent Improvements Across Model Scales

The authors conduct experiments on dense and MoE models ranging from 0.5B to 16B parameters, covering both GQA and MLA attention mechanisms. Results show that DepthGraft brings consistent improvements on downstream tasks. On a 16B MoE model, DepthGraft achieves a 7.0% absolute improvement on MMLU-Pro, 6.6% on CMMLU, and 6.4% on C3. These results validate the potential of DepthGraft as a scalable design for large-scale pre-training.

One detail: these improvements are achieved while reducing KV cache by about one third, indicating that the method strikes a good balance between efficiency and performance.

GQA and MLA: GQA (Grouped-Query Attention) allows multiple query heads to share a set of key-value heads, significantly compressing KV cache during inference, and has been adopted by mainstream models such as Llama; MLA (Multi-head Latent Attention) is a low-rank latent attention proposed by DeepSeek^Ad, which compresses KV projections into low-dimensional latent vectors and expands them on demand, further reducing cache usage. The two compress KV representations differently, and DepthGraft is validated under both mechanisms, demonstrating that its regularization effect from hierarchical reconstruction does not depend on a specific attention implementation and has architectural generality.

Analysis: Why Hierarchical Sharing Is More Effective

The authors compare adjacent-layer sharing and global anchor designs, finding that they fail to fully activate the regularization effect of KV sharing. DepthGraft, through stride-aligned source pair selection, disperses gradient flow to shallow and middle layers, allowing deep-layer KV representations to benefit from features at different abstraction levels simultaneously. This hierarchical structure avoids the potential information bottleneck of global anchors while providing richer gradient paths than adjacent sharing.

From a representation learning perspective, sharpened key retrieval directions help the model match queries and keys more precisely, while rich value content subspaces enhance information aggregation. Together, these improve the effectiveness of deep networks.

Conclusion and Outlook

DepthGraft redefines cross-layer KV sharing as a structural regularization mechanism and achieves non-local distribution of gradient flow through hierarchical reconstruction. Experiments show that the method is effective across various model scales and attention mechanisms, and can reduce KV cache. The project page has been open-sourced, providing a foundation for future research. In the future, how to adaptively select sharing topologies and balance regularization strength with model capacity remains a direction worth exploring.

Background Supplement

Early prefill completion property: In autoregressive inference, the prefill stage needs to compute and cache KV for all input tokens. If deep-layer KV can be reconstructed from KV already computed in shallow and middle layers, then deep layers do not need to wait for their own projections to complete before obtaining usable representations, allowing prefill to finish earlier and shortening first-token latency. DepthGraft preserves this property while reducing KV cache by about one third, meaning its benefits come not only from storage savings but also from inference latency.

One-line Summary

DepthGraft uses hierarchical cross-layer KV reconstruction to turn KV sharing into implicit regularization, improving deep model performance while reducing cache.