混合LLM中的注意力失忆：CoT微调破坏长距离召回及修复方法

周新宇

LARK，香港科技大学（广州）

朱博宇

伦敦大学学院

徐毅

Mistral AI

李志伟

LARK，香港科技大学（广州）

陈英发

王慧明

新加坡科技设计大学

郭志江

LARK，香港科技大学（广州）

摘要

思维链监督微调被广泛用于提升推理能力，然而我们发现，在混合线性注意力模型中，它会系统性地损害长上下文召回能力。在包括 HypeNet 和 Jet-Nemotron 在内的多种架构上，经过思维链监督微调后，模型在"大海捞针"测试中的检索性能显著下降，并且在更困难的检索设置和更长的上下文窗口下，退化程度更为严重。例如，HypeNet-9B 在 NIAH-S2@256K 上的性能从某个值下降到了另一个值。我们将此归因于思维链监督微调使注意力梯度偏向短距离模式，从而破坏了负责长距离路由的查询-键投影。基于这一观察，我们提出了 QK-Restore，一种无需训练的方法，该方法仅从监督微调前的检查点恢复查询投影和键投影，同时保留所有其他监督微调后的参数。我们进一步引入了一种 Procrustes 变体，以平衡路由保留与推理适应。在各种架构上，QK-Restore 在零训练成本下持续恢复长上下文能力，同时保持推理性能；例如，在 HypeNet-5B 上，它将 S3@256K 从某个值提升到另一个值，同时保持了强大的推理性能。

摘要

思维链（CoT）监督微调（SFT）被广泛用于提升推理能力，然而我们发现它会系统性地降低混合线性注意力模型的长上下文召回能力。在包括 HypeNet 和 Jet-Nemotron 在内的多种架构上，经过 CoT-SFT 后，大海捞针（NIAH）测试中的检索性能显著下降，并且在更困难的检索设置和更长的上下文窗口下，这种退化变得更加严重。例如，HypeNet-9B 在 NIAH-S2@256K 上的表现从 [原文此处数值缺失] 下降至 [原文此处数值缺失]。我们将此归因于 CoT-SFT 使注意力梯度偏向短程模式，从而破坏了负责长程路由的查询-键投影（query-key projections）。基于这一发现，我们提出了 QK-Restore，一种无需训练的方法，该方法仅从预 SFT 检查点恢复 [原文此处变量缺失] 和 [原文此处变量缺失]，同时保留所有其他后 SFT 参数。我们还引入了一种 Procrustes 变体，以平衡路由保留与推理适应。在多种架构上，QK-Restore 在零训练成本下一致地恢复了长上下文能力，同时保持了推理性能；例如，在 HypeNet-5B 上，它将 S3@256K 从 [原文此处数值缺失] 提升至 [原文此处数值缺失]，同时保持了强大的推理性能。

1 引言

高效序列模型，例如线性注意力（katharopoulos2020transformersrnnsfastautoregressive; qin2024variouslengthsconstantspeed; yang2024gatedlinearattentiontransformers）和状态空间模型（gu2024mambalineartimesequencemodeling; dao2024transformersssmsgeneralizedmodels），已成为长上下文处理中 softmax 注意力（vaswani2023attentionneed）的极具吸引力的替代方案，通过紧凑的循环或结构化状态将二次复杂度降低为线性复杂度。然而，这种压缩在长上下文理解（zhang2024hedgehogporcupineexpressive）所依赖的强召回任务（例如“大海捞针”NIAH; hsieh2024ruler）上引入了信息瓶颈。混合线性注意力模型通过保留少量 softmax 注意力层用于全局召回，同时将其余层转换为高效的线性注意力层（nvidia2025nvidianemotron3efficient），来缓解这种权衡。尽管 Qwen3.5（qwen35blog）和 Kimi Linear（kimiteam2025kimilinearexpressiveefficient）等原生混合模型实现了强大的性能和效率，但从头开始预训练它们仍然非常耗费资源。这一成本推动了近期的工作，即通过知识蒸馏，仅使用一小部分预训练 token（chen2026hybridlinearattentionright; li2025distilling; hoshino2025radredundancyawaredistillationhybrid），将预训练的 softmax 注意力模型转换为混合线性注意力模型。现有的转换工作主要关注在架构转换过程中保留预训练能力，而下游推理后训练对长上下文检索的影响仍未得到充分探索。

虽然蒸馏混合模型主要作为高效的基础模型使用，但现代大语言模型通常依赖后训练来获得更强的指令遵循和推理能力（Guo_2025；qwen2025qwen25technicalreport；ouyang2022traininglanguagemodelsfollow）。特别是，链式思维监督微调被广泛用于增强数学和多步推理（wei2023chainofthoughtpromptingelicitsreasoning；li2025tl）。我们观察到，在数学推理数据上应用链式思维监督微调可以提升推理性能，但会显著降低蒸馏混合模型在预训练或架构转换过程中获得的长上下文召回能力，尤其是在更具挑战性的 NIAH 设置和扩展上下文长度中（图 1）。与一般的灾难性遗忘不同，这种退化具有高度结构性：它主要影响由保留的 softmax 注意力层介导的长距离检索行为，而预期的推理增益则基本保持不变。这揭示了混合模型中的一个根本性矛盾：链式思维监督微调在强化局部推理的同时，可能会破坏长距离检索所需的路由机制。

为了研究这一现象，我们将链式思维推理轨迹建模为中间推理状态上的潜在马尔可夫过程，以反映数学推导中局部的逐步结构（wang2026doeschainofthoughthelpmarkovian；prystawski2023thinkstepstepreasoning）。在此链式思维-马尔可夫假设下，我们推导出注意力 logits 上的期望梯度幅度随 token 距离呈指数衰减。我们通过测量 token 自相关性和注意力 logits 的梯度，从经验上验证了这一预测。这些结果表明，链式思维监督微调可以改善局部多步推理，同时会侵蚀检索所需的长距离路由行为。

为了更好地隔离召回能力下降的根源，我们通过路由-提取分解来分析保留的 softmax 注意力层。我们发现，CoT-SFT 会在查询-键投影中引发局部性偏差漂移，而查询-键投影决定了信息检索的来源；相比之下，值侧提取仍能从 SFT 后的适应中受益。由于蒸馏混合模型中的长上下文召回高度依赖于少量保留的 softmax 注意力层，这种查询-键漂移会不成比例地破坏检索能力。受此不对称性的启发，我们提出了 QK-Restore，这是一种无需训练的方法，仅从 SFT 前的检查点恢复这些层中的查询和键投影，同时保留所有其他 SFT 后的参数，从而在恢复长距离路由能力的同时，保留 CoT-SFT 带来的大部分推理改进。我们的贡献可总结如下：

• •

  我们识别出 CoT-SFT 引发的召回能力下降是蒸馏混合模型的一种结构化后训练失效模式。

• •

  我们提供了理论和实证分析，表明 CoT-SFT 将训练信号集中在局部 token 交互上，并在保留的 softmax 注意力层的查询-键路由几何结构中引发局部性偏差漂移。

• •

  我们引入了 QK-Restore，这是一种无需训练的方法，仅恢复查询和键投影，从而在基本保持推理性能的同时，恢复长上下文召回能力。

2 相关工作

高效长上下文模型与混合注意力机制。近期混合注意力模型将 softmax 注意力层与高效循环或线性注意力层交错结合，在常识推理和需要高召回率的长上下文任务上均取得了有竞争力的性能，同时提升了推理效率（qwen35blog; kimiteam2025kimilinearexpressiveefficient; nvidia2025nvidianemotron3efficient）。从零开始大规模预训练强大的混合模型仍然成本过高，这促使近期研究转向通过 Transformer 到混合模型的转换或知识蒸馏来获得混合模型（chen2026hybridlinearattentionright; li2025distilling）。然而，现有工作主要聚焦于获取高性能的转换后混合模型检查点，而它们在后续后训练阶段的行为仍未被充分探索。

将 Transformer 蒸馏为混合模型。Transformer 到混合模型的蒸馏将选定的 softmax 注意力层转换为线性或循环混合器，其中层选择对长上下文检索性能至关重要（goldstein2026radladsrapidattentiondistillation; chen2026hybridlinearattentionright; li2025distilling; gu2026jet）。然而，转换后强大的召回性能并不一定意味着在面向推理的后训练后仍能保持稳定；具体而言，在多个具有强长上下文检索能力的蒸馏混合模型中，我们观察到 CoT-SFT 会显著降低召回能力，这揭示了一个现有蒸馏方法尚未解决的关键失效模式。

思维链动态与注意力路由。近期研究将思维链推理视为一种在中间推理状态上的结构化生成过程（wang2026doeschainofthoughthelpmarkovian; prystawski2023thinkstepstepreasoning）。马尔可夫视角通过将推理建模为潜在状态之间的局部转移来形式化这一结构，表明思维链监督主要强化了短程的步骤间依赖关系（wang2026doeschainofthoughthelpmarkovian）。与先前利用这种局部性来解释思维链在推理中有效性的工作不同，我们将其与高效长上下文模型中的一种失败模式联系起来：偏向局部性的思维链监督微调会侵蚀召回所需的长程注意力路由。这一观点与将注意力视为信息路由机制的分析相一致，其中查询-键交互定义了用于选择位置的路由逻辑值，而值投影则决定了从这些位置提取的内容（vaswani2023attentionneed）。

3 背景与问题设定

3.1 混合模型

在本工作中，我们主要关注混合模型。为了在提升长上下文效率的同时保留 softmax 注意力的表达能力，混合架构将少量 softmax 注意力层与高效的线性注意力层交错排列（yang2025zebra; chen2026hybridlinearattentionright）。

Softmax 注意力层 对于 Transformer 的第 层、第 个注意力头，每个头的投影矩阵（）对隐藏状态 进行变换：

其中标量逻辑值 和注意力权重 为

每个头的输出向量 和该层的输出为

线性注意力层 对于第 层，一个循环状态矩阵 通过以下方式演化：

其中 、 是特定于层的更新和读取算子。以 Lightning Attention（qin2024lightningattention2freelunch）为例，状态更新和读取过程为：

我们将 softmax 注意力层的集合记为 。实证研究表明，混合模型中长距离召回能力不成比例地依赖于一组 softmax 注意力层，而其余大部分层可以用循环机制替代，且性能下降极小（wang2025systematicanalysishybridlinear; chen2026hybridlinearattentionright; jelassi2024repeat）。因此，保留 的路由行为对于长上下文召回至关重要。

3.2 注意力路由

我们通过 softmax 之前的 logit 来实现注意力路由的操作化，该 logit 决定了位置 对位置 的注意力强度。由于路由模式完全由 决定，对这些矩阵的改动会直接改变长距离检索行为。为了理解 CoT-SFT 训练如何影响检索行为，我们考察了关于 QK 指标的梯度。 的梯度更新为：

如果 在小距离时很大，在大距离时可忽略不计，那么梯度下降会系统地推动 趋向局部模式，无论上下文长度如何。此外，我们有：

其中 是梯度向量。因此，理解 的距离依赖性成为理解 CoT-SFT 如何影响检索的核心。

4 为什么 CoT-SFT 会破坏路由？

4.1 CoT 数据假设

我们首先通过一个潜在马尔可夫模型（wang2026doeschainofthoughthelpmarkovian; prystawski2023thinkstepstepreasoning）来刻画 CoT 数据结构，这为后续的梯度分析提供了动机。

假设 4.1（CoT-马尔可夫结构）。

存在潜在推理状态 ， ，使得：

1. 马尔可夫转移： 。

2. 观测模型： 。位置 处的 token 由当前推理状态决定。

3. 遍历性：该链是不可约且非周期的，具有平稳分布 ， 。

4. 可逆性：细致平衡条件成立： 对所有 成立。

5. 谱间隙：转移矩阵的第二大特征值幅度为 。

我们在附录 A 中讨论该假设有效性的相关论述。

4.2 梯度局部性理论

Softmax 注意力层是混合模型中长距离召回的唯一场所。因此，我们分析 CoT-SFT 如何通过梯度局部性侵蚀这些层中的路由能力。

4.2.1 定义与假设

定义 4.2（距离条件梯度幅度）。

假设 4.3（范数边界）。

存在常数，使得在训练分布上几乎必然地，

假设 4.4（得分函数恒等式）。

设 表示损失函数关于位置 处每个注意力头输出的梯度。在交叉熵损失下的近似最优性条件下，模型满足：对于所有 以及 上几乎必然地，

上述假设的论证包含在附录 B 和附录 C 中。

4.2.2 梯度局部性定理

引理 4.5（谱相关性衰减）。

对于一条具有谱间隙 的平稳可逆遍历马尔可夫链，以及任意均值为零的 ：

定理 4.6（梯度局部性）。

设 为假设 4.3 中的常数， 为假设 4.1 中的谱半径。在假设 4.3–4.4 以及 CoT-马尔可夫模型（假设 4.1）下，对于所有 ：

常数 仅取决于模型的范数边界；衰减率完全由潜在转移矩阵的谱间隙决定。完整证明包含在附录 F 中。

4.3 实证验证

定理 4.6 表明，CoT 训练会优先强化邻近 token 之间的交互，导致远距离路由信号随距离增加而衰减。这一理论链条依赖于关于数据和模型的两个实证主张：

1. 引理 4.5 的假设：控制 token 生成的潜在马尔可夫链具有谱间隙 。然而，由于潜在状态不可观测，我们转而考虑观测到的 token 序列，并使用 token 自相关性作为 的代理指标。具体而言， 衡量的是某个 token 在距离 处重新出现的概率比随机情况高出多少，该值按所有 token 类型以其平稳频率 加权平均后得到。如附录 G 所证明，在隐马尔可夫模型下成立。

2. 定理 4.6 的预测：期望梯度幅度随 呈指数衰减，且 CoT 数据的有效窗口比通用文本更小。我们直接从模型中测量每个距离 处的 ：

3. 曲线拟合。两者均采用形式 ，其中 作为衰减的自然尺度。因此，我们通过对数线性回归来拟合模型。取对数可将模型线性化：

4. 解释。如图 3 所示，左子图展示了模型 token 自相关衰减参数 ，这是一个语料库级别的统计量，用于衡量文本保持结构化、非随机自相似性的特征距离。由于数学推理会反复引用相同的符号、变量和公式片段，思维链文本在更长距离上（ token）维持着比普通散文（ token）更高的自相似性。右子图展示了注意力梯度衰减参数 ，它通过衡量梯度信号 的有效作用范围来刻画训练动态。两条衰减曲线均能很好地用指数拟合来描述，这与定理的预测一致。

当 时，存在一个距离区间，在该区间内数据需要长程注意力，但训练梯度不再对其进行强化。关于预训练阶段也观察到的注意力梯度衰减的讨论，请参见附录 H。

5 方法

定理 4.6 已经确定，注意力 logit 上的梯度随距离呈几何级数衰减。现在我们研究这种衰减会传播到哪些权重指标中。

5.1 路由-提取梯度解耦

定理 5.1（路由-提取梯度解耦）

在思维链监督微调条件下，当 token 生成过程被建模为具有谱间隙 的平稳可逆遍历马尔可夫链时，对 softmax 注意力层参数的梯度更新表现出以下不对称性：

[路由] 来自距离为 的位置对的 的逐对贡献满足：

其中 为常数。 也满足相同界限。

[提取] 在任意位置 的值向量上的梯度满足：

其中 对所有 成立，且 对所有 成立。

详细证明见附录 I。因此，思维链监督微调对路由和知识提取的影响不同。 仅从短程位置对接收有意义的梯度信号，而到达 的梯度在所有位置上都有统一的下界，与上下文长度无关。

5.2 QK-恢复

定理 5.1 表明，CoT-SFT 以根本不同的方式影响 softmax 注意力层的两个功能组件。

路由参数（控制模型关注的位置）接收到的梯度信号随距离呈指数衰减：它们仅由短距离配对塑造，并逐渐丧失长距离路由能力。提取参数（控制检索内容）在每个位置接收到的梯度均具有一致的下界：它们会累积价值处理方面的改进。

换句话说，CoT-SFT 同时具有破坏性和改进性。这些效应被分离到不相交的参数集中。我们提出了 QK-Restore，该方法通过移植 SFT 前检查点的参数，同时保留 SFT 后的参数（算法 1），从而消除路由破坏并保留提取改进。

6 实验

6.1 实验设置

模型。HypeNet (chen2026hybridlinearattentionright) 近期展现了强大的长上下文召回性能，因此我们主要关注该模型在 2B 到 9B 规模上的表现。为确保分析的全面性，我们还包含了 Jet-Nemotron-2B111https://huggingface.co/collections/jet-ai/jet-nemotron。

训练。为研究 CoT-SFT 的影响，我们主要聚焦于数学领域，该领域研究广泛且拥有充足的高质量数据集 (ultradata-math; mitra2024orcamath; yu2025dapoopensourcellmreinforcement)。我们在预训练和 SFT 阶段自行训练了不同规模的 HypeNet，并基于 Jet-Nemotron 的预训练检查点对其应用 SFT。数据集和训练配置的更多细节见附录 J。

评估。我们评估模型在长上下文召回和数学推理任务上的表现。对于长上下文召回，我们报告在 NIAH 上的准确率。为衡量数学推理能力，我们在 GSM8K (cobbe2021gsm8k) 和 MATH500 (lightman2023letsverifystepstep) 上对模型进行测试。

6.2 主要结果

性能对比结果见表 1。

CoT-SFT 导致长上下文性能下降。一个关键发现是，CoT-SFT 通常会削弱模型在预训练阶段获得的长上下文检索能力，尤其是在更具挑战性的检索设置和更长的上下文窗口下。虽然像 NIAH-Single-1 这样的简单任务几乎仍能保持饱和性能，但在 NIAH-Single-2 和 NIAH-Single-3 上出现了显著的性能下降。

例如在表 1 中，HypeNet-2B 在 128K 上下文下的 NIAH-Single-3 任务上，经过 SFT 后从 [原文此处缺失数值] 下降至 [原文此处缺失数值]；而 HypeNet-9B 在 256K 上下文下的 NIAH-Single-2 任务上，从 [原文此处缺失数值] 下降至 [原文此处缺失数值]。此外，随着上下文长度增加，性能下降持续加剧，这表明 SFT 对负责长距离检索的机制造成了不成比例的破坏。

QK-Restore 恢复长上下文能力。QK-Restore 能够持续缓解监督微调（SFT）过程中产生的长上下文性能损失，在退化最严重的配置中观察到的提升幅度最大。

具体在表 1 中，对于 HypeNet-2B，在 256K 长度下的 NIAH-Single-3 任务上，性能从提升至（+19.6），而 HypeNet-9B 则从提升至（+19.8）。在 OpenThoughts-3 的链式思维监督微调（CoT-SFT）下，HypeNet-9B 在 256K 长度下的 NIAH-Single-2 任务上从提升至，在 NIAH-Single-3 任务上从提升至。在若干配置中，QK-Restore 甚至超越了原始的预训练基线。例如，HypeNet-2B 在 64K 长度下的 NIAH-Single-3 任务上，从预训练结果提升至。这表明我们的方法并非简单地恢复预训练行为；相反，它能够在保留预训练阶段检索结构的同时，仍然受益于监督微调（SFT）过程中引入的任务特定表示。

QK-Restore 保留下游推理性能。除了检索能力之外，我们还评估了恢复路由参数对下游推理能力的影响。在各项基准测试中，QK-Restore 在显著提升长上下文检索能力的同时，很大程度上保留了从链式思维监督微调（CoT-SFT）中获得的性能提升。

在数学推理任务上，性能保持接近监督微调（SFT）后的模型水平。例如，对于 HypeNet-5B，QK-Restore 在 MATH500 和 GSM8K 上分别仅带来 和 个百分点的微小变化，同时恢复了长上下文检索性能。在其他模型规模上也观察到了类似的趋势。这些结果表明，可以选择性地恢复路由参数以恢复长上下文检索能力，而不会显著影响后训练阶段获得的任务特定知识和推理能力。

检索难度加剧长上下文退化。我们进一步观察到，检索的复杂性对鲁棒性有强烈影响。NIAH-Single-1 在各训练阶段几乎保持饱和状态，而在 Single-2 和 Single-3 任务中则出现了显著更大的差距。这一趋势表明，难度更高的检索设置对精确的长距离 token 交互提出了更高要求，因此对退化现象更为敏感。

7 分析

7.1 非 CoT SFT 分析

我们的理论分析将长上下文性能下降归因于 CoT 监督所引发的局部优化动态。一个自然的问题是：这种现象是后训练阶段的普遍现象，还是 CoT 风格推理轨迹所特有的？为了回答这个问题，我们使用 Tulu-3 数据集，将 CoT-SFT 与传统的指令遵循 SFT 进行了比较。

如表 2 所示，标准的指令微调并未表现出 CoT-SFT 下观察到的相同性能下降模式。在 HypeNet-5B 和 HypeNet-9B 上，经过 Tulu-3 训练后，长上下文检索性能基本得以保持，并且经常得到提升。例如，HypeNet-5B 在 NIAH-Single-3 的 256K 长度上从 [原文此处缺少数值] 提升至 [原文此处缺少数值]，而 HypeNet-9B 则从 [原文此处缺少数值] 提升至 [原文此处缺少数值]。在多个上下文长度下的 NIAH-Single-2 上也观察到了类似的提升。

此外，根据 IFEval 的衡量，经过 Tulu-3 微调后，指令遵循能力显著提升。尽管有这些改进，长上下文检索能力仍然保持完好。这与 CoT-SFT 的行为形成了鲜明对比——在相同的架构和上下文长度下，CoT-SFT 的检索性能出现了下降。

这些发现表明，长上下文性能下降并非后训练阶段的固有结果。相反，它与 CoT 风格的监督密切相关，这支持了我们的假设：推理轨迹的顺序依赖结构将梯度更新集中在短程交互上，并对长程路由行为产生了不成比例的影响。

7.2 超越数学推理的分析

我们的主要实验聚焦于数学推理，因为该领域容易获得高质量的 CoT 监督信号。为了评估观察到的检索-推理权衡是否也存在于数学领域之外，我们利用 OpenThoughts-3 进行了额外实验，并在 LiveCodeBench-V5 上评估了下游编码性能。

如表 3 所示，在编码领域也观察到了相同的定性行为。CoT-SFT 显著提升了编码能力，使 HypeNet-5B 在 LiveCodeBench-V5 上的得分从 [原文缺失数值] 提升至 [原文缺失数值]，HypeNet-9B 从 [原文缺失数值] 提升至 [原文缺失数值]。然而，这些提升伴随着长上下文检索能力的明显下降，尤其是在具有挑战性的 NIAH 设置和较长上下文长度上。

应用 QK-Restore 可以在保持下游编码性能的同时，恢复大部分丢失的检索能力。例如，在 HypeNet-5B 上，LiveCodeBench-V5 得分进一步从 [原文缺失数值] 提升至 [原文缺失数值]，而 NIAH-Single-3 在 256K 长度下的得分从 [原文缺失数值] 恢复至 [原文缺失数值]。类似地，HypeNet-9B 在 NIAH-Single-3 的 256K 长度下得分从 [原文缺失数值] 提升至 [原文缺失数值]，同时保持了有竞争力的编码性能（[原文缺失数值]）。

这些结果表明，CoT-SFT 导致的检索能力下降并非数学推理所特有。相反，它似乎源于与长推理链相关的一种更普遍的优化现象。QK-Restore 在数学和编码两个领域均有效，这进一步支持了我们的核心假设：长距离检索主要由路由参数控制，而任务特定的推理能力则主要编码在网络的其它部分。

7.3 注意力路由恢复分析

在本节中，我们对注意力图进行可视化，以描述 CoT 监督微调如何改变 HypeNet-5B 中 softmax 注意力层的长程路由行为。我们基于在 Miromind 数据集上训练的模型进行分析。

具体而言，我们比较了 HypeNet-5B 在 SFT 前和 CoT-SFT 后两个检查点之间所有 softmax 注意力层的注意力矩阵。然后，我们通过每个注意力头的平均注意力距离指标来识别受影响最大的层：

其中 索引注意力头， 是从查询 到键 的 softmax 后权重。根据 SFT 后 值下降的头数量对层进行排序，第 33 层成为受影响最大的层，32 个头中有 25 个转向更短的注意力距离。因此，我们将分析重点放在这一层。

我们在一个包含 8,192 个 token 的合成 NIAH 提示词上可视化召回率最高的注意力头，目标事实嵌入在上下文深度中。对于每个注意力头，我们展示一个四面板图：SFT 前注意力矩阵（蓝色）、SFT 后矩阵（红色）、QK-Restore 矩阵（绿色）以及逐元素对数比率（面板 4），其中正值表示 QK-Restore 比 SFT 后恢复了更多长程注意力的位置。

分析。图4展示了HypeNet-5B中第33层第8个注意力头的模式。从图4中我们观察到，非对角线的下三角区域主要呈红色，这证实了QK-Restore部分恢复了被CoT微调抑制的长距离路由：对于大多数查询位置，QK-Restore对远距离键的关注程度高于微调后的模型。其次，右下角区域（查询和键位置均位于"针"之后）沿主对角线呈现蓝色带状，这意味着微调后模型在该区域对邻近token赋予了更高权重，这与我们的假设相符。这些观察结果证明了QK-Restore在缓解路由崩溃和恢复长距离注意力行为特征方面的有效性。

7.4 QK-Restore消融实验

QK-Restore将预微调检查点的完整矩阵移植到微调后模型中，大幅恢复了长距离召回能力，但导致数学性能下降。微调后模型可分解为：

其中 编码了有害的局部性漂移， 编码了有益的数学能力适应。由于这两个分量在 中纠缠在一起，QK-Restore将它们一并丢弃，在恢复路由的同时牺牲了数学收益。如图5和图6所示， 中出现了不可忽略的漂移，其中 。

因此，我们将其建模为一个Procrustes问题：

我们将其称为QK-Pro变体（更多细节见附录N）。

此外，我们评估了注意力中各组件（即 ）的作用，并将CoT-SFT与冻结 的预防性方案（记为QK-Frozen）进行对比，以回答一个问题：是在训练过程中预防路由漂移更好，还是在事后进行修正？

分析。表 4 验证了三个关键主张。V-Restore 无法恢复 NIAH，同时导致 MATH500 性能下降，证实其无助于路由恢复。Q-Restore 和 K-Restore 各自实现了部分召回恢复，但在 128K 长度下显著不足（例如，K-Restore 在 NIAH-Single-3 上崩溃，得分 13.2 对比 30.8），表明孤立恢复引入了不匹配，且联合恢复对于连贯的路由几何结构是必要的。QK-Pro 在 MATH500 上取得了最高性能，但以牺牲长程路由恢复来换取数学能力保留。QK-Frozen 在 NIAH 和数学任务上始终不如 QK-Restore，表明无约束的 SFT 需要更好的适应性。

8 结论

在这项工作中，我们将 CoT-SFT 引发的长上下文召回退化识别为蒸馏混合模型中的一个关键后训练失败模式。我们表明，面向推理的 SFT 会使优化偏向短程注意力模式，扰乱保留的 softmax 注意力层的查询-键路由几何结构，并削弱长程检索能力。受此观察启发，我们提出了 QK-Restore，一种无需训练的方法，该方法仅从 SFT 前的检查点恢复查询和键投影，同时保留所有其他参数。跨多种架构的实验表明，QK-Restore 在基本保留 SFT 推理增益的同时，显著恢复了长上下文能力。总体而言，我们的发现强调了路由稳定性是构建同时支持检索与推理的高效长上下文模型的关键因素。

参考文献

附录 A 关于 CoT-马尔可夫结构假设的讨论

讨论 1。1 步马尔可夫条件是为了分析而做的理想化处理。任意步的马尔可夫链等价于在状态乘积空间上的 1 步链，产生相同的指数衰减，但推理范围更大。

讨论 2. 可逆性在引理 4.5 的证明中被用于其中一步：它意味着马尔可夫算子在 上是自伴的。我们承认，因果生成的序列（例如 CoT 推理链）通常不满足细致平衡条件：其潜在转移概率在构造上就是不对称的。然而，可逆性是一个充分条件而非必要条件。对于一般的遍历链（可能不可逆），相同的界仍然成立，只需将 替换为 的第二大奇异值，该值通过其奇异值分解而非特征分解获得。由于奇异值始终是实数且非负，证明过程无需修改即可进行，最终界的形式也保持不变。因此，我们在引理 4.5 中假设可逆性是为了表述清晰，同时理解所有后续结果通过将特征值替换为奇异值后，对一般遍历情况同样成立。

附录 B 范数界假设的论证

对 的界。值向量为 ，其中 表示层归一化。层归一化对每个坐标进行标准化，使得 ，其中 和 是 各分量的经验均值和标准差。归一化后，仿射变换前的向量满足 ，因此得到

其中 是一个仅依赖于学习到的缩放参数 和偏置 的常数。投影矩阵 是一个固定的训练权重；其算子范数是一个有限的、依赖于模型的常数。综合以上，我们得到：

对 的界。梯度通过最后一层层归一化（其雅可比矩阵为 ）和语言模型头 反向传播，得到

其中 是预测分布， 是独热目标。根据算子范数的次可乘性：

由于 是概率向量的差，因此 。对于一个固定的训练模型， 和 都是有限常数。因此：

附录 C 得分函数恒等式假设的论证

在最优解 处，关于注意力输出的交叉熵损失对于任意固定上下文 满足：

由于交叉熵梯度在其最小值处消失。应用全期望公式可得。更精确地说，因为梯度是在给定上下文下计算的，然后对从条件分布中抽取的目标取平均，该等式有条件地成立：

附录 D 预备定义与引理

我们列出证明所需的以下定义和引理。

定义 D.1（空间）

对于具有有限状态空间和平稳分布的马尔可夫链，定义加权函数空间：

带有内积

以及范数。

此外，均值为零意味着，即该函数在下的期望为零。

定义 D.2（作用于函数的算子）

我们定义作用于函数的马尔可夫算子：

引理 D.3

对于可逆链，细致平衡条件成立：，这使得在中是对称的。

证明

根据 D.2 中马尔可夫算子的定义，我们有：

由于，自伴算子的有限维谱定理给出了一个实的标准正交特征基，其特征值为，谱间隙为，其中：

进一步，对于均值为零的：

这意味着沿方向没有分量。其展开仅使用：

附录 E 谱相关衰减引理的证明

证明

在平稳状态下（），塔性质表明：

注意内部期望是步预测：对于固定的起始状态，经过步后的期望值恰好是：

因此：

然后我们通过标准正交特征基展开：

利用标准正交性，我们进一步简化：

然后对其进行约束：

应用柯西-施瓦茨不等式：

现在将所有结果合并：

附录 F 梯度局部性定理的证明

证明

为清晰起见，我们在以下推导中省略层/头的上标。首先，对梯度进行约束：

然后我们考虑如何约束。应用塔性质：

其中我们定义：

根据假设 4.4：

类似地，中心化不会改变协方差（因为）。和都是上的均值为零的函数。

然后应用引理 4.5：

对项应用相同的约束，该项满足作为凸组合：

其中且。∎

附录 G：模型 token 自相关衰减的证明

定义中心化指示特征以及滞后 处的 -加权自相关：

对于分子 ，我们将其展开为：

考虑 的期望：

因此我们可以进一步简化：

在隐马尔可夫模型（HMM）中，模型 token 在给定其潜在状态的情况下是条件独立的： 。根据塔性质：

其中 是 的一个均值为零的函数：

由于 和 都是均值为零的，我们可以直接应用引理 4.5：

其中 是潜在转移矩阵 的第二大特征值。现在我们可以限定自相关的范围：

结论。在语料库中观察到 且 证实了潜在链具有谱间隙 ，这恰好是引理 4.5 的假设条件。

附录 H：预训练阶段注意力梯度衰减的讨论

定理 4.6 基于引理 4.5，该引理适用于任何具有谱间隙的遍历马尔可夫链：这是数据生成过程的一个属性，并非 CoT 结构所特有，在一般的预训练数据语料库中也能观察到。

因此，梯度局部性界限适用于任何训练语料库，包括一般的预训练数据。路由坍缩的严重程度由失配程度 决定：即数据要求模型关注的距离与梯度实际强化该关注的距离之间的差异。对于一般的预训练语料库，这种失配程度小于 CoT 语料库，因此路由退化得更慢，但机制是相同的。

然而，对于混合架构来说，即使是预训练阶段较慢的侵蚀，也会在微调开始之前显著降低可用的召回容量，使模型更接近阈值。因此，混合架构的脆弱性源于预训练阶段，并会因 CoT 微调而进一步加剧。

附录 I：路由-提取梯度解耦定理的证明

证明。

根据公式 7，我们有：

其中 是对应的 KV 头。然后考虑其 Frobenius 范数：

通过定理 4.6 以及有界范数 、 。

对于提取参数，值向量通过 进入损失函数。累积所有下游依赖关系：

现在我们假设对于所有 和所有 ：

附录 J：训练数据与配置详情

对于 HypeNet，我们同时进行了预训练和 SFT。

预训练数据。由于我们主要关注长上下文召回和数学推理能力，在预训练阶段，我们混合使用了 FineWeb-Edu [lozhkov2024fineweb-edu] 和 UltraData-Math [ultradata-math]。根据近期预训练研究的观察结果——即保持网络语料库为主要来源，同时增加结构化推理数据 [allal2025smollm2smolgoesbig]——我们构建了一个包含通用网络数据和数学专注数据的混合数据集。先前的证据表明，网络规模的语料库能保留广泛的语言和事实能力，而专门的数学语料库则能提升推理和 STEM 表现。

预训练配置。我们对 HypeNet 采用三阶段训练，保持与 chen2026hybridlinearattentionright 中相同的设置，如表 5、表 6 和表 7 所示。

SFT 数据。为了进一步提升模型的数学推理性能，我们使用 MiroMind-M1 数据集 [li2025miromindm1opensourceadvancementmathematical] 进行 CoT-SFT，该数据集收集自 OpenR1 [openr1]、Open-thoughts [guha2025openthoughts]、Light-R1 [wen2025light] 和 Synthetic-1 [2025synthetic1]。SFT 配置。详细配置如表 8 和表 9 所示。

附录 K 评估细节

在评估 NIAH 时，我们使用 LM Evaluation Harness222https://github.com/EleutherAI/lm-evaluation-harness [eval-harness] 进行官方测试。在评估 MATH500 和 GSM8K 时，我们报告每个问题在 16 次独立生成上的平均 pass@1 以及 Maj@16（即多数投票作为最终预测）。具体来说，对于 MATH500，生成长度为 8,192；对于 GSM8K，我们将其设置为 2,048，因为该任务相对简单。对于 LiveCodeBench，我们报告在 16,384 响应长度下 8 次独立生成的平均 pass@1。为了衡量模型在 Tulu3 上微调后的性能，我们报告 IFEval 任务 [zhou2023instructionfollowingevaluationlargelanguage] 上的提示词级别严格准确率。

附录 L 纯 Softmax 注意力模型分析

我们的主要实验集中在混合模型上，在本节中，我们探讨 SFT 是否也会降低纯 Softmax 注意力模型的性能。我们选取了 Qwen2.5-3B、Qwen2.5-7B [qwen2025qwen25technicalreport] 和 Mistral-7B-Instruct-v0.3 [jiang2023mistral7b] 进行研究。关于如何识别用于检索的顶层，我们在附录 M 中提供了详细信息。

分析。表 10 显示，对于 Mistral-7B，在较长上下文中，它对 SFT 表现出强烈的敏感性。SFT 导致 NIAH-Single-3 在 64K 处严重退化，从 下降到 ，表明长距离检索能力崩溃。QK-Restore 对这种退化的恢复效果有限，表明预训练模型中的长上下文表示较弱。

对于 Qwen2.5-7B 和 Qwen2.5-3B，长上下文性能在 NIAH-Single-2 和 Single-3 任务上高度稳定。SFT 仅引起微小波动（例如，Qwen2.5-7B 在 64K 处从某值下降至某值），这表明这些模型在预训练期间已经学会了稳健的长上下文检索。因此，QK-Restore 仅带来边际收益，这与能力几乎未损失的接近饱和状态一致。为了探究其底层机制，我们测量了 Qwen2.5-3B 上的注意力梯度衰减，如图 7 所示。与 HypeNet-2B（其全注意力层呈现拟合良好的指数衰减）不同，该模型具有明显的近平坦尾部。我们假设，在纯 softmax 注意力模型中，长上下文召回是通过丰富的全注意力层在结构上得以维持的。

附录 M 纯 Softmax 注意力模型中检索层的识别方法

为了识别哪些 softmax 注意力层主要负责长距离检索，我们在集合上进行了留一法消融研究。对于每一层，我们通过设置来构建一个消融模型。每个消融模型在多个检索基准（Ruler-QA-SQuAD [hsieh2024ruler]、FDA [arora2023language]、SWDE [lockard2019openceres]）上进行评估。层的重要性得分定义为：

其中表示各基准上的平均检索准确率。较大的值表明该层对检索贡献不成比例地大，且其路由权重至关重要；接近零的得分则表示冗余。各层按降序排列，选择前个层作为权重恢复的目标集合。在我们的实验中，我们保持与 HypeNet 中相同。

附录 N QK-Pro 推导的详细过程

双线性约束没有封闭形式的联合解，因此我们通过将其中一个因子固定为其 SFT 前的值来进行线性化，使剩余问题成为可通过拉格朗日乘数法求解的标准约束最小二乘系统。

我们固定，求最小偏差：

然后应用拉格朗日乘数法：

求导，得到：