Depthgraft improves large language model learning with efficient cross-layer sharing
DepthGraft: Structural Regularization through Hierarchical Cross-Layer KV Reconstruction
Computation and Language
Summary
Large language models get better when they have many layers, but information can get lost as it moves through these layers. The authors studied a way to share information more efficiently across layers by reusing certain data structures, called key-value pairs, which helps training and performance. They designed a method called DepthGraft that shares information in a hierarchical way, improving results and reducing memory needs. Tests on models of various sizes showed consistent improvements, especially on complex benchmarks.
What this means in practice
- •For ai model engineers: Enhance the training efficiency and accuracy of large transformer-based language models by using DepthGraft for better layer interaction and memory usage.
- •For cloud ai service operators: Reduce memory requirements and improve throughput of large language models deployed in the cloud by integrating DepthGraft’s hierarchical key-value sharing.
Authors
Xiaohan Qin, Xiangdong Zhang, Yu Wang, Huaijin Wu, Zhuo Xia, Yebin Yang, Junchi Yan
Abstract
Scaling depth is a key route to stronger large language models, but Pre-LN architectures often suffer from depth-wise information dilution, limiting the effective use of deep layers. While cross-layer connectivity has proven effective at mitigating this problem, existing designs usually introduce additional computation or memory access. We revisit cross-layer KV sharing as an efficient form of such connectivity. Beyond cache compression, each KV reuse edge creates a direct gradient path from deep layers to source KV projections, turning KV sharing into an implicit structural regularizer for KV representation geometry. We provide a systematic analysis of this mechanism, showing that proper sharing topologies sharpen key retrieval directions while enriching the value content subspace, whereas existing adjacent or global-anchor designs do not fully exploit this effect. Motivated by this, we propose DepthGraft, a hierarchical KV-sharing architecture that reconstructs deep-layer KV from stride-aligned shallow-middle source pairs. DepthGraft distributes non-local gradient flow across source layers, while preserving prefill early completion and reducing KV cache by roughly one third. Experiments on dense and MoE models from 0.5B to 16B parameters show consistent downstream improvements under both GQA and MLA. Notably, on the 16B MoE model, DepthGraft achieves a 7.0% absolute improvement on MMLU-Pro, along with gains of 6.6% on CMMLU and 6.4% on C3. These results validate the potential of DepthGraft as a scalable design for large-scale pretraining. The project page is available at https://github.com/jianke0604/DepthGraft.