DepthGraft:通过分层跨层 KV 重建实现结构正则化
DepthGraft: Structural Regularization through Hierarchical Cross-Layer KV Reconstruction
Xiaohan Qin · Xiangdong Zhang · Yu Wang · Huaijin Wu · Zhuo Xia · Yebin Yang · et al.
中文摘要
扩展深度是增强大语言模型的关键路径,但 Pre-LN 架构常出现深度方向的信息稀释,限制深层有效利用。跨层连接虽能缓解该问题,现有设计却通常引入额外计算或内存访问。重新审视跨层 KV 共享作为一种高效连接形式:除缓存压缩外,每条 KV 复用边为深层到源 KV 投影创建直接梯度路径,使 KV 共享成为 KV 表示几何的隐式结构正则化器。系统分析表明,合适的共享拓扑能锐化键检索方向并丰富值内容子空间,而现有相邻或全局锚点设计未充分利用该效应。据此提出 DepthGraft,一种分层 KV 共享架构,从步幅对齐的浅层-中层源对重建深层 KV。DepthGraft 将非局部梯度流分布到源层,同时保留预填充提前完成并将 KV 缓存减少约三分之一。在 0.5B 至 16B 参数的稠密与 MoE 模型上实验,GQA 与 MLA 下下游任务均一致提升。16B MoE 模型上,MMLU-Pro 绝对提升 7.0%,CMMLU 提升 6.6%,C3 提升 6.4%。结果验证 DepthGraft 作为大规模预训练可扩展设计的潜力。项目页面:https://github.com/jianke0604/DepthGraft。
关键要点
- 01Pre-LN 架构存在深度方向信息稀释,深层利用率受限
- 02跨层 KV 共享形成梯度路径,可作 KV 表示几何的隐式结构正则化器
- 03DepthGraft 从步幅对齐的浅层-中层源对重建深层 KV,缓存减少约三分之一
- 040.5B 至 16B 稠密与 MoE 模型在 GQA 和 MLA 下均获一致提升
- 0516B MoE 上 MMLU-Pro 绝对提升 7.0%,CMMLU 6.6%,C3 6.4%
解读
在 GitHub 查看内容由 AI 生成,仅供参考。如发现错误,欢迎反馈。
原始英文摘要
arXiv:2611.03142v1 Announce Type: new 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.
1. 问题与价值
这篇论文解决的是 Pre-LN Transformer 随深度增加出现的信息稀释问题。具体来说,在 Pre-LN 架构中,残差流以固定加法更新方式聚合所有前层计算,当网络层数变深时,早期层形成的特征会逐渐被稀释,导致深层难以有效检索这些特征。结果是增加层数带来的收益低于其表示容量所暗示的水平。
这个问题重要,因为深度扩展是提升大语言模型能力的主要路线之一。如果深层不能被有效利用,那么单纯堆层数就是在浪费算力和参数。论文指出,现有跨层连接方法虽然有效,但通常引入额外计算或内存访问开销。
受益方是做大模型预训练和推理部署的团队。如果 DepthGraft 成立,训练方可以在不增加训练 FLOPs 的前提下获得更好的下游性能,推理方可以同时获得约三分之一的 KV cache 缩减和 prefill 加速。论文在 16B MoE 模型上报告 MMLU-Pro 绝对提升 7.0%,CMMLU 提升 6.6%,C3 提升 6.4%,同时 TTFT 在 128K 上下文下降低 28.8%。
2. 已有工作的不足
在 DepthGraft 之前,跨层 KV 共享主要有三类设计。
第一类是相邻层共享,代表方法是 CLA 和 MLKV。它们把连续层分组,组内只有第一层计算 KV,其余层复用。论文指出其不足在于查询空间分离不足:相邻层的隐藏状态高度相关,源层和复用层的 query 空间几乎重叠,跨层梯度无法提供足够的非局部压力来驱动 K/V 几何特化。实验上,CLA/MLKV 在 16B A3B MoE 上平均分只有 38.05,低于 GQA 基线的 42.38。
第二类是全局锚点共享,代表方法是 YOCO 和 FusedKV。YOCO 用单一锚点层服务所有深层,FusedKV 用两个全局锚点做通道级融合。论文指出其不足在于梯度拥塞:少量锚点层要同时满足大量异构深层 query 空间的目标,形成高扇入多目标压力,源表示变成过载的折中。同时正则化在结构上不均匀,只有锚点层获得直接跨层 KV 梯度,其余浅中层不参与。实验上 FusedKV 平均分 45.64,虽然超过基线但明显低于 DepthGraft 的 47.56。
第三类是更广泛的跨层连接方法,如 Attention Residuals、MoDA、DenseFormer 等。它们有效缓解了信息稀释,但论文指出这些密集跨层访问引入额外计算、内存流量和系统复杂度,需要仔细的性能-效率权衡。
已有工作的共同盲点是:没有系统研究 KV 共享的跨层梯度动力学,特别是共享拓扑如何重塑 KV 表示几何。它们把 KV 共享主要当作 cache 压缩手段,而不是训练期结构正则化器。
3. 重建作者的思考路径
从旧方法出发。YOCO 和 FusedKV 已经证明全局锚点 KV 共享可以省 cache 且不崩性能,FusedKV 甚至能小幅超过基线。但 FusedKV 的收益有限,且锚点层梯度集中。如果我是作者,会先问:为什么 FusedKV 能超过基线?仅仅是因为省了参数、让深层更「轻」吗?
一个自然的猜测是:深层复用浅层 KV 时,梯度从深层通过共享 KV 回传到源层的投影矩阵。这意味着源层的 KV 投影同时被自己的 query 和深层的 query 训练。这应该会改变源层 KV 的表示性质。但怎么验证?可以看有效秩。如果这个猜测对,那么源层 K 的有效秩应该比基线更低(更集中),V 的有效秩应该更高(更丰富)。这是可以直接做实验验证的。
撞上的第一个矛盾:如果跨层梯度真的有用,为什么 CLA/MLKV 反而比基线差?它们也有跨层梯度。可能的解释是相邻层的 query 太像,梯度是冗余的,没有提供新的「视角」。这提示跨层距离是关键变量。
第二个矛盾:如果跨层梯度有用,为什么 YOCO 不如 FusedKV?YOCO 只有一个锚点,FusedKV 有两个。两个锚点给了深层两个抽象层次的记忆基。但 FusedKV 的锚点只有两个,所有深层的梯度都压在这两个锚点上。这会不会造成「梯度拥塞」——太多异构目标压在同一组参数上?
此时浮现两条设计原则:源层和复用层要有足够的深度距离;梯度流要分散到多个源层而不是集中在少数锚点。但还要满足效率约束:所有源 KV 必须在复用层之前产生,才能保持 prefill 提前完成。
怎么同时满足「距离足够远」和「源层分散」?一个直接的想法是把网络分成三段:浅、中、深。深层不复用自己的 KV,而是从浅层和中层各取一个源。这样每个深层有两个抽象层次的记忆基,且源层数量是深层的两倍,梯度被分散。
但「取哪个浅层、哪个中层」?如果所有深层都取同一个浅层锚点和同一个中层锚点,那就退化成 FusedKV。要分散,就得让不同深层取不同的源。一个工程上简单的方案是步幅对齐:第 ℓ 个深层取第 ℓ−2L/3 个浅层和第 ℓ−L/3 个中层。这样每个源层只服务一个或少数几个深层,梯度自然分散,而且实现上只是索引偏移,不引入复杂路由。
改完发现的新问题:如果直接复制源 KV,深层失去了自己的 KV 投影,表达能力可能受限。解决方式是学习通道级系数,对两个源的 K 和 V 分别做逐通道加权求和。但 K 要过 RoPE,如果系数在 RoPE 对内不一致,RoPE 和重建就不对易,cache 逻辑会乱。所以 K 的系数要在 RoPE 对内绑定。
与论文实际路径的差异:我猜的路径是从「为什么 FusedKV 有效」出发,先假设跨层梯度改变表示几何,再用有效秩验证。论文的实际路径确实如此,但论文比我猜的更早地引入了有效秩作为核心诊断工具,并且在第 3.2 节先建立了标准注意力下 K/V 几何不对称性的理论框架(K 通过 softmax 竞争自然集中,V 通过线性读出自然扩展),然后才论证 KV 共享放大这种不对称性。我猜的路径里没有先建立这个「基线不对称性」的铺垫,而是直接跳到跨层梯度的作用。另外,论文对「相邻共享为什么失败」的分析比我猜的更定量,明确用了 query 空间分离不足和表示约束主导弱正则化收益的论证。
4. 核心 intuition
用一个日常类比:把 KV 共享想象成「两个老师教同一个学生」。
标准 Transformer 里,每一层的 KV 投影只被自己这一层的 query 训练,相当于一个老师只教一个学生。这个老师的教法会偏向学生当前的需求,K 会变得越来越「专一」——只保留能区分当前 query 的检索方向;V 会变得越来越「丰富」——因为要传递的内容越来越多。
跨层 KV 共享相当于让两个不同年级的老师共用同一个学生。浅层老师教基础词汇,深层老师教抽象语义。两个老师的需求有重叠但不完全一样。对于 K(检索索引),两个老师都需要的那些区分方向会得到一致强化,而各自特有的方向会互相干扰,所以 K 变得更集中、更锐利。对于 V(内容通道),两个老师要传递的内容不同,学生必须同时覆盖两边的需求,所以 V 的覆盖范围变大、更丰富。
这个类比能对上机制:K 通过 softmax 竞争,两个 query 空间的共享方向得到相干强化,特有方向互相干扰,所以有效秩下降;V 通过线性读出,没有 softmax 竞争,两个读出模式把互补的梯度路由到同一个 V 矩阵,所以有效秩上升。相邻层共享失败的原因是两个老师教的内容几乎一样,没有新的信息;全局锚点失败的原因是太多老师挤在一个学生身上,学生变成谁都不满意的折中。
5. 方法 pipeline 举例
假设有一个 9 层模型,L=9,分成三段,每段 3 层。浅层是第 0、1、2 层,中层是第 3、4、5 层,深层是第 6、7、8 层。输入是一个 4 token 的序列,隐藏维度 D 简化为 4,KV 头数 1,头维度 d=4。
第一步,浅层正常计算。第 0 层输入隐藏状态 X0 是 4×4 矩阵,经过 W_K0 和 W_V0 投影得到 K0 和 V0,各是 4×4。第 1、2 层同理,各自产生 K1、V1、K2、V2。这些 KV 被缓存。
第二步,中层正常计算。第 3、4、5 层各自产生 K3、V3、K4、V4、K5、V5,也被缓存。
第三步,深层重建 KV。第 6 层的步幅对齐源是:浅层源 s(6)=6−2×9/3=6−6=0,中层源 m(6)=6−9/3=6−3=3。所以第 6 层从第 0 层和第 3 层重建 KV。重建公式是 K6 = a6,1⊙K0 + a6,2⊙K3,V6 = b6,1⊙V0 + b6,2⊙V3。其中 a6,1、a6,2、b6,1、b6,2 都是 4 维通道系数,初始化为 0.5。所以初始时 K6 就是 K0 和 K3 的逐元素平均。
第四步,第 6 层用重建的 K6、V6 和自己计算的 Q6 做注意力。Q6 由 X6 经过 W_Q6 投影得到。注意力输出是 softmax(Q6·K6^T/√d)·V6,再经过 W_O6 输出。
第五步,第 7 层同理,源是 s(7)=7−6=1,m(7)=7−3=4,从第 1 层和第 4 层重建。第 8 层从第 2 层和第 5 层重建。
第六步,反向传播。第 6 层的注意力损失梯度通过 K6 和 V6 回传到第 0 层和第 3 层的 KV 投影。第 0 层的 W_K0 同时收到来自第 0 层自身注意力的梯度和来自第 6 层的梯度。第 3 层的 W_K3 同理。第 1、2、4、5 层也各自收到来自对应深层的梯度。
第七步,推理时。prefill 阶段只需要跑完第 0 到 5 层,所有 KV cache 就齐了。第 6 到 8 层只需要对最后一个 prompt token 做前向,用已缓存的浅中层 KV 计算注意力,产生第一个输出 token。KV cache 只存 6 层的量,是基线 9 层的 2/3。
6. 核心数学推导
先解释符号。L 是层数,ℓ 是层索引。Xℓ 是第 ℓ 层的隐藏状态,维度 T×D,T 是序列长度,D 是隐藏维度。Qℓ、Kℓ、Vℓ 分别是第 ℓ 层的 query、key、value,由 Xℓ 分别乘以投影矩阵 W_Qℓ、W_Kℓ、W_Vℓ 得到。d 是每个头的维度。H_q 是 query 头数,H_k 是 KV 头数。在 GQA 下,query 头被分配到 KV 组,φ(h) 表示第 h 个 query 头所属的 KV 组。
注意力输出公式:Attnℓ = Concat over h of softmax(Qℓ,h · Kℓ,φ(h)^T / √d) · Vℓ,φ(h),再乘以 W_Oℓ。这是标准注意力,原文公式可辨认。
现在推导标准注意力下 K 和 V 的梯度。定义 Sℓ = Qℓ·Kℓ^T/√d 是 softmax 前的 logits,Pℓ = softmax(Sℓ) 是注意力权重,Oℓ = Pℓ·Vℓ 是注意力输出。设 Gℓ = ∂L/∂Sℓ 是损失对 logits 的梯度,Eℓ = ∂L/∂Oℓ 是损失对输出的梯度。
K 的梯度:因为 Sℓ = Qℓ·Kℓ^T/√d,所以 ∂Sℓ/∂Kℓ 涉及 Qℓ,系数是 1/√d。链式法则给出 ∇Kℓ = (1/√d)·Gℓ^T·Qℓ。这个公式说明 K 的更新方向由 query 和 logits 梯度共同决定。由于 softmax 的竞争性归一化,只有那些能产生有用 logits 分离的 key 方向才会被持续强化。
V 的梯度:因为 Oℓ = Pℓ·Vℓ,所以 ∂Oℓ/∂Vℓ 涉及 Pℓ。链式法则给出 ∇Vℓ = Pℓ^T·Eℓ。这个公式说明 V 的更新方向由注意力权重和输出梯度共同决定。一旦 Pℓ 固定,输出对 V 是线性读出,没有 softmax 竞争,所以 V 可以被推向覆盖更多独立内容方向。
现在加入跨层共享。假设源层 s 的 KV 被远处的层 ℓ 复用。源层 K 的梯度变成两项之和:∇share K_s = (1/√d)·G_s^T·Q_s + (1/√d)·G_ℓ^T·Q_ℓ。第一项是源层自身的梯度,第二项是复用层的梯度。源层 V 的梯度类似:∇share V_s = P_s^T·E_s + P_ℓ^T·E_ℓ。
关键观察:对于 K,两个 query 空间共享的判别方向会得到相干强化,而层特有的方向会通过 softmax 竞争互相干扰,所以 K 的有效秩下降。对于 V,两个读出项不通过 softmax 竞争,不同注意力模式把互补梯度路由到同一个 V 矩阵,所以 V 的有效秩上升。
有效秩的定义:对中心化的表示矩阵 M(N×d),计算特征协方差 C_M = M^T·M/N,其特征值为 λ_1 到 λ_d。归一化概率 p_i = λ_i / Σ_j λ_j。有效秩 erank(C_M) = exp(−Σ_i p_i·log p_i)。这是 Shannon 熵的指数形式,衡量特征能量在方向上的均匀程度。有效秩高表示能量均匀分布,有效秩低表示能量集中在少数方向。
这个推导在算什么:它在算跨层 KV 共享如何通过梯度路径改变源层 K 和 V 的表示几何。核心结论是 K 的有效秩下降(检索更锐利),V 的有效秩上升(内容更丰富),且这种不对称性来自 softmax 竞争和线性读出的机制差异。
7. 实验设计
数据集方面,预训练语料是私有的,包含 web text、code、math、多语言来源,经过精确/模糊去重、规则过滤和 LLM 质量过滤。下游评估用 OpenCompass 框架,覆盖四组能力:知识与语言理解(MMLU 5-shot、MMLU-Pro 5-shot、LAMBADA 0-shot)、推理与数学(ARC-Challenge 25-shot、SIQA 5-shot、GSM8K 4-shot)、代码生成(HumanEval 0-shot、MBPP 3-shot、LCBench 5-shot)、中文理解(C-Eval 5-shot、CMMLU 5-shot、C3 3-shot)。长上下文评估用 RULER 的 10 个任务,上下文长度从 4K 到 32K。
基线方面,主对比是 GQA 基线和 MLA 基线,即同规模同训练配方但不做 KV 共享的标准模型。消融中对比了 CLA/MLKV(相邻共享,3 层一组)、YOCO(全局单锚点)、FusedKV(全局双锚点)。这些基线在 16B A3B MoE 上用相同训练协议复现。
指标方面,下游任务用准确率或 pass rate,汇总用所有 benchmark 的平均分。效率指标用 TTFT(time-to-first-token)和 TPOT(time-per-output-token),KV cache 用量用 MiB。表示几何用有效秩追踪。
消融设计分两个轴。第一个轴是梯度连通性:Shallow-Copy (detach) 阻断跨层梯度,Shallow-Copy 只允许浅层到深层的单源复制梯度,DepthGraft 用浅中层双源重建。第二个轴是阶段布局:2-stage、3-stage、4-stage 三源、4-stage 两源。消融在 16B A3B MoE 上进行。
实验能否支撑结论:主结果表 1 和表 2 覆盖 0.5B 到 16B 共 6 个模型配置,dense 和 MoE 都有,GQA 和 MLA 都有,一致性较好。消融中 Shallow-Copy (detach) 对比 Shallow-Copy 是关键的因果检验,能支撑「跨层梯度是收益来源」的论断。有效秩曲线(图 2)和梯度范数分布(图 3b)为机制解释提供了直接证据。
缺什么:第一,所有训练用的是私有数据,外部无法完全复现数据分布的影响。第二,没有报告训练 token 总量,无法判断收益是否在更长训练后会消失或扩大。第三,SFT 后收益保留的实验只在 16B 一个规模上做了。第四,没有做超过 32K 上下文的长程评估。第五,没有报告多次随机种子下的方差,无法判断小规模模型上 +0.75 的平均分提升是否显著。
8. take-aways
9. 最脆弱的假设
最脆弱的假设是:有效秩的变化是性能提升的因果机制,而不是伴随现象。论文的核心论证链条是「跨层梯度 → K 有效秩下降、V 有效秩上升 → 更好的下游性能」。但论文提供的证据主要是相关性:DepthGraft 同时改善了有效秩和下游性能,Shallow-Copy (detach) 同时失去了两者。这不能排除第三种因素同时驱动有效秩和性能,例如深层 KV 投影的移除本身改变了优化景观,有效秩只是这个变化的被动反映。
论文给出的支持证据:第一,Shallow-Copy (detach) 与 Shallow-Copy 的对比显示梯度路径是关键,这是较有力的因果证据。第二,附录 C 的命题 1 在显式假设下证明了有效秩变化方向的条件性结论,但假设本身(跨层 key 梯度共享低维任务相关子空间、value 梯度有互补内容方向)没有被直接验证。第三,图 2 显示有效秩差异在训练早期就出现,与性能差异的时间线一致。
如果这个假设倒塌,即有效秩只是副现象,那么论文的「KV 共享作为结构正则化器」的理论框架就退化为一种事后解释,DepthGraft 的实际价值将主要来自工程层面的 cache 节省和 prefill 加速,而非论文强调的表示几何改善。这会削弱论文的学术贡献,但不完全否定方法的实用价值。
10. 一周最小复现
目标:在 1 张 24GB 消费级显卡上复现最核心的结论——跨层梯度是 DepthGraft 收益的主要来源。
步骤:
预期指标:由于训练 token 量远小于论文,下游分数的绝对提升会小于论文报告的 +0.75,但三个变体之间的相对排序应该保持。如果一周内看不到 DepthGraft 与 detach 变体的差距,说明梯度正则化效应在极短训练下不明显,需要延长训练或增大模型。
11. 反例设计
核心主张是:跨层 KV 共享的梯度路径通过正则化 KV 表示几何来提升下游性能,且这种正则化需要足够的查询空间分离和分散的梯度流。
反例实验:构造一个「梯度路径存在但查询空间不分离」的变体,如果它的性能与 DepthGraft 相当,就推翻了「查询空间分离是必要条件」的论断。
具体设计:在 16B A3B MoE 上,保持 DepthGraft 的三阶段布局和双源重建,但把源层选择从步幅对齐改为「最近邻源」。即每个深层 ℓ 从浅层和中层各选距离最近的源,例如第 6 层从第 5 层(中层最近)和第 2 层(浅层最近)重建,第 7 层从第 5 层和第 2 层重建,第 8 层也从第 5 层和第 2 层重建。这样梯度路径仍然存在,但浅中层源与深层的距离被最小化,查询空间分离不足。同时梯度集中在少数几个源上,类似于全局锚点。
判据:如果这个变体的平均分达到或超过 DepthGraft 的 47.56(允许 0.5 分的噪声区间),则论文关于「充分跨层距离」和「分散梯度流」的设计原则被推翻。如果它显著低于 DepthGraft 但高于 GQA 基线 42.38,则部分支持论文但说明距离和分散的贡献可以分离。如果它低于基线,则进一步支持论文。
更强的反例:构造一个「梯度路径存在但有效秩不变化」的变体。在 DepthGraft 的梯度回传路径上加入一个正交性约束或谱归一化,强制源层 K 和 V 的有效秩在训练中保持不变。如果这个变体的下游性能仍然与 DepthGraft 相当,则有效秩变化不是性能提升的必要条件,论文的机制解释被推翻。判据是平均分是否在 47.56 的 0.5 分范围内。
12. 落地判断
用于「抢占式 GPU 自托管 27B 模型做多智能体脚手架」是可行的,但收益和风险需要分开看。
怎么用:如果 27B 模型是 GQA 架构,可以按论文的三阶段布局改造。假设 27B 模型有约 40 层,分成 13+13+14 或类似的三段。深层 14 层不独立计算 KV,从步幅对齐的浅中层源重建。实现上需要修改模型前向逻辑,在深层用重建的 KV 替代独立投影。如果模型是 MLA 架构,论文在 9B MLA 上验证过,但 MLA 的 KV 结构不同,需要按 MLA 的 latent KV 做重建,实现复杂度更高。
成本量级:KV cache 从 27B 模型全量降到约三分之二。以 27B 模型、GQA、8 个 KV 头、头维度 128、40 层计算,每 token 的 KV cache 约为 40 层 × 2(K 和 V)× 8 头 × 128 维 × 2 字节 = 约 163,840 字节,即约 160 KiB/token。DepthGraft 后约 107 KiB/token。在 32K 上下文下,单序列 KV cache 从约 5 GiB 降到约 3.3 GiB。对于多智能体脚手架同时跑多个序列的场景,这直接增加了可并发的序列数。prefill 时间在长上下文下可降低约四分之一到三分之一。
最大的坑:第一,27B 模型如果是已有预训练模型,不能直接套用 DepthGraft——论文的所有收益都是从零预训练得到的。对已训练模型做结构改造需要继续预训练,成本可能超过收益。第二,多智能体脚手架通常需要频繁的长上下文 prefill,DepthGraft 的 prefill 加速依赖「深层只处理最后一个 token」的实现,这需要推理框架支持非标准的 prefill 调度,在 vLLM 或 TensorRT-LLM 上可能需要自定义 kernel。第三,论文的 TPOT 开销低于 1% 是在其测试条件下,自托管场景的 batch size 和序列长度分布不同,解码开销可能放大。第四,如果 27B 模型用了 sliding window attention 或其他注意力变体,需要按论文附录 D.2 的方式限制 KV 复用在同类型注意力层之间,否则梯度范数不匹配可能导致训练不稳定。