DepthGraft: Structural Regularization through Hierarchical Cross-Layer KV Reconstruction

📄 arXiv: 2611.03142v1 📥 PDF

作者: Xiaohan Qin, Xiangdong Zhang, Yu Wang, Huaijin Wu, Zhuo Xia, Yebin Yang, Junchi Yan

分类: cs.CL

发布日期: 2026-11-02

🔗 代码/项目: GITHUB

💡 一句话要点

提出DepthGraft以解决深层信息稀释问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 深度学习 大型语言模型 KV共享 结构正则化 梯度流动 模型优化 预训练

📋 核心要点

  1. 现有的Pre-LN架构在深层信息利用上存在稀释问题，限制了模型的表现。
  2. DepthGraft通过分层KV共享重建深层KV，优化了梯度流动并减少了内存使用。
  3. 在多个参数规模的模型上，DepthGraft在下游任务上实现了显著的性能提升，尤其是在16B MoE模型上取得了7.0%的绝对提升。

📝 摘要（中文）

深度扩展是增强大型语言模型的关键途径，但Pre-LN架构常常面临深度信息稀释的问题，限制了深层的有效利用。尽管跨层连接被证明能有效缓解这一问题，但现有设计通常会引入额外的计算或内存访问。本文重新审视了跨层KV共享作为一种高效的连接形式，提出DepthGraft，一个分层KV共享架构，通过从步幅对齐的浅中层源对重建深层KV，分配非局部梯度流，同时保持预填充的早期完成，减少约三分之一的KV缓存。实验结果表明，DepthGraft在多个模型上均实现了显著的下游性能提升。

🔬 方法详解

问题定义：本文旨在解决Pre-LN架构中深层信息稀释的问题，现有方法在跨层连接上引入了额外的计算和内存开销，影响了模型的效率和性能。

核心思路：DepthGraft通过跨层KV共享的方式，重建深层KV，形成直接的梯度路径，从而优化深层信息的利用，减少计算和内存负担。

技术框架：DepthGraft的整体架构包括多个层次的KV共享模块，利用步幅对齐的浅中层源对来重建深层KV，确保非局部梯度流动的同时，保持预填充的早期完成。

关键创新：DepthGraft的主要创新在于将KV共享转化为隐式结构正则化，优化了KV表示几何，提升了关键检索方向和价值内容子空间的丰富性，超越了现有的相邻或全局锚设计。

关键设计：在设计中，DepthGraft通过合理的共享拓扑结构来提高关键检索方向的清晰度，具体参数设置和损失函数的设计确保了模型的有效性和稳定性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，DepthGraft在多个模型上均表现出色，尤其是在16B MoE模型上，MMLU-Pro任务实现了7.0%的绝对提升，CMMLU和C3任务分别提升6.6%和6.4%。这些结果验证了DepthGraft作为大规模预训练设计的潜力。

🎯 应用场景

DepthGraft的研究成果具有广泛的应用潜力，尤其在大型语言模型的预训练和微调过程中，可以有效提升模型的性能和效率。其设计理念可推广至其他深度学习任务，推动更高效的模型架构发展，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.