深层KV不再各层自建:跨层重建把16B MoE的MMLU-Pro抬高7.0个百分点
核心概要
DepthGraft把跨层KV共享从省缓存的推理技巧改造成一种训练期的结构正则:深层不再独立计算KV,而是从浅层与中层按相同相对深度取一对源层,用可学习的逐通道系数重建。这样深层的梯度会沿共享边回流到源层KV投影,使键的检索方向更集中、值的内容子空间更丰富。在0.5B到16B的稠密与MoE模型上,下游平均分一致提高,16B MoE的MMLU-Pro绝对提升7.0个百分点,同时KV缓存减少约三分之一。
深度剖析
跨层KV共享被当作一种跨层连接来用,而不只是压缩推理缓存的手段。深层直接读取浅层与中层的记忆表示,绕开残差流里逐层稀释的信息通道。 此前的相邻层共享与全局锚点共享主要从缓存压缩和推理加速角度设计,没有系统考察共享边带来的跨层梯度如何改变KV表示。 在16B A3B MoE上,同一三阶段布局下,切断复用KV的反向梯度后平均分为42.86,与GQA基线42.38相当;打开反向路径后升到45.30。
深层KV由一对按步长对齐的浅层与中层源重建,每个深层因此同时拿到两个抽象层级的记忆基。 单源复制只把中层或浅层之一接入共享图,全局锚点方案则把全部深层目标压在一两个源层上。 同一16B A3B MoE设置下,DepthGraft平均分47.56,高于单源浅层复制的45.30和两阶段单源复制的43.58。
深层不再物化独立的KV张量,KV缓存与预填充成本随之下降,而解码开销基本不变。 标准Transformer为每一层缓存各自的KV,预填充必须把整段提示走完全部层才能开始生成。 在9B MLA与16B GQA上以固定2K生成长度测试,128K上下文时首 token 延迟分别降低26.8%和28.8%,逐 token 延迟在所有测试长度上变化低于1%。
启示与展望
讨论范围限于0.5B至16B参数的稠密与DeepSeekMoE式模型,全部从零预训练、使用8192上下文与同一套OpenCompass少样本评测协议;长上下文结果来自16B A3B模型在32K上下文继续训练后的RULER评测。适合关注大模型架构与推理效率的研究者和工程师阅读。
理论部分明确是有条件的谱分析,并不主张任意训练动态下有效秩都单调变化;训练损失的相对下降幅度被描述为温和,下游收益大于损失差异所暗示的程度。作者把更严格的理论解释、更多共享拓扑与更大预训练规模列为后续工作。
