DepthGraft：分层KV重建如何成为深度模型的结构正则化利器

深度扩展是提升大语言模型能力的关键路径，但Pre-LN架构常受困于深度方向的信息稀释，导致深层网络难以有效利用。跨层连接被证明能缓解这一问题，然而现有设计往往引入额外的计算或内存访问开销。最近一篇题为《DepthGraft: Structural Regularization through Hierarchical Cross-Layer KV Reconstruction》的工作重新审视了跨层KV共享，揭示其除了缓存压缩之外的隐式正则化作用，并提出了一种分层共享架构DepthGraft。

跨层KV共享：不只是缓存压缩

KV共享在推理加速中常被用于减少KV缓存，但这项研究指出，每条KV重用边都会在深层与源层KV投影之间建立直接的梯度路径。这意味着KV共享实际上充当了KV表示几何的隐式结构正则化器。作者系统分析了这一机制，发现恰当的共享拓扑能够锐化键的检索方向，同时丰富值的内容子空间。而现有的相邻层共享或全局锚点设计并未充分利用这一效应。

这一视角将KV共享从单纯的效率优化提升为一种结构正则化手段，为深度模型的设计提供了新思路。

Pre-LN架构：Pre-LN（Pre-Layer Normalization）指在Transformer每个子层（自注意力、前馈网络）之前施加LayerNorm，而非原始Post-LN在残差相加之后归一化。Pre-LN使残差路径保持恒等映射，训练更稳定、可扩展至更深网络，是当前大模型的主流选择。但其残差流在深层会因反复叠加而归一化尺度趋同，导致不同层表示差异缩小，即文中所述“深度方向的信息稀释”——深层难以从浅层获得充分、有区分度的信号，这也是跨层连接被引入的结构性动机。

DepthGraft架构：分层重建与梯度分布

受上述分析启发，作者提出了DepthGraft，一种分层KV共享架构。它从步幅对齐的浅层-中层源对中重建深层KV，从而将非局部梯度流分布到多个源层。同时，该设计保留了预填充早期完成特性，并将KV缓存减少约三分之一。

DepthGraft的核心在于通过分层重建，让深层网络能够接收到来自不同深度源层的梯度信号，从而缓解信息稀释问题。这种设计不引入额外计算或内存访问，却实现了跨层连接的正则化效果。

步幅对齐（stride-aligned）源对选择：DepthGraft并非简单令相邻层共享，而是按固定步幅从浅层与中层各取一层组成源对，再由该对重建深层KV。步幅对齐确保源层在深度上分布均匀，避免梯度路径过度集中；浅层源提供低阶、局部特征，中层源提供较抽象语义，使深层KV表示同时吸收不同抽象层级的梯度信号。与相邻共享相比，梯度可跨越多个层回传，路径更丰富；与全局锚点相比，又避免所有深层依赖单一源层造成的信息瓶颈。

实验验证：多模型规模下的一致提升

作者在0.5B到16B参数的密集模型和MoE模型上进行了实验，覆盖GQA和MLA两种注意力机制。结果显示，DepthGraft在 downstream 任务上带来一致改进。在16B MoE模型上，DepthGraft在MMLU-Pro上取得7.0%的绝对提升，CMMLU提升6.6%，C3提升6.4%。这些结果验证了DepthGraft作为大规模预训练可扩展设计的潜力。

说个细节，这些提升是在KV缓存减少约三分之一的情况下实现的，表明该方法在效率与性能之间取得了良好平衡。

GQA与MLA：GQA（Grouped-Query Attention）让多个查询头共享一组键值头，在推理时显著压缩KV缓存，已被Llama等主流模型采用；MLA（Multi-head Latent Attention）是DeepSeek^广告提出的低秩潜在注意力，将KV投影压缩到低维潜在向量再按需展开，进一步降低缓存占用。两者对KV表示的压缩方式不同，DepthGraft在两种机制下均验证有效，说明其分层重建的正则化效应不依赖特定注意力实现，具备架构通用性。

分析：为何分层共享更有效

作者对比了相邻层共享和全局锚点设计，发现它们未能充分激发KV共享的正则化效应。DepthGraft通过步幅对齐的源对选择，将梯度流分散到浅层和中层，使得深层KV表示能够同时受益于不同抽象层级的特征。这种分层结构既避免了全局锚点可能带来的信息瓶颈，又比相邻共享具有更丰富的梯度路径。

从表示学习角度看，锐化的键检索方向有助于模型更精准地匹配查询与键，而丰富的值内容子空间则增强了信息聚合能力。这两者共同作用，提升了深层网络的有效性。

总结与展望

DepthGraft将跨层KV共享重新定义为结构正则化机制，并通过分层重建实现了梯度流的非局部分布。实验表明，该方法在多种模型规模和注意力机制下均有效，且能减少KV缓存。项目页面已开源，为后续研究提供了基础。未来，如何自适应地选择共享拓扑、平衡正则化强度与模型容量，仍是值得探索的方向。

背景补充

预填充早期完成特性：在自回归推理中，预填充阶段需为全部输入token计算并缓存KV。若深层KV可由浅层、中层已算出的KV重建，则深层不必等待自身投影完成即可获得可用表示，从而让预填充更早结束、缩短首token延迟。DepthGraft在减少约三分之一KV缓存的同时保留该特性，意味着其收益不仅来自存储节省，也体现在推理时延上。